···
+
from graphlib import TopologicalSorter
+
from pathlib import Path
+
from typing import Any, Generator, Literal
···
+
Order = Literal["arbitrary", "reverse-topological", "topological"]
class CalledProcessError(Exception):
···
+
async def nix_instantiate(attr_path: str) -> Path:
+
out = await check_subprocess_output(
+
stdout=asyncio.subprocess.PIPE,
+
stderr=asyncio.subprocess.PIPE,
+
drv = out.decode("utf-8").strip().split("!", 1)[0]
+
async def nix_query_requisites(drv: Path) -> list[Path]:
+
requisites = await check_subprocess_output(
+
stdout=asyncio.subprocess.PIPE,
+
stderr=asyncio.subprocess.PIPE,
+
for requisite in requisites.decode("utf-8").splitlines()
+
if requisite != drv_str
+
async def attr_instantiation_worker(
+
semaphore: asyncio.Semaphore,
+
eprint(f"Instantiating {attr_path}…")
+
return (await nix_instantiate(attr_path), attr_path)
+
async def requisites_worker(
+
semaphore: asyncio.Semaphore,
+
) -> tuple[Path, list[Path]]:
+
eprint(f"Obtaining requisites for {drv}…")
+
return (drv, await nix_query_requisites(drv))
+
def requisites_to_attrs(
+
drv_attr_paths: dict[Path, str],
+
requisites: list[Path],
+
Converts a set of requisite `.drv`s to a set of attribute paths.
+
Derivations that do not correspond to any of the packages we want to update will be discarded.
+
drv_attr_paths[requisite]
+
for requisite in requisites
+
if requisite in drv_attr_paths
+
def reverse_edges(graph: dict[str, set[str]]) -> dict[str, set[str]]:
+
Flips the edges of a directed graph.
+
reversed_graph: dict[str, set[str]] = {}
+
for dependent, dependencies in graph.items():
+
for dependency in dependencies:
+
reversed_graph.setdefault(dependency, set()).add(dependent)
+
def get_independent_sorter(
+
) -> TopologicalSorter[str]:
+
Returns a sorter which treats all packages as independent,
+
which will allow them to be updated in parallel.
+
attr_deps: dict[str, set[str]] = {
+
package["attrPath"]: set() for package in packages
+
sorter = TopologicalSorter(attr_deps)
+
async def get_topological_sorter(
+
) -> tuple[TopologicalSorter[str], list[dict]]:
+
Returns a sorter which returns packages in topological or reverse topological order,
+
which will ensure a package is updated before or after its dependencies, respectively.
+
semaphore = asyncio.Semaphore(max_workers)
+
attr_instantiation_worker(semaphore, package["attrPath"])
+
for package in packages
+
drv_requisites = await asyncio.gather(
+
*(requisites_worker(semaphore, drv) for drv in drv_attr_paths.keys())
+
drv_attr_paths[drv]: requisites_to_attrs(drv_attr_paths, requisites)
+
for drv, requisites in drv_requisites
+
attr_deps = reverse_edges(attr_deps)
+
# Adjust packages order based on the topological one
+
ordered = list(TopologicalSorter(attr_deps).static_order())
+
packages = sorted(packages, key=lambda package: ordered.index(package["attrPath"]))
+
sorter = TopologicalSorter(attr_deps)
+
return sorter, packages
async def run_update_script(
merge_lock: asyncio.Lock,
···
packages_to_update.task_done()
+
async def populate_queue(
+
attr_packages: dict[str, dict],
+
sorter: TopologicalSorter[str],
+
packages_to_update: asyncio.Queue[dict | None],
+
Keeps populating the queue with packages that can be updated
+
according to ordering requirements. If topological order
+
is used, the packages will appear in waves, as packages with
+
no dependencies are processed and removed from the sorter.
+
With `order="none"`, all packages will be enqueued simultaneously.
+
# Fill up an update queue,
+
while sorter.is_active():
+
ready_packages = list(sorter.get_ready())
+
eprint(f"Enqueuing group of {len(ready_packages)} packages")
+
for package in ready_packages:
+
await packages_to_update.put(attr_packages[package])
+
await packages_to_update.join()
+
sorter.done(*ready_packages)
+
# Add sentinels, one for each worker.
+
# A worker will terminate when it gets a sentinel from the queue.
+
for i in range(num_workers):
+
await packages_to_update.put(None)
+
attr_packages: dict[str, dict],
+
sorter: TopologicalSorter[str],
merge_lock = asyncio.Lock()
packages_to_update: asyncio.Queue[dict | None] = asyncio.Queue()
···
temp_dirs: list[tuple[str, str] | None] = []
# Do not create more workers than there are packages.
+
num_workers = min(max_workers, len(attr_packages))
nixpkgs_root_output = await check_subprocess_output(
···
temp_dir = stack.enter_context(make_worktree()) if commit else None
temp_dirs.append(temp_dir)
+
queue_task = populate_queue(
# Prepare updater workers for each temp_dir directory.
# At most `num_workers` instances of `run_update_script` will be running at one time.
···
···
with open(packages_path) as f:
+
if order != "arbitrary":
+
eprint("Sorting packages…")
+
reverse_order = order == "reverse-topological"
+
sorter, packages = await get_topological_sorter(
+
sorter = get_independent_sorter(packages)
+
attr_packages = {package["attrPath"]: package for package in packages}
eprint("Going to be running update for following packages:")
···
eprint("Running update for:")
+
await start_updates(max_workers, keep_going, commit, attr_packages, sorter)
eprint("Packages updated!")
···
help="Commit the changes",
+
choices=["arbitrary", "reverse-topological", "topological"],
+
help="Sort the packages based on dependency relation",
help="JSON file containing the list of package names and their update scripts",
···
except KeyboardInterrupt as e:
pyrox.dev