Add as_packages parameter to find_shortest_chains

AUTHORS.rst

@@ -8,4 +8,5 @@ Authors
 * NetworkX developers - The shortest path algorithm was adapted from the NetworkX library https://networkx.org/.
 * Peter Byfield - https://github.com/Peter554
 * Shane Smiskol - https://github.com/sshane
-* Andreas Rammhold - https://github.com/andir
+* Andreas Rammhold - https://github.com/andir
+* Nicholas Bunn - https://github.com/NicholasBunn

CHANGELOG.rst

@@ -2,6 +2,11 @@
 Changelog
 =========
 
+latest
+------
+
+* Added `as_packages` option to the `find_shortest_chains` method.
+
 3.4.1 (2024-07-12)
 ------------------
 

docs/usage.rst

@@ -226,6 +226,9 @@ Methods for analysing import chains
 
     :param str importer: A module or subpackage within the graph.
     :param str imported: Another module or subpackage within the graph.
+    :param bool as_packages: Whether or not to treat the imported and importer as an individual module,
+                            or as a package (including any descendants, if there are any). If treating them as packages, all descendants
+                            of ``importer`` and ``imported`` will be checked too. Defaults to True.
     :return: The shortest import chains that exist between the ``importer`` and ``imported``, and between any modules
              contained within them. Only one chain per upstream/downstream pair will be included. Any chains that are
              contained within other chains in the result set will be excluded.

src/grimp/adaptors/graph.py

