From 7bc1e5cecb3456e20379bba7203e1ce547f53e0b Mon Sep 17 00:00:00 2001 From: Peter Byfield Date: Tue, 10 Nov 2026 15:57:40 +0100 Subject: [PATCH] Speed up find_matching_direct_imports Previously this checked every import in the graph against both expressions, which made each call proportional to the size of the graph. This is slow when it is called many times, as Import Linter does for each of the expressions in a contract's ignore_imports. Instead, find the modules that match each expression, then only look at the imports of the modules that match the importer expression. Assistant-model: Claude Code --- CHANGELOG.rst | 4 +- docs/usage.rst | 12 +- rust/src/graph/direct_import_queries.rs | 212 ++++++++++++++++-- rust/src/graph/graph_manipulation.rs | 16 +- rust/src/graph/mod.rs | 18 +- src/grimp/application/graph.py | 22 +- .../application/graph/test_direct_imports.py | 102 ++++++++- 7 files changed, 338 insertions(+), 48 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ea3e8ab7..d2484c47 100644 --- a/CHANGELOG.rst +++ b/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) ----------------- diff --git a/docs/usage.rst b/docs/usage.rst index df1c516e..447dd96f 100644 --- a/docs/usage.rst +++ b/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. diff --git a/rust/src/graph/direct_import_queries.rs b/rust/src/graph/direct_import_queries.rs index b3534d23..d7a50ceb 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; @@ -67,14 +67,58 @@ impl Graph { Ok(all_imports) } - pub fn modules_directly_imported_by(&self, importer: ModuleToken) -> &FxHashSet { - 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 { + 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 { - 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 { + 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 = 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() } @@ -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()); + } } diff --git a/rust/src/graph/graph_manipulation.rs b/rust/src/graph/graph_manipulation.rs index b6c762c3..a780ab3d 100644 --- a/rust/src/graph/graph_manipulation.rs +++ b/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. diff --git a/rust/src/graph/mod.rs b/rust/src/graph/mod.rs index cf745368..40ff48cb 100644 --- a/rust/src/graph/mod.rs +++ b/rust/src/graph/mod.rs @@ -394,11 +394,16 @@ impl GraphWrapper { .direct_import_exists(importer, imported, as_packages)?) } - pub fn find_modules_directly_imported_by(&self, module: &str) -> PyResult> { + #[pyo3(signature = (module, as_package=false))] + pub fn find_modules_directly_imported_by( + &self, + module: &str, + as_package: bool, + ) -> PyResult> { 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() @@ -406,11 +411,16 @@ impl GraphWrapper { .collect()) } - pub fn find_modules_that_directly_import(&self, module: &str) -> PyResult> { + #[pyo3(signature = (module, as_package=false))] + pub fn find_modules_that_directly_import( + &self, + module: &str, + as_package: bool, + ) -> PyResult> { 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() diff --git a/src/grimp/application/graph.py b/src/grimp/application/graph.py index 7e53481a..3c679621 100644 --- a/src/grimp/application/graph.py +++ b/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]: """ diff --git a/tests/unit/application/graph/test_direct_imports.py b/tests/unit/application/graph/test_direct_imports.py index d7d724a3..ac65fabd 100644 --- a/tests/unit/application/graph/test_direct_imports.py +++ b/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( @@ -344,6 +386,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(