Fix handling of within-descendant-imports when squashing modules

CHANGELOG.rst

@@ -2,6 +2,11 @@
 Changelog
 =========
 
+Unreleased
+----------
+
+* Fixed handling of within-descendant-imports when squashing modules (see `Issue 195 <https://github.com/seddonym/grimp/issues/195>`_).
+
 3.7 (2025-03-07)
 ----------------
 

rust/src/graph/graph_manipulation.rs

@@ -178,11 +178,6 @@ impl Graph {
             })
             .collect();
 
-        // Remove any descendants.
-        for descendant in descendants {
-            self.remove_module(descendant);
-        }
-
         // Add descendants and imports to parent module.
         for imported in modules_imported_by_descendants {
             self.add_import(module, imported);
@@ -192,6 +187,11 @@ impl Graph {
             self.add_import(importer, module);
         }
 
+        // Remove any descendants.
+        for descendant in descendants {
+            self.remove_module(descendant);
+        }
+
         self.mark_module_squashed(module);
     }
 }

tests/unit/adaptors/graph/test_manipulation.py

@@ -447,3 +447,22 @@ def test_raises_module_not_present_if_no_module(self):
 
         with pytest.raises(ModuleNotPresent):
             graph.squash_module("foo")
+
+    def test_correctly_handles_imports_within_descendants(self):
+        graph = ImportGraph()
+
+        graph.add_module("animals")
+        graph.add_module("food")
+        graph.add_import(importer="animals.dog", imported="food.chicken")
+        graph.add_import(importer="app.cli", imported="animals.dog")
+        # We want to check that this import within the descendants of `animals` does
+        # not cause problems. If this import is not properly removed then the imports map can
+        # become corrupted, since it will contain an import to modules that no longer exist.
+        # See https://github.com/seddonym/grimp/issues/195 for more details.
+        graph.add_import(importer="animals.dog", imported="animals.base")
+
+        graph.squash_module("animals")
+
+        # If the `animals` module was squashed correctly then the following calls should not panic.
+        assert graph.find_modules_directly_imported_by("animals") == {"food.chicken"}
+        assert graph.find_modules_that_directly_import("animals") == {"app.cli"}