@@ -287,20 +287,29 @@ def find_shortest_chain(self, importer: str, imported: str) -> Optional[Tuple[st
 
         return self._find_shortest_chain(importer=importer, imported=imported)
 
-    def find_shortest_chains(self, importer: str, imported: str) -> Set[Tuple[str, ...]]:
+    def find_shortest_chains(
+        self, importer: str, imported: str, as_packages: bool = True
+    ) -> Set[Tuple[str, ...]]:
         """
         Find the shortest import chains that exist between the importer and imported, and
-        between any modules contained within them. Only one chain per upstream/downstream pair
-        will be included. Any chains that are contained within other chains in the result set
-        will be excluded.
+        between any modules contained within them if as_packages is True. Only one chain per
+        upstream/downstream pair will be included. Any chains that are contained within other
+        chains in the result set will be excluded.
+
+        The default behavior is to treat the import and imported as packages, however, if
+        as_packages is False, both the importer and imported will be treated as modules instead.
 
         Returns:
             A set of tuples of strings. Each tuple is ordered from importer to imported modules.
         """
         shortest_chains = set()
 
-        upstream_modules = self._all_modules_in_package(imported)
-        downstream_modules = self._all_modules_in_package(importer)
+        upstream_modules = (
+            {imported} if not as_packages else self._all_modules_in_package(imported)
+        )
+        downstream_modules = (
+            {importer} if not as_packages else self._all_modules_in_package(importer)
+        )
 
         if upstream_modules & downstream_modules:
             # If there are shared modules between the two, one of the modules is a descendant

src/grimp/application/ports/graph.py

@@ -231,12 +231,19 @@ def find_shortest_chain(self, importer: str, imported: str) -> tuple[str, ...] |
         raise NotImplementedError
 
     @abc.abstractmethod
-    def find_shortest_chains(self, importer: str, imported: str) -> Set[Tuple[str, ...]]:
+    def find_shortest_chains(
+        self, importer: str, imported: str, as_packages: bool = True
+    ) -> Set[Tuple[str, ...]]:
         """
         Find the shortest import chains that exist between the importer and imported, and
-        between any modules contained within them. Only one chain per upstream/downstream pair
-        will be included. Any chains that are contained within other chains in the result set
-        will be excluded.
+        between any modules contained within them if as_packages is True. Only one chain per
+        upstream/downstream pair will be included. Any chains that are contained within other
+        chains in the result set will be excluded.
+
+        The default behavior is to treat the import and imported as packages, however, if
+        as_packages is False, both the importer and imported will be treated as modules instead.
+
+
 
         Returns:
             A set of tuples of strings. Each tuple is ordered from importer to imported modules.

tests/assets/testpackage/three/__init__.py

@@ -0,0 +1 @@
+from ..one import alpha

tests/assets/testpackage/three/alpha.py

@@ -0,0 +1,3 @@
+from ..one import alpha
+
+BAR = alpha.BAR

tests/assets/testpackage/three/beta.py

@@ -0,0 +1,5 @@
+from testpackage.one import alpha
+
+
+def foo():
+    return alpha.BAR

tests/assets/testpackage/three/gamma.py

@@ -0,0 +1,2 @@
+from .beta import foo
+from .. import utils

tests/functional/test_build_and_use_graph.py

@@ -1,4 +1,6 @@
 from grimp import build_graph
+from typing import Set, Tuple, Optional
+import pytest
 
 """
 For ease of reference, these are the imports of all the files:
@@ -15,6 +17,12 @@
 testpackage.two.beta: testpackage.one.alpha
 testpackage.two.gamma: testpackage.two.beta, testpackage.utils
 testpackage.utils: testpackage.one, testpackage.two.alpha
+testpackage.three: testpackage.one.alpha
+testpackage.three.alpha: testpackage.one.alpha
+testpackage.three.beta: testpackage.one.alpha
+testpackage.three.gamma: testpackage.two.beta, testpackage.utils
+
+
 
 """
 
@@ -38,6 +46,10 @@ def test_modules():
         "testpackage.two.beta",
         "testpackage.two.gamma",
         "testpackage.utils",
+        "testpackage.three",
+        "testpackage.three.beta",
+        "testpackage.three.gamma",
+        "testpackage.three.alpha",
     }
 
 
@@ -123,6 +135,9 @@ def test_find_modules_that_directly_import():
         "testpackage.one.beta",
         "testpackage.two.alpha",
         "testpackage.two.beta",
+        "testpackage.three",
+        "testpackage.three.alpha",
+        "testpackage.three.beta",
     } == result
 
 
@@ -189,14 +204,62 @@ def test_find_shortest_chain():
     ) == graph.find_shortest_chain(importer="testpackage.utils", imported="testpackage.one.alpha")
 
 
-def test_find_shortest_chains():
+@pytest.mark.parametrize(
+    "as_packages, expected_result",
+    [
+        pytest.param(
+            True,
+            {
+                ("testpackage.three", "testpackage.one.alpha"),
+                ("testpackage.three.alpha", "testpackage.one.alpha"),
+                ("testpackage.three.beta", "testpackage.one.alpha"),
+                (
+                    "testpackage.three.gamma",
+                    "testpackage.utils",
+                    "testpackage.two.alpha",
+                    "testpackage.one.alpha",
+                ),
+            },
+            id="Treating imports as packages",
+        ),
+        pytest.param(
+            False,
+            {
+                ("testpackage.three", "testpackage.one.alpha"),
+            },
+            id="Treating imports as modules (explicit)",
+        ),
+        pytest.param(
+            None,
+            {
+                ("testpackage.three", "testpackage.one.alpha"),
+                ("testpackage.three.alpha", "testpackage.one.alpha"),
+                ("testpackage.three.beta", "testpackage.one.alpha"),
+                (
+                    "testpackage.three.gamma",
+                    "testpackage.utils",
+                    "testpackage.two.alpha",
+                    "testpackage.one.alpha",
+                ),
+            },
+            id="Treating imports as modules (implicit)",
+        ),
+    ],
+)
+def test_find_shortest_chains(as_packages: Optional[bool], expected_result: Set[Tuple]):
+    importer = "testpackage.three"
+    imported = "testpackage.one.alpha"
+
     graph = build_graph("testpackage", cache_dir=None)
 
-    assert {
-        ("testpackage.two.alpha", "testpackage.one.alpha"),
-        ("testpackage.two.beta", "testpackage.one.alpha"),
-        ("testpackage.two.gamma", "testpackage.utils", "testpackage.one"),
-    } == graph.find_shortest_chains(importer="testpackage.two", imported="testpackage.one")
+    if as_packages is not None:
+        shortest_chain = graph.find_shortest_chains(
+            importer=importer, imported=imported, as_packages=as_packages
+        )
+    else:
+        shortest_chain = graph.find_shortest_chains(importer=importer, imported=imported)
+
+    assert expected_result == shortest_chain
 
 
 class TestFindDownstreamModules:
@@ -212,6 +275,10 @@ def test_as_package_false(self):
             "testpackage.two.beta",
             "testpackage.two.gamma",
             "testpackage.utils",
+            "testpackage.three.beta",
+            "testpackage.three",
+            "testpackage.three.alpha",
+            "testpackage.three.gamma",
         } == result
 
     def test_as_package_true(self):
@@ -224,6 +291,10 @@ def test_as_package_true(self):
             "testpackage.two.beta",
             "testpackage.two.gamma",
             "testpackage.utils",
+            "testpackage.three.beta",
+            "testpackage.three",
+            "testpackage.three.alpha",
+            "testpackage.three.gamma",
         } == result