Add as_packages parameter to find_shortest_chain

CHANGELOG.rst

@@ -2,6 +2,11 @@
 Changelog
 =========
 
+Unreleased
+----------
+
+* Added `as_packages` option to the `find_shortest_chain` method.
+
 3.6 (2025-02-07)
 ----------------
 

docs/usage.rst

@@ -216,15 +216,19 @@ Methods for analysing import chains
     :return: All the modules that are imported (even indirectly) by the supplied module.
     :rtype: A set of strings.
 
-.. py:function:: ImportGraph.find_shortest_chain(importer, imported)
+.. py:function:: ImportGraph.find_shortest_chain(importer, imported, as_packages=False)
 
     :param str importer: The module at the start of a potential chain of imports between ``importer`` and ``imported``
         (i.e. the module that potentially imports ``imported``, even indirectly).
     :param str imported: The module at the end of the potential chain of imports.
+    :param bool as_packages: Whether to treat the supplied modules as individual modules,
+         or as packages (including any descendants, if there are any). If
+         treating them as packages, all descendants of ``importer`` and
+         ``imported`` will be checked too.
     :return: The shortest chain of imports between the supplied modules, or None if no chain exists.
     :rtype: A tuple of strings, ordered from importer to imported modules, or None.
 
-.. py:function:: ImportGraph.find_shortest_chains(importer, imported)
+.. py:function:: ImportGraph.find_shortest_chains(importer, imported, as_packages=True)
 
     :param str importer: A module or subpackage within the graph.
     :param str imported: Another module or subpackage within the graph.

rust/src/lib.rs

@@ -296,17 +296,18 @@ impl GraphWrapper {
             .collect())
     }
 
-    // TODO(peter) Add `as_packages` argument here? The implementation already supports it!
+    #[pyo3(signature = (importer, imported, as_packages=false))]
     pub fn find_shortest_chain(
         &self,
         importer: &str,
         imported: &str,
+        as_packages: bool,
     ) -> PyResult<Option<Vec<String>>> {
         let importer = self.get_visible_module_by_name(importer)?.token();
         let imported = self.get_visible_module_by_name(imported)?.token();
         Ok(self
             ._graph
-            .find_shortest_chain(importer, imported, false)?
+            .find_shortest_chain(importer, imported, as_packages)?
             .map(|chain| {
                 chain
                     .iter()

src/grimp/adaptors/graph.py

@@ -108,12 +108,14 @@ def find_downstream_modules(self, module: str, as_package: bool = False) -> Set[
     def find_upstream_modules(self, module: str, as_package: bool = False) -> Set[str]:
         return self._rustgraph.find_upstream_modules(module, as_package)
 
-    def find_shortest_chain(self, importer: str, imported: str) -> tuple[str, ...] | None:
+    def find_shortest_chain(
+        self, importer: str, imported: str, as_packages: bool = False
+    ) -> tuple[str, ...] | None:
         for module in (importer, imported):
             if not self._rustgraph.contains_module(module):
                 raise ValueError(f"Module {module} is not present in the graph.")
 
-        chain = self._rustgraph.find_shortest_chain(importer, imported)
+        chain = self._rustgraph.find_shortest_chain(importer, imported, as_packages)
         return tuple(chain) if chain else None
 
     def find_shortest_chains(

src/grimp/application/ports/graph.py

@@ -220,11 +220,19 @@ def find_upstream_modules(self, module: str, as_package: bool = False) -> Set[st
         raise NotImplementedError
 
     @abc.abstractmethod
-    def find_shortest_chain(self, importer: str, imported: str) -> tuple[str, ...] | None:
+    def find_shortest_chain(
+        self, importer: str, imported: str, as_packages: bool = False
+    ) -> tuple[str, ...] | None:
         """
         Attempt to find the shortest chain of imports between two modules, in the direction
         of importer to imported.
 
+        Optional args:
+            as_packages: Whether to treat the supplied modules as individual modules,
+                         or as packages (including any descendants, if there are any). If
+                         treating them as subpackages, all descendants of the supplied modules
+                         will be checked too.
+
         Returns:
             Tuple of module names, from importer to imported, or None if no chain exists.
         """

tests/unit/adaptors/graph/test_chains.py

@@ -168,6 +168,25 @@ def test_demonstrate_nondeterminism_of_equal_chains(self):
         other_chain = (source, d, e, f, destination)
         assert (result == one_chain) or (result == other_chain)
 
+    @pytest.mark.parametrize(
+        "as_packages, expected_result",
+        (
+            (False, None),
+            (True, ("green.foo", "blue.bar")),
+        ),
+    )
+    def test_as_packages(self, as_packages: bool, expected_result: Set[Tuple]):
+        graph = ImportGraph()
+        graph.add_module("green")
+        graph.add_module("blue")
+        graph.add_import(importer="green.foo", imported="blue.bar")
+
+        result = graph.find_shortest_chain(
+            importer="green", imported="blue", as_packages=as_packages
+        )
+
+        assert result == expected_result
+
 
 class TestFindShortestChains:
     @pytest.mark.parametrize(