Add as_package parameter to direct import methods

CHANGELOG.rst

@@ -5,7 +5,9 @@ Changelog
 latest
 ------
 
-* Speed up ``ImportGraph.find_matching_modules`` for expressions that don't contain wildcards.
+* 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

@@ -1,6 +1,6 @@
 use crate::errors::{GrimpError, GrimpResult};
 use crate::graph::{
-    EMPTY_IMPORT_DETAILS, EMPTY_MODULE_TOKENS, ExtendWithDescendants, Graph, MODULE_NAMES,
+    EMPTY_IMPORT_DETAILS, EMPTY_MODULE_TOKENS, ExtendWithDescendants, Graph, ModuleIterator,
     ModuleToken, PyImportDetails,
 };
 use crate::module_expressions::ModuleExpression;
@@ -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(
@@ -93,22 +137,22 @@ impl Graph {
         importer_expression: &ModuleExpression,
         imported_expression: &ModuleExpression,
     ) -> FxHashSet<(ModuleToken, ModuleToken)> {
-        let interner = MODULE_NAMES.read().unwrap();
-        self.imports
+        let importers: FxHashSet<ModuleToken> = self
+            .find_matching_modules(importer_expression)
+            .tokens()
+            .collect();
+        let importeds: FxHashSet<ModuleToken> = self
+            .find_matching_modules(imported_expression)
+            .tokens()
+            .collect();
+
+        importers
             .iter()
-            .flat_map(|(importer, imports)| {
-                imports
+            .flat_map(|importer| {
+                self.modules_directly_imported_by(*importer)
                     .iter()
-                    .cloned()
-                    .map(move |imported| (importer, imported))
-            })
-            .filter(|(importer, imported)| {
-                let importer = self.get_module(*importer).unwrap();
-                let importer_name = interner.resolve(importer.interned_name).unwrap();
-                let imported = self.get_module(*imported).unwrap();
-                let imported_name = interner.resolve(imported.interned_name).unwrap();
-                importer_expression.is_match(importer_name)
-                    && imported_expression.is_match(imported_name)
+                    .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();
@@ -155,4 +269,60 @@ mod test {
             FxHashSet::from_iter([(dog, chicken), (cat, fish)])
         );
     }
+
+    #[test]
+    fn test_find_matching_direct_imports_with_and_without_wildcards() {
+        let mut graph = Graph::default();
+
+        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 fish = graph.get_or_add_module("pkg.food.fish").token;
+        let golden = graph.get_or_add_module("pkg.colors.golden").token;
+        let tesco = graph.get_or_add_module("pkg.shops.tesco").token;
+        graph.get_or_add_module("pkg.shops.coop");
+
+        graph.add_import(dog, chicken);
+        graph.add_import(dog, golden);
+        graph.add_import(cat, fish);
+        graph.add_import(tesco, chicken);
+
+        let find = |importer_expression: &str, imported_expression: &str| {
+            graph.find_matching_direct_imports(
+                &importer_expression.parse().unwrap(),
+                &imported_expression.parse().unwrap(),
+            )
+        };
+
+        // Neither side has wildcards.
+        assert_eq!(
+            find("pkg.animals.dog", "pkg.food.chicken"),
+            FxHashSet::from_iter([(dog, chicken)])
+        );
+        // Only the importer is a single module.
+        assert_eq!(
+            find("pkg.animals.dog", "pkg.*.*"),
+            FxHashSet::from_iter([(dog, chicken), (dog, golden)])
+        );
+        // Only the imported is a single module.
+        assert_eq!(
+            find("pkg.*.*", "pkg.food.chicken"),
+            FxHashSet::from_iter([(dog, chicken), (tesco, chicken)])
+        );
+        // Both sides have wildcards.
+        assert_eq!(
+            find("pkg.shops.*", "pkg.food.*"),
+            FxHashSet::from_iter([(tesco, chicken)])
+        );
+        // A module with no imports.
+        assert!(find("pkg.shops.coop", "pkg.*.*").is_empty());
+        // A module that has no importers.
+        assert!(find("pkg.*.*", "pkg.shops.coop").is_empty());
+        // Modules that aren't in the graph.
+        assert!(find("pkg.animals.missing", "pkg.*.*").is_empty());
+        assert!(find("pkg.*.*", "pkg.food.missing").is_empty());
+        assert!(find("pkg.animals.missing", "pkg.food.missing").is_empty());
+        // Wildcards that match nothing.
+        assert!(find("pkg.nothing.*", "pkg.*.*").is_empty());
+    }
 }

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()