Add as_package parameter to direct import methods (retry branch)

CHANGELOG.rst

@@ -5,9 +5,10 @@ Changelog
 latest
 ------
 
-* Speed up ``ImportGraph.find_matching_modules`` for expressions that don't contain wildcards.
-* Speed up ``ImportGraph.find_matching_direct_imports`` by only looking at the imports of the modules that match
-  one of the expressions.
+* Speed up ``find_matching_modules`` for expressions that don't contain wildcards.
+* Speed up ``find_matching_direct_imports``.
+* Add ``as_package`` parameter to ``find_modules_directly_imported_by`` and
+  ``find_modules_that_directly_import``.
 
 3.19 (2026-11-06)
 -----------------

docs/usage.rst

@@ -145,15 +145,23 @@ Methods for analysing direct imports
     :return: Whether or not the importer directly imports the imported module.
     :rtype: ``True`` or ``False``.
 
-.. py:function:: ImportGraph.find_modules_directly_imported_by(module)
+.. py:function:: ImportGraph.find_modules_directly_imported_by(module, as_package=False)
 
     :param str module: A module name.
+    :param bool as_package: Whether or not to treat the supplied module as an individual module,
+                           or as a package (i.e. including any descendants, if there are any). If
+                           treating it as a package, the result will include modules *external* to the
+                           package, and won't include modules within it.
     :return: Set of all modules in the graph are imported by the supplied module.
     :rtype: A set of strings.
 
-.. py:function:: ImportGraph.find_modules_that_directly_import(module)
+.. py:function:: ImportGraph.find_modules_that_directly_import(module, as_package=False)
 
     :param str module: A module name.
+    :param bool as_package: Whether or not to treat the supplied module as an individual module,
+                           or as a package (i.e. including any descendants, if there are any). If
+                           treating it as a package, the result will include modules *external* to the
+                           package, and won't include modules within it.
     :return: Set of all modules in the graph that directly import the supplied module.
     :rtype: A set of strings.
 

rust/src/graph/direct_import_queries.rs

@@ -67,14 +67,58 @@ impl Graph {
         Ok(all_imports)
     }
 
