nixos-rebuild-ng: improve developer experience (#356468)

pkgs/by-name/ni/nixos-rebuild-ng/README.md

···
 
 And use `nixos-rebuild-ng` instead of `nixos-rebuild`.
 
+## Development
+
+Run:
+
+```console
+nix-build -A nixos-rebuild-ng.tests.ci
+```
+
+The command above will run the unit tests and linters, and also check if the
+code is formatted. However, sometimes is more convenient to run just a few
+tests to debug, in this case you can run:
+
+```console
+nix-shell -A nixos-rebuild-ng.devShell
+```
+
+The command above should automatically put you inside `src` directory, and you
+can run:
+
+```console
+# run program
+python -m nixos_rebuild
+# run tests
+python -m pytest
+# check types
+mypy .
+# fix lint issues
+ruff check --fix .
+# format code
+ruff format .
+```
+
 ## Current caveats
 
 - For now we will install it in `nixos-rebuild-ng` path by default, to avoid
···
 - [ ] Improve documentation
 - [ ] `nixos-rebuild repl` (calling old `nixos-rebuild` for now)
 - [ ] `nix` build/bootstrap
-- [ ] Reduce build closure
+- [ ] Generate tab completion via [`shtab`](https://docs.iterative.ai/shtab/)
+- [x] Reduce build closure
+
+## TODON'T
+
+- Reimplement `systemd-run` logic (will be moved to the new
+  [`apply`](https://github.com/NixOS/nixpkgs/pull/344407) script)

pkgs/by-name/ni/nixos-rebuild-ng/package.nix

···
 {
   lib,
   installShellFiles,
+  mkShell,
   nix,
   nixos-rebuild,
   python3,
+  python3Packages,
+  runCommand,
   withNgSuffix ? true,
 }:
-python3.pkgs.buildPythonApplication {
+python3Packages.buildPythonApplication rec {
   pname = "nixos-rebuild-ng";
   version = "0.0.0";
   src = ./src;
   pyproject = true;
 
-  build-system = with python3.pkgs; [
+  build-system = with python3Packages; [
     setuptools
   ];
 
-  dependencies = with python3.pkgs; [
+  dependencies = with python3Packages; [
     tabulate
-    types-tabulate
   ];
 
   nativeBuildInputs = [
···
       mv $out/bin/nixos-rebuild $out/bin/nixos-rebuild-ng
     '';
 
-  nativeCheckInputs = with python3.pkgs; [
+  nativeCheckInputs = with python3Packages; [
     pytestCheckHook
-    mypy
-    ruff
   ];
 
   pytestFlagsArray = [ "-vv" ];
 
-  postCheck = ''
-    echo -e "\x1b[32m## run mypy\x1b[0m"
-    mypy nixos_rebuild tests
-    echo -e "\x1b[32m## run ruff\x1b[0m"
-    ruff check nixos_rebuild tests
-    echo -e "\x1b[32m## run ruff format\x1b[0m"
-    ruff format --check nixos_rebuild tests
-  '';
+  passthru =
+    let
+      python-with-pkgs = python3.withPackages (
+        ps: with ps; [
+          mypy
+          pytest
+          ruff
+          types-tabulate
+          # dependencies
+          tabulate
+        ]
+      );
+    in
+    {
+      devShell = mkShell {
+        packages = [ python-with-pkgs ];
+        shellHook = ''
+          cd pkgs/by-name/ni/nixos-rebuild-ng/src || true
+        '';
+      };
+
+      # NOTE: this is a passthru test rather than a build-time test because we
+      # want to keep the build closures small
+      tests.ci = runCommand "${pname}-ci" { nativeBuildInputs = [ python-with-pkgs ]; } ''
+        export RUFF_CACHE_DIR="$(mktemp -d)"
+
+        echo -e "\x1b[32m## run mypy\x1b[0m"
+        mypy ${src}
+        echo -e "\x1b[32m## run ruff\x1b[0m"
+        ruff check ${src}
+        echo -e "\x1b[32m## run ruff format\x1b[0m"
+        ruff format --check ${src}
+
+        touch $out
+      '';
+    };
 
   meta = {
     description = "Rebuild your NixOS configuration and switch to it, on local hosts and remote";

pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__init__.py

···
 from subprocess import run
 from typing import assert_never
 
-from tabulate import tabulate
-
 from .models import Action, Flake, NRError, Profile
 from .nix import (
     edit,
+    find_file,
+    get_nixpkgs_rev,
     list_generations,
     nixos_build,
     nixos_build_flake,
···
     if args.upgrade or args.upgrade_all:
         upgrade_channels(bool(args.upgrade_all))
 
-    match action := Action(args.action):
+    action = Action(args.action)
+    # Only run shell scripts from the Nixpkgs tree if the action is
+    # "switch", "boot", or "test". With other actions (such as "build"),
+    # the user may reasonably expect that no code from the Nixpkgs tree is
+    # executed, so it's safe to run nixos-rebuild against a potentially
+    # untrusted tree.
+    can_run = action in (Action.SWITCH, Action.BOOT, Action.TEST)
+    if can_run and not flake:
+        nixpkgs_path = find_file("nixpkgs", nix_flags)
+        rev = get_nixpkgs_rev(nixpkgs_path)
+        if nixpkgs_path and rev:
+            (nixpkgs_path / ".version-suffix").write_text(rev)
+
+    match action:
         case Action.SWITCH | Action.BOOT:
             info("building the system configuration...")
             if args.rollback:
···
             if args.json:
                 print(json.dumps(generations, indent=2))
             else:
+                from tabulate import tabulate
+
                 headers = {
                     "generation": "Generation",
                     "date": "Build-date",
···
             raise ex
         else:
             sys.exit(str(ex))
-
-
-if __name__ == "__main__":
-    main()

pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/__main__.py

···
+from . import main
+
+if __name__ == "__main__":
+    main()

pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/models.py

···
 class Flake:
     path: Path
     attr: str
-    _re: ClassVar[re.Pattern[str]] = re.compile(
-        r"^(?P<path>[^\#]*)\#?(?P<attr>[^\#\"]*)$"
-    )
+    _re: ClassVar = re.compile(r"^(?P<path>[^\#]*)\#?(?P<attr>[^\#\"]*)$")
 
     @override
     def __str__(self) -> str:
···
         m = cls._re.match(flake_str)
         assert m is not None, f"got no matches for {flake_str}"
         attr = m.group("attr")
-        if not attr:
-            attr = f"nixosConfigurations.{hostname or "default"}"
-        else:
-            attr = f"nixosConfigurations.{attr}"
-        return Flake(Path(m.group("path")), attr)
+        nixos_attr = f"nixosConfigurations.{attr or hostname or "default"}"
+        return Flake(Path(m.group("path")), nixos_attr)
 
     @classmethod
     def from_arg(cls, flake_arg: Any) -> Flake | None:

pkgs/by-name/ni/nixos-rebuild-ng/src/nixos_rebuild/nix.py

···
 from __future__ import annotations
 
 import os
+import shutil
 from datetime import datetime
 from pathlib import Path
 from subprocess import PIPE, CalledProcessError, run
···
     NRError,
     Profile,
 )
-from .utils import dict_to_flags
+from .utils import dict_to_flags, info
 
 FLAKE_FLAGS: Final = ["--extra-experimental-features", "nix-command flakes"]
 
···
             run([os.getenv("EDITOR", "nano"), nixos_config], check=False)
         else:
             raise NRError("cannot find NixOS config file")
+
+
+def find_file(file: str, nix_flags: list[str] | None = None) -> Path | None:
+    "Find classic Nixpkgs location."
+    r = run(
+        ["nix-instantiate", "--find-file", file, *(nix_flags or [])],
+        stdout=PIPE,
+        check=False,
+        text=True,
+    )
+    if r.returncode:
+        return None
+    return Path(r.stdout.strip())
+
+
+def get_nixpkgs_rev(nixpkgs_path: Path | None) -> str | None:
+    """Get Nixpkgs path as a Git revision.
+
+    Can be used to generate `.version-suffix` file."""
+    if not nixpkgs_path:
+        return None
+
+    # Git is not included in the closure for nixos-rebuild so we need to check
+    if not shutil.which("git"):
+        info(f"warning: Git not found; cannot figure out revision of '{nixpkgs_path}'")
+        return None
+
+    # Get current revision
+    r = run(
+        ["git", "-C", nixpkgs_path, "rev-parse", "--short", "HEAD"],
+        check=False,
+        stdout=PIPE,
+        text=True,
+    )
+    rev = r.stdout.strip()
+
+    if rev:
+        # Check if repo is dirty
+        if run(["git", "-C", nixpkgs_path, "diff", "--quiet"], check=False).returncode:
+            rev += "M"
+        return f".git.{rev}"
+    else:
+        return None
 
 
 def _parse_generation_from_nix_store(path: Path, profile: Profile) -> Generation:

pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_main.py

···
 
 
 @patch(get_qualified_name(nr.nix.run, nr.nix), autospec=True)
-def test_execute_nix_boot(mock_run: Any, tmp_path: Path) -> None:
+@patch(get_qualified_name(nr.nix.shutil.which), autospec=True, return_value="/bin/git")
+def test_execute_nix_boot(mock_which: Any, mock_run: Any, tmp_path: Path) -> None:
+    nixpkgs_path = tmp_path / "nixpkgs"
+    nixpkgs_path.mkdir()
     config_path = tmp_path / "test"
     config_path.touch()
     mock_run.side_effect = [
+        # update_nixpkgs_rev
+        CompletedProcess([], 0, str(nixpkgs_path)),
+        CompletedProcess([], 0, "nixpkgs-rev"),
+        CompletedProcess([], 0),
         # nixos_build
         CompletedProcess([], 0, str(config_path)),
         # set_profile
···
     nr.execute(["nixos-rebuild", "boot", "--no-flake", "-vvv"])
 
     assert nr.VERBOSE is True
-    assert mock_run.call_count == 3
+    assert mock_run.call_count == 6
     mock_run.assert_has_calls(
         [
+            call(
+                ["nix-instantiate", "--find-file", "nixpkgs", "-vvv"],
+                stdout=PIPE,
+                check=False,
+                text=True,
+            ),
+            call(
+                ["git", "-C", nixpkgs_path, "rev-parse", "--short", "HEAD"],
+                check=False,
+                stdout=PIPE,
+                text=True,
+            ),
+            call(
+                ["git", "-C", nixpkgs_path, "diff", "--quiet"],
+                check=False,
+            ),
             call(
                 [
                     "nix-build",
···
 
 
 @patch(get_qualified_name(nr.nix.run, nr.nix), autospec=True)
-def test_execute_switch_rollback(mock_run: Any) -> None:
+def test_execute_switch_rollback(mock_run: Any, tmp_path: Path) -> None:
+    nixpkgs_path = tmp_path / "nixpkgs"
+    nixpkgs_path.touch()
+
     nr.execute(["nixos-rebuild", "switch", "--rollback", "--install-bootloader"])
 
     assert nr.VERBOSE is False
-    assert mock_run.call_count == 2
+    assert mock_run.call_count == 3
+    # ignoring update_nixpkgs_rev calls
     mock_run.assert_has_calls(
         [
             call(

pkgs/by-name/ni/nixos-rebuild-ng/src/tests/test_nix.py

···
         mock_run.assert_called_with(["editor", default_nix], check=False)
 
 
+@patch(get_qualified_name(n.shutil.which), autospec=True, return_value="/bin/git")
+def test_get_nixpkgs_rev(mock_which: Any) -> None:
+    assert n.get_nixpkgs_rev(None) is None
+
+    path = Path("/path/to/nix")
+
+    with patch(
+        get_qualified_name(n.run, n),
+        autospec=True,
+        side_effect=[CompletedProcess([], 0, "")],
+    ) as mock_run:
+        assert n.get_nixpkgs_rev(path) is None
+        mock_run.assert_called_with(
+            ["git", "-C", path, "rev-parse", "--short", "HEAD"],
+            check=False,
+            stdout=PIPE,
+            text=True,
+        )
+
+    expected_calls = [
+        call(
+            ["git", "-C", path, "rev-parse", "--short", "HEAD"],
+            check=False,
+            stdout=PIPE,
+            text=True,
+        ),
+        call(
+            ["git", "-C", path, "diff", "--quiet"],
+            check=False,
+        ),
+    ]
+
+    with patch(
+        get_qualified_name(n.run, n),
+        autospec=True,
+        side_effect=[
+            CompletedProcess([], 0, "0f7c82403fd6"),
+            CompletedProcess([], returncode=0),
+        ],
+    ) as mock_run:
+        assert n.get_nixpkgs_rev(path) == ".git.0f7c82403fd6"
+        mock_run.assert_has_calls(expected_calls)
+
+    with patch(
+        get_qualified_name(n.run, n),
+        autospec=True,
+        side_effect=[
+            CompletedProcess([], 0, "0f7c82403fd6"),
+            CompletedProcess([], returncode=1),
+        ],
+    ) as mock_run:
+        assert n.get_nixpkgs_rev(path) == ".git.0f7c82403fd6M"
+        mock_run.assert_has_calls(expected_calls)
+
+
 def test_get_generations_from_nix_store(tmp_path: Path) -> None:
     nixos_path = tmp_path / "nixos-system"
     nixos_path.mkdir()