nixos-rebuild-ng: refactors (#368619)

Author: thiagokokada

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

@@ -43,19 +43,19 @@ def get_parser() -> tuple[argparse.ArgumentParser, dict[str, argparse.ArgumentPa
     common_build_flags.add_argument("--print-build-logs", "-L", action="store_true")
     common_build_flags.add_argument("--show-trace", action="store_true")
 
-    flake_build_flags = argparse.ArgumentParser(add_help=False)
-    flake_build_flags.add_argument("--accept-flake-config", action="store_true")
-    flake_build_flags.add_argument("--refresh", action="store_true")
-    flake_build_flags.add_argument("--impure", action="store_true")
-    flake_build_flags.add_argument("--offline", action="store_true")
-    flake_build_flags.add_argument("--no-net", action="store_true")
-    flake_build_flags.add_argument("--recreate-lock-file", action="store_true")
-    flake_build_flags.add_argument("--no-update-lock-file", action="store_true")
-    flake_build_flags.add_argument("--no-write-lock-file", action="store_true")
-    flake_build_flags.add_argument("--no-registries", action="store_true")
-    flake_build_flags.add_argument("--commit-lock-file", action="store_true")
-    flake_build_flags.add_argument("--update-input")
-    flake_build_flags.add_argument("--override-input", nargs=2)
+    flake_common_flags = argparse.ArgumentParser(add_help=False)
+    flake_common_flags.add_argument("--accept-flake-config", action="store_true")
+    flake_common_flags.add_argument("--refresh", action="store_true")
+    flake_common_flags.add_argument("--impure", action="store_true")
+    flake_common_flags.add_argument("--offline", action="store_true")
+    flake_common_flags.add_argument("--no-net", action="store_true")
+    flake_common_flags.add_argument("--recreate-lock-file", action="store_true")
+    flake_common_flags.add_argument("--no-update-lock-file", action="store_true")
+    flake_common_flags.add_argument("--no-write-lock-file", action="store_true")
+    flake_common_flags.add_argument("--no-registries", action="store_true")
+    flake_common_flags.add_argument("--commit-lock-file", action="store_true")
+    flake_common_flags.add_argument("--update-input")
+    flake_common_flags.add_argument("--override-input", nargs=2)
 
     classic_build_flags = argparse.ArgumentParser(add_help=False)
     classic_build_flags.add_argument("--no-build-output", "-Q", action="store_true")
@@ -74,7 +74,7 @@ def get_parser() -> tuple[argparse.ArgumentParser, dict[str, argparse.ArgumentPa
     sub_parsers = {
         "common_flags": common_flags,
         "common_build_flags": common_build_flags,
-        "flake_build_flags": flake_build_flags,
+        "flake_common_flags": flake_common_flags,
         "classic_build_flags": classic_build_flags,
         "copy_flags": copy_flags,
     }
@@ -262,19 +262,27 @@ def parser_warn(msg: str) -> None:
 def reexec(
     argv: list[str],
     args: argparse.Namespace,
-    build_flags: dict[str, Args],
-    flake_build_flags: dict[str, Args],
+    build_flags: Args,
+    flake_build_flags: Args,
 ) -> None:
     drv = None
     attr = "config.system.build.nixos-rebuild"
     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(attr, flake, **flake_build_flags, no_link=True)
+            drv = nix.build_flake(
+                attr,
+                flake,
+                flake_build_flags | {"no_link": True},
+            )
         else:
             build_attr = BuildAttr.from_arg(args.attr, args.file)
-            drv = nix.build(attr, build_attr, **build_flags, no_out_link=True)
+            drv = nix.build(
+                attr,
+                build_attr,
+                build_flags | {"no_out_link": True},
+            )
     except CalledProcessError:
         logger.warning(
             "could not build a newer version of nixos-rebuild, "
@@ -319,7 +327,8 @@ def execute(argv: list[str]) -> None:
     common_flags = vars(args_groups["common_flags"])
     common_build_flags = common_flags | vars(args_groups["common_build_flags"])
     build_flags = common_build_flags | vars(args_groups["classic_build_flags"])
-    flake_build_flags = common_build_flags | vars(args_groups["flake_build_flags"])
+    flake_common_flags = common_flags | vars(args_groups["flake_common_flags"])
+    flake_build_flags = common_build_flags | flake_common_flags
     copy_flags = common_flags | vars(args_groups["copy_flags"])
 
     if args.upgrade or args.upgrade_all:
@@ -350,7 +359,7 @@ def execute(argv: list[str]) -> None:
     flake = Flake.from_arg(args.flake, target_host)
 
     if can_run and not flake:
-        nixpkgs_path = nix.find_file("nixpkgs", **build_flags)
+        nixpkgs_path = nix.find_file("nixpkgs", build_flags)
         rev = nix.get_nixpkgs_rev(nixpkgs_path)
         if nixpkgs_path and rev:
             (nixpkgs_path / ".version-suffix").write_text(rev)
@@ -370,7 +379,10 @@ def execute(argv: list[str]) -> None:
 
             dry_run = action == Action.DRY_BUILD
             no_link = action in (Action.SWITCH, Action.BOOT)
+            build_flags |= {"no_out_link": no_link, "dry_run": dry_run}
+            flake_build_flags |= {"no_link": no_link, "dry_run": dry_run}
             rollback = bool(args.rollback)
+
             match action:
                 case Action.BUILD_VM:
                     attr = "config.system.build.vm"
@@ -395,24 +407,22 @@ def execute(argv: list[str]) -> None:
                 case (_, True, _, _):
                     raise NRError(f"--rollback is incompatible with '{action}'")
                 case (_, False, Remote(_), Flake(_)):
-                    path_to_config = nix.remote_build_flake(
+                    path_to_config = nix.build_remote_flake(
                         attr,
                         flake,
                         build_host,
+                        eval_flags=flake_common_flags,
                         flake_build_flags=flake_build_flags,
                         copy_flags=copy_flags,
-                        build_flags=build_flags,
                     )
                 case (_, False, None, Flake(_)):
                     path_to_config = nix.build_flake(
                         attr,
                         flake,
-                        no_link=no_link,
-                        dry_run=dry_run,
-                        **flake_build_flags,
+                        flake_build_flags=flake_build_flags,
                     )
                 case (_, False, Remote(_), None):
-                    path_to_config = nix.remote_build(
+                    path_to_config = nix.build_remote(
                         attr,
                         build_attr,
                         build_host,
@@ -424,9 +434,7 @@ def execute(argv: list[str]) -> None:
                     path_to_config = nix.build(
                         attr,
                         build_attr,
-                        no_out_link=no_link,
-                        dry_run=dry_run,
-                        **build_flags,
+                        build_flags=build_flags,
                     )
                 case never:
                     # should never happen, but mypy is not smart enough to
@@ -442,7 +450,7 @@ def execute(argv: list[str]) -> None:
                     path_to_config,
                     to_host=target_host,
                     from_host=build_host,
-                    **copy_flags,
+                    copy_flags=copy_flags,
                 )
                 if action in (Action.SWITCH, Action.BOOT):
                     nix.set_profile(
@@ -468,7 +476,7 @@ def execute(argv: list[str]) -> None:
                     f"Done. The virtual machine can be started by running '{vm_path}'"
                 )
         case Action.EDIT:
-            nix.edit(flake, **flake_build_flags)
+            nix.edit(flake, flake_build_flags)
         case Action.DRY_RUN:
             assert False, "DRY_RUN should be a DRY_BUILD alias"
         case Action.LIST_GENERATIONS:
@@ -488,9 +496,9 @@ def execute(argv: list[str]) -> None:
                 print(tabulate(generations, headers=headers))
         case Action.REPL:
             if flake:
-                nix.repl_flake("toplevel", flake, **flake_build_flags)
+                nix.repl_flake("toplevel", flake, flake_build_flags)
             else:
-                nix.repl("system", build_attr, **build_flags)
+                nix.repl("system", build_attr, build_flags)
         case _:
             assert_never(action)
 

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

@@ -32,7 +32,7 @@
 def build(
     attr: str,
     build_attr: BuildAttr,
-    **build_flags: Args,
+    build_flags: Args | None = None,
 ) -> Path:
     """Build NixOS attribute using classic Nix.
 
@@ -52,7 +52,7 @@ def build(
 def build_flake(
     attr: str,
     flake: Flake,
-    **flake_build_flags: Args,
+    flake_build_flags: Args | None = None,
 ) -> Path:
     """Build NixOS attribute using Flakes.
 
@@ -70,13 +70,13 @@ def build_flake(
     return Path(r.stdout.strip())
 
 
-def remote_build(
+def build_remote(
     attr: str,
     build_attr: BuildAttr,
     build_host: Remote | None,
-    build_flags: dict[str, Args] | None = None,
-    instantiate_flags: dict[str, Args] | None = None,
-    copy_flags: dict[str, Args] | None = None,
+    build_flags: Args | None = None,
+    instantiate_flags: Args | None = None,
+    copy_flags: Args | None = None,
 ) -> Path:
     # We need to use `--add-root` otherwise Nix will print this warning:
     # > warning: you did not specify '--add-root'; the result might be removed
@@ -89,12 +89,12 @@ def remote_build(
             build_attr.to_attr(attr),
             "--add-root",
             tmpdir.TMPDIR_PATH / uuid4().hex,
-            *dict_to_flags(instantiate_flags or {}),
+            *dict_to_flags(instantiate_flags),
         ],
         stdout=PIPE,
     )
     drv = Path(r.stdout.strip()).resolve()
-    copy_closure(drv, to_host=build_host, from_host=None, **(copy_flags or {}))
+    copy_closure(drv, to_host=build_host, from_host=None, copy_flags=copy_flags)
 
     # Need a temporary directory in remote to use in `nix-store --add-root`
     r = run_wrapper(
@@ -109,7 +109,7 @@ def remote_build(
                 drv,
                 "--add-root",
                 remote_tmpdir / uuid4().hex,
-                *dict_to_flags(build_flags or {}),
+                *dict_to_flags(build_flags),
             ],
             remote=build_host,
             stdout=PIPE,
@@ -124,13 +124,13 @@ def remote_build(
         run_wrapper(["rm", "-rf", remote_tmpdir], remote=build_host, check=False)
 
 
-def remote_build_flake(
+def build_remote_flake(
     attr: str,
     flake: Flake,
     build_host: Remote,
-    flake_build_flags: dict[str, Args] | None = None,
-    copy_flags: dict[str, Args] | None = None,
-    build_flags: dict[str, Args] | None = None,
+    eval_flags: Args | None = None,
+    copy_flags: Args | None = None,
+    flake_build_flags: Args | None = None,
 ) -> Path:
     r = run_wrapper(
         [
@@ -139,20 +139,20 @@ def remote_build_flake(
             "eval",
             "--raw",
             flake.to_attr(attr, "drvPath"),
-            *dict_to_flags(flake_build_flags or {}),
+            *dict_to_flags(eval_flags),
         ],
         stdout=PIPE,
     )
     drv = Path(r.stdout.strip())
-    copy_closure(drv, to_host=build_host, from_host=None, **(copy_flags or {}))
+    copy_closure(drv, to_host=build_host, from_host=None, copy_flags=copy_flags)
     r = run_wrapper(
         [
             "nix",
             *FLAKE_FLAGS,
             "build",
             f"{drv}^*",
             "--print-out-paths",
-            *dict_to_flags(build_flags or {}),
+            *dict_to_flags(flake_build_flags),
         ],
         remote=build_host,
         stdout=PIPE,
@@ -164,7 +164,7 @@ def copy_closure(
     closure: Path,
     to_host: Remote | None,
     from_host: Remote | None = None,
-    **copy_flags: Args,
+    copy_flags: Args | None = None,
 ) -> None:
     """Copy a nix closure to or from host to localhost.
 
@@ -192,6 +192,7 @@ def nix_copy(to_host: Remote, from_host: Remote) -> None:
             [
                 "nix",
                 "copy",
+                *dict_to_flags(copy_flags),
                 "--from",
                 f"ssh://{from_host.host}",
                 "--to",
@@ -221,7 +222,7 @@ def nix_copy(to_host: Remote, from_host: Remote) -> None:
                 nix_copy_closure(to_host, to=True)
 
 
-def edit(flake: Flake | None, **flake_flags: Args) -> None:
+def edit(flake: Flake | None, flake_flags: Args | None = None) -> None:
     "Try to find and open NixOS configuration file in editor."
     if flake:
         run_wrapper(
@@ -250,7 +251,7 @@ def edit(flake: Flake | None, **flake_flags: Args) -> None:
             raise NRError("cannot find NixOS config file")
 
 
-def find_file(file: str, **nix_flags: Args) -> Path | None:
+def find_file(file: str, nix_flags: Args | None = None) -> Path | None:
     "Find classic Nix file location."
     r = run_wrapper(
         ["nix-instantiate", "--find-file", file, *dict_to_flags(nix_flags)],
@@ -420,14 +421,14 @@ def get_generation_info(generation: Generation) -> GenerationJson:
         )
 
 
-def repl(attr: str, build_attr: BuildAttr, **nix_flags: Args) -> None:
+def repl(attr: str, build_attr: BuildAttr, nix_flags: Args | None = None) -> None:
     run_args = ["nix", "repl", "--file", build_attr.path]
     if build_attr.attr:
         run_args.append(build_attr.attr)
     run_wrapper([*run_args, *dict_to_flags(nix_flags)])
 
 
-def repl_flake(attr: str, flake: Flake, **flake_flags: Args) -> None:
+def repl_flake(attr: str, flake: Flake, flake_flags: Args | None = None) -> None:
     expr = Template(
         files(__package__).joinpath(FLAKE_REPL_TEMPLATE).read_text()
     ).substitute(

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

@@ -1,8 +1,9 @@
 import logging
 from collections.abc import Mapping, Sequence
-from typing import Any, TypeAlias, assert_never, override
+from typing import Any, assert_never, override
 
-Args: TypeAlias = bool | str | list[str] | int | None
+type Arg = bool | str | list[str] | int | None
+type Args = dict[str, Arg]
 
 
 class LogFormatter(logging.Formatter):
@@ -19,7 +20,10 @@ def format(self, record: logging.LogRecord) -> str:
         return formatter.format(record)
 
 
-def dict_to_flags(d: Mapping[str, Args]) -> list[str]:
+def dict_to_flags(d: Args | None) -> list[str]:
+    if not d:
+        return []
+
     flags = []
     for key, value in d.items():
         flag = f"--{'-'.join(key.split('_'))}"

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

@@ -123,8 +123,8 @@ def run_side_effect(args: list[str], **kwargs: Any) -> CompletedProcess[str]:
                     "<nixpkgs/nixos>",
                     "--attr",
                     "config.system.build.toplevel",
-                    "--no-out-link",
                     "-vvv",
+                    "--no-out-link",
                 ],
                 check=True,
                 stdout=PIPE,
@@ -188,8 +188,8 @@ def run_side_effect(args: list[str], **kwargs: Any) -> CompletedProcess[str]:
                     "build",
                     "--print-out-paths",
                     "/path/to/config#nixosConfigurations.hostname.config.system.build.toplevel",
-                    "--no-link",
                     "-v",
+                    "--no-link",
                 ],
                 check=True,
                 stdout=PIPE,
@@ -371,6 +371,7 @@ def run_side_effect(args: list[str], **kwargs: Any) -> CompletedProcess[str]:
                     "build",
                     f"'{config_path}^*'",
                     "--print-out-paths",
+                    "--no-link",
                 ],
                 check=True,
                 stdout=PIPE,