-    pub fn modules_directly_imported_by(&self, importer: ModuleToken) -> &FxHashSet<ModuleToken> {
-        self.imports.get(importer).unwrap_or(&EMPTY_MODULE_TOKENS)
+    /// Find the modules directly imported by the module.
+    ///
+    /// If `as_package` is true, the module is treated as a package: the result is the modules
+    /// directly imported by it or any of its descendants, excluding modules within the package.
+    pub fn modules_directly_imported_by(
+        &self,
+        importer: ModuleToken,
+        as_package: bool,
+    ) -> FxHashSet<ModuleToken> {
+        let mut importers: FxHashSet<_> = importer.into();
+        if as_package {
+            importers.extend_with_descendants(self);
+        }
+
+        let mut imported_modules: FxHashSet<_> = importers
+            .iter()
+            .flat_map(|importer| self.imports.get(*importer).unwrap_or(&EMPTY_MODULE_TOKENS))
+            .cloned()
+            .collect();
+        if as_package {
+            imported_modules.retain(|imported| !importers.contains(imported));
+        }
+        imported_modules
     }
 
-    pub fn modules_that_directly_import(&self, imported: ModuleToken) -> &FxHashSet<ModuleToken> {
-        self.reverse_imports
-            .get(imported)
-            .unwrap_or(&EMPTY_MODULE_TOKENS)
+    /// Find the modules that directly import the module.
+    ///
+    /// If `as_package` is true, the module is treated as a package: the result is the modules
+    /// that directly import it or any of its descendants, excluding modules within the package.
+    pub fn modules_that_directly_import(
+        &self,
+        imported: ModuleToken,
+        as_package: bool,
+    ) -> FxHashSet<ModuleToken> {
+        let mut importeds: FxHashSet<_> = imported.into();
+        if as_package {
+            importeds.extend_with_descendants(self);
+        }
+
+        let mut importers: FxHashSet<_> = importeds
+            .iter()
+            .flat_map(|imported| {
+                self.reverse_imports
+                    .get(*imported)
+                    .unwrap_or(&EMPTY_MODULE_TOKENS)
+            })
+            .cloned()
+            .collect();
+        if as_package {
+            importers.retain(|importer| !importeds.contains(importer));
+        }
+        importers
     }
 
     pub fn get_import_details(
@@ -105,10 +149,10 @@ impl Graph {
         importers
             .iter()
             .flat_map(|importer| {
-                self.modules_directly_imported_by(*importer)
-                    .iter()
-                    .filter(|imported| importeds.contains(*imported))
-                    .map(move |imported| (*importer, *imported))
+                self.modules_directly_imported_by(*importer, false)
+                    .into_iter()
+                    .filter(|imported| importeds.contains(imported))
+                    .map(move |imported| (*importer, imported))
             })
             .collect()
     }
@@ -118,6 +162,76 @@ impl Graph {
 mod test {
     use super::*;
 
+    #[test]
+    fn test_modules_directly_imported_by() {
+        let mut graph = Graph::default();
+
+        let animals = graph.get_or_add_module("pkg.animals").token;
+        let dog = graph.get_or_add_module("pkg.animals.dog").token;
+        let cat = graph.get_or_add_module("pkg.animals.cat").token;
+        let food = graph.get_or_add_module("pkg.food").token;
+        let chicken = graph.get_or_add_module("pkg.food.chicken").token;
+        let tesco = graph.get_or_add_module("pkg.shops.tesco").token;
+
+        graph.add_import(dog, chicken);
+        graph.add_import(cat, food);
+        // An import within the package.
+        graph.add_import(dog, cat);
+        // An import of the package, not by it.
+        graph.add_import(tesco, animals);
+
+        // Without as_package, only the module itself is looked at.
+        assert_eq!(
+            graph.modules_directly_imported_by(animals, false),
+            FxHashSet::default()
+        );
+        assert_eq!(
+            graph.modules_directly_imported_by(dog, false),
+            FxHashSet::from_iter([chicken, cat])
+        );
+        // With as_package, descendants are looked at too, but modules within the package
+        // are excluded.
+        assert_eq!(
+            graph.modules_directly_imported_by(animals, true),
+            FxHashSet::from_iter([chicken, food])
+        );
+    }
+
+    #[test]
+    fn test_modules_that_directly_import() {
+        let mut graph = Graph::default();
+
+        let animals = graph.get_or_add_module("pkg.animals").token;
+        let dog = graph.get_or_add_module("pkg.animals.dog").token;
+        let cat = graph.get_or_add_module("pkg.animals.cat").token;
+        let chicken = graph.get_or_add_module("pkg.food.chicken").token;
+        let shops = graph.get_or_add_module("pkg.shops").token;
+        let tesco = graph.get_or_add_module("pkg.shops.tesco").token;
+
+        graph.add_import(tesco, animals);
+        graph.add_import(shops, cat);
+        // An import within the package.
+        graph.add_import(dog, cat);
+        // An import by the package, not of it.
+        graph.add_import(dog, chicken);
+
+        // Without as_package, only the module itself is looked at.
+        assert_eq!(
+            graph.modules_that_directly_import(animals, false),
+            FxHashSet::from_iter([tesco])
+        );
+        assert_eq!(
+            graph.modules_that_directly_import(cat, false),
+            FxHashSet::from_iter([shops, dog])
+        );
+        // With as_package, descendants are looked at too, but modules within the package
+        // are excluded.
+        assert_eq!(
+            graph.modules_that_directly_import(animals, true),
+            FxHashSet::from_iter([tesco, shops])
+        );
+    }
+
     #[test]
     fn test_find_matching_direct_imports() {
         let mut graph = Graph::default();

rust/src/graph/graph_manipulation.rs

@@ -91,10 +91,10 @@ impl Graph {
         self.module_children.remove(module);
 
         // Update imports.
-        for imported in self.modules_directly_imported_by(module).clone() {
+        for imported in self.modules_directly_imported_by(module, false) {
             self.remove_import(module, imported);
         }
-        for importer in self.modules_that_directly_import(module).clone() {
+        for importer in self.modules_that_directly_import(module, false) {
             self.remove_import(importer, module);
         }
         self.imports.remove(module);
@@ -164,19 +164,11 @@ impl Graph {
 
         let modules_imported_by_descendants: FxHashSet<_> = descendants
             .iter()
-            .flat_map(|descendant| {
-                self.modules_directly_imported_by(*descendant)
-                    .iter()
-                    .cloned()
-            })
+            .flat_map(|descendant| self.modules_directly_imported_by(*descendant, false))
             .collect();
         let modules_that_import_descendants: FxHashSet<_> = descendants
             .iter()
-            .flat_map(|descendant| {
-                self.modules_that_directly_import(*descendant)
-                    .iter()
-                    .cloned()
-            })
+            .flat_map(|descendant| self.modules_that_directly_import(*descendant, false))
             .collect();
 
         // Add descendants and imports to parent module.

rust/src/graph/mod.rs

@@ -394,23 +394,33 @@ impl GraphWrapper {
             .direct_import_exists(importer, imported, as_packages)?)
     }
 
-    pub fn find_modules_directly_imported_by(&self, module: &str) -> PyResult<HashSet<String>> {
+    #[pyo3(signature = (module, as_package=false))]
+    pub fn find_modules_directly_imported_by(
+        &self,
+        module: &str,
+        as_package: bool,
+    ) -> PyResult<HashSet<String>> {
         let module = self.get_visible_module_by_name(module)?.token();
         Ok(self
             ._graph
-            .modules_directly_imported_by(module)
+            .modules_directly_imported_by(module, as_package)
             .iter()
             .into_module_iterator(&self._graph)
             .visible()
             .names()
             .collect())
     }
 
-    pub fn find_modules_that_directly_import(&self, module: &str) -> PyResult<HashSet<String>> {
+    #[pyo3(signature = (module, as_package=false))]
+    pub fn find_modules_that_directly_import(
+        &self,
+        module: &str,
+        as_package: bool,
+    ) -> PyResult<HashSet<String>> {
         let module = self.get_visible_module_by_name(module)?.token();
         Ok(self
             ._graph
-            .modules_that_directly_import(module)
+            .modules_that_directly_import(module, as_package)
             .iter()
             .into_module_iterator(&self._graph)
             .visible()

src/grimp/application/graph.py

@@ -223,27 +223,41 @@ def direct_import_exists(
             importer=importer, imported=imported, as_packages=as_packages
         )
 
-    def find_modules_directly_imported_by(self, module: str) -> set[str]:
+    def find_modules_directly_imported_by(self, module: str, as_package: bool = False) -> set[str]:
         """
         Find all modules that are directly imported by the supplied module.
 
+        Args:
+            module:     The name of the importing module.
+            as_package: Whether or not to treat the supplied module as an individual module,
+                        or as a package (i.e. including any descendants, if there are any). If
+                        treating it as a package, the result will include modules *external* to
+                        the package, and won't include modules within it.
+
         If the module is not present in the graph, grimp.exceptions.ModuleNotPresent will be raised.
         """
         if not self._rustgraph.contains_module(module):
             msg = f'"{module}" not present in the graph.'
             raise ModuleNotPresent(msg)
-        return self._rustgraph.find_modules_directly_imported_by(module)
+        return self._rustgraph.find_modules_directly_imported_by(module, as_package)
 
-    def find_modules_that_directly_import(self, module: str) -> set[str]:
+    def find_modules_that_directly_import(self, module: str, as_package: bool = False) -> set[str]:
         """
         Find all modules that directly import the supplied module.
 
+        Args:
+            module:     The name of the imported module.
+            as_package: Whether or not to treat the supplied module as an individual module,
+                        or as a package (i.e. including any descendants, if there are any). If
+                        treating it as a package, the result will include modules *external* to
+                        the package, and won't include modules within it.
+
         If the module is not present in the graph, grimp.exceptions.ModuleNotPresent will be raised.
         """
         if not self._rustgraph.contains_module(module):
             msg = f'"{module}" not present in the graph.'
             raise ModuleNotPresent(msg)
-        return self._rustgraph.find_modules_that_directly_import(module)
+        return self._rustgraph.find_modules_that_directly_import(module, as_package)
 
     def get_import_details(self, *, importer: str, imported: str) -> list[DetailedImport]:
         """

tests/unit/application/graph/test_direct_imports.py

@@ -34,20 +34,62 @@ def test_find_modules_that_directly_import():
     assert {a, f} == graph.find_modules_that_directly_import("bar")
 
 
-def test_find_modules_directly_imported_by_raises_module_not_present():
+def test_find_modules_directly_imported_by_as_package():
+    graph = ImportGraph()
+    graph.add_import(importer="foo", imported="bar")
+    graph.add_import(importer="foo.one", imported="baz")
+    graph.add_import(importer="foo.one.alpha", imported="baz.one")
+    # Imports within the package.
+    graph.add_import(importer="foo.one", imported="foo.two")
+    graph.add_import(importer="foo.two", imported="foo")
+    # Imports of the package are not imports by the package.
+    graph.add_import(importer="qux", imported="foo")
+
+    assert {"bar"} == graph.find_modules_directly_imported_by("foo")
+    assert {"bar", "baz", "baz.one"} == graph.find_modules_directly_imported_by(
+        "foo", as_package=True
+    )
+    assert {"baz", "baz.one", "foo.two"} == graph.find_modules_directly_imported_by(
+        "foo.one", as_package=True
+    )
+
+
+def test_find_modules_that_directly_import_as_package():
+    graph = ImportGraph()
+    graph.add_import(importer="bar", imported="foo")
+    graph.add_import(importer="baz", imported="foo.one")
+    graph.add_import(importer="baz.one", imported="foo.one.alpha")
+    # Imports within the package.
+    graph.add_import(importer="foo.two", imported="foo.one")
+    graph.add_import(importer="foo", imported="foo.two")
+    # Imports by the package are not imports of the package.
+    graph.add_import(importer="foo", imported="qux")
+
+    assert {"bar"} == graph.find_modules_that_directly_import("foo")
+    assert {"bar", "baz", "baz.one"} == graph.find_modules_that_directly_import(
+        "foo", as_package=True
+    )
+    assert {"baz", "baz.one", "foo.two"} == graph.find_modules_that_directly_import(
+        "foo.one", as_package=True
+    )
+
+
+@pytest.mark.parametrize("as_package", (False, True))
+def test_find_modules_directly_imported_by_raises_module_not_present(as_package):
     graph = ImportGraph()
     graph.add_module("foo")
 
     with pytest.raises(ModuleNotPresent, match=re.escape('"bar" not present in the graph.')):
-        graph.find_modules_directly_imported_by("bar")
+        graph.find_modules_directly_imported_by("bar", as_package=as_package)
 
 
-def test_find_modules_that_directly_import_raises_module_not_present():
+@pytest.mark.parametrize("as_package", (False, True))
+def test_find_modules_that_directly_import_raises_module_not_present(as_package):
     graph = ImportGraph()
     graph.add_module("foo")
 
     with pytest.raises(ModuleNotPresent, match=re.escape('"bar" not present in the graph.')):
-        graph.find_modules_that_directly_import("bar")
+        graph.find_modules_that_directly_import("bar", as_package=as_package)
 
 
 @pytest.mark.parametrize(