refactor: reduce duplication in gather_transitive_closure util

aspect-build · Dec 12, 2023 · 81af4c2 · 81af4c2
1 parent f076364
commit 81af4c2
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 25 deletions.
diff --git a/npm/private/test/transitive_closure_tests.bzl b/npm/private/test/transitive_closure_tests.bzl
@@ -40,26 +40,16 @@ TEST_PACKAGES = {
 # buildifier: disable=function-docstring
 def test_walk_deps(ctx):
     env = unittest.begin(ctx)
-    closure = {}
     no_optional = True
     not_no_optional = False
 
-    # Test empty case
-    gather_transitive_closure({}, no_optional, {}, closure)
-    asserts.equals(env, 0, len(closure))
-
     # Walk the example tree above
-    gather_transitive_closure(TEST_PACKAGES, not_no_optional, {
-        "@aspect-test/a": "5.0.0",
-    }, closure)
+    closure = gather_transitive_closure(TEST_PACKAGES, "@aspect-test/a/5.0.0", not_no_optional)
     expected = {"@aspect-test/a": ["5.0.0"], "@aspect-test/b": ["5.0.0"], "@aspect-test/c": ["2.0.0", "1.0.0"], "@aspect-test/d": ["2.0.0_@aspect-test+c@1.0.0"]}
     asserts.equals(env, expected, closure)
 
     # Run again with no_optional set, this means we shouldn't walk the dep from @aspect-test/b/5.0.0 -> @aspect-test/c/2.0.0
-    closure = {}
-    gather_transitive_closure(TEST_PACKAGES, no_optional, {
-        "@aspect-test/a": "5.0.0",
-    }, closure)
+    closure = gather_transitive_closure(TEST_PACKAGES, "@aspect-test/a/5.0.0", no_optional)
     expected = {"@aspect-test/a": ["5.0.0"], "@aspect-test/b": ["5.0.0"], "@aspect-test/c": ["1.0.0"], "@aspect-test/d": ["2.0.0_@aspect-test+c@1.0.0"]}
     asserts.equals(env, expected, closure)
 

diff --git a/npm/private/transitive_closure.bzl b/npm/private/transitive_closure.bzl
@@ -3,7 +3,7 @@
 load("@bazel_skylib//lib:dicts.bzl", "dicts")
 load(":utils.bzl", "utils")
 
-def gather_transitive_closure(packages, no_optional, direct_deps, transitive_closure):
+def gather_transitive_closure(packages, package, no_optional):
     """Walk the dependency tree, collecting the transitive closure of dependencies and their versions.
 
     This is needed to resolve npm dependency cycles.
@@ -12,11 +12,18 @@ def gather_transitive_closure(packages, no_optional, direct_deps, transitive_clo
 
     Args:
         packages: dictionary from pnpm lock
+        package: the package to collect deps for
         no_optional: whether to exclude optionalDependencies
-        direct_deps: the immediate dependencies of a given package
-        transitive_closure: a dictionary which is mutated as the return value
+
+    Returns:
+        A dictionary of transitive dependencies, mapping package names to dependent versions.
     """
-    stack = [direct_deps]
+    root_package = packages[package]
+
+    transitive_closure = {}
+    transitive_closure[root_package["name"]] = [root_package["version"]]
+
+    stack = [_get_package_info_deps(root_package, no_optional)]
     iteration_max = 999999
     for i in range(0, iteration_max + 1):
         if not len(stack):
@@ -46,7 +53,12 @@ def gather_transitive_closure(packages, no_optional, direct_deps, transitive_clo
                 continue
             else:
                 package_info = packages[package_key]
-            stack.append(package_info["dependencies"] if no_optional else dicts.add(package_info["dependencies"], package_info["optional_dependencies"]))
+            stack.append(_get_package_info_deps(package_info, no_optional))
+
+    return transitive_closure
+
+def _get_package_info_deps(package_info, no_optional):
+    return package_info["dependencies"] if no_optional else dicts.add(package_info["dependencies"], package_info["optional_dependencies"])
 
 def _gather_package_info(package_path, package_snapshot):
     if package_path.startswith("/"):
@@ -166,17 +178,11 @@ def translate_to_transitive_closure(lock_importers, lock_packages, prod = False,
     # Collect transitive dependencies for each package
     for package in packages.keys():
         package_info = packages[package]
-        transitive_closure = {}
-        transitive_closure[package_info["name"]] = [package_info["version"]]
-        dependencies = package_info["dependencies"] if no_optional else dicts.add(package_info["dependencies"], package_info["optional_dependencies"])
 
-        gather_transitive_closure(
+        package_info["transitive_closure"] = gather_transitive_closure(
             packages,
+            package,
             no_optional,
-            dependencies,
-            transitive_closure,
         )
 
-        package_info["transitive_closure"] = transitive_closure
-
     return (importers, packages)