diff --git a/pkgs/clan-cli/clan_cli/tests/test_vars.py b/pkgs/clan-cli/clan_cli/tests/test_vars.py index 21b431eb9..98482e17e 100644 --- a/pkgs/clan-cli/clan_cli/tests/test_vars.py +++ b/pkgs/clan-cli/clan_cli/tests/test_vars.py @@ -10,6 +10,7 @@ from clan_cli.tests.helpers import cli from clan_cli.vars.check import check_vars from clan_cli.vars.generate import ( Generator, + GeneratorKey, create_machine_vars_interactive, get_generators, run_generators, @@ -53,39 +54,53 @@ def test_dependencies_as_files(temp_dir: Path) -> None: def test_required_generators() -> None: - gen_1 = Generator(name="gen_1", dependencies=[]) - gen_2 = Generator(name="gen_2", dependencies=["gen_1"]) - gen_2a = Generator(name="gen_2a", dependencies=["gen_2"]) - gen_2b = Generator(name="gen_2b", dependencies=["gen_2"]) + # Create generators with proper machine context + machine_name = "test_machine" + gen_1 = Generator(name="gen_1", dependencies=[], machine=machine_name) + gen_2 = Generator( + name="gen_2", + dependencies=[gen_1.key], + machine=machine_name, + ) + gen_2a = Generator( + name="gen_2a", + dependencies=[gen_2.key], + machine=machine_name, + ) + gen_2b = Generator( + name="gen_2b", + dependencies=[gen_2.key], + machine=machine_name, + ) gen_1.exists = True gen_2.exists = False gen_2a.exists = False gen_2b.exists = True - generators = { - generator.name: generator for generator in [gen_1, gen_2, gen_2a, gen_2b] + generators: dict[GeneratorKey, Generator] = { + generator.key: generator for generator in [gen_1, gen_2, gen_2a, gen_2b] } def generator_names(generator: list[Generator]) -> list[str]: return [gen.name for gen in generator] - assert generator_names(requested_closure(["gen_1"], generators)) == [ + assert generator_names(requested_closure([gen_1.key], generators)) == [ "gen_1", "gen_2", "gen_2a", "gen_2b", ] - assert generator_names(requested_closure(["gen_2"], generators)) == [ + assert generator_names(requested_closure([gen_2.key], generators)) == [ "gen_2", "gen_2a", "gen_2b", ] - assert generator_names(requested_closure(["gen_2a"], generators)) == [ + assert generator_names(requested_closure([gen_2a.key], generators)) == [ "gen_2", "gen_2a", "gen_2b", ] - assert generator_names(requested_closure(["gen_2b"], generators)) == [ + assert generator_names(requested_closure([gen_2b.key], generators)) == [ "gen_2", "gen_2a", "gen_2b", diff --git a/pkgs/clan-cli/clan_cli/vars/generate.py b/pkgs/clan-cli/clan_cli/vars/generate.py index eb91dd58c..db358c9ad 100644 --- a/pkgs/clan-cli/clan_cli/vars/generate.py +++ b/pkgs/clan-cli/clan_cli/vars/generate.py @@ -26,10 +26,7 @@ from clan_lib.machines.list import list_full_machines from clan_lib.nix import nix_config, nix_shell, nix_test_store from .check import check_vars -from .graph import ( - minimal_closure, - requested_closure, -) +from .graph import minimal_closure, requested_closure from .prompt import Prompt, ask from .var import Var @@ -40,19 +37,34 @@ if TYPE_CHECKING: from clan_lib.machines.machines import Machine +@dataclass(frozen=True) +class GeneratorKey: + """A key uniquely identifying a generator within a clan.""" + + machine: str | None + name: str + + @dataclass class Generator: name: str files: list[Var] = field(default_factory=list) share: bool = False prompts: list[Prompt] = field(default_factory=list) - dependencies: list[str] = field(default_factory=list) + dependencies: list[GeneratorKey] = field(default_factory=list) migrate_fact: str | None = None machine: str | None = None _flake: "Flake | None" = None + @property + def key(self) -> GeneratorKey: + return GeneratorKey(machine=self.machine, name=self.name) + + def __hash__(self) -> int: + return hash(self.key) + @cached_property def exists(self) -> bool: assert self.machine is not None @@ -124,7 +136,10 @@ class Generator: name=gen_name, share=gen_data["share"], files=files, - dependencies=gen_data["dependencies"], + dependencies=[ + GeneratorKey(machine=machine_name, name=dep) + for dep in gen_data["dependencies"] + ], migrate_fact=gen_data.get("migrateFact"), prompts=prompts, machine=machine_name, @@ -217,22 +232,27 @@ def decrypt_dependencies( result: dict[str, dict[str, bytes]] = {} - for generator_name in set(generator.dependencies): - result[generator_name] = {} + for dep_key in set(generator.dependencies): + # For now, we only support dependencies from the same machine + if dep_key.machine != machine.name: + msg = f"Cross-machine dependencies are not supported. Generator {generator.name} depends on {dep_key.name} from machine {dep_key.machine}" + raise ClanError(msg) - dep_generator = next(g for g in generators if g.name == generator_name) + result[dep_key.name] = {} + + dep_generator = next((g for g in generators if g.name == dep_key.name), None) if dep_generator is None: - msg = f"Generator {generator_name} not found in machine {machine.name}" + msg = f"Generator {dep_key.name} not found in machine {machine.name}" raise ClanError(msg) dep_files = dep_generator.files for file in dep_files: if file.secret: - result[generator_name][file.name] = secret_vars_store.get( + result[dep_key.name][file.name] = secret_vars_store.get( dep_generator, file.name ) else: - result[generator_name][file.name] = public_vars_store.get( + result[dep_key.name][file.name] = public_vars_store.get( dep_generator, file.name ) return result @@ -411,8 +431,8 @@ def _get_closure( from . import graph vars_generators = Generator.get_machine_generators(machine.name, machine.flake) - generators: dict[str, Generator] = { - generator.name: generator for generator in vars_generators + generators: dict[GeneratorKey, Generator] = { + generator.key: generator for generator in vars_generators } result_closure = [] @@ -423,9 +443,11 @@ def _get_closure( result_closure = graph.all_missing_closure(generators) # specific generator selected elif full_closure: - result_closure = requested_closure([generator_name], generators) + gen_key = GeneratorKey(machine=machine.name, name=generator_name) + result_closure = requested_closure([gen_key], generators) else: - result_closure = minimal_closure([generator_name], generators) + gen_key = GeneratorKey(machine=machine.name, name=generator_name) + result_closure = minimal_closure([gen_key], generators) if include_previous_values: for generator in result_closure: diff --git a/pkgs/clan-cli/clan_cli/vars/graph.py b/pkgs/clan-cli/clan_cli/vars/graph.py index 32c9591c4..626db646c 100644 --- a/pkgs/clan-cli/clan_cli/vars/graph.py +++ b/pkgs/clan-cli/clan_cli/vars/graph.py @@ -7,7 +7,7 @@ from typing import TYPE_CHECKING from clan_lib.errors import ClanError if TYPE_CHECKING: - from .generate import Generator + from .generate import Generator, GeneratorKey class GeneratorNotFoundError(ClanError): @@ -15,71 +15,78 @@ class GeneratorNotFoundError(ClanError): def missing_dependency_closure( - requested_generators: Iterable[str], generators: dict[str, Generator] -) -> set[str]: + requested_generators: Iterable[GeneratorKey], + generators: dict[GeneratorKey, Generator], +) -> set[GeneratorKey]: closure = set(requested_generators) # extend the graph to include all dependencies which are not on disk dep_closure = set() queue = list(closure) while queue: - gen_name = queue.pop(0) + gen_key = queue.pop(0) - if gen_name not in generators: - msg = f"Requested generator {gen_name} not found" + if gen_key not in generators: + msg = f"Requested generator {gen_key.name} not found" raise GeneratorNotFoundError(msg) - for dep in generators[gen_name].dependencies: - if dep not in closure and not generators[dep].exists: - dep_closure.add(dep) - queue.append(dep) + for dep_key in generators[gen_key].dependencies: + if ( + dep_key not in closure + and dep_key in generators + and not generators[dep_key].exists + ): + dep_closure.add(dep_key) + queue.append(dep_key) return dep_closure def add_missing_dependencies( - requested_generators: Iterable[str], generators: dict -) -> set[str]: + requested_generators: Iterable[GeneratorKey], + generators: dict[GeneratorKey, Generator], +) -> set[GeneratorKey]: closure = set(requested_generators) return missing_dependency_closure(closure, generators) | closure def add_dependents( - requested_generators: Iterable[str], generators: dict[str, Generator] -) -> set[str]: + requested_generators: Iterable[GeneratorKey], + generators: dict[GeneratorKey, Generator], +) -> set[GeneratorKey]: closure = set(requested_generators) # build reverse dependency graph (graph of dependents) - dependents_graph: dict[str, set[str]] = {} - for gen_name, gen in generators.items(): - for dep in gen.dependencies: - if dep not in dependents_graph: - dependents_graph[dep] = set() - dependents_graph[dep].add(gen_name) + dependents_graph: dict[GeneratorKey, set[GeneratorKey]] = {} + for gen_key, gen in generators.items(): + for dep_key in gen.dependencies: + if dep_key not in dependents_graph: + dependents_graph[dep_key] = set() + dependents_graph[dep_key].add(gen_key) # extend the graph to include all dependents of the current closure queue = list(closure) while queue: - gen_name = queue.pop(0) - for dep in dependents_graph.get(gen_name, []): - if dep not in closure: - closure.add(dep) - queue.append(dep) + gen_key = queue.pop(0) + for dep_key in dependents_graph.get(gen_key, []): + if dep_key not in closure: + closure.add(dep_key) + queue.append(dep_key) return closure def toposort_closure( - _closure: Iterable[str], generators: dict[str, Generator] + _closure: Iterable[GeneratorKey], generators: dict[GeneratorKey, Generator] ) -> list[Generator]: closure = set(_closure) # return the topological sorted list of generators to execute final_dep_graph = {} - for gen_name in sorted(closure): - deps = set(generators[gen_name].dependencies) & closure - final_dep_graph[gen_name] = deps + for gen_key in sorted(closure, key=lambda k: (k.machine or "", k.name)): + deps = set(generators[gen_key].dependencies) & closure + final_dep_graph[gen_key] = deps sorter = TopologicalSorter(final_dep_graph) result = list(sorter.static_order()) - return [generators[gen_name] for gen_name in result] + return [generators[gen_key] for gen_key in result] # all generators in topological order -def full_closure(generators: dict[str, Generator]) -> list[Generator]: +def full_closure(generators: dict[GeneratorKey, Generator]) -> list[Generator]: """ From a set of generators, return all generators in topological order. This includes all dependencies and dependents of the generators. @@ -89,7 +96,7 @@ def full_closure(generators: dict[str, Generator]) -> list[Generator]: # just the missing generators including their dependents -def all_missing_closure(generators: dict[str, Generator]) -> list[Generator]: +def all_missing_closure(generators: dict[GeneratorKey, Generator]) -> list[Generator]: """ From a set of generators, return all incomplete generators in topological order. @@ -97,14 +104,14 @@ def all_missing_closure(generators: dict[str, Generator]) -> list[Generator]: : A generator is missing if at least one of its files is missing. """ # collect all generators that are missing from disk - closure = {gen_name for gen_name, gen in generators.items() if not gen.exists} + closure = {gen_key for gen_key, gen in generators.items() if not gen.exists} closure = add_dependents(closure, generators) return toposort_closure(closure, generators) # only a selected list of generators including their missing dependencies and their dependents def requested_closure( - requested_generators: list[str], generators: dict[str, Generator] + requested_generators: list[GeneratorKey], generators: dict[GeneratorKey, Generator] ) -> list[Generator]: closure = set(requested_generators) # extend the graph to include all dependencies which are not on disk @@ -116,12 +123,12 @@ def requested_closure( # just enough to ensure that the list of selected generators are in a consistent state. # empty if nothing is missing. def minimal_closure( - requested_generators: list[str], generators: dict[str, Generator] + requested_generators: list[GeneratorKey], generators: dict[GeneratorKey, Generator] ) -> list[Generator]: closure = set(requested_generators) final_closure = missing_dependency_closure(closure, generators) # add requested generators if not already exist - for gen_name in closure: - if not generators[gen_name].exists: - final_closure.add(gen_name) + for gen_key in closure: + if not generators[gen_key].exists: + final_closure.add(gen_key) return toposort_closure(final_closure, generators)