diff --git a/doc/changelog.rst b/doc/changelog.rst index 4d03dbf..4bfab95 100644 --- a/doc/changelog.rst +++ b/doc/changelog.rst @@ -4,6 +4,10 @@ Release history 2.3 --- +* Add `modulegraph2.Modulegraph.add_dependencies_for_source`` + that will add the imports in a python code fragment + to a graph as roots. + * Add ``modulegraph2.stdlib_module_names`` and add ``--exclude-stdlib`` to the standard interface to exclude the stdlib from the graph. @@ -19,6 +23,13 @@ Release history * Fix incompatibility with Python 3.11 when implict namespace packages are used. +* Add ``ModuleGraph.import_package`` which will add all + submodules of a package to the graph, without adding + the package to the graph roots. + + This currently only supports packages found in the + filesystem (excluding zipfiles). + 2.2.1 ----- diff --git a/modulegraph2/_modulegraph.py b/modulegraph2/_modulegraph.py index e1b469e..77e15ce 100644 --- a/modulegraph2/_modulegraph.py +++ b/modulegraph2/_modulegraph.py @@ -2,9 +2,11 @@ This module contains the definition of the ModuleGraph class. """ import ast +import contextlib import importlib import operator import os +import pathlib import sys from types import ModuleType from typing import ( @@ -58,8 +60,9 @@ class ModuleGraph(ObjectGraph[Union[BaseNode, PyPIDistribution], DependencyInfo] of python modules and scripts. The roots of the graph are those nodes that are added to the - graph using :meth:`add_script() <ModuleGraph.add_script>` and - :meth:`add_module() <ModuleGraph.add_module>`. + graph using :meth:`add_script() <ModuleGraph.add_script>`, + :meth:`add_module() <ModuleGraph.add_module>` and + :meth:`add_dependencies_for_source() <ModuleGraph.add_dependencies_for_source>`. Args: * use_stdlib_implies: Use the built-in implied actions for the stdlib. @@ -110,7 +113,7 @@ def distributions(self, reachable: bool = True) -> Iterator[PyPIDistribution]: unless they are also the *distribution* attribute of a node. Args: - reacable: IF true only report on nodes that are reachable from + reachable: IF true only report on nodes that are reachable from a graph root, otherwise report on all nodes. """ seen: Set[str] = set() @@ -192,6 +195,26 @@ def add_module(self, module_name: str) -> BaseNode: self._run_stack() return node + def add_dependencies_for_source(self, source_code: str) -> None: + """ + Add the modules imported by the Python code in *source_code* + to the graph as roots. + + Args: + source_code: Source code for a python script or module + """ + ast_node = compile( + source_code, + "-source-code-", + "exec", + flags=ast.PyCF_ONLY_AST, + dont_inherit=True, + ) + for info in extract_ast_info(ast_node): + self.add_module(info.import_module) + for name in info.import_names: + self.add_module(f"{info.import_module}.{name}") + def add_distribution(self, distribution: Union[PyPIDistribution, str]): """ Add a package distribution to the graph, with references @@ -229,6 +252,17 @@ def add_distribution(self, distribution: Union[PyPIDistribution, str]): # Hooks # + @contextlib.contextmanager + def hook_context(self): + """ + Contextmanger that allows using hook APIs outside of a hook + callback. + """ + try: + yield + finally: + self._run_stack() + def import_module(self, importing_module: BaseNode, import_name: str) -> BaseNode: """ Import 'import_name' and add an edge from 'module' to 'import_name' @@ -251,6 +285,53 @@ def import_module(self, importing_module: BaseNode, import_name: str) -> BaseNod self.add_edge(importing_module, node, DEFAULT_DEPENDENCY) return node + def import_package(self, importing_module: BaseNode, package_name: str) -> BaseNode: + """ + Add all modules in a package to the graph and process imports. + + Will not raise an exception for non-existing packages or when + the *package_name* refers to a module instead of a package. + + This is an API to be used by hooks. The graph is not fully updated + after calling this method. + + This method only supports packages that are found in the + filesystem. + Args: + importing_module: The module that triggers this import + + package_name: Name of the package to import + """ + node = self._find_module(package_name) + if node is None: + node = self._find_or_load_module(None, package_name) + + self._run_stack() + + if not isinstance(node, Package): + return node + + # This is something not natively supported by + # importlib, because of this the function tries + # to find submodules by looking in the filesystem. + # + spec = importlib.util.find_spec(package_name) + if spec is not None: + for path_name in spec.submodule_search_locations or []: + path = pathlib.Path(path_name) + for fn in path.glob("*.py"): + module = fn.stem + if not module.isidentifier(): + continue + + abs_name = f"{package_name}.{module}" + module_node = self._find_or_load_module(node, abs_name) + self.add_edge(node, module_node, DEFAULT_DEPENDENCY) + if isinstance(module_node, Package): + self._work_stack.append((self.import_package, (abs_name,))) + + return node + def add_post_processing_hook(self, hook: ProcessingCallback) -> None: """ Add a hook function to be ran whenever a node is fully processed.