nixos-rebuild-ng: refactor part of __init__.py code in services.py

There are a bunch of logic that were happening in `__init__.py` file
that was not really fit in other files in the project but also didn't
seems fit to the `__init__.py`.

While moving this to `services.py` doesn't really solve the problem, at
least it reduces the size of the main `execute()` function to a
reasonable size, and I hope this will make it easier for other people to
understand what is happening.

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

@@ -3,17 +3,15 @@ import json
 import logging
 import os
 import sys
-from pathlib import Path
 from subprocess import CalledProcessError, run
 from typing import Final, assert_never
 
-from . import nix, tmpdir
+from . import nix
 from .constants import EXECUTABLE, WITH_NIX_2_18, WITH_REEXEC, WITH_SHELL_FILES
-from .models import Action, BuildAttr, Flake, ImageVariants, NixOSRebuildError, Profile
+from .models import Action, BuildAttr, Flake, Profile
-from .process import Remote, cleanup_ssh
+from .process import Remote
-from .utils import Args, LogFormatter, tabulate
+from .services import build_and_activate_system, reexec
-
+from .utils import LogFormatter, tabulate
-NIXOS_REBUILD_ATTR: Final = "config.system.build.nixos-rebuild"
 
 logger: Final = logging.getLogger(__name__)
 logger.setLevel(logging.INFO)
@@ -271,72 +269,6 @@ def parse_args(
     return args, args_groups
 
 
-def reexec(
-    argv: list[str],
-    args: argparse.Namespace,
-    build_flags: Args,
-    flake_build_flags: Args,
-) -> None:
-    drv = None
-    try:
-        # Parsing the args here but ignore ask_sudo_password since it is not
-        # needed and we would end up asking sudo password twice
-        if flake := Flake.from_arg(args.flake, Remote.from_arg(args.target_host, None)):
-            drv = nix.build_flake(
-                NIXOS_REBUILD_ATTR,
-                flake,
-                flake_build_flags | {"no_link": True},
-            )
-        else:
-            build_attr = BuildAttr.from_arg(args.attr, args.file)
-            drv = nix.build(
-                NIXOS_REBUILD_ATTR,
-                build_attr,
-                build_flags | {"no_out_link": True},
-            )
-    except CalledProcessError:
-        logger.warning(
-            "could not build a newer version of nixos-rebuild, using current version",
-            exc_info=logger.isEnabledFor(logging.DEBUG),
-        )
-
-    if drv:
-        new = drv / f"bin/{EXECUTABLE}"
-        current = Path(argv[0])
-        if new != current:
-            logger.debug(
-                "detected newer version of script, re-exec'ing, current=%s, new=%s",
-                current,
-                new,
-            )
-            # Manually call clean-up functions since os.execve() will replace
-            # the process immediately
-            cleanup_ssh()
-            tmpdir.TMPDIR.cleanup()
-            try:
-                os.execve(new, argv, os.environ | {"_NIXOS_REBUILD_REEXEC": "1"})
-            except Exception:
-                # Possible errors that we can have here:
-                # - Missing the binary
-                # - Exec format error (e.g.: another OS/CPU arch)
-                logger.warning(
-                    "could not re-exec in a newer version of nixos-rebuild, "
-                    "using current version",
-                    exc_info=logger.isEnabledFor(logging.DEBUG),
-                )
-                # We already run clean-up, let's re-exec in the current version
-                # to avoid issues
-                os.execve(current, argv, os.environ | {"_NIXOS_REBUILD_REEXEC": "1"})
-
-
-def validate_image_variant(image_variant: str, variants: ImageVariants) -> None:
-    if image_variant not in variants:
-        raise NixOSRebuildError(
-            "please specify one of the following supported image variants via "
-            "--image-variant:\n" + "\n".join(f"- {v}" for v in variants)
-        )
-
-
 def execute(argv: list[str]) -> None:
     args, args_groups = parse_args(argv)
 
@@ -395,147 +327,20 @@ def execute(argv: list[str]) -> None:
             | Action.BUILD_VM
             | Action.BUILD_VM_WITH_BOOTLOADER
         ):
