diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 957e128a..3e1d66aa 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -57,3 +57,7 @@ Changelog 1.0b10 (2019-5-15) ------------------ * Fix Windows incompatibility. + +latest +------ +* Add remove_module method. diff --git a/docs/usage.rst b/docs/usage.rst index ad7e0dde..ecd750b6 100644 --- a/docs/usage.rst +++ b/docs/usage.rst @@ -202,6 +202,15 @@ Methods for manipulating the graph :param bool is_squashed: If True, the module should be treated as a 'squashed module' (see `Terminology`_ above). :return: None +.. py:function:: ImportGraph.remove_module(module) + + Remove a module from the graph. + + If the module is not present in the graph, no exception will be raised. + + :param str module: The name of a module, for example ``'mypackage.foo'``. + :return: None + .. py:function:: ImportGraph.add_import(importer, imported, line_number=None, line_contents=None) Add a direct import between two modules to the graph. If the modules are not already diff --git a/src/grimp/adaptors/graph.py b/src/grimp/adaptors/graph.py index aaa5878a..2339aeaa 100644 --- a/src/grimp/adaptors/graph.py +++ b/src/grimp/adaptors/graph.py @@ -41,6 +41,10 @@ def add_module(self, module: str, is_squashed: bool = False) -> None: if is_squashed: self._mark_module_as_squashed(module) + def remove_module(self, module: str) -> None: + if module in self.modules: + self._networkx_graph.remove_node(module) + def add_import( self, *, importer: str, diff --git a/src/grimp/application/ports/graph.py b/src/grimp/application/ports/graph.py index 74a58c98..86f93508 100644 --- a/src/grimp/application/ports/graph.py +++ b/src/grimp/application/ports/graph.py @@ -30,6 +30,15 @@ def add_module(self, module: str, is_squashed: bool = False) -> None: """ raise NotImplementedError + @abc.abstractmethod + def remove_module(self, module: str) -> None: + """ + Remove a module from the graph, if it exists. + + If the module is not present in the graph, no exception will be raised. + """ + raise NotImplementedError + @abc.abstractmethod def add_import( self, *, diff --git a/tests/functional/test_build_and_use_graph.py b/tests/functional/test_build_and_use_graph.py index b73581ce..d5c16a7a 100644 --- a/tests/functional/test_build_and_use_graph.py +++ b/tests/functional/test_build_and_use_graph.py @@ -50,6 +50,15 @@ def test_add_module(): assert number_of_modules + 1 == len(graph.modules) +def test_remove_module(): + graph = build_graph('testpackage') + number_of_modules = len(graph.modules) + + graph.remove_module('testpackage.two.alpha') + assert 'testpackage.two.alpha' not in graph.modules + assert number_of_modules - 1 == len(graph.modules) + + def test_add_and_remove_import(): graph = build_graph('testpackage') a = 'testpackage.one.delta.blue' diff --git a/tests/unit/adaptors/test_graph.py b/tests/unit/adaptors/test_graph.py index a1407d88..5be593e8 100644 --- a/tests/unit/adaptors/test_graph.py +++ b/tests/unit/adaptors/test_graph.py @@ -522,6 +522,21 @@ def test_add_module(): assert graph.modules == {module} +def test_remove_module(): + graph = ImportGraph() + a, b = {'mypackage.blue', 'mypackage.green'} + + graph.add_module(a) + graph.add_module(b) + graph.add_import(importer=a, imported=b) + + graph.remove_module(b) + assert {a} == graph.modules + + # Removing a non-existent module doesn't cause an error. + graph.remove_module('mypackage.yellow') + + class TestAddSquashedModule: def test_can_repeatedly_add_same_squashed_module(self): graph = ImportGraph()