Speed up find_matching_direct_imports

CHANGELOG.rst

@@ -2,6 +2,13 @@
 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.
+
 3.19 (2026-11-06)
 -----------------
 

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;
@@ -93,22 +93,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()
     }
@@ -155,4 +155,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/hierarchy_queries.rs

@@ -54,12 +54,19 @@ impl Graph {
         &self,
         expression: &ModuleExpression,
     ) -> impl ModuleIterator<'_> + use<'_> {
-        let interner = MODULE_NAMES.read().unwrap();
-        let modules: FxHashSet<_> = self
-            .modules
-            .values()
-            .filter(|m| expression.is_match(interner.resolve(m.interned_name).unwrap()))
-            .collect();
+        let modules: FxHashSet<_> = if expression.as_ref().contains('*') {
+            let interner = MODULE_NAMES.read().unwrap();
+            self.modules
+                .values()
+                .filter(|m| expression.is_match(interner.resolve(m.interned_name).unwrap()))
+                .collect()
+        } else {
+            // Without wildcards the expression can only match the module with that name,
+            // so there's no need to check every module.
+            self.get_module_by_name(expression.as_ref())
+                .into_iter()
+                .collect()
+        };
         modules.into_iter()
     }
 }
@@ -80,6 +87,8 @@ mod tests {
 
     #[parameterized(case = {
         FindMatchingModulesTestCase::new("foo", ["foo"]),
+        FindMatchingModulesTestCase::new("foo.bar.baz", ["foo.bar.baz"]),
+        FindMatchingModulesTestCase::new("foo.missing", Vec::<&str>::new()),
         FindMatchingModulesTestCase::new("foo.*", ["foo.bar"]),
         FindMatchingModulesTestCase::new("foo.**", ["foo.bar", "foo.bar.baz"]),
     })]

tests/unit/application/graph/test_direct_imports.py

@@ -344,6 +344,58 @@ def test_finds_matching_direct_imports(self, import_line_number, import_line_con
             {"importer": "pkg.animals.dog", "imported": "pkg.food.chicken"},
         ]
 
+    @pytest.mark.parametrize(
+        "expression, expected_imports",
+        [
+            # Neither side has wildcards.
+            (
+                "pkg.animals.dog -> pkg.food.chicken",
+                [("pkg.animals.dog", "pkg.food.chicken")],
+            ),
+            # Only the importer has no wildcards.
+            (
+                "pkg.animals.dog -> pkg.*.*",
+                [
+                    ("pkg.animals.dog", "pkg.colors.golden"),
+                    ("pkg.animals.dog", "pkg.food.chicken"),
+                ],
+            ),
+            # Only the imported has no wildcards.
+            (
+                "pkg.*.* -> pkg.food.chicken",
+                [("pkg.animals.dog", "pkg.food.chicken"), ("pkg.shops.tesco", "pkg.food.chicken")],
+            ),
+            # Both sides have wildcards.
+            (
+                "pkg.shops.* -> pkg.food.*",
+                [("pkg.shops.tesco", "pkg.food.chicken")],
+            ),
+            # A module that doesn't import anything.
+            ("pkg.shops.coop -> pkg.*.*", []),
+            # A module that isn't imported by anything.
+            ("pkg.*.* -> pkg.shops.coop", []),
+            # Modules that aren't in the graph.
+            ("pkg.animals.missing -> pkg.*.*", []),
+            ("pkg.*.* -> pkg.food.missing", []),
+            ("pkg.animals.missing -> pkg.food.missing", []),
+            # Wildcards that match nothing.
+            ("pkg.nothing.* -> pkg.*.*", []),
+        ],
+    )
+    def test_finds_matching_direct_imports_with_and_without_wildcards(
+        self, expression, expected_imports
+    ):
+        graph = ImportGraph()
+        graph.add_module("pkg.shops.coop")
+        graph.add_import(importer="pkg.animals.dog", imported="pkg.food.chicken")
+        graph.add_import(importer="pkg.animals.dog", imported="pkg.colors.golden")
+        graph.add_import(importer="pkg.animals.cat", imported="pkg.food.fish")
+        graph.add_import(importer="pkg.shops.tesco", imported="pkg.food.chicken")
+
+        assert graph.find_matching_direct_imports(expression) == [
+            {"importer": importer, "imported": imported} for importer, imported in expected_imports
+        ]
+
     def test_deduplicates_imports(self):
         graph = ImportGraph()
         graph.add_import(