-            logger.info("building the system configuration...")
+            build_and_activate_system(
-
+                action=action,
-            dry_run = action == Action.DRY_BUILD
+                args=args,
-            no_link = action in (Action.SWITCH, Action.BOOT)
+                build_host=build_host,
-            rollback = bool(args.rollback)
+                target_host=target_host,
-
+                profile=profile,
-            match action:
+                flake=flake,
-                case Action.BUILD_IMAGE if flake:
+                build_attr=build_attr,
-                    variants = nix.get_build_image_variants_flake(
+                build_flags=build_flags,
-                        flake,
+                common_flags=common_flags,
-                        eval_flags=flake_common_flags,
+                copy_flags=copy_flags,
-                    )
+                flake_build_flags=flake_build_flags,
-                    validate_image_variant(args.image_variant, variants)
+                flake_common_flags=flake_common_flags,
-                    attr = f"config.system.build.images.{args.image_variant}"
+            )
-                case Action.BUILD_IMAGE:
-                    variants = nix.get_build_image_variants(
-                        build_attr,
-                        instantiate_flags=common_flags,
-                    )
-                    validate_image_variant(args.image_variant, variants)
-                    attr = f"config.system.build.images.{args.image_variant}"
-                case Action.BUILD_VM:
-                    attr = "config.system.build.vm"
-                case Action.BUILD_VM_WITH_BOOTLOADER:
-                    attr = "config.system.build.vmWithBootLoader"
-                case _:
-                    attr = "config.system.build.toplevel"
-
-            match (action, rollback, build_host, flake):
-                case (Action.SWITCH | Action.BOOT, True, _, _):
-                    path_to_config = nix.rollback(profile, target_host, sudo=args.sudo)
-                case (Action.TEST | Action.BUILD, True, _, _):
-                    maybe_path_to_config = nix.rollback_temporary_profile(
-                        profile,
-                        target_host,
-                        sudo=args.sudo,
-                    )
-                    if maybe_path_to_config:  # kinda silly but this makes mypy happy
-                        path_to_config = maybe_path_to_config
-                    else:
-                        raise NixOSRebuildError("could not find previous generation")
-                case (_, True, _, _):
-                    raise NixOSRebuildError(
-                        f"--rollback is incompatible with '{action}'"
-                    )
-                case (_, False, Remote(_), Flake(_)):
-                    path_to_config = nix.build_remote_flake(
-                        attr,
-                        flake,
-                        build_host,
-                        eval_flags=flake_common_flags,
-                        flake_build_flags=flake_build_flags
-                        | {"no_link": no_link, "dry_run": dry_run},
-                        copy_flags=copy_flags,
-                    )
-                case (_, False, None, Flake(_)):
-                    path_to_config = nix.build_flake(
-                        attr,
-                        flake,
-                        flake_build_flags=flake_build_flags
-                        | {"no_link": no_link, "dry_run": dry_run},
-                    )
-                case (_, False, Remote(_), None):
-                    path_to_config = nix.build_remote(
-                        attr,
-                        build_attr,
-                        build_host,
-                        realise_flags=common_flags,
-                        instantiate_flags=build_flags,
-                        copy_flags=copy_flags,
-                    )
-                case (_, False, None, None):
-                    path_to_config = nix.build(
-                        attr,
-                        build_attr,
-                        build_flags=build_flags
-                        | {"no_out_link": no_link, "dry_run": dry_run},
-                    )
-                case never:
-                    # should never happen, but mypy is not smart enough to
-                    # handle this with assert_never
-                    # https://github.com/python/mypy/issues/16650
-                    # https://github.com/python/mypy/issues/16722
-                    raise AssertionError(
-                        f"expected code to be unreachable, but got: {never}"
-                    )
-
-            if not rollback:
-                nix.copy_closure(
-                    path_to_config,
-                    to_host=target_host,
-                    from_host=build_host,
-                    copy_flags=copy_flags,
-                )
-                if action in (Action.SWITCH, Action.BOOT):
-                    nix.set_profile(
-                        profile,
-                        path_to_config,
-                        target_host=target_host,
-                        sudo=args.sudo,
-                    )
-
-            # Print only the result to stdout to make it easier to script
-            def print_result(msg: str, result: str | Path) -> None:
-                print(msg, end=" ", file=sys.stderr, flush=True)
-                print(result, flush=True)
-
-            match action:
-                case Action.SWITCH | Action.BOOT | Action.TEST | Action.DRY_ACTIVATE:
-                    nix.switch_to_configuration(
-                        path_to_config,
-                        action,
-                        target_host=target_host,
-                        sudo=args.sudo,
-                        specialisation=args.specialisation,
-                        install_bootloader=args.install_bootloader,
-                    )
-                    print_result("Done. The new configuration is", path_to_config)
-                case Action.BUILD:
-                    print_result("Done. The new configuration is", path_to_config)
-                case Action.BUILD_VM | Action.BUILD_VM_WITH_BOOTLOADER:
-                    # If you get `not-found`, please open an issue
-                    vm_path = next(path_to_config.glob("bin/run-*-vm"), "not-found")
-                    print_result(
-                        "Done. The virtual machine can be started by running", vm_path
-                    )
-                case Action.BUILD_IMAGE:
-                    if flake:
-                        image_name = nix.get_build_image_name_flake(
-                            flake,
-                            args.image_variant,
-                            eval_flags=flake_common_flags,
-                        )
-                    else:
-                        image_name = nix.get_build_image_name(
-                            build_attr,
-                            args.image_variant,
-                            instantiate_flags=flake_common_flags,
-                        )
-                    disk_path = path_to_config / image_name
-                    print_result("Done. The disk image can be found in", disk_path)
 
         case Action.EDIT:
             nix.edit(flake, flake_build_flags)

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

@@ -0,0 +1,319 @@
+import argparse
+import logging
+import os
+import sys
+from pathlib import Path
+from subprocess import CalledProcessError
+from typing import Final
+
+from . import nix, tmpdir
+from .constants import EXECUTABLE
+from .models import Action, BuildAttr, Flake, ImageVariants, NixOSRebuildError, Profile
+from .process import Remote, cleanup_ssh
+from .utils import Args
+
+NIXOS_REBUILD_ATTR: Final = "config.system.build.nixos-rebuild"
+
+logger: Final = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+def reexec(
+    argv: list[str],
+    args: argparse.Namespace,
+    build_flags: Args,
+    flake_build_flags: Args,
+) -> None:
+    drv = None
+    try:
+        # Parsing the args here but ignore ask_sudo_password since it is not
+        # needed and we would end up asking sudo password twice
+        if flake := Flake.from_arg(args.flake, Remote.from_arg(args.target_host, None)):
+            drv = nix.build_flake(
+                NIXOS_REBUILD_ATTR,
+                flake,
+                flake_build_flags | {"no_link": True},
+            )
+        else:
+            build_attr = BuildAttr.from_arg(args.attr, args.file)
+            drv = nix.build(
+                NIXOS_REBUILD_ATTR,
+                build_attr,
+                build_flags | {"no_out_link": True},
+            )
+    except CalledProcessError:
+        logger.warning(
+            "could not build a newer version of nixos-rebuild, using current version",
+            exc_info=logger.isEnabledFor(logging.DEBUG),
+        )
+
+    if drv:
+        new = drv / f"bin/{EXECUTABLE}"
+        current = Path(argv[0])
+        if new != current:
+            logger.debug(
+                "detected newer version of script, re-exec'ing, current=%s, new=%s",
+                current,
+                new,
+            )
+            # Manually call clean-up functions since os.execve() will replace
+            # the process immediately
+            cleanup_ssh()
+            tmpdir.TMPDIR.cleanup()
+            try:
+                os.execve(new, argv, os.environ | {"_NIXOS_REBUILD_REEXEC": "1"})
+            except Exception:
+                # Possible errors that we can have here:
+                # - Missing the binary
+                # - Exec format error (e.g.: another OS/CPU arch)
+                logger.warning(
+                    "could not re-exec in a newer version of nixos-rebuild, "
+                    "using current version",
+                    exc_info=logger.isEnabledFor(logging.DEBUG),
+                )
+                # We already run clean-up, let's re-exec in the current version
+                # to avoid issues
+                os.execve(current, argv, os.environ | {"_NIXOS_REBUILD_REEXEC": "1"})
+
+
+def _validate_image_variant(image_variant: str, variants: ImageVariants) -> None:
+    if image_variant not in variants:
+        raise NixOSRebuildError(
+            "please specify one of the following supported image variants via "
+            "--image-variant:\n" + "\n".join(f"- {v}" for v in variants)
+        )
+
+
+def _get_system_attr(
+    action: Action,
+    args: argparse.Namespace,
+    flake: Flake | None,
+    build_attr: BuildAttr,
+    common_flags: Args,
+    flake_common_flags: Args,
+) -> str:
+    match action:
+        case Action.BUILD_IMAGE if flake:
+            variants = nix.get_build_image_variants_flake(
+                flake,
+                eval_flags=flake_common_flags,
+            )
+            _validate_image_variant(args.image_variant, variants)
+            attr = f"config.system.build.images.{args.image_variant}"
+        case Action.BUILD_IMAGE:
+            variants = nix.get_build_image_variants(
+                build_attr,
+                instantiate_flags=common_flags,
+            )
+            _validate_image_variant(args.image_variant, variants)
+            attr = f"config.system.build.images.{args.image_variant}"
+        case Action.BUILD_VM:
+            attr = "config.system.build.vm"
+        case Action.BUILD_VM_WITH_BOOTLOADER:
+            attr = "config.system.build.vmWithBootLoader"
+        case _:
+            attr = "config.system.build.toplevel"
+
+    return attr
+
+
+def _build_system(
+    attr: str,
+    action: Action,
+    args: argparse.Namespace,
+    build_host: Remote | None,
+    target_host: Remote | None,
+    profile: Profile,
+    flake: Flake | None,
+    build_attr: BuildAttr,
+    build_flags: Args,
+    common_flags: Args,
+    copy_flags: Args,
+    flake_build_flags: Args,
+    flake_common_flags: Args,
+) -> Path:
+    dry_run = action == Action.DRY_BUILD
+    no_link = action in (Action.SWITCH, Action.BOOT)
+
+    match (action, args.rollback, build_host, flake):
+        case (Action.SWITCH | Action.BOOT, True, _, _):
+            path_to_config = nix.rollback(profile, target_host, sudo=args.sudo)
+        case (Action.TEST | Action.BUILD, True, _, _):
+            maybe_path_to_config = nix.rollback_temporary_profile(
+                profile,
+                target_host,
+                sudo=args.sudo,
+            )
+            if maybe_path_to_config:  # kinda silly but this makes mypy happy
+                path_to_config = maybe_path_to_config
+            else:
+                raise NixOSRebuildError("could not find previous generation")
+        case (_, True, _, _):
+            raise NixOSRebuildError(f"--rollback is incompatible with '{action}'")
+        case (_, False, Remote(_), Flake(_)):
+            path_to_config = nix.build_remote_flake(
+                attr,
+                flake,
+                build_host,
+                eval_flags=flake_common_flags,
+                flake_build_flags=flake_build_flags
+                | {"no_link": no_link, "dry_run": dry_run},
+                copy_flags=copy_flags,
+            )
+        case (_, False, None, Flake(_)):
+            path_to_config = nix.build_flake(
+                attr,
+                flake,
+                flake_build_flags=flake_build_flags
+                | {"no_link": no_link, "dry_run": dry_run},
+            )
+        case (_, False, Remote(_), None):
+            path_to_config = nix.build_remote(
+                attr,
+                build_attr,
+                build_host,
+                realise_flags=common_flags,
+                instantiate_flags=build_flags,
+                copy_flags=copy_flags,
+            )
+        case (_, False, None, None):
+            path_to_config = nix.build(
+                attr,
+                build_attr,
+                build_flags=build_flags | {"no_out_link": no_link, "dry_run": dry_run},
+            )
+        case never:
+            # should never happen, but mypy is not smart enough to
+            # handle this with assert_never
+            # https://github.com/python/mypy/issues/16650
+            # https://github.com/python/mypy/issues/16722
+            raise AssertionError(f"expected code to be unreachable, but got: {never}")
+
+    if not args.rollback:
+        nix.copy_closure(
+            path_to_config,
+            to_host=target_host,
+            from_host=build_host,
+            copy_flags=copy_flags,
+        )
+
+    return path_to_config
+
+
+def _activate_system(
+    path_to_config: Path,
+    action: Action,
+    args: argparse.Namespace,
+    build_host: Remote | None,
+    target_host: Remote | None,
+    profile: Profile,
+    flake: Flake | None,
+    build_attr: BuildAttr,
+    flake_common_flags: Args,
+    common_flags: Args,
+) -> None:
+    # Print only the result to stdout to make it easier to script
+    def print_result(msg: str, result: str | Path) -> None:
+        print(msg, end=" ", file=sys.stderr, flush=True)
+        print(result, flush=True)
+
+    match action:
+        case Action.SWITCH | Action.BOOT if not args.rollback:
+            nix.set_profile(
+                profile,
+                path_to_config,
+                target_host=target_host,
+                sudo=args.sudo,
+            )
+            nix.switch_to_configuration(
+                path_to_config,
+                action,
+                target_host=target_host,
+                sudo=args.sudo,
+                specialisation=args.specialisation,
+                install_bootloader=args.install_bootloader,
+            )
+            print_result("Done. The new configuration is", path_to_config)
+        case Action.SWITCH | Action.BOOT | Action.TEST | Action.DRY_ACTIVATE:
+            nix.switch_to_configuration(
+                path_to_config,
+                action,
+                target_host=target_host,
+                sudo=args.sudo,
+                specialisation=args.specialisation,
+                install_bootloader=args.install_bootloader,
+            )
+            print_result("Done. The new configuration is", path_to_config)
+        case Action.BUILD:
+            print_result("Done. The new configuration is", path_to_config)
+        case Action.BUILD_VM | Action.BUILD_VM_WITH_BOOTLOADER:
+            # If you get `not-found`, please open an issue
+            vm_path = next(path_to_config.glob("bin/run-*-vm"), "not-found")
+            print_result("Done. The virtual machine can be started by running", vm_path)
+        case Action.BUILD_IMAGE:
+            if flake:
+                image_name = nix.get_build_image_name_flake(
+                    flake,
+                    args.image_variant,
+                    eval_flags=flake_common_flags,
+                )
+            else:
+                image_name = nix.get_build_image_name(
+                    build_attr,
+                    args.image_variant,
+                    instantiate_flags=common_flags,
+                )
+            disk_path = path_to_config / image_name
+            print_result("Done. The disk image can be found in", disk_path)
+
+
+def build_and_activate_system(
+    action: Action,
+    args: argparse.Namespace,
+    build_host: Remote | None,
+    target_host: Remote | None,
+    profile: Profile,
+    flake: Flake | None,
+    build_attr: BuildAttr,
+    build_flags: Args,
+    common_flags: Args,
+    copy_flags: Args,
+    flake_build_flags: Args,
+    flake_common_flags: Args,
+) -> None:
+    logger.info("building the system configuration...")
+    attr = _get_system_attr(
+        action=action,
+        args=args,
+        flake=flake,
+        build_attr=build_attr,
+        common_flags=common_flags,
+        flake_common_flags=flake_common_flags,
+    )
+    path_to_config = _build_system(
+        attr,
+        action=action,
+        args=args,
+        build_host=build_host,
+        target_host=target_host,
+        profile=profile,
+        flake=flake,
+        build_attr=build_attr,
+        build_flags=build_flags,
+        common_flags=common_flags,
+        copy_flags=copy_flags,
+        flake_build_flags=flake_build_flags,
+        flake_common_flags=flake_common_flags,
+    )
+    _activate_system(
+        path_to_config,
+        action=action,
+        args=args,
+        build_host=build_host,
+        target_host=target_host,
+        profile=profile,
+        flake=flake,
+        build_attr=build_attr,
+        common_flags=common_flags,
+        flake_common_flags=flake_common_flags,
+    )

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

@@ -141,7 +141,7 @@ def test_reexec(mock_build: Mock, mock_execve: Mock, monkeypatch: MonkeyPatch) -
     mock_build.assert_has_calls(
         [
             call(
-                nr.NIXOS_REBUILD_ATTR,
+                nr.services.NIXOS_REBUILD_ATTR,
                 nr.models.BuildAttr(ANY, ANY),
                 {"build": True, "no_out_link": True},
             )
@@ -185,7 +185,7 @@ def test_reexec_flake(
 
     nr.reexec(argv, args, {"build": True}, {"flake": True})
     mock_build.assert_called_once_with(
-        nr.NIXOS_REBUILD_ATTR,
+        nr.services.NIXOS_REBUILD_ATTR,
         nr.models.Flake(ANY, ANY),
         {"flake": True, "no_link": True},
     )
@@ -536,7 +536,7 @@ def test_execute_nix_switch_flake(mock_run: Mock, tmp_path: Path) -> None:
 )
 @patch("subprocess.run", autospec=True)
 @patch("uuid.uuid4", autospec=True)
-@patch(get_qualified_name(nr.cleanup_ssh), autospec=True)
+@patch(get_qualified_name(nr.services.cleanup_ssh), autospec=True)
 @pytest.mark.skipif(
     not WITH_NIX_2_18,
     reason="Tests internal logic based on the assumption that Nix >= 2.18",
@@ -755,7 +755,7 @@ def test_execute_nix_switch_build_target_host(
     clear=True,
 )
 @patch("subprocess.run", autospec=True)
-@patch(get_qualified_name(nr.cleanup_ssh), autospec=True)
+@patch(get_qualified_name(nr.services.cleanup_ssh), autospec=True)
 def test_execute_nix_switch_flake_target_host(
     mock_cleanup_ssh: Mock,
     mock_run: Mock,
@@ -862,7 +862,7 @@ def test_execute_nix_switch_flake_target_host(
     clear=True,
 )
 @patch("subprocess.run", autospec=True)
-@patch(get_qualified_name(nr.cleanup_ssh), autospec=True)
+@patch(get_qualified_name(nr.services.cleanup_ssh), autospec=True)
 def test_execute_nix_switch_flake_build_host(
     mock_cleanup_ssh: Mock,
     mock_run: Mock,