diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 2c4fd490..33d2f5c2 100644 --- a/CHANGELOG.rst +++ b/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) ----------------- diff --git a/rust/src/graph/direct_import_queries.rs b/rust/src/graph/direct_import_queries.rs index b3534d23..1785a717 100644 --- a/rust/src/graph/direct_import_queries.rs +++ b/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 = self + .find_matching_modules(importer_expression) + .tokens() + .collect(); + let importeds: FxHashSet = 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()); + } } diff --git a/rust/src/graph/hierarchy_queries.rs b/rust/src/graph/hierarchy_queries.rs index 08379e44..443e2997 100644 --- a/rust/src/graph/hierarchy_queries.rs +++ b/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"]), })] diff --git a/tests/unit/application/graph/test_direct_imports.py b/tests/unit/application/graph/test_direct_imports.py index d7d724a3..71270363 100644 --- a/tests/unit/application/graph/test_direct_imports.py +++ b/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(