From 247819c34623c6f402c7b40bf326cff232015c43 Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Thu, 23 Jul 2026 03:40:07 +0200 Subject: [PATCH 01/12] Retrieve nested imports with `ProfmodExtractor` kernprof.py Updated call to `line_profiler.autoprofile.autoprofile.run()` to pass the `config` parameter line_profiler/autoprofile/ast_tree_profiler.py::AstTreeProfiler .__init__() - Loosened type of parameter `prof_mod` to `Sequence[str]` - Added parameter `config` .profile() Updated instantiation of `line_profiler.autoprofile.profmod_extractor.ProfmodExtractor` to pass the `config` parameter line_profiler/autoprofile/autoprofile.py::run() Updated instantiation of `line_profiler.autoprofile.ast_tree_profiler.AstTreeProfiler` to pass the `config` parameter line_profiler/autoprofile/profmod_extractor.py _ImportFinder[.find()] New `ast.NodeVisitor` subclass (and method) for locating imports nested inside other statements (e.g. conditionals and definitions) ProfmodExtractor .__init__() - Loosened type of parameter `prof_mod` to `Sequence[str]` - Added parameter `config` ._get_modnames_to_profile_from_prof_mod() Loosened type of parameter `prof_mod` to `Sequence[str]` ._ast_get_imports_from_tree() Now an alias for `_ImportFinder.find()` ._find_modnames_in_tree_imports() Loosened types of parameters to `Sequence[...]` .extract_all() Now handling nested imports in accordance to the `config` line_profiler/autoprofile/run_module.py ::AstTreeModuleProfiler._check_profile_full_script() Loosened type of parameter `prof_mod` to `Sequence[str]` line_profiler/rc/line_profiler.toml ::[tool.line_profiler.prof_mod_import_discovery] New boolean config table for controlling what kind of statements to descend into to check for imports; keys: [`conditionals`, 'try_except', 'contexts', 'loops', 'definitions'] tests/test_autoprofile.py::test_nested_import_discovery() New test checking that the above toggles work --- kernprof.py | 1 + .../autoprofile/ast_tree_profiler.py | 17 +- line_profiler/autoprofile/autoprofile.py | 23 +- .../autoprofile/profmod_extractor.py | 257 +++++++++++++++--- line_profiler/autoprofile/run_module.py | 3 +- line_profiler/rc/line_profiler.toml | 35 +++ tests/test_autoprofile.py | 118 +++++++- 7 files changed, 397 insertions(+), 57 deletions(-) diff --git a/kernprof.py b/kernprof.py index 590fd48c..988ee7f3 100755 --- a/kernprof.py +++ b/kernprof.py @@ -1242,6 +1242,7 @@ def _main_profile(options, module=False, exit_on_error=True): prof_mod=options.prof_mod, profile_imports=options.prof_imports, as_module=module is not None, + config=options.config, ) else: # Note: to reduce complications (e.g. whenever something diff --git a/line_profiler/autoprofile/ast_tree_profiler.py b/line_profiler/autoprofile/ast_tree_profiler.py index e88f1100..d20bc9d3 100644 --- a/line_profiler/autoprofile/ast_tree_profiler.py +++ b/line_profiler/autoprofile/ast_tree_profiler.py @@ -6,6 +6,7 @@ from typing import Any from ._import_targets import ImportTarget +from ..toml_config import ConfigSource from .ast_profile_transformer import ( AstProfileTransformer, ast_create_profile_node, @@ -29,7 +30,7 @@ class AstTreeProfiler: def __init__( self, script_file: str, - prof_mod: list[str], + prof_mod: Sequence[str], profile_imports: bool, ast_transformer_class_handler: ( type[AstProfileTransformer] @@ -37,6 +38,7 @@ def __init__( profmod_extractor_class_handler: ( type[ProfmodExtractor] ) = ProfmodExtractor, + config: ConfigSource | None = None, ) -> None: """Initializes the AST tree profiler instance with the script file path @@ -57,16 +59,21 @@ def __init__( profmod_extractor_class_handler (type[ProfmodExtractor]): the ProfmodExtractor class that handles mapping prof_mod to objects in the script. + + config (ConfigSource | None): + optional :py:class:`.ConfigSource` to load additional + configurations from. """ self._script_file = script_file self._prof_mod = prof_mod self._profile_imports = profile_imports self._ast_transformer_class_handler = ast_transformer_class_handler self._profmod_extractor_class_handler = profmod_extractor_class_handler + self._config = config @staticmethod def _check_profile_full_script( - script_file: str, prof_mod: list[str] + script_file: str, prof_mod: Sequence[str], ) -> bool: """Check whether whole script should be profiled. @@ -77,13 +84,13 @@ def _check_profile_full_script( script_file (str): path to script being profiled. - prof_mod (List[str]): + prof_mod (Sequence[str]): list of imports to profile in script. passing the path to script will profile the whole script. the objects can be specified using its dotted path or full path (if applicable). Returns: - (bool): profile_full_script + profile_full_script (bool): if True, profile whole script. """ script_file_realpath = os.path.realpath(script_file) @@ -207,7 +214,7 @@ def profile(self) -> ast.Module: tree = self._get_script_ast_tree(self._script_file) tree_imports_to_profile_dict = self._profmod_extractor_class_handler( - tree, self._script_file, self._prof_mod + tree, self._script_file, self._prof_mod, self._config, ).extract_all() tree_profiled = self._profile_ast_tree( tree, diff --git a/line_profiler/autoprofile/autoprofile.py b/line_profiler/autoprofile/autoprofile.py index 4471a778..3e98a92c 100644 --- a/line_profiler/autoprofile/autoprofile.py +++ b/line_profiler/autoprofile/autoprofile.py @@ -47,11 +47,13 @@ def main(): from __future__ import annotations import importlib.util +import os import sys import types from collections.abc import MutableMapping -from typing import Any, cast, Dict, Mapping -from typing import ContextManager +from typing import Any, cast + +from ..toml_config import ConfigSource from ..line_profiler_utils import restore from .ast_tree_profiler import AstTreeProfiler from .run_module import AstTreeModuleProfiler @@ -84,6 +86,7 @@ def run( prof_mod: list[str], profile_imports: bool = False, as_module: bool = False, + config: os.PathLike[str] | str | None = None, ) -> None: """Automatically profile a script and run it. @@ -106,7 +109,10 @@ def run( if True, when auto-profiling whole script, profile all imports aswell. as_module (bool): - Whether we're running script_file as a module + whether we're running script_file as a module + + config (os.PathLike[str] | str | None): + optional path to load the session config from """ Profiler: type[AstTreeModuleProfiler] | type[AstTreeProfiler] @@ -128,7 +134,10 @@ def run( namespace: MutableMapping[str, Any] = vars(module_obj) namespace.update(ns) - profiler = Profiler(script_file, prof_mod, profile_imports) + profiler = Profiler( + script_file, prof_mod, profile_imports, + config=ConfigSource.from_config(config), + ) tree_profiled = profiler.profile() _extend_line_profiler_for_profiling_imports(ns[PROFILER_LOCALS_NAME]) @@ -138,4 +147,8 @@ def run( # then restore it via the context manager, so that the executed # code is run as `__main__` sys.modules['__main__'] = module_obj - exec(code_obj, cast(Dict[str, Any], namespace), namespace) # type: ignore[redundant-cast] + exec( + code_obj, + cast('dict[str, Any]', namespace), # type: ignore[ty:redundant-cast] + namespace, + ) diff --git a/line_profiler/autoprofile/profmod_extractor.py b/line_profiler/autoprofile/profmod_extractor.py index 4a49bedf..030680bb 100644 --- a/line_profiler/autoprofile/profmod_extractor.py +++ b/line_profiler/autoprofile/profmod_extractor.py @@ -3,9 +3,11 @@ import ast import os import sys -from typing import cast +from collections.abc import Sequence +from typing import TYPE_CHECKING, ClassVar, Literal, cast, get_args from warnings import warn +from ..toml_config import ConfigSource from .util_static import ( modname_to_modpath, modpath_to_modname, @@ -15,17 +17,185 @@ from ._import_targets import ImportTarget -class ProfmodExtractor: - """Map prof_mod to imports in an abstract syntax tree. - - Takes the paths and dotted paths in prod_mod and finds their respective imports in an - abstract syntax tree. +# Node types where code blocks can be found +_BodiedNodeType = Literal[ + # Basic top-level nodes + 'Module', 'Interactive', + # Definition nodes + 'FunctionDef', 'AsyncFunctionDef', 'ClassDef', + # Loop nodes + 'For', 'AsyncFor', 'While', + # Conditional nodes + 'If', 'match_case', + # Context nodes + 'With', 'AsyncWith', + # `try-except` nodes + 'Try', 'TryStar', 'ExceptHandler', +] + + +class _ImportFinder(ast.NodeVisitor): """ + Locate all the imports inside an AST, including those nested inside + other nodes. + """ + _bodied_node_types: ClassVar[set[_BodiedNodeType]] = cast( + set[_BodiedNodeType], set(get_args(_BodiedNodeType)), + ) def __init__( - self, tree: ast.Module, script_file: str, prof_mod: list[str] + self, + node_types: dict[_BodiedNodeType, bool], + found_imports: ( + dict[tuple[str | int, ...], list[ImportTarget]] | None + ) = None, + ) -> None: + self._current_loc: list[str | int] = [] + self._should_visit = node_types + if found_imports is None: + found_imports = {} + self.found_imports = found_imports + + @classmethod + def find( + cls, + node: ast.AST, + *, + collect_from_conditionals: bool = True, + collect_from_try_except: bool = True, + collect_from_contexts: bool = True, + collect_from_loops: bool = True, + collect_from_definitions: bool = False, + ) -> dict[tuple[str | int, ...], list[ImportTarget]]: + """ + Parameters: + node (ast.AST): + AST node + collect_from_conditionals (bool): + Whether to collect imports inside :py:class:`ast.If` and + :py:class:`ast.match_case` nodes (and their children). + collect_from_try_except (bool): + Whether to collect imports inside + :py:class:`ast.ExceptHandler`, :py:class:`ast.Try`, and + :py:class:`ast.TryStar` nodes (and their children). + collect_from_contexts (bool): + Whether to collect imports inside + :py:class:`ast.AsyncWith` and :py:class:`ast.With` nodes + (and their children). + collect_from_loops (bool): + Whether to collect imports inside + :py:class:`ast.AsyncFor`, :py:class:`ast.For`, and + :py:class:`ast.While` nodes (and their children). + collect_from_definitions (bool): + Whether to collect imports inside + :py:class:`AsyncFunctionDef`, :py:class:`ast.ClassDef` + and :py:class:`ast.FunctionDef` nodes (and their + children). + + Returns: + found_imports \ +(dict[tuple[str | int, ...], list[ImportTarget]]): + The import targets and their locations in the AST + """ + node_types = cls.filter_node_types( + collect_from_conditionals=collect_from_conditionals, + collect_from_try_except=collect_from_try_except, + collect_from_contexts=collect_from_contexts, + collect_from_loops=collect_from_loops, + collect_from_definitions=collect_from_definitions, + ) + visitor = cls(node_types) + visitor.visit(node) + return visitor.found_imports + + @classmethod + def filter_node_types( + cls, + *, + collect_from_conditionals: bool = True, + collect_from_try_except: bool = True, + collect_from_contexts: bool = True, + collect_from_loops: bool = True, + collect_from_definitions: bool = False, + ) -> dict[_BodiedNodeType, bool]: + """ + Select the which of the "bodied" node types (i.e. those than can + contain nexted code blocks, e.g. try-except statements) to + visit. + """ + allowed: set[_BodiedNodeType] = {'Module', 'Interactive'} + if collect_from_conditionals: + allowed.update({'If', 'match_case'}) + if collect_from_try_except: + allowed.update({'Try', 'TryStar', 'ExceptHandler'}) + if collect_from_contexts: + allowed.update({'AsyncWith', 'With'}) + if collect_from_loops: + allowed.update({'AsyncFor', 'For', 'While'}) + if collect_from_definitions: + allowed.update({ + 'AsyncFunctionDef', 'ClassDef', 'FunctionDef', + }) + return { + node_type: node_type in allowed + for node_type in cls._bodied_node_types + } + + def generic_visit(self, node: ast.AST) -> None: + """ + For the "bodied" node types, only visit their children if said + type is selected via the init args; for other types, visit by + default. + """ + node_type = type(node).__name__ + if not self._should_visit.get(cast(_BodiedNodeType, node_type), True): + return # Deselected types + for field, value in ast.iter_fields(node): + if isinstance(value, ast.AST): + self._current_loc.append(field) + try: + self.visit(value) + finally: + self._current_loc.pop() + elif isinstance(value, Sequence): # Node with "body" + if not all(isinstance(item, ast.AST) for item in value): + continue + if TYPE_CHECKING: + value = cast(Sequence[ast.AST], value) + self._current_loc.append(field) + try: + # Parse import targets + imports = ImportTarget._from_ast_nodes(value) + # Descend into children nodes + if imports: + loc = tuple(self._current_loc) + self.found_imports[loc] = imports + for i, item in enumerate(value): + self._current_loc.append(i) + try: + self.visit(item) + finally: + self._current_loc.pop() + finally: + self._current_loc.pop() + + +class ProfmodExtractor: + """ + Map ``prof_mod`` to imports in an abstract syntax tree. Takes the + paths and dotted paths in ``prof_mod`` and finds their respective + imports in an abstract syntax tree. + """ + def __init__( + self, + tree: ast.Module, + script_file: str, + prof_mod: Sequence[str], + config: ConfigSource | None = None, ) -> None: - """Initializes the AST tree profiler instance with the AST, script file path and prof_mod + """ + Initializes the AST tree profiler instance with the AST, + script-file path and ``prof_mod`` Args: tree (_ast.Module): @@ -34,14 +204,27 @@ def __init__( script_file (str): path to script being profiled. - prof_mod (list[str]): - list of imports to profile in script. - passing the path to script will profile the whole script. - the objects can be specified using its dotted path or full path (if applicable). + prof_mod (Sequence[str]): + optional list of imports to profile in script. + passing the path to script will profile the whole + script. + the objects can be specified using its dotted path or + full path (if applicable). + + config (ConfigSource | None): + optional :py:class:`.ConfigSource` to load additional + configurations from. """ self._tree = tree self._script_file = script_file self._prof_mod = prof_mod + if config is None: + config = ConfigSource.from_default() + self._config = ( + config + .get_subconfig('prof_mod_import_discovery') + .conf_dict + ) @staticmethod def _is_path(text: str) -> bool: @@ -62,7 +245,7 @@ def _is_path(text: str) -> bool: @classmethod def _get_modnames_to_profile_from_prof_mod( - cls, script_file: str, prof_mod: list[str] + cls, script_file: str, prof_mod: Sequence[str] ) -> list[str]: """Grab the valid paths and all dotted paths in prof_mod and their subpackages and submodules, in the form of dotted paths. @@ -81,7 +264,7 @@ def _get_modnames_to_profile_from_prof_mod( script_file (str): path to script being profiled. - prof_mod (list[str]): + prof_mod (Sequence[str]): list of imports to profile in script. passing the path to script will profile the whole script. the objects can be specified using its dotted path or full path (if applicable). @@ -141,32 +324,12 @@ def _get_modnames_to_profile_from_prof_mod( return modnames_to_profile - @staticmethod - def _ast_get_imports_from_tree( - tree: ast.Module, - ) -> dict[tuple[str | int, ...], list[ImportTarget]]: - """Get all top-level imports in an abstract syntax tree. - - Args: - tree (_ast.Module): - abstract syntax tree to fetch imports from. - - Returns: - import_targets (dict[tuple[str | int, ...], list[ImportTarget]]) - - Note: - Imports nested in e.g. try-except or if statements are not - currently returned. - - See also: - :py:meth:`.ImportTarget._from_ast_nodes` - """ - # TODO: descend into bodied statements (e.g. try-except) - return {('body',): ImportTarget._from_ast_nodes(tree.body)} + _ast_get_imports_from_tree = _ImportFinder.find @staticmethod def _find_modnames_in_tree_imports( - modnames_to_profile: list[str], import_targets: list[ImportTarget], + modnames_to_profile: Sequence[str], + import_targets: Sequence[ImportTarget], ) -> dict[int, list[ImportTarget]]: """Map modnames to imports from an abstract sytax tree. @@ -179,10 +342,10 @@ def _find_modnames_in_tree_imports( The import's alias is stored in the output dict. Args: - modnames_to_profile (list[str]): + modnames_to_profile (Sequence[str]): list of dotted paths to profile. - import_targets (list[ImportTarget]): + import_targets (Sequence[ImportTarget]): list of all import targets in the tree Returns: @@ -249,16 +412,20 @@ def extract_all(self) -> dict[tuple[str | int, ...], list[ImportTarget]]: :py:const`None` for non-star-imports) Notes: - - As of now, ``from import *`` is not supported, - and will result in a :py:class:`UserWarning`. - - - Nested imports (e.g. imports in try-except/if blocks) are - not currently retrieved. + As of now, ``from import *`` is not supported, and + will result in a :py:class:`UserWarning`. """ modnames_to_profile = self._get_modnames_to_profile_from_prof_mod( self._script_file, self._prof_mod ) - import_targets = self._ast_get_imports_from_tree(self._tree) + import_targets = self._ast_get_imports_from_tree( + self._tree, + collect_from_conditionals=self._config['conditionals'], + collect_from_try_except=self._config['try_except'], + collect_from_contexts=self._config['contexts'], + collect_from_loops=self._config['loops'], + collect_from_definitions=self._config['definitions'], + ) raw: dict[tuple[str | int, ...], list[ImportTarget]] = { (*loc, index): filtered_imports for loc, imports in import_targets.items() diff --git a/line_profiler/autoprofile/run_module.py b/line_profiler/autoprofile/run_module.py index 9cde259c..76d8da43 100644 --- a/line_profiler/autoprofile/run_module.py +++ b/line_profiler/autoprofile/run_module.py @@ -2,6 +2,7 @@ import ast import os +from collections.abc import Sequence from typing import cast from .ast_tree_profiler import AstTreeProfiler @@ -96,7 +97,7 @@ def _is_main(fname: str) -> bool: @classmethod def _check_profile_full_script( - cls, script_file: str, prof_mod: list[str] + cls, script_file: str, prof_mod: Sequence[str], ) -> bool: rp = os.path.realpath paths_to_check = {rp(script_file)} diff --git a/line_profiler/rc/line_profiler.toml b/line_profiler/rc/line_profiler.toml index 6680c06a..10151798 100644 --- a/line_profiler/rc/line_profiler.toml +++ b/line_profiler/rc/line_profiler.toml @@ -206,3 +206,38 @@ hits = 9 time = 12 perhit = 8 percent = 8 + +# `line_profiler.autoprofile.profmod_extractor.ProfmodExtractor` +# selective import-profiling options + +# Note: +# - When using `preimports = true`, the `prof-mod` names are always +# profiled. +# - Otherwise, if the file-/module-path of the profiled code is itself +# among the `prof-mod` names, its imports are ALWAYS profiled. +# - Else, imports are selectively discovered from the AST of the +# profiled code, and import targets specifically matching the +# `prof-mod` names are profiled. Import discovery is dictated by the +# following options: +# - If the option corresponding to a language construction is `false`, +# we won't descend therein to look for imports. +# - If all these options are set to false, only raw (from-)import +# statements on the top level of the profiled code will be checked. + +[tool.line_profiler.prof_mod_import_discovery] + +# - `conditionals` (bool): +# Whether to look into `if-elif-else` and `match-case` statements. +conditionals = true +# - `try_except` (bool): +# Whether to look into `try-except-else-finally` statements. +try_except = true +# - `contexts` (bool): +# Whether to look into `[async] with` statements. +contexts = true +# - `loops` (bool): +# Whether to look into `[async] for-else` and `while-else` statements. +loops = true +# - `definitions` (bool): +# Whether to look into `[async] def` and `class` statements. +definitions = false diff --git a/tests/test_autoprofile.py b/tests/test_autoprofile.py index bf2ae917..5e95270b 100644 --- a/tests/test_autoprofile.py +++ b/tests/test_autoprofile.py @@ -4,9 +4,10 @@ import contextlib import os import re +import shlex import subprocess import sys -import shlex +import textwrap import tempfile from collections.abc import Collection, Sequence from typing import Any, ClassVar, Literal @@ -14,6 +15,7 @@ import pytest import ubelt as ub +from line_profiler.toml_config import ConfigSource from line_profiler.autoprofile.ast_tree_profiler import AstTreeProfiler from line_profiler.autoprofile.profmod_extractor import ProfmodExtractor @@ -1550,3 +1552,117 @@ def func() -> None: output_module = ast.unparse(module_ast) for pattern, expected in re_checks: assert bool(re.search(pattern, output_module)) == expected + + +@pytest.mark.parametrize( + ('prof_mod', 'expected', 'options'), + [(['qux', 'quux'], {'qux.jam', 'spam', 'ham', 'eggs'}, {'conditionals'}), + (['os', 'qux'], {'qux.jam'}, {'conditionals'}), + (['os', 'qux'], {'fork', 'register_at_fork'}, {'try_except'}), + (['foobar'], {'baz'}, {'try_except', 'contexts'}), + (['ersatz_foobar.my_baz'], {'baz'}, {'try_except', 'contexts'}), + (['os.fork', 'foobar.bar'], {'fork'}, {'try_except', 'contexts'}), + (['backup_fred', 'operator'], {'methodcaller', 'setitem'}, + {'definitions'}), + (['backup_fred', 'operator'], {'fred'}, {'loops'})], +) +def test_nested_import_discovery( + prof_mod: list[str], + expected: set[str], + options: set[Literal[ + 'conditionals', 'try_except', 'contexts', 'loops', 'definitions', + ]], +) -> None: + """ + Check the source code transformed by :py:class:`.AstTreeProfiler` to + see if the import-discovery selection options in the TOML file + (``[tool.line_profiler.prof_mod_import_discovery]``) are handled + correctly. + """ + test_module = ub.codeblock(""" + from collections.abc import Generator + from contextlib import contextmanager + from functools import partial + from importlib import import_module + from sys import path, version_info + + notify_fork = partial(print, 'Forking...') + try: + from os import fork + except Exception: # Windows + pass + else: + from os import register_at_fork + + register_at_fork(before=notify_fork) + + import foo, bar + + if version_info > (3, 14): + import qux.jam + from quux import spam, ham, eggs + else: + qux = ham = spam = eggs = None + + + @contextmanager + def _restore_sys_path() -> Generator[None, None, None]: + from operator import methodcaller, setitem + + old = methodcaller('copy')(path) + try: + yield + finally: + setitem(path, slice(None), old) + + + with _restore_sys_path(): + try: + from foobar import baz + except ImportError: + from ersatz_foobar import my_baz as baz + + + for _fred in 'fred', 'some_fred', 'other_fred': + try: + fred = import_module(_fred) + except ImportError: + continue + else: + del _fred + break + else: # Fallback + import backup_fred as fred + """).strip('\n') + + config_file_lines = ['[tool.line_profiler.prof_mod_import_discovery]'] + for option in [ + 'conditionals', 'try_except', 'contexts', 'loops', 'definitions', + ]: + line = f'{option} = {str(option in options).lower()}' + config_file_lines.append(line) + config_file = '\n'.join(config_file_lines) + + with tempfile.TemporaryDirectory() as tmp: + mod_fname = os.path.join(tmp, 'test_module.py') + with open(mod_fname, 'w') as fobj: + print(test_module, file=fobj) + + cfg_fname = os.path.join(tmp, 'config.toml') + with open(cfg_fname, 'w') as fobj: + print(config_file, file=fobj) + + config = ConfigSource.from_config(cfg_fname) + atp = AstTreeProfiler(mod_fname, prof_mod, False, config=config) + output_module = ast.unparse(atp.profile()) + + for label, module_text in [ + ('input', test_module), ('output', output_module), + ]: + print(f'{label.capitalize()}:\n{textwrap.indent(module_text, " ")}\n') + + prof_pattern = ( + r'\badd_imported_function_or_module\((\w+(?:\.\w+)*)\)' + ) + profiled_names = re.findall(prof_pattern, output_module) + assert set(profiled_names) == expected From cd33423e47658852637ace921c9a801c6bb9353e Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Tue, 21 Jul 2026 02:07:21 +0200 Subject: [PATCH 02/12] Minor refactoring + new, more granular test line_profiler/autoprofile/profmod_extractor.py Updated wordings of several docstrings and some private names tests/test_autoprofile.py test_nested_import_discovery() Refactored internals out to be shared with other tests test_import_discovery_in_all_compound_statements() New test for exhaustively testing each of the language constructions which create compound statements containing other statements, that our profiling of imports therein are correctly controlled by the respective config options --- .../autoprofile/profmod_extractor.py | 26 +- tests/test_autoprofile.py | 249 ++++++++++++++++-- 2 files changed, 244 insertions(+), 31 deletions(-) diff --git a/line_profiler/autoprofile/profmod_extractor.py b/line_profiler/autoprofile/profmod_extractor.py index 030680bb..15f1fd21 100644 --- a/line_profiler/autoprofile/profmod_extractor.py +++ b/line_profiler/autoprofile/profmod_extractor.py @@ -18,7 +18,7 @@ # Node types where code blocks can be found -_BodiedNodeType = Literal[ +_CompoundNodeType = Literal[ # Basic top-level nodes 'Module', 'Interactive', # Definition nodes @@ -39,13 +39,13 @@ class _ImportFinder(ast.NodeVisitor): Locate all the imports inside an AST, including those nested inside other nodes. """ - _bodied_node_types: ClassVar[set[_BodiedNodeType]] = cast( - set[_BodiedNodeType], set(get_args(_BodiedNodeType)), + _compound_node_types: ClassVar[set[_CompoundNodeType]] = cast( + set[_CompoundNodeType], set(get_args(_CompoundNodeType)), ) def __init__( self, - node_types: dict[_BodiedNodeType, bool], + node_types: dict[_CompoundNodeType, bool], found_imports: ( dict[tuple[str | int, ...], list[ImportTarget]] | None ) = None, @@ -117,13 +117,13 @@ def filter_node_types( collect_from_contexts: bool = True, collect_from_loops: bool = True, collect_from_definitions: bool = False, - ) -> dict[_BodiedNodeType, bool]: + ) -> dict[_CompoundNodeType, bool]: """ - Select the which of the "bodied" node types (i.e. those than can - contain nexted code blocks, e.g. try-except statements) to + Select the which of the compound node types (i.e. those than can + contain nested code blocks, e.g. try-except statements) to visit. """ - allowed: set[_BodiedNodeType] = {'Module', 'Interactive'} + allowed: set[_CompoundNodeType] = {'Module', 'Interactive'} if collect_from_conditionals: allowed.update({'If', 'match_case'}) if collect_from_try_except: @@ -138,17 +138,19 @@ def filter_node_types( }) return { node_type: node_type in allowed - for node_type in cls._bodied_node_types + for node_type in cls._compound_node_types } def generic_visit(self, node: ast.AST) -> None: """ - For the "bodied" node types, only visit their children if said + For the compound node types, only visit their children if said type is selected via the init args; for other types, visit by default. """ node_type = type(node).__name__ - if not self._should_visit.get(cast(_BodiedNodeType, node_type), True): + if not self._should_visit.get( + cast(_CompoundNodeType, node_type), True, + ): return # Deselected types for field, value in ast.iter_fields(node): if isinstance(value, ast.AST): @@ -157,7 +159,7 @@ def generic_visit(self, node: ast.AST) -> None: self.visit(value) finally: self._current_loc.pop() - elif isinstance(value, Sequence): # Node with "body" + elif isinstance(value, Sequence): # Compound node if not all(isinstance(item, ast.AST) for item in value): continue if TYPE_CHECKING: diff --git a/tests/test_autoprofile.py b/tests/test_autoprofile.py index 5e95270b..45d2c75f 100644 --- a/tests/test_autoprofile.py +++ b/tests/test_autoprofile.py @@ -1485,6 +1485,44 @@ def add_imported_function_or_module(cls, obj) -> None: assert namespace['Parser'] is XMLParser +_ImportDiscoveryOption = Literal[ + 'conditionals', 'try_except', 'contexts', 'loops', 'definitions', +] +_CompoundStatement = Literal[ + 'function-def', + 'async-function-def', # 3.5+ + 'class-def', + 'for-else', + 'async-for-else', # 3.5+ + 'while-else', + 'if-elif-else', + 'match-case', # 3.10+ + 'with', + 'async-with', # 3.5+ + 'try-except-else-finally', + 'try-except*-else-finally', # 3.11+ +] + + +def _get_toml_import_discovery_section( + options: set[_ImportDiscoveryOption], +) -> str: + config_file_lines = ['[tool.line_profiler.prof_mod_import_discovery]'] + for option in [ + 'conditionals', 'try_except', 'contexts', 'loops', 'definitions', + ]: + line = f'{option} = {str(option in options).lower()}' + config_file_lines.append(line) + return '\n'.join(config_file_lines) + + +def _grep_profiled_names(module_text: str) -> list[str]: + prof_pattern = ( + r'\badd_imported_function_or_module\((\w+(?:\.\w+)*)\)' + ) + return re.findall(prof_pattern, module_text) + + @pytest.mark.parametrize( ('prof_mod', 'expected_targets', 'profile_imports', 'profile_whole_file', 'expect_warnings'), @@ -1564,20 +1602,19 @@ def func() -> None: (['os.fork', 'foobar.bar'], {'fork'}, {'try_except', 'contexts'}), (['backup_fred', 'operator'], {'methodcaller', 'setitem'}, {'definitions'}), - (['backup_fred', 'operator'], {'fred'}, {'loops'})], -) + (['backup_fred', 'operator'], {'fred'}, {'loops'})]) def test_nested_import_discovery( prof_mod: list[str], expected: set[str], - options: set[Literal[ - 'conditionals', 'try_except', 'contexts', 'loops', 'definitions', - ]], + options: set[_ImportDiscoveryOption], ) -> None: """ Check the source code transformed by :py:class:`.AstTreeProfiler` to see if the import-discovery selection options in the TOML file (``[tool.line_profiler.prof_mod_import_discovery]``) are handled - correctly. + correctly in a real-ish script, with some of the compound statements + hosting the import statements nested inside other coumpound + statements. """ test_module = ub.codeblock(""" from collections.abc import Generator @@ -1635,14 +1672,6 @@ def _restore_sys_path() -> Generator[None, None, None]: import backup_fred as fred """).strip('\n') - config_file_lines = ['[tool.line_profiler.prof_mod_import_discovery]'] - for option in [ - 'conditionals', 'try_except', 'contexts', 'loops', 'definitions', - ]: - line = f'{option} = {str(option in options).lower()}' - config_file_lines.append(line) - config_file = '\n'.join(config_file_lines) - with tempfile.TemporaryDirectory() as tmp: mod_fname = os.path.join(tmp, 'test_module.py') with open(mod_fname, 'w') as fobj: @@ -1650,7 +1679,7 @@ def _restore_sys_path() -> Generator[None, None, None]: cfg_fname = os.path.join(tmp, 'config.toml') with open(cfg_fname, 'w') as fobj: - print(config_file, file=fobj) + print(_get_toml_import_discovery_section(options), file=fobj) config = ConfigSource.from_config(cfg_fname) atp = AstTreeProfiler(mod_fname, prof_mod, False, config=config) @@ -1661,8 +1690,190 @@ def _restore_sys_path() -> Generator[None, None, None]: ]: print(f'{label.capitalize()}:\n{textwrap.indent(module_text, " ")}\n') - prof_pattern = ( - r'\badd_imported_function_or_module\((\w+(?:\.\w+)*)\)' + assert set(_grep_profiled_names(output_module)) == expected + + +@pytest.mark.parametrize( + ('compound_statement', 'options', 'should_be_profiled'), + [('function-def', set(), False), + ('function-def', {'definitions'}, True), + ('async-function-def', set(), False), + ('async-function-def', {'definitions'}, True), + ('class-def', set(), False), + ('class-def', {'definitions'}, True), + ('for-else', set(), False), + ('for-else', {'loops'}, True), + ('async-for-else', {'definitions'}, False), + ('async-for-else', {'definitions', 'loops'}, True), + ('while-else', set(), False), + ('while-else', {'loops'}, True), + ('if-elif-else', set(), False), + ('if-elif-else', {'conditionals'}, True), + ('match-case', set(), False), + ('match-case', {'conditionals'}, True), + ('with', set(), False), + ('with', {'contexts'}, True), + ('async-with', {'definitions'}, False), + ('async-with', {'definitions', 'contexts'}, True), + ('try-except-else-finally', set(), False), + ('try-except-else-finally', {'try_except'}, True), + ('try-except*-else-finally', set(), False), + ('try-except*-else-finally', {'try_except'}, True)]) +def test_import_discovery_in_all_compound_statements( + compound_statement: _CompoundStatement, + options: set[_ImportDiscoveryOption], + should_be_profiled: bool, +) -> None: + """ + Exhaustive "unit" test for imports nested in all the kwown + compound-statement language constructions, and all their respective + config-level switches. + + Notes: + - If a construction is not valid in the current Python version, + the subtest is skipped. + + - Some ``async`` constructions are nested inside a coroutine + definition by necessity. + """ + test_cases = { + 'function-def': """ + def func(): + import foo, bar + import baz + import foobar + + ... + """, + 'async-function-def': """ + async def coroutine(awaitable): + import foo + import bar + import baz, foobar + + await awaitable + """, + 'class-def': """ + class Class: + import foo + import bar, baz + import foobar + + ... + """, + 'for-else': """ + for _ in range(5): + import foo + import bar + + ... + else: + import baz + import foobar + + ... + """, + 'async-for-else': """ + async def agen(awaitable): + async for x in (await awaitable): + import foo, bar + + yield x + else: + import baz, foobar + ... + """, + 'while-else': """ + while True: + import foo, bar, baz + ... + else: + import foobar + """, + 'if-elif-else': """ + if True: + import foo + elif False: + import bar, baz + else: + import foobar + """, + 'match-case': """ + match [1, 2, 3]: + case [1, *a, 2]: + import foo + ... + case [1, 2, 3, b]: + import bar + ... + case [1, *c]: + import baz + ... + case _: + import foobar + """, + 'with': """ + with ctx: + import foo, bar, baz, foobar + ... + """, + 'async-with': """ + async def afunc(): + async with actx: + import foo + import bar, baz, foobar + ... + """, + 'try-except-else-finally': """ + try: + import foo + except ImportError: + import bar + else: + import baz + finally: + import foobar + """, + } + test_cases['try-except*-else-finally'] = ( + test_cases['try-except-else-finally'].replace('except', 'except*') ) - profiled_names = re.findall(prof_pattern, output_module) - assert set(profiled_names) == expected + version_bounds = { + 'async-function-def': (3, 5), + 'async-def': (3, 5), + 'async-for-else': (3, 5), + 'async-with': (3, 5), + 'match-case': (3, 10), + 'try-except*-else-finally': (3, 11), + } + all_names = {'foo', 'bar', 'baz', 'foobar'} + + test_case = ub.codeblock(test_cases[compound_statement]).strip('\n') + version_bound: tuple[int, ...] = version_bounds.get(compound_statement, ()) + if sys.version_info < version_bound: + pytest.skip( + reason=f'cannot test {compound_statement} on {sys.version_info}', + ) + + with tempfile.TemporaryDirectory() as tmp: + case_fname = os.path.join(tmp, 'test_case.py') + with open(case_fname, 'w') as fobj: + print(test_case, file=fobj) + + cfg_fname = os.path.join(tmp, 'config.toml') + with open(cfg_fname, 'w') as fobj: + print(_get_toml_import_discovery_section(options), file=fobj) + + config = ConfigSource.from_config(cfg_fname) + atp = AstTreeProfiler( + case_fname, list(all_names), False, config=config, + ) + output = ast.unparse(atp.profile()) + + for label, module_text in [ + ('input', test_case), ('output', output), + ]: + print(f'{label.capitalize()}:\n{textwrap.indent(module_text, " ")}\n') + + expected = all_names if should_be_profiled else set() + assert set(_grep_profiled_names(output)) == expected From 951caaa86568d11cff8ebd6be134cfcbbf18045e Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Tue, 21 Jul 2026 04:04:50 +0200 Subject: [PATCH 03/12] More granular control of import discovery line_profiler/autoprofile/profmod_extractor.py _ImportFinder.find() - Replaced param `collect_from_definitions` with `collect_from_func_defs` and `collect_from_class_defs` - The `collect_from_*` params now default to and accept `None`, in which case the values are resolved from the default config file _ImportFinder.filter_node_types() Replaced param `collect_from_definitions` with `collect_from_func_defs` and `collect_from_class_defs` ProfmodExtractor._config Chaged typing (`dict[str, bool]` -> `ConfigSource`) ProfmodExtractor._ast_get_imports_from_tree() Now a thin wrapper around `_ImportFinder.find()` so that instead of passing the `collect_from_*` arguments we can just pass the `._config` ProfmodExtractor.extract_all() Updated call to `._ast_get_imports_from_tree()` line_profiler/rc/line_profiler.toml ::[tool.line_profiler.prof_mod_import_discovery] - Replaced key-value pair `definitions = false` with `func_defs = false` and `class_defs = true`; because while class definitions are typically once-and-done, if we were to `add_imported_function_or_module()` on every function call it may disproportionately impact performance - Default for `loops` now false, again because of performance concerns tests/test_autoprofile.py test_multitarget_import_transformation_executes() Refactored internals outside so that other tests can reuse them test_nested_import_discovery() Updated parametrization and test-module body because of the split of `definitions` into `func_defs` and `class_defs` test_import_discovery_in_all_compound_statements() Updated parametrization because of the split of `definitions` into `func_defs` and `class_defs` test_nested_imports_correct_deduplication_across_scopes() New test ensuring that deduplication of the import target only happens in the same scope; e.g. importing the same function in the function body of different functions should result in a `add_imported_function_or_module()` call being interpolated once in each of the functions --- .../autoprofile/profmod_extractor.py | 117 ++++++++------ line_profiler/rc/line_profiler.toml | 15 +- tests/test_autoprofile.py | 146 ++++++++++++++---- 3 files changed, 204 insertions(+), 74 deletions(-) diff --git a/line_profiler/autoprofile/profmod_extractor.py b/line_profiler/autoprofile/profmod_extractor.py index 15f1fd21..0711ba46 100644 --- a/line_profiler/autoprofile/profmod_extractor.py +++ b/line_profiler/autoprofile/profmod_extractor.py @@ -21,8 +21,10 @@ _CompoundNodeType = Literal[ # Basic top-level nodes 'Module', 'Interactive', - # Definition nodes - 'FunctionDef', 'AsyncFunctionDef', 'ClassDef', + # Function-definition nodes + 'FunctionDef', 'AsyncFunctionDef', + # Class-definition nodes + 'ClassDef', # Loop nodes 'For', 'AsyncFor', 'While', # Conditional nodes @@ -61,49 +63,63 @@ def find( cls, node: ast.AST, *, - collect_from_conditionals: bool = True, - collect_from_try_except: bool = True, - collect_from_contexts: bool = True, - collect_from_loops: bool = True, - collect_from_definitions: bool = False, + collect_from_conditionals: bool | None = None, + collect_from_try_except: bool | None = None, + collect_from_contexts: bool | None = None, + collect_from_loops: bool | None = None, + collect_from_func_defs: bool | None = None, + collect_from_class_defs: bool | None = None, ) -> dict[tuple[str | int, ...], list[ImportTarget]]: """ Parameters: node (ast.AST): AST node - collect_from_conditionals (bool): + collect_from_conditionals (bool | None): Whether to collect imports inside :py:class:`ast.If` and :py:class:`ast.match_case` nodes (and their children). - collect_from_try_except (bool): + collect_from_try_except (bool | None): Whether to collect imports inside :py:class:`ast.ExceptHandler`, :py:class:`ast.Try`, and :py:class:`ast.TryStar` nodes (and their children). - collect_from_contexts (bool): + collect_from_contexts (bool | None): Whether to collect imports inside :py:class:`ast.AsyncWith` and :py:class:`ast.With` nodes (and their children). - collect_from_loops (bool): + collect_from_loops (bool | None): Whether to collect imports inside :py:class:`ast.AsyncFor`, :py:class:`ast.For`, and :py:class:`ast.While` nodes (and their children). - collect_from_definitions (bool): + collect_from_func_defs (bool | None): Whether to collect imports inside - :py:class:`AsyncFunctionDef`, :py:class:`ast.ClassDef` - and :py:class:`ast.FunctionDef` nodes (and their - children). + :py:class:`AsyncFunctionDef` and + :py:class:`ast.FunctionDef` nodes (and their children). + collect_from_class_defs (bool | None): + Whether to collect imports inside + :py:class:`ast.ClassDef` nodes (and their children). Returns: found_imports \ (dict[tuple[str | int, ...], list[ImportTarget]]): The import targets and their locations in the AST + + Notes: + If a ``collect_from_*`` option is set to :py:const:`None`, + the value is taken from the default configs. """ - node_types = cls.filter_node_types( - collect_from_conditionals=collect_from_conditionals, - collect_from_try_except=collect_from_try_except, - collect_from_contexts=collect_from_contexts, - collect_from_loops=collect_from_loops, - collect_from_definitions=collect_from_definitions, - ) + default_config = ConfigSource.from_default() + filter_kwargs = { + 'collect_from_conditionals': collect_from_conditionals, + 'collect_from_try_except': collect_from_try_except, + 'collect_from_contexts': collect_from_contexts, + 'collect_from_loops': collect_from_loops, + 'collect_from_func_defs': collect_from_func_defs, + 'collect_from_class_defs': collect_from_class_defs, + } + consolidated_filter_kwargs: dict[str, bool] = { + k: v if filter_kwargs[k] is None else cast(bool, filter_kwargs[k]) + for k, v in cls._get_filter_args(default_config).items() + } + node_types = cls.filter_node_types(**consolidated_filter_kwargs) visitor = cls(node_types) visitor.visit(node) return visitor.found_imports @@ -112,11 +128,12 @@ def find( def filter_node_types( cls, *, - collect_from_conditionals: bool = True, - collect_from_try_except: bool = True, - collect_from_contexts: bool = True, - collect_from_loops: bool = True, - collect_from_definitions: bool = False, + collect_from_conditionals: bool = False, + collect_from_try_except: bool = False, + collect_from_contexts: bool = False, + collect_from_loops: bool = False, + collect_from_func_defs: bool = False, + collect_from_class_defs: bool = False, ) -> dict[_CompoundNodeType, bool]: """ Select the which of the compound node types (i.e. those than can @@ -132,15 +149,31 @@ def filter_node_types( allowed.update({'AsyncWith', 'With'}) if collect_from_loops: allowed.update({'AsyncFor', 'For', 'While'}) - if collect_from_definitions: - allowed.update({ - 'AsyncFunctionDef', 'ClassDef', 'FunctionDef', - }) + if collect_from_func_defs: + allowed.update({'AsyncFunctionDef', 'FunctionDef'}) + if collect_from_class_defs: + allowed.update({'ClassDef'}) return { node_type: node_type in allowed for node_type in cls._compound_node_types } + @staticmethod + def _get_filter_args(config: ConfigSource) -> dict[str, bool]: + cfg = ( + config + .get_subconfig('prof_mod_import_discovery') + .conf_dict + ) + return { + 'collect_from_conditionals': cfg['conditionals'], + 'collect_from_try_except': cfg['try_except'], + 'collect_from_contexts': cfg['contexts'], + 'collect_from_loops': cfg['loops'], + 'collect_from_func_defs': cfg['func_defs'], + 'collect_from_class_defs': cfg['class_defs'], + } + def generic_visit(self, node: ast.AST) -> None: """ For the compound node types, only visit their children if said @@ -222,11 +255,7 @@ def __init__( self._prof_mod = prof_mod if config is None: config = ConfigSource.from_default() - self._config = ( - config - .get_subconfig('prof_mod_import_discovery') - .conf_dict - ) + self._config = config @staticmethod def _is_path(text: str) -> bool: @@ -326,7 +355,14 @@ def _get_modnames_to_profile_from_prof_mod( return modnames_to_profile - _ast_get_imports_from_tree = _ImportFinder.find + @staticmethod + def _ast_get_imports_from_tree( + node: ast.AST, config: ConfigSource | None = None, + ) -> dict[tuple[str | int, ...], list[ImportTarget]]: + if config is None: + config = ConfigSource.from_default() + kwargs = _ImportFinder._get_filter_args(config) + return _ImportFinder.find(node, **kwargs) @staticmethod def _find_modnames_in_tree_imports( @@ -421,12 +457,7 @@ def extract_all(self) -> dict[tuple[str | int, ...], list[ImportTarget]]: self._script_file, self._prof_mod ) import_targets = self._ast_get_imports_from_tree( - self._tree, - collect_from_conditionals=self._config['conditionals'], - collect_from_try_except=self._config['try_except'], - collect_from_contexts=self._config['contexts'], - collect_from_loops=self._config['loops'], - collect_from_definitions=self._config['definitions'], + self._tree, self._config, ) raw: dict[tuple[str | int, ...], list[ImportTarget]] = { (*loc, index): filtered_imports diff --git a/line_profiler/rc/line_profiler.toml b/line_profiler/rc/line_profiler.toml index 10151798..4de354fe 100644 --- a/line_profiler/rc/line_profiler.toml +++ b/line_profiler/rc/line_profiler.toml @@ -223,6 +223,10 @@ percent = 8 # we won't descend therein to look for imports. # - If all these options are set to false, only raw (from-)import # statements on the top level of the profiled code will be checked. +# - For performance reasons (each generated call to +# `add_imported_function_or_module()` can result in arbitrarily deep +# descent into the obejc tto be profiled), import discovery in loops +# and function/method definitions are disable dby default. [tool.line_profiler.prof_mod_import_discovery] @@ -237,7 +241,10 @@ try_except = true contexts = true # - `loops` (bool): # Whether to look into `[async] for-else` and `while-else` statements. -loops = true -# - `definitions` (bool): -# Whether to look into `[async] def` and `class` statements. -definitions = false +loops = false +# - `func_defs` (bool): +# Whether to look into `[async] def` statements. +func_defs = false +# - `class_defs` (bool): +# Whether to look into `class` statements. +class_defs = true diff --git a/tests/test_autoprofile.py b/tests/test_autoprofile.py index 45d2c75f..b100024b 100644 --- a/tests/test_autoprofile.py +++ b/tests/test_autoprofile.py @@ -10,7 +10,7 @@ import textwrap import tempfile from collections.abc import Collection, Sequence -from typing import Any, ClassVar, Literal +from typing import Any, Literal, get_args from warnings import catch_warnings, WarningMessage import pytest @@ -1424,6 +1424,17 @@ def _check_warnings( ) +class _RecordingProfiler: + """ + Mock :py:class:`line_profiler.LineProfiler` object. + """ + def __init__(self) -> None: + self.profiled_objects: list[Any] = [] + + def add_imported_function_or_module(self, obj) -> None: + self.profiled_objects.append(obj) + + def test_multitarget_import_transformation_executes() -> None: """ Test the runtime behavior of the transformed AST, including: @@ -1438,22 +1449,13 @@ def test_multitarget_import_transformation_executes() -> None: """ from xml.etree.ElementTree import Element, dump, XMLParser - class RecordingProfiler: - """ - Mock :py:class:`line_profiler.LineProfiler` object. - """ - @classmethod - def add_imported_function_or_module(cls, obj) -> None: - cls.profiled_objects.append(obj) - - profiled_objects: ClassVar[list[Any]] = [] - input_module = ub.codeblock(""" import os, sys as system from xml.etree.ElementTree import ( # `xml_dump` not profiled Element, dump as xml_dump, XMLParser as Parser, ) """) + mock_prof = _RecordingProfiler() with tempfile.TemporaryDirectory() as tmp: fpath = ub.Path(tmp) / 'script.py' fpath.write_text(input_module) @@ -1467,11 +1469,11 @@ def add_imported_function_or_module(cls, obj) -> None: ], False, ).profile() - namespace = {'profile': RecordingProfiler()} + namespace = {'profile': mock_prof} code = compile(module_ast, str(fpath), 'exec') exec(code, namespace) - assert RecordingProfiler.profiled_objects == [ + assert mock_prof.profiled_objects == [ os, sys, Element, @@ -1486,7 +1488,8 @@ def add_imported_function_or_module(cls, obj) -> None: _ImportDiscoveryOption = Literal[ - 'conditionals', 'try_except', 'contexts', 'loops', 'definitions', + 'conditionals', 'try_except', 'contexts', 'loops', + 'func_defs', 'class_defs', ] _CompoundStatement = Literal[ 'function-def', @@ -1505,12 +1508,13 @@ def add_imported_function_or_module(cls, obj) -> None: def _get_toml_import_discovery_section( - options: set[_ImportDiscoveryOption], + options: set[_ImportDiscoveryOption] | None = None, ) -> str: + all_options = set(get_args(_ImportDiscoveryOption)) + if options is None: + options = all_options config_file_lines = ['[tool.line_profiler.prof_mod_import_discovery]'] - for option in [ - 'conditionals', 'try_except', 'contexts', 'loops', 'definitions', - ]: + for option in all_options: line = f'{option} = {str(option in options).lower()}' config_file_lines.append(line) return '\n'.join(config_file_lines) @@ -1601,7 +1605,9 @@ def func() -> None: (['ersatz_foobar.my_baz'], {'baz'}, {'try_except', 'contexts'}), (['os.fork', 'foobar.bar'], {'fork'}, {'try_except', 'contexts'}), (['backup_fred', 'operator'], {'methodcaller', 'setitem'}, - {'definitions'}), + {'func_defs'}), + (['backup_fred', 'operator'], {'__getattr__'}, + {'class_defs'}), (['backup_fred', 'operator'], {'fred'}, {'loops'})]) def test_nested_import_discovery( prof_mod: list[str], @@ -1617,11 +1623,12 @@ def test_nested_import_discovery( statements. """ test_module = ub.codeblock(""" - from collections.abc import Generator + from collections.abc import Generator, Iterable, Mapping from contextlib import contextmanager from functools import partial from importlib import import_module from sys import path, version_info + from typing import Any notify_fork = partial(print, 'Forking...') try: @@ -1652,6 +1659,18 @@ def _restore_sys_path() -> Generator[None, None, None]: finally: setitem(path, slice(None), old) + class MyMapping(Mapping[str, Any]): + from operator import getitem as __getattr__ + + def __getitem__(self, key: str) -> Any: + ... + + def __iter__(self) -> Iterable[str]: + ... + + def __len__(self) -> int: + ... + with _restore_sys_path(): try: @@ -1696,15 +1715,15 @@ def _restore_sys_path() -> Generator[None, None, None]: @pytest.mark.parametrize( ('compound_statement', 'options', 'should_be_profiled'), [('function-def', set(), False), - ('function-def', {'definitions'}, True), + ('function-def', {'func_defs'}, True), ('async-function-def', set(), False), - ('async-function-def', {'definitions'}, True), + ('async-function-def', {'func_defs'}, True), ('class-def', set(), False), - ('class-def', {'definitions'}, True), + ('class-def', {'class_defs'}, True), ('for-else', set(), False), ('for-else', {'loops'}, True), - ('async-for-else', {'definitions'}, False), - ('async-for-else', {'definitions', 'loops'}, True), + ('async-for-else', {'func_defs'}, False), + ('async-for-else', {'func_defs', 'loops'}, True), ('while-else', set(), False), ('while-else', {'loops'}, True), ('if-elif-else', set(), False), @@ -1713,8 +1732,8 @@ def _restore_sys_path() -> Generator[None, None, None]: ('match-case', {'conditionals'}, True), ('with', set(), False), ('with', {'contexts'}, True), - ('async-with', {'definitions'}, False), - ('async-with', {'definitions', 'contexts'}, True), + ('async-with', {'func_defs'}, False), + ('async-with', {'func_defs', 'contexts'}, True), ('try-except-else-finally', set(), False), ('try-except-else-finally', {'try_except'}, True), ('try-except*-else-finally', set(), False), @@ -1877,3 +1896,76 @@ async def afunc(): expected = all_names if should_be_profiled else set() assert set(_grep_profiled_names(output)) == expected + + +@pytest.mark.parametrize('call', ['first', 'second', 'third']) +def test_nested_imports_correct_deduplication_across_scopes( + call: Literal['first', 'second', 'third'], +) -> None: + """ + Test that there is no aliasing in the check we have against + inserting duplicate ``profile.add_imported_function_or_module(...)`` + statements: duplicates should only be counted within the same scope. + + Note: + - At runtime, duplicates don't really matter in terms of + CORRECTNESS, because ultimately + :py:class:`line_profiler.LineProfiler.add_callable` is + idempotent. + + - However, since calls to + :py:func:`line_profiler.autoprofile.line_profiler_utils\ +.add_imported_function_or_module` + can result in arbitrary deep descent into the profiled object, + these interpolated calls can have an impact on the + PERFORMANCE, especially when inserted into function/method + bodies. For this reason, import discovery in function bodies + is off by default. + """ + from textwrap import indent + + test_module = ub.codeblock(""" + def first() -> str: + from textwrap import indent + + return indent('first', ' ') + + + def second() -> str: + from textwrap import indent as ind + + return ind('second', ' ') + + + def third() -> str: + from textwrap import indent, dedent + from textwrap import indent as _indent # Duplicate + + return _indent('third', ' ') + """).strip('\n') + + mock_prof = _RecordingProfiler() + with tempfile.TemporaryDirectory() as tmp: + mod_fname = os.path.join(tmp, 'test_module.py') + with open(mod_fname, 'w') as fobj: + print(test_module, file=fobj) + + cfg_fname = os.path.join(tmp, 'config.toml') + with open(cfg_fname, 'w') as fobj: + print(_get_toml_import_discovery_section(), file=fobj) + + mod_ast = AstTreeProfiler( + mod_fname, ['textwrap.indent'], False, + config=ConfigSource.from_config(cfg_fname), + ).profile() + print(ast.unparse(mod_ast)) + + namespace: dict[str, Any] = {'profile': mock_prof} + code = compile(mod_ast, mod_fname, 'exec') + exec(code, namespace) + + # Make the call; regardless of which of the functions is called, + # `textwrap.indent()` should be presented to the profiler exactly + # once + assert namespace[call]() == ' ' + call + assert mock_prof.profiled_objects == [indent] From 6c8937156647f27be790c47b932f71cc5370fb60 Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Wed, 22 Jul 2026 01:57:02 +0200 Subject: [PATCH 04/12] Handle star-imports line_profiler/autoprofile/ast_profile_transformer.py Updated type annotations and docstring formatting ast_create_star_import_node() New function parallel to `ast_create_profile_node()` which creates an additional AST node for profiling star-imports AstProfileTransformer.__init__() - Relaxed type of `profiled_imports` (`list[str] | None` -> `Collection[str] | None`) - Added param `profile_star_imports` line_profiler/autoprofile/ast_tree_profiler.py::AstTreeProfiler ._profile_ast_tree() - Added optional param `profile_star_imports` for toggling whether to handle star-imports - Added optional param `modnames_to_profile` to help with calls to `ast_create_star_import_node()` .profile() Added optional param `profile_star_imports` for toggling whether to handle star-imports line_profiler/autoprofile/line_profiler_utils.py::add_star_import() New function/pseudo-method for `LineProfiler` so that it can retrieve the objects imported by a star-import and profile them line_profiler/autoprofile/profmod_extractor.py::ProfmodExtractor ._find_modnames_in_tree_imports() - Loosened typing for param `modnames_to_profile` - Minor refactoring .extract_all() Added argument `filter_star_imports` for toggling whether to return star-imports (false) or dropping them with a warning (true) ._modnames_to_profile New cached property for use by `AstTreeProfiler` --- .../autoprofile/ast_profile_transformer.py | 203 +++++++++++++----- .../autoprofile/ast_tree_profiler.py | 91 ++++++-- .../autoprofile/line_profiler_utils.py | 87 +++++++- .../autoprofile/profmod_extractor.py | 64 +++--- 4 files changed, 342 insertions(+), 103 deletions(-) diff --git a/line_profiler/autoprofile/ast_profile_transformer.py b/line_profiler/autoprofile/ast_profile_transformer.py index 4d1f600e..5a089dbc 100644 --- a/line_profiler/autoprofile/ast_profile_transformer.py +++ b/line_profiler/autoprofile/ast_profile_transformer.py @@ -1,7 +1,7 @@ from __future__ import annotations import ast -from collections.abc import Callable, Sequence +from collections.abc import Callable, Collection, Sequence from functools import partial from os import PathLike from typing import cast, TypeVar @@ -17,12 +17,15 @@ def ast_create_profile_node( profiler_name: str = 'profile', attr: str = 'add_imported_function_or_module', ) -> ast.Expr: - """Create an abstract syntax tree node that adds an object to the profiler to be profiled. + """ + Create an abstract syntax tree node that adds an object to the + profiler to be profiled, by calling the ``attr`` method of + ``profile`` and passing ``modname`` to it. + + At runtime, this adds the object to the profiler so it can be + profiled. This node must be added after the first instance of + ``modname`` in the AST and before it is used. - An abstract syntax tree node is created which calls the attr method from profile and - passes modname to it. - At runtime, this adds the object to the profiler so it can be profiled. - This node must be added after the first instance of modname in the AST and before it is used. The node will look like: >>> # xdoctest: +SKIP >>> import foo.bar @@ -33,14 +36,15 @@ def ast_create_profile_node( name of the imported module. profiler_name (str): - name of the LineProfiler object. + name of the :py:class:`line_profiler.LineProfiler` object. attr (str): - name of the method of the LineProfiler object to call on the imported module. + name of the method of the :py:class:`LineProfiler` object to + call on the imported module. Returns: (_ast.Expr): expr - AST node that adds modname to profiler. + AST node that adds ``modname`` to profiler. """ func = ast.Attribute( value=ast.Name(id=profiler_name, ctx=ast.Load()), @@ -55,6 +59,83 @@ def ast_create_profile_node( return expr +def ast_create_star_import_node( + modname: str, + targets: Collection[str] | None, + profiler_name: str = 'profile', + attr: str = 'add_star_import', +) -> ast.Expr: + """ + AST node similar to that created by + :py:func:`.ast_create_profile_node`, except that it handles + star-imports (``from ... import *``), like: + + >>> # doctest: +SKIP + >>> from foo.bar import * + >>> profile.add_star_import( + ... 'foo.bar', ['foo.bar', 'spam.ham'], locals(), + ... ) + + Args: + modname (str): + name of the imported module. + + targets (Collection[str] | None): + profile-on-import module targets; if :py:const:`None`, all + the names imported by the star-import will be added to the + profiler. + + profiler_name (str): + name of the :py:class:`line_profiler.LineProfiler` object. + + attr (str): + name of the method of the + :py:class:`line_profiler.LineProfiler` object to call on the + imported module. + + Returns: + (_ast.Expr): expr + AST node that adds ``modname`` to profiler. + """ + func_node = ast.Attribute( + value=ast.Name(id=profiler_name, ctx=ast.Load()), + attr=attr, + ctx=ast.Load(), + ) + modname_node = ast.Constant(value=modname) + if targets is None: + targets_node: ast.Constant | ast.List = ast.Constant(value=None) + else: + targets_node = ast.List( + elts=[ast.Constant(value=t) for t in targets], + ctx=ast.Load(), + ) + namespace_node = ast.Call( + func=ast.Name(id='locals', ctx=ast.Load()), args=[], keywords=[], + ) + expr = ast.Expr(value=ast.Call( + func=func_node, + args=[modname_node, targets_node, namespace_node], + keywords=[], + )) + return expr + + +def _ast_create_node_from_import_target( + target: ImportTarget, + modnames_to_profile: Collection[str] | None = None, + profile_star_imports: bool = False, +) -> ast.Expr | None: + if target.resolved_name is None: # Star-imports + if not profile_star_imports: + return None + assert target.name.endswith('.*') + return ast_create_star_import_node( + target.name[:-2], modnames_to_profile, + ) + return ast_create_profile_node(target.resolved_name) + + class AstProfileTransformer(ast.NodeTransformer): """Transform an abstract syntax tree adding profiling to all of its objects. @@ -67,27 +148,31 @@ class AstProfileTransformer(ast.NodeTransformer): def __init__( self, profile_imports: bool = False, - profiled_imports: list[str] | None = None, + profiled_imports: Collection[str] | None = None, profiler_name: str = 'profile', + profile_star_imports: bool = False, ) -> None: """Initializes the AST transformer with the profiler name. Args: profile_imports (bool): - If True, profile all imports. + if True, profile all concrete (non-star) imports. - profiled_imports (List[str]): + profiled_imports (Collection[str]): list of dotted paths of imports to skip that have already been added to profiler. profiler_name (str): the profiler name used as decorator and for the method call to add to the object to the profiler. + + profile_star_imports (bool): + if this and ``profile_imports`` are True, also profile + star-imports. """ self._profile_imports = bool(profile_imports) - self._profiled_imports = ( - profiled_imports if profiled_imports is not None else [] - ) + self._profiled_imports = set(profiled_imports or ()) self._profiler_name = profiler_name + self._profile_star_imports = profile_star_imports self._dropped_star_imports: set[ImportTarget] = set() def _visit_func_def( @@ -101,11 +186,11 @@ def _visit_func_def( e.g. @staticmethod. Args: - node (Union[_ast.FunctionDef, _ast.AsyncFunctionDef]): + node (_ast.FunctionDef | _ast.AsyncFunctionDef): function/method in the AST Returns: - (Union[_ast.FunctionDef, _ast.AsyncFunctionDef]): node + node (_ast.FunctionDef | _ast.AsyncFunctionDef): function/method with profiling decorator """ decor_ids = set() @@ -126,58 +211,65 @@ def _visit_import( node: _Import, get_import_targets: Callable[[_Import], Sequence[ImportTarget]], ) -> _Import | list[_Import | ast.Expr]: - """Add a node that profiles an import - - If profile_imports is True and the import is not in profiled_imports, - a node which calls the profiler method, which adds the object to the profiler, - is added immediately after the import. + """ + Add a node that profiles an import. If ``profile_imports`` is + true and an import target is not in ``profiled_imports``, a node + which calls the profiler method adding the object to the + profiler is added immediately after the import. Args: - node (Union[_ast.Import,_ast.ImportFrom]): - import in the AST + node (_Import): + import[-from] node in the AST + get_import_targets \ +(Callable[[_Import], Sequence[ImportTarget]]): + helper callable for analyzing the node Returns: - (Union[Union[_ast.Import,_ast.ImportFrom],List[Union[_ast.Import,_ast.ImportFrom,_ast.Expr]]]): node - if profile_imports is False: - returns the import node - if profile_imports is True: - returns list containing the import node and the profiling node + node (_Import | list[_Import | _ast.Expr]): + if ``profile_imports`` is False: + the import node + if ``profile_imports`` is True: + a list containing the import node and the profiling + node(s) """ if not self._profile_imports: self.generic_visit(node) return node + this_visit = cast(_Import, self.generic_visit(node)) visited: list[_Import | ast.Expr] = [this_visit] - for name, target in zip( - node.names, get_import_targets(node), strict=True, - ): - node_name = name.name if name.asname is None else name.asname - if target.resolved_name is None: - # TODO: handle starred imports - self._dropped_star_imports.add(target) + for target in get_import_targets(node): + name = target.name + if name in self._profiled_imports: continue - if node_name in self._profiled_imports: - continue - self._profiled_imports.append(node_name) - expr = ast_create_profile_node(node_name) - visited.append(expr) + expr = _ast_create_node_from_import_target( + target, profile_star_imports=self._profile_star_imports, + ) + if expr is None: # Bookkeeping + self._dropped_star_imports.add(target) + else: + self._profiled_imports.add(name) + visited.append(expr) return visited def visit_Import( - self, node: ast.Import + self, node: ast.Import, ) -> ast.Import | list[ast.Import | ast.Expr]: - """Add a node that profiles an object imported using the "import foo" sytanx + """ + Add nodes that profile objects imported using the + ``import foo`` syntax. Args: node (_ast.Import): import in the AST Returns: - (Union[_ast.Import,List[Union[_ast.Import,_ast.Expr]]]): node - if profile_imports is False: - returns the import node - if profile_imports is True: - returns list containing the import node and the profiling node + node (_ast.Import | list[_ast.Import | _ast.Expr]): + if ``profile_imports`` is False: + the import node + if ``profile_imports`` is True: + a list containing the import node and the + profiling node(s) """ # Note: we don't actually care about the `ImportTarget.index` # here; in fact, we're just reusing the name-resolution @@ -188,18 +280,21 @@ def visit_Import( def visit_ImportFrom( self, node: ast.ImportFrom ) -> ast.ImportFrom | list[ast.ImportFrom | ast.Expr]: - """Add a node that profiles an object imported using the "from foo import bar" syntax + """ + Add nodes that profile objects imported using the + ``from foo import bar`` syntax. Args: node (_ast.ImportFrom): import in the AST Returns: - (Union[_ast.ImportFrom,List[Union[_ast.ImportFrom,_ast.Expr]]]): node - if profile_imports is False: - returns the import node - if profile_imports is True: - returns list containing the import node and the profiling node + node (_ast.Import | list[_ast.Import | _ast.Expr]): + if ``profile_imports`` is False: + the import node + if ``profile_imports`` is True: + a list containing the import node and the + profiling node(s) """ get_targets = partial(ImportTarget._from_import_from_node, 0) return self._visit_import(node, get_targets) diff --git a/line_profiler/autoprofile/ast_tree_profiler.py b/line_profiler/autoprofile/ast_tree_profiler.py index d20bc9d3..a6cf3a6d 100644 --- a/line_profiler/autoprofile/ast_tree_profiler.py +++ b/line_profiler/autoprofile/ast_tree_profiler.py @@ -2,13 +2,15 @@ import ast import os -from collections.abc import MutableSequence, Sequence -from typing import Any +from collections.abc import Collection, MutableSequence, Sequence +from typing import Any, cast from ._import_targets import ImportTarget from ..toml_config import ConfigSource -from .ast_profile_transformer import ( +from .ast_profile_transformer import ( # noqa: F401 AstProfileTransformer, + _ast_create_node_from_import_target, + # Keep import below for compatibility ast_create_profile_node, ) from .profmod_extractor import ProfmodExtractor @@ -124,6 +126,8 @@ def _profile_ast_tree( ], profile_full_script: bool = False, profile_imports: bool = False, + modnames_to_profile: Collection[str] = (), + profile_star_imports: bool = False, ) -> ast.Module: """ Add profiling to an abstract syntax tree by adding nodes to the @@ -160,6 +164,17 @@ def _profile_ast_tree( if True, and ``profile_full_script`` is True, profile all imports as well. + modnames_to_profile (Collection[str]): + module names to be profiled; needed for processing + star-imports (``from ... import *``). + + profile_star_imports (bool): + if True, profile star-imports (those included in + ``tree_imports_to_profile_dict``, or retrieved by + :py:class:`.AstProfileTransformer` if both + ``profile_full_script`` and ``profile_imports`` are + true) + Returns: (_ast.Module): tree abstract syntax tree with profiling. @@ -172,36 +187,51 @@ def _profile_ast_tree( imports = tree_imports_to_profile_dict[tree_loc] *loc, tree_index = tree_loc assert isinstance(tree_index, int) - body = self._descend(tree, loc) + body = cast(MutableSequence[ast.AST], self._descend(tree, loc)) assert isinstance(body, MutableSequence) for imp in reversed(imports): # Reversing keeps the order of the inserted nodes # consistent with the imports - name = imp.resolved_name - if name is None: # Star-imports; TODO: handle this - continue - expr = ast_create_profile_node(name) - body.insert(tree_index + 1, expr) - profiled_imports.append(name) + expr = _ast_create_node_from_import_target( + imp, modnames_to_profile, profile_star_imports, + ) + if expr is not None: + body.insert(tree_index + 1, expr) + profiled_imports.append(imp.name) if profile_full_script: tree = self._ast_transformer_class_handler._transform( tree, self._script_file, profile_imports=profile_imports, profiled_imports=profiled_imports, + profile_star_imports=profile_star_imports, ) ast.fix_missing_locations(tree) return tree - def profile(self) -> ast.Module: - """Create an abstract syntax tree of a script and add profiling to it. + def profile(self, profile_star_imports: bool = False) -> ast.Module: + """ + Create an abstract syntax tree of a script and add profiling to + it: + + - Read a script file and generates an abstract syntax tree. + + - Then matches imports in the script's AST with the names in + ``prof_mod``. + + - The matched imports are added to the profiler for profiling. + + - If the path to the script is found in ``prof_mod``, all + functions/methods, classes & modules are added to the + profiler. - Reads a script file and generates an abstract syntax tree. - Then matches imports in the script's AST with the names in prof_mod. - The matched imports are added to the profiler for profiling. - If path to script is found in prof_mod, all functions/methods, classes & modules are - added to the profiler. - If profile_imports is True as well as path to script in prof_mod, all the imports - in the script are added to the profiler. + - If ``profile_imports`` is True as well as path to script in + ``prof_mod``, all the imports in the script are added to the + profiler. + + Args: + profile_star_imports (bool): + if True, add targets imported by ``from ... import *`` + statements to the profiler. Returns: (_ast.Module): tree @@ -213,14 +243,33 @@ def profile(self) -> ast.Module: tree = self._get_script_ast_tree(self._script_file) - tree_imports_to_profile_dict = self._profmod_extractor_class_handler( + # Note: warnings about dropped star-imports can be issued from 2 + # places: + # - `ProfmodExtractor.extract_all(filter_star_imports=True)` + # - `._profile_ast_tree(...), where both + # `profile_full_script=True` and `profile_imports=True` + # So take care to enture that we don't have duplicate warnings + extractor = self._profmod_extractor_class_handler( tree, self._script_file, self._prof_mod, self._config, - ).extract_all() + ) + if profile_star_imports: + # Star imports recovered -> nothing to warn either way + filter_star_imports_in_extract_all = False + else: + filter_star_imports_in_extract_all = not ( + profile_full_script and self._profile_imports + ) + tree_imports_to_profile_dict = extractor.extract_all( + filter_star_imports=filter_star_imports_in_extract_all, + ) + tree_profiled = self._profile_ast_tree( tree, tree_imports_to_profile_dict, profile_full_script=profile_full_script, profile_imports=self._profile_imports, + modnames_to_profile=extractor._modnames_to_profile, + profile_star_imports=profile_star_imports, ) return tree_profiled diff --git a/line_profiler/autoprofile/line_profiler_utils.py b/line_profiler/autoprofile/line_profiler_utils.py index 31d95c7e..90efac42 100644 --- a/line_profiler/autoprofile/line_profiler_utils.py +++ b/line_profiler/autoprofile/line_profiler_utils.py @@ -1,9 +1,14 @@ from __future__ import annotations import inspect +import operator +from collections.abc import Callable, Collection, MutableMapping from functools import cached_property, partial, partialmethod +from importlib import import_module from types import FunctionType, MethodType, ModuleType -from typing import TYPE_CHECKING, Any, Literal, overload +from typing import TYPE_CHECKING, Any, Literal, cast, overload + +from .profmod_extractor import _should_profile if TYPE_CHECKING: # pragma: no cover from ..profiler_mixin import CLevelCallable, CythonCallable @@ -113,3 +118,83 @@ def add_imported_function_or_module( # they're called self.enable_by_count() return 1 if count else 0 + + +def add_star_import( + self, + import_from: str, + targets: Collection[str] | None, + namespace: MutableMapping[str, Any], + **kwargs +) -> Literal[0, 1]: + """ + Helper method for a :py:class:`~.line_profiler.LineProfiler` to + handle star-imports (``from import *``). + + Args: + import_from (str): + Module name to star-import from. + targets (Collection[str] | None): + Profile-on-import module targets; if :py:const:`None`, all + the names imported by the star-import will be added to the + profiler. + namespace (MutableMapping[str, Any]) + Namespace into which the names from ``import_from`` should + be imported. + **kwargs + Passed to :py:func:`.add_imported_function_or_module`. + + Returns: + 1 if any function is added to the profiler, 0 otherwise. + """ + # Dynamically inspect the module to see which names would've been + # inserted by a star-import + module = import_module(import_from) # TODO + try: + all_names: set[str] | None = set(cast( + Collection[str], getattr(module, '__all__', None), + )) + except Exception: + # Could be for whatever reason, there's no guarantee that a + # module declares the `.__all__` CORRECTLY + all_names = None + + check_attr: Callable[[str], bool] + if all_names is None: # Default behavior: take all public names + check_attr = lambda attr: ( # noqa: E731 + not attr.startswith('_') + ) + else: # If we have a valid `.__all__`, take names therefrom + check_attr = partial(operator.contains, all_names) + imported_names: dict[str, Any] = { + attr: value + for attr, value in inspect.getmembers(module) + if check_attr(attr) + } + + # Decide on which of the names to pass to the profiler + # Note: :the name shoul've already been inserted into the namespace + # by the import statement itself, so this is just post-hoc + # bookkeeping + add: Callable[[Any], int] + if hasattr(self, 'add_imported_function_or_module'): + # Pseudo-method inserted by `.autoprofile.run()` + add = partial(self.add_imported_function_or_module, **kwargs) + else: + add = partial(add_imported_function_or_module, self, **kwargs) + count = 0 + sentinel = object() + for name, value in imported_names.items(): + if not ( + targets is None + or _should_profile(targets, f'{import_from}.{name}') + ): + # Check that the name should be profiled (if we have + # constrained `targets`) + continue + if namespace.get(name, sentinel) is not value: + # Check that the actual object in the namespace is + # consistent with what is imported + continue + count += add(value) + return 1 if count else 0 diff --git a/line_profiler/autoprofile/profmod_extractor.py b/line_profiler/autoprofile/profmod_extractor.py index 0711ba46..5a8833dd 100644 --- a/line_profiler/autoprofile/profmod_extractor.py +++ b/line_profiler/autoprofile/profmod_extractor.py @@ -3,7 +3,8 @@ import ast import os import sys -from collections.abc import Sequence +from collections.abc import Collection, Sequence +from functools import cached_property from typing import TYPE_CHECKING, ClassVar, Literal, cast, get_args from warnings import warn @@ -366,7 +367,7 @@ def _ast_get_imports_from_tree( @staticmethod def _find_modnames_in_tree_imports( - modnames_to_profile: Sequence[str], + modnames_to_profile: Collection[str], import_targets: Sequence[ImportTarget], ) -> dict[int, list[ImportTarget]]: """Map modnames to imports from an abstract sytax tree. @@ -380,7 +381,7 @@ def _find_modnames_in_tree_imports( The import's alias is stored in the output dict. Args: - modnames_to_profile (Sequence[str]): + modnames_to_profile (Collection[str]): list of dotted paths to profile. import_targets (Sequence[ImportTarget]): @@ -400,12 +401,7 @@ def _find_modnames_in_tree_imports( modname = import_target.name if modname in modname_added_list: continue - # Check if either the parent module or submodule are in - # `modnames_to_profile` - if ( - modname not in modnames_to_profile - and modname.rsplit('.', 1)[0] not in modnames_to_profile - ): + if not _should_profile(modnames_to_profile, modname): continue modname_added_list.append(modname) try: @@ -414,13 +410,20 @@ def _find_modnames_in_tree_imports( filtered_imports[import_target.index] = [import_target] return filtered_imports - def extract_all(self) -> dict[tuple[str | int, ...], list[ImportTarget]]: + def extract_all( + self, filter_star_imports: bool = True, + ) -> dict[tuple[str | int, ...], list[ImportTarget]]: """ Map ``prof_mod`` to imports in an abstract syntax tree. Takes the paths and dotted paths in ``prof_mod`` and finds their respective imports in an abstract syntax tree, returning their aliases and the location they appear in the AST. + Args: + filter_star_imports (bool): + If true, filter out star imports + (``from import *``) with a warning. + Returns: tree_imports_to_profile_dict \ (dict[tuple[str | int, ...], list[ImportTarget]]); @@ -448,14 +451,7 @@ def extract_all(self) -> dict[tuple[str | int, ...], list[ImportTarget]]: Name under which the import is inserted into the namespace (should never be :py:const`None` for non-star-imports) - - Notes: - As of now, ``from import *`` is not supported, and - will result in a :py:class:`UserWarning`. """ - modnames_to_profile = self._get_modnames_to_profile_from_prof_mod( - self._script_file, self._prof_mod - ) import_targets = self._ast_get_imports_from_tree( self._tree, self._config, ) @@ -463,14 +459,12 @@ def extract_all(self) -> dict[tuple[str | int, ...], list[ImportTarget]]: (*loc, index): filtered_imports for loc, imports in import_targets.items() for index, filtered_imports in self._find_modnames_in_tree_imports( - modnames_to_profile, imports + self._modnames_to_profile, imports, ).items() } filtered: dict[tuple[str | int, ...], list[ImportTarget]] = {} star_imports: set[ImportTarget] = set() for loc, imports in raw.items(): - # TODO: runtime introspection of imports to handle - # star-imports # Notes: # - We don't issue the warning in # `._find_modnames_in_tree_imports()` because that is a @@ -482,15 +476,16 @@ def extract_all(self) -> dict[tuple[str | int, ...], list[ImportTarget]]: # which would be the sole target thereof (so # `indices_to_drop` should either be `[]` or `[0]`); # but it doesn't hurt to be cautious - indices_to_drop = [ - i for i, imp in enumerate(imports) - if imp.resolved_name is None # Star-imports - ] - for i in reversed(indices_to_drop): - imp = imports.pop(i) - star_imports.add(imp) + if filter_star_imports: + indices_to_drop = [ + i for i, imp in enumerate(imports) + if imp.resolved_name is None # Star-imports + ] + for i in reversed(indices_to_drop): + star_imports.add(imports.pop(i)) if imports: filtered[loc] = imports + # Attribute the warning to the caller ImportTarget._check_and_warn_dropped_imports( star_imports, "we don't currently handle `from ... import *` statements", @@ -563,3 +558,18 @@ def run(self) -> dict[int, str]: stacklevel=2, # Attribute warning to caller ) return result + + @cached_property + def _modnames_to_profile(self) -> frozenset[str]: + return frozenset(self._get_modnames_to_profile_from_prof_mod( + self._script_file, self._prof_mod, + )) + + +def _should_profile(targets: Collection[str], modname: str) -> bool: + """ + Check if either the parent module or submodule are in + `targets` + """ + names = {modname, modname.rsplit('.', 1)[0]} + return bool(names.intersection(targets)) From b3f2b619293bf4482dd5ba78ed796bbd7c1a193f Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Thu, 23 Jul 2026 03:54:45 +0200 Subject: [PATCH 05/12] Restructured configs line_profiler/autoprofile/ast_tree_profiler.py ::AstTreeProfiler.profile() Param `profile_star_imports` now takes `None`, where the value is then resolved from the `config` supplied at initialization line_profiler/autoprofile/profmod_extractor.py _ImportFinder._get_filter_args() Updated config loc ProfmodExtractor.extract_all() Param `filter_star_imports` now takes `None`, where the value is then resolved from the `config` supplied at initialization line_profiler/rc/line_profiler.toml::[tool.line_profiler.autoprofile] New subtable for config options related to `line_profiler.autoprofile`: - `prof_star_imports`: New boolean value for the default of: - `AstTreeProfiler.profile(profile_star_imports=...)` - `ProfmodExtractor.extract_all(filter_star_imports=...)` (negated) - `import_discovery`: Migrated from `tool.line_profiler.prof_mod_import_discovery` --- .../autoprofile/ast_tree_profiler.py | 13 +++-- .../autoprofile/profmod_extractor.py | 21 ++++++-- line_profiler/rc/line_profiler.toml | 54 ++++++++++++------- tests/test_autoprofile.py | 4 +- 4 files changed, 62 insertions(+), 30 deletions(-) diff --git a/line_profiler/autoprofile/ast_tree_profiler.py b/line_profiler/autoprofile/ast_tree_profiler.py index a6cf3a6d..bc2b434d 100644 --- a/line_profiler/autoprofile/ast_tree_profiler.py +++ b/line_profiler/autoprofile/ast_tree_profiler.py @@ -13,7 +13,7 @@ # Keep import below for compatibility ast_create_profile_node, ) -from .profmod_extractor import ProfmodExtractor +from .profmod_extractor import ProfmodExtractor, _should_profile_star_imports __docstubs__ = """ from .ast_profile_transformer import AstProfileTransformer @@ -208,7 +208,7 @@ def _profile_ast_tree( ast.fix_missing_locations(tree) return tree - def profile(self, profile_star_imports: bool = False) -> ast.Module: + def profile(self, profile_star_imports: bool | None = None) -> ast.Module: """ Create an abstract syntax tree of a script and add profiling to it: @@ -229,14 +229,19 @@ def profile(self, profile_star_imports: bool = False) -> ast.Module: profiler. Args: - profile_star_imports (bool): + profile_star_imports (bool | None): if True, add targets imported by ``from ... import *`` - statements to the profiler. + statements to the profiler; + if :py:const:`None`, it is loaded from the ``config`` + (from `autoprofile.prof_star_imports`). Returns: (_ast.Module): tree abstract syntax tree with profiling. """ + if profile_star_imports is None: + profile_star_imports = _should_profile_star_imports(self._config) + profile_full_script = self._check_profile_full_script( self._script_file, self._prof_mod ) diff --git a/line_profiler/autoprofile/profmod_extractor.py b/line_profiler/autoprofile/profmod_extractor.py index 5a8833dd..e286a0ab 100644 --- a/line_profiler/autoprofile/profmod_extractor.py +++ b/line_profiler/autoprofile/profmod_extractor.py @@ -163,7 +163,7 @@ def filter_node_types( def _get_filter_args(config: ConfigSource) -> dict[str, bool]: cfg = ( config - .get_subconfig('prof_mod_import_discovery') + .get_subconfig('autoprofile', 'import_discovery') .conf_dict ) return { @@ -411,7 +411,7 @@ def _find_modnames_in_tree_imports( return filtered_imports def extract_all( - self, filter_star_imports: bool = True, + self, filter_star_imports: bool | None = None, ) -> dict[tuple[str | int, ...], list[ImportTarget]]: """ Map ``prof_mod`` to imports in an abstract syntax tree. @@ -420,9 +420,11 @@ def extract_all( aliases and the location they appear in the AST. Args: - filter_star_imports (bool): + filter_star_imports (bool | None): If true, filter out star imports - (``from import *``) with a warning. + (``from import *``) with a warning; + if :py:const:`None`, it is loaded from the ``config`` + (as the negation of `autoprofile.prof_star_imports`). Returns: tree_imports_to_profile_dict \ @@ -452,6 +454,10 @@ def extract_all( the namespace (should never be :py:const`None` for non-star-imports) """ + if filter_star_imports is None: + filter_star_imports = not _should_profile_star_imports( + self._config, + ) import_targets = self._ast_get_imports_from_tree( self._tree, self._config, ) @@ -573,3 +579,10 @@ def _should_profile(targets: Collection[str], modname: str) -> bool: """ names = {modname, modname.rsplit('.', 1)[0]} return bool(names.intersection(targets)) + + +def _should_profile_star_imports(config: ConfigSource | None) -> bool: + if config is None: + config = ConfigSource.from_default() + kvps = config.get_subconfig('autoprofile').conf_dict + return bool(kvps['prof_star_imports']) diff --git a/line_profiler/rc/line_profiler.toml b/line_profiler/rc/line_profiler.toml index 4de354fe..5cf11cf5 100644 --- a/line_profiler/rc/line_profiler.toml +++ b/line_profiler/rc/line_profiler.toml @@ -207,28 +207,42 @@ time = 12 perhit = 8 percent = 8 -# `line_profiler.autoprofile.profmod_extractor.ProfmodExtractor` -# selective import-profiling options +# `line_profiler.autoprofile` import-profiling options # Note: -# - When using `preimports = true`, the `prof-mod` names are always -# profiled. -# - Otherwise, if the file-/module-path of the profiled code is itself -# among the `prof-mod` names, its imports are ALWAYS profiled. -# - Else, imports are selectively discovered from the AST of the -# profiled code, and import targets specifically matching the -# `prof-mod` names are profiled. Import discovery is dictated by the -# following options: -# - If the option corresponding to a language construction is `false`, -# we won't descend therein to look for imports. -# - If all these options are set to false, only raw (from-)import -# statements on the top level of the profiled code will be checked. -# - For performance reasons (each generated call to -# `add_imported_function_or_module()` can result in arbitrarily deep -# descent into the obejc tto be profiled), import discovery in loops -# and function/method definitions are disable dby default. - -[tool.line_profiler.prof_mod_import_discovery] +# 1. When using `preimports = true`, the `prof-mod` names are always +# profiled, regardless of whether they are directly or indirectly +# imported in the profiled code. +# 2, If (i) the file-/module-path of the profiled code is itself among +# the `prof-mod` names, and (ii) `prof-imports = true`, its imports +# are ALWAYS profiled via module-AST rewriting. +# 3. Else, imports are selectively discovered from the AST of the +# profiled code, and import targets specifically matching the +# `prof-mod` names are profiled. Import discovery is dictated by the +# `import_discovery` options: +# - If the option corresponding to a language construction is `false`, +# we won't descend therein to look for imports. +# - If all these options are set to false, only raw (from-)import +# statements on the top level of the profiled code will be checked. +# 4. For performance reasons (each generated call to +# `add_imported_function_or_module()` can result in arbitrarily deep +# descent into the object to be profiled), import discovery in loops +# and function/method definitions are disabled by default. +# 5. Star-imports (`from ... import *`) are ignored (with a warning) by +# default, due to the extra complexity in retrieving the import +# targets; to profile those (either when individual import targets +# match the `prof-mod` names, or when the profiled code goes through +# whole-module rewriting (see point 2) with `prof-imports` true), set +# `prof_star_imports` to true. + +[tool.line_profiler.autoprofile] + +# - `prof_star_imports` (bool): +# Whether to profile targets implicitly imported via +# `from ... import *` statements +prof_star_imports = false + +[tool.line_profiler.autoprofile.import_discovery] # - `conditionals` (bool): # Whether to look into `if-elif-else` and `match-case` statements. diff --git a/tests/test_autoprofile.py b/tests/test_autoprofile.py index b100024b..5ae1acba 100644 --- a/tests/test_autoprofile.py +++ b/tests/test_autoprofile.py @@ -1513,7 +1513,7 @@ def _get_toml_import_discovery_section( all_options = set(get_args(_ImportDiscoveryOption)) if options is None: options = all_options - config_file_lines = ['[tool.line_profiler.prof_mod_import_discovery]'] + config_file_lines = ['[tool.line_profiler.autoprofile.import_discovery]'] for option in all_options: line = f'{option} = {str(option in options).lower()}' config_file_lines.append(line) @@ -1617,7 +1617,7 @@ def test_nested_import_discovery( """ Check the source code transformed by :py:class:`.AstTreeProfiler` to see if the import-discovery selection options in the TOML file - (``[tool.line_profiler.prof_mod_import_discovery]``) are handled + (``[tool.line_profiler.autoprofile.import_discovery]``) are handled correctly in a real-ish script, with some of the compound statements hosting the import statements nested inside other coumpound statements. From 46ef76609a8b199c50e4ceb20aa44a81cb28661e Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Thu, 23 Jul 2026 16:44:22 +0200 Subject: [PATCH 06/12] `import_dicovery` options for `ast_profile_transformer` line_profiler/autoprofile/ast_profile_transformer.py ::AstProfileTransformer .__init__() Added optional param `profile_imports_in` and attributes `._should_visit_imports` and `._current_loc` for controlling whether an import statement should be profiled ._visit_import() Added check to skip import statements nested inside deselected compound statements .visit() New method wrapping around `NodeTransformer.visit()` and do bookkeeping on the current location (i.e. what kinds of nodes are we nested in) ._transform() Added optional param `config` to read in the `tool.line_profiler.autoprofile.import_dicovery` table and decide on which import statements to profile line_profiler/autoprofile/ast_tree_profiler.py ::AstTreeProfiler.__init__() Param `config` now keyword-only to prevent cluttering the signature line_profiler/autoprofile/autoprofile.py _extend_line_profiler_for_profiling_imports() Added missing `.add_star_import()` method to the `prof` run() Param `config` now keyword-only to prevent cluttering the signature line_profiler/autoprofile/profmod_extractor.py _ImportFinder.__init__() Loosened type annotation on param `node_types` (`dict[...]` -> `Mapping[...]`) ProfmodExtractor.__init__() Param `config` now keyword-only to prevent cluttering the signature tests/test_autoprofile.py test_import_discovery_in_all_compound_statements() Added parametrization to test the two methods of import discovery and profiling (`ProfmodExtractor` via `AstTreeProfiler`, and `AstProfileTransformer`) test_nested_imports_correct_deduplication_across_scopes() Ditto above (FIXME: currently failing) --- .../autoprofile/ast_profile_transformer.py | 99 ++++++++++++++++--- .../autoprofile/ast_tree_profiler.py | 4 +- line_profiler/autoprofile/autoprofile.py | 31 ++++-- .../autoprofile/profmod_extractor.py | 5 +- tests/test_autoprofile.py | 74 +++++++++++--- 5 files changed, 177 insertions(+), 36 deletions(-) diff --git a/line_profiler/autoprofile/ast_profile_transformer.py b/line_profiler/autoprofile/ast_profile_transformer.py index 5a089dbc..a67380e2 100644 --- a/line_profiler/autoprofile/ast_profile_transformer.py +++ b/line_profiler/autoprofile/ast_profile_transformer.py @@ -1,16 +1,24 @@ from __future__ import annotations import ast -from collections.abc import Callable, Collection, Sequence +from collections.abc import Callable, Collection, Mapping, Sequence from functools import partial from os import PathLike -from typing import cast, TypeVar +from types import MappingProxyType +from typing import TypeVar, cast, get_args +from ..toml_config import ConfigSource from ._import_targets import ImportTarget +from .profmod_extractor import _CompoundNodeType, _ImportFinder _Import = TypeVar('_Import', ast.Import, ast.ImportFrom) +_PROFILE_IMPORTS_IN_DEFAULT: MappingProxyType[_CompoundNodeType, bool] +_PROFILE_IMPORTS_IN_DEFAULT = MappingProxyType(dict.fromkeys( + get_args(_CompoundNodeType), True, +)) + def ast_create_profile_node( modname: str, @@ -137,12 +145,19 @@ def _ast_create_node_from_import_target( class AstProfileTransformer(ast.NodeTransformer): - """Transform an abstract syntax tree adding profiling to all of its objects. - - Adds profiler decorators on all functions & methods that are not already decorated with - the profiler. - If profile_imports is True, a profiler method call to profile is added to all imports - immediately after the import. + """ + Transform an abstract syntax tree adding profiling to all of its + objects, by: + + - Adding profiler decorators on all functions & methods that are not + already decorated with the profiler. + - If ``profile_imports`` is True, a profiler method call (see + :py:func:`line_profiler.autoprofile.line_profiler_utils\ +.add_imported_function_or_module` + and + :py:func:`line_profiler.autoprofile.line_profiler_utils\ +.add_star_import`) + is added to all imports immediately after the import. """ def __init__( @@ -150,7 +165,11 @@ def __init__( profile_imports: bool = False, profiled_imports: Collection[str] | None = None, profiler_name: str = 'profile', + *, profile_star_imports: bool = False, + profile_imports_in: Mapping[ + _CompoundNodeType, bool + ] = _PROFILE_IMPORTS_IN_DEFAULT, ) -> None: """Initializes the AST transformer with the profiler name. @@ -168,12 +187,22 @@ def __init__( profile_star_imports (bool): if this and ``profile_imports`` are True, also profile star-imports. + + profile_imports_in \ +(Mapping[Literal['Module', 'Interactive', \ +'FunctionDef', 'AsyncFunctionDef', 'ClassDef', \ +'For', `AsyncFor`, `While`, 'If', 'match_case', \ +'With', 'AsyncWith'. 'Try', 'TryStar', 'ExceptHandler'], bool]): + for each of the compound-statemnt node type, whether to + profile import statements residing therein. """ self._profile_imports = bool(profile_imports) self._profiled_imports = set(profiled_imports or ()) self._profiler_name = profiler_name self._profile_star_imports = profile_star_imports + self._should_visit_imports = dict(profile_imports_in) self._dropped_star_imports: set[ImportTarget] = set() + self._current_loc: list[str] = [] def _visit_func_def( self, node: ast.FunctionDef | ast.AsyncFunctionDef @@ -212,9 +241,16 @@ def _visit_import( get_import_targets: Callable[[_Import], Sequence[ImportTarget]], ) -> _Import | list[_Import | ast.Expr]: """ - Add a node that profiles an import. If ``profile_imports`` is - true and an import target is not in ``profiled_imports``, a node - which calls the profiler method adding the object to the + Add a node that profiles an import. If: + + - ``profile_imports`` is true, + + - The import statement isn't nested in a compound-statement node + type explicitly excluded via ``profile_imports_in``, and + + - The import target is not in ``profiled_imports``, + + a node which calls the profiler method adding the object to the profiler is added immediately after the import. Args: @@ -233,6 +269,18 @@ def _visit_import( node(s) """ if not self._profile_imports: + should_profile = False + else: + # Check if this node is nested inside compound statements + # that we shouldn't look for imports in + *ancestry, _ = self._current_loc + svi = self._should_visit_imports + should_profile = all( + svi.get(cast(_CompoundNodeType, a_type), True) + for a_type in ancestry + ) + + if not should_profile: self.generic_visit(node) return node @@ -277,6 +325,15 @@ def visit_Import( get_targets = partial(ImportTarget._from_import_node, 0) return self._visit_import(node, get_targets) + def visit(self, node: ast.AST) -> ast.AST | list[ast.AST]: + # Bookkeeping + loc = self._current_loc + loc.append(type(node).__name__) + try: + return super().visit(node) + finally: + loc.pop() + def visit_ImportFrom( self, node: ast.ImportFrom ) -> ast.ImportFrom | list[ast.ImportFrom | ast.Expr]: @@ -299,11 +356,23 @@ def visit_ImportFrom( get_targets = partial(ImportTarget._from_import_from_node, 0) return self._visit_import(node, get_targets) + @staticmethod + def _get_profile_imports_in( + config: ConfigSource | None = None, + ) -> dict[_CompoundNodeType, bool]: + if config is None: + config = ConfigSource.from_default() + return _ImportFinder.filter_node_types( + **_ImportFinder._get_filter_args(config), + ) + @classmethod def _transform( cls, node: ast.Module, filename: PathLike[str] | str | None = None, + *, + config: ConfigSource | None = None, **kwargs, ) -> ast.Module: """ @@ -314,6 +383,9 @@ def _transform( AST module node filename (PathLike[str] | str | None): Optional filename to be used in error/warning messages + config (ConfigSource | None): + Optional :py:class:`.ConfigSource` to load options from, + controlling whether an import should be profiled **kwargs Passed to the initializer @@ -321,12 +393,15 @@ def _transform( node (ast.Module): Input module node """ + kwargs.setdefault( + 'profile_imports_in', cls._get_profile_imports_in(config), + ) transformer = cls(**kwargs) dropped_star_imports = transformer._dropped_star_imports if filename is None: filename = '???' try: - return transformer.visit(node) + return cast(ast.Module, transformer.visit(node)) finally: ImportTarget._check_and_warn_dropped_imports( dropped_star_imports, diff --git a/line_profiler/autoprofile/ast_tree_profiler.py b/line_profiler/autoprofile/ast_tree_profiler.py index bc2b434d..750cd149 100644 --- a/line_profiler/autoprofile/ast_tree_profiler.py +++ b/line_profiler/autoprofile/ast_tree_profiler.py @@ -40,6 +40,7 @@ def __init__( profmod_extractor_class_handler: ( type[ProfmodExtractor] ) = ProfmodExtractor, + *, config: ConfigSource | None = None, ) -> None: """Initializes the AST tree profiler instance with the script file path @@ -204,6 +205,7 @@ def _profile_ast_tree( profile_imports=profile_imports, profiled_imports=profiled_imports, profile_star_imports=profile_star_imports, + config=self._config, ) ast.fix_missing_locations(tree) return tree @@ -255,7 +257,7 @@ def profile(self, profile_star_imports: bool | None = None) -> ast.Module: # `profile_full_script=True` and `profile_imports=True` # So take care to enture that we don't have duplicate warnings extractor = self._profmod_extractor_class_handler( - tree, self._script_file, self._prof_mod, self._config, + tree, self._script_file, self._prof_mod, config=self._config, ) if profile_star_imports: # Star imports recovered -> nothing to warn either way diff --git a/line_profiler/autoprofile/autoprofile.py b/line_profiler/autoprofile/autoprofile.py index 3e98a92c..a53dce51 100644 --- a/line_profiler/autoprofile/autoprofile.py +++ b/line_profiler/autoprofile/autoprofile.py @@ -57,27 +57,39 @@ def main(): from ..line_profiler_utils import restore from .ast_tree_profiler import AstTreeProfiler from .run_module import AstTreeModuleProfiler -from .line_profiler_utils import add_imported_function_or_module +from .line_profiler_utils import ( + add_imported_function_or_module, add_star_import, +) from .util_static import modpath_to_modname PROFILER_LOCALS_NAME = 'prof' def _extend_line_profiler_for_profiling_imports(prof: Any) -> None: - """Allow profiler to handle functions/methods, classes & modules with a single call. + """ + Allow profiler to handle imported functions/methods, classes and + modules, and also star-import targets, with a single call. This + adds to a :py:class:`line_profiler.LineProfiler` instance: + + - A method that can identify whether the object is a + function/method, class, or module, and handle it's profiling + accordingly; and + + - A method that can retrieve the names imported via a star-import + and use the above handling to profile them. - Add a method to LineProfiler that can identify whether the object is a - function/method, class or module and handle it's profiling accordingly. Mainly used for profiling objects that are imported. - (Workaround to keep changes needed by autoprofile separate from base LineProfiler) Args: prof (LineProfiler): - instance of LineProfiler. + instance of :py:class:`line_profiler.LineProfiler`. + + Notes: + This is a workaround to keep changes needed by autoprofile + separate from the base :py:class:`line_profiler.LineProfiler`. """ - prof.add_imported_function_or_module = types.MethodType( - add_imported_function_or_module, prof - ) + for func in add_imported_function_or_module, add_star_import: + setattr(prof, func.__name__, types.MethodType(func, prof)) def run( @@ -86,6 +98,7 @@ def run( prof_mod: list[str], profile_imports: bool = False, as_module: bool = False, + *, config: os.PathLike[str] | str | None = None, ) -> None: """Automatically profile a script and run it. diff --git a/line_profiler/autoprofile/profmod_extractor.py b/line_profiler/autoprofile/profmod_extractor.py index e286a0ab..3caab517 100644 --- a/line_profiler/autoprofile/profmod_extractor.py +++ b/line_profiler/autoprofile/profmod_extractor.py @@ -3,7 +3,7 @@ import ast import os import sys -from collections.abc import Collection, Sequence +from collections.abc import Collection, Mapping, Sequence from functools import cached_property from typing import TYPE_CHECKING, ClassVar, Literal, cast, get_args from warnings import warn @@ -48,7 +48,7 @@ class _ImportFinder(ast.NodeVisitor): def __init__( self, - node_types: dict[_CompoundNodeType, bool], + node_types: Mapping[_CompoundNodeType, bool], found_imports: ( dict[tuple[str | int, ...], list[ImportTarget]] | None ) = None, @@ -227,6 +227,7 @@ def __init__( tree: ast.Module, script_file: str, prof_mod: Sequence[str], + *, config: ConfigSource | None = None, ) -> None: """ diff --git a/tests/test_autoprofile.py b/tests/test_autoprofile.py index 5ae1acba..c6632f46 100644 --- a/tests/test_autoprofile.py +++ b/tests/test_autoprofile.py @@ -16,6 +16,9 @@ import pytest import ubelt as ub from line_profiler.toml_config import ConfigSource +from line_profiler.autoprofile.ast_profile_transformer import ( + AstProfileTransformer, +) from line_profiler.autoprofile.ast_tree_profiler import AstTreeProfiler from line_profiler.autoprofile.profmod_extractor import ProfmodExtractor @@ -1431,6 +1434,9 @@ class _RecordingProfiler: def __init__(self) -> None: self.profiled_objects: list[Any] = [] + def __call__(self, x: Any) -> Any: + return x + def add_imported_function_or_module(self, obj) -> None: self.profiled_objects.append(obj) @@ -1712,6 +1718,8 @@ def __len__(self) -> int: assert set(_grep_profiled_names(output_module)) == expected +@pytest.mark.parametrize('use_component', + ['ast_tree_profiler', 'ast_profile_transformer']) @pytest.mark.parametrize( ('compound_statement', 'options', 'should_be_profiled'), [('function-def', set(), False), @@ -1741,6 +1749,7 @@ def __len__(self) -> int: def test_import_discovery_in_all_compound_statements( compound_statement: _CompoundStatement, options: set[_ImportDiscoveryOption], + use_component: Literal['ast_tree_profiler', 'ast_profile_transformer'], should_be_profiled: bool, ) -> None: """ @@ -1884,10 +1893,23 @@ async def afunc(): print(_get_toml_import_discovery_section(options), file=fobj) config = ConfigSource.from_config(cfg_fname) - atp = AstTreeProfiler( - case_fname, list(all_names), False, config=config, - ) - output = ast.unparse(atp.profile()) + if use_component == 'ast_tree_profiler': + atp = AstTreeProfiler( + # `profile_imports=False` prevents + # `AstProfileTransformer` from rewriting the imports, so + # we're really testing `ProfmodExtractor` here + case_fname, list(all_names), False, config=config, + ) + module_ast = atp.profile() + else: + module_ast = AstProfileTransformer._transform( + ast.parse(test_case), + case_fname, + profile_imports=True, + profiled_imports=[], + config=config, + ) + output = ast.unparse(module_ast) for label, module_text in [ ('input', test_case), ('output', output), @@ -1898,9 +1920,25 @@ async def afunc(): assert set(_grep_profiled_names(output)) == expected -@pytest.mark.parametrize('call', ['first', 'second', 'third']) +@pytest.mark.parametrize( + ('call', 'use_component', 'expected_profiled_objects'), + [ + # Nothing special happens when calling `first()` and `second()`, + # there's only a single import target (`textwrap.indent`) inside + # the function + ('first', 'ast_tree_profiler', ['indent']), + ('first', 'ast_profile_transformer', ['indent']), + ('second', 'ast_tree_profiler', ['indent']), + ('second', 'ast_profile_transformer', ['indent']), + # With `third()`, because `textwrap.dedent()` is also imported, + # it is also profiled when using `AstProfileTransformer` + ('third', 'ast_tree_profiler', ['indent']), + ('third', 'ast_profile_transformer', ['indent', 'dedent']), + ]) def test_nested_imports_correct_deduplication_across_scopes( call: Literal['first', 'second', 'third'], + use_component: Literal['ast_tree_profiler', 'ast_profile_transformer'], + expected_profiled_objects: Sequence[Literal['indent', 'dedent']], ) -> None: """ Test that there is no aliasing in the check we have against @@ -1922,8 +1960,6 @@ def test_nested_imports_correct_deduplication_across_scopes( bodies. For this reason, import discovery in function bodies is off by default. """ - from textwrap import indent - test_module = ub.codeblock(""" def first() -> str: from textwrap import indent @@ -1954,10 +1990,23 @@ def third() -> str: with open(cfg_fname, 'w') as fobj: print(_get_toml_import_discovery_section(), file=fobj) - mod_ast = AstTreeProfiler( - mod_fname, ['textwrap.indent'], False, - config=ConfigSource.from_config(cfg_fname), - ).profile() + config = ConfigSource.from_config(cfg_fname) + if use_component == 'ast_tree_profiler': + # Ditto comment in + # `test_import_discovery_in_all_compound_statements()` + mod_ast = AstTreeProfiler( + mod_fname, ['textwrap.indent'], False, + config=config, + ).profile() + else: + mod_ast = AstProfileTransformer._transform( + ast.parse(test_module), + mod_fname, + profile_imports=True, + profiled_imports=[], + config=config, + ) + mod_ast = ast.fix_missing_locations(mod_ast) print(ast.unparse(mod_ast)) namespace: dict[str, Any] = {'profile': mock_prof} @@ -1968,4 +2017,5 @@ def third() -> str: # `textwrap.indent()` should be presented to the profiler exactly # once assert namespace[call]() == ' ' + call - assert mock_prof.profiled_objects == [indent] + profiled_objects = [func.__name__ for func in mock_prof.profiled_objects] + assert profiled_objects == list(expected_profiled_objects) From ec9923922054b71814b9337be2be3d80281190fa Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Fri, 24 Jul 2026 00:22:07 +0200 Subject: [PATCH 07/12] Fix `ast_profile_transformer` profiling dedup line_profiler/autoprofile/ast_profile_transformer.py _DuplicateChecker New object used by `AstProfileTransformer` to figure out whether to insert a post-import profiling node for an import target, and to keep track of the profiled targets _ContextAwareDuplicateChecker Concrete implmentation of `_DuplicateChecker` which handles contextualization (e.g. the name being profiled on-import in a function shouldn't cause it to not be in another) _LegacyDuplicateChecker Deprecated implementation of `_DuplicateChecker` that keeps the old behavior, working only on a collection of already-profiled names AstProfileTransformer .__init__() Now also taking a mapping for `profiled_imports`, allowing for context-aware on-import profiling deduplication ._visit_import() - Updated signature so that the location of the `ImportTarget` can be correctly resolved - Updated implementation to use a `_DuplicateChecker` instance for bookkeeping .visit_Import(), .visit_ImportFrom() Updated calls to `._visit_import()` .generic_visit() New method which does what `NodeTransformer.generic_visit()` does but with extra bookkeeping, allowing the `_DuplicateChecker` to check the current context ._visit_generic_child(), ._visit_generic_children() New helper methods used by `.generic_visit()` line_profiler/autoprofile/ast_tree_profiler.py ::AstTreeProfiler._profile_ast_tree() - Loosened type hints on parameter `tree_imports_to_profile_dict` (`dict[..., list[...]]` -> `Mapping[..., Sequence[...]]`) - Updated implementation to: - Maintain an updated copy of `tree_imports_to_profile_dict`, taking into account the profiling nodes inserted - Pass said copy to `AstProfileTransformer._transform(profiled_imports=...)` so as to use the new context-aware deduplication tests/test_autoprofile.py test_import_discovery_in_all_compound_statements() Updated call to `AstProfileTransformer._transform()` test_nested_imports_correct_deduplication_across_scopes() Updated parametrization and implementation to allow for testing how `ProfmodExtractor` and `AstProfileTransformer` interact test_ast_profile_transformer_deprecated_profiled_imports() New test that `AstProfileTransformer(profiled_imports=...)` still supports passing a `Collection[str]`, which results in (more or less) the old behavior and the issuance of a `DeprecationWarning` --- .../autoprofile/ast_profile_transformer.py | 308 +++++++++++++++--- .../autoprofile/ast_tree_profiler.py | 32 +- tests/test_autoprofile.py | 118 ++++++- 3 files changed, 400 insertions(+), 58 deletions(-) diff --git a/line_profiler/autoprofile/ast_profile_transformer.py b/line_profiler/autoprofile/ast_profile_transformer.py index a67380e2..6b8d773c 100644 --- a/line_profiler/autoprofile/ast_profile_transformer.py +++ b/line_profiler/autoprofile/ast_profile_transformer.py @@ -1,12 +1,15 @@ from __future__ import annotations import ast -from collections.abc import Callable, Collection, Mapping, Sequence -from functools import partial +from collections.abc import ( + Callable, Collection, Mapping, MutableSequence, Sequence, +) from os import PathLike from types import MappingProxyType -from typing import TypeVar, cast, get_args +from typing import Any, Protocol, TypeVar, cast, get_args +from warnings import warn +from .. import _diagnostics as diagnostics from ..toml_config import ConfigSource from ._import_targets import ImportTarget from .profmod_extractor import _CompoundNodeType, _ImportFinder @@ -144,6 +147,132 @@ def _ast_create_node_from_import_target( return ast_create_profile_node(target.resolved_name) +class _DuplicateChecker(Protocol): + """ + Protocol for objects which helps with deduplication. + """ + def should_profile_import( + self, target: ImportTarget, context: Sequence[str | int], /, + ) -> bool: + ... + + def record_profiled_import( + self, target: ImportTarget, context: Sequence[str | int], /, + ) -> Any: + ... + + +class _ContextAwareDuplicateChecker: + """ + This checker is context-aware and only deduplicates imports of the + same object in the same scope. + """ + def __init__( + self, + profiled_imports: ( + Mapping[Sequence[str | int], Sequence[ImportTarget]] | None + ) = None, + ) -> None: + pi: dict[ + tuple[str | int, ...], dict[int, list[ImportTarget]] + ] + self._profiled_imports = pi = {} + for context, imports in (profiled_imports or {}).items(): + ctx, index = self._check_context(context) + pi.setdefault(ctx, {})[index] = list(imports) + + def should_profile_import( + self, target: ImportTarget, context: Sequence[str | int], + ) -> bool: + # Note: the `index` shouldn't be needed since we're already + # going through the import targets in order + ctx, _ = self._check_context(context) + if ctx not in self._profiled_imports: + return True + ctx_profiled_names = { + imp.name + for imports in self._profiled_imports[ctx].values() + for imp in imports + } + return target.name not in ctx_profiled_names + + def record_profiled_import( + self, target: ImportTarget, context: Sequence[str | int], + ) -> None: + ctx, index = self._check_context(context) + ( + self._profiled_imports + .setdefault(ctx, {}) + .setdefault(index, []) + .append(target) + ) + + @staticmethod + def _check_context( + context: Sequence[str | int], + ) -> tuple[tuple[str | int, ...], int]: + *ctx, index = context + if not isinstance(index, int): + raise TypeError( + f'context[-1] = {context[-1]!r}: expected an integer', + ) + return tuple(ctx), index + + +class _LegacyDuplicateChecker: + """ + This checker replicates legacy behavior: all imports in the same + module are treated on equal footing regardless of scoping, and each + import target is only passed once to the profiler. + + Notes: + Using this results in a :py:class:`DeprecationWarning`. This is + because such deduplication can result in the profiler never + getting passed an intended target. Consider the following + example: + + >>> # doctest: +SKIP + >>> + >>> + >>> def foo(): + ... from spam import ham + ... + ... return ham() + ... + >>> + >>> def bar(): + ... from spam import ham + ... + ... return ham() + ... + >>> + >>> if __name__ == '__main__': + ... bar() + + In the above example, ``spam.ham()`` is never profiled even + after AST rewrite, because only the import in ``foo()`` has a + post-import profiling statement inserted. + """ + def __init__(self, profiled_imports: Collection[str]) -> None: + msg = ( + 'AstProfileTransformer(profiled_imports=) ' + 'is deprecated because of erroneous resolution of duplicate ' + 'imports; future code should use either ' + '`Mapping[Sequence[str | int], ImportTarget]` ' + '(e.g. the return value of `ProfmodExtractor.extract_all()`) ' + 'or `None`' + ) + diagnostics.log.warning(f'DeprecationWarning: {msg}') + warn(msg, category=DeprecationWarning, stacklevel=2) # Caller + self._profiled_imports = set(profiled_imports) + + def should_profile_import(self, target: ImportTarget, _) -> bool: + return target.name not in self._profiled_imports + + def record_profiled_import(self, target: ImportTarget, _) -> None: + self._profiled_imports.add(target.name) + + class AstProfileTransformer(ast.NodeTransformer): """ Transform an abstract syntax tree adding profiling to all of its @@ -163,7 +292,11 @@ class AstProfileTransformer(ast.NodeTransformer): def __init__( self, profile_imports: bool = False, - profiled_imports: Collection[str] | None = None, + profiled_imports: ( + Mapping[Sequence[str | int], Sequence[ImportTarget]] + | Collection[str] + | None + ) = None, profiler_name: str = 'profile', *, profile_star_imports: bool = False, @@ -171,18 +304,32 @@ def __init__( _CompoundNodeType, bool ] = _PROFILE_IMPORTS_IN_DEFAULT, ) -> None: - """Initializes the AST transformer with the profiler name. + """ + Initializes the AST transformer. Args: profile_imports (bool): if True, profile all concrete (non-star) imports. - profiled_imports (Collection[str]): - list of dotted paths of imports to skip that have already been added to profiler. + profiled_imports \ +(Mapping[Sequence[str | int], Sequence[ImportTarget]] \ +| Collection[str] | None): + ``Mapping[Sequence[str | int], Sequence[ImportTarget]]`` + mapping from the locations of already-profiled + import targets to those import targets themselves, + in the same format as the return value of + :py:meth:`line_profiler.autoprofile\ +.ProfmodExtractor.extract_all`. + ``Collection[str]`` + DEPRECATED; dotted paths of imports to skip that + have already been added to profiler. + :py:const:`None` + equivalent to ``{}``, i.e. no prior profiled + imports. profiler_name (str): - the profiler name used as decorator and for the method call to add to the object - to the profiler. + the profiler name used as decorator and for the method + call to add to the object to the profiler. profile_star_imports (bool): if this and ``profile_imports`` are True, also profile @@ -197,12 +344,39 @@ def __init__( profile import statements residing therein. """ self._profile_imports = bool(profile_imports) - self._profiled_imports = set(profiled_imports or ()) + if profiled_imports is None: + self._duplicate_checker: _DuplicateChecker + self._duplicate_checker = _ContextAwareDuplicateChecker() + elif ( + isinstance(profiled_imports, Mapping) + and all( + isinstance(imports, Sequence) + for imports in profiled_imports.values() + ) + ): + self._duplicate_checker = _ContextAwareDuplicateChecker(cast( + Mapping[Sequence[str | int], Sequence[ImportTarget]], + profiled_imports, + )) + elif ( + isinstance(profiled_imports, Collection) + and all(isinstance(imp, str) for imp in profiled_imports) + ): + profiled_imports = cast(Collection[str], profiled_imports) + self._duplicate_checker = _LegacyDuplicateChecker(profiled_imports) + else: # nocover + raise TypeError( + f'profiled_imports = {profiled_imports!r}: ' + 'expected `Collection[str]`, ' + '`Mapping[Sequence[str | int], Sequence[ImportTarget]]`, ' + 'or `None`', + ) self._profiler_name = profiler_name self._profile_star_imports = profile_star_imports self._should_visit_imports = dict(profile_imports_in) self._dropped_star_imports: set[ImportTarget] = set() - self._current_loc: list[str] = [] + self._current_loc: list[str | int] = [] + self._current_node_ancestry: list[str] = [] def _visit_func_def( self, node: ast.FunctionDef | ast.AsyncFunctionDef @@ -238,7 +412,7 @@ def _visit_func_def( def _visit_import( self, node: _Import, - get_import_targets: Callable[[_Import], Sequence[ImportTarget]], + get_import_targets: Callable[[int, _Import], Sequence[ImportTarget]], ) -> _Import | list[_Import | ast.Expr]: """ Add a node that profiles an import. If: @@ -257,7 +431,7 @@ def _visit_import( node (_Import): import[-from] node in the AST get_import_targets \ -(Callable[[_Import], Sequence[ImportTarget]]): +(Callable[[int, _Import], Sequence[ImportTarget]]): helper callable for analyzing the node Returns: @@ -273,7 +447,7 @@ def _visit_import( else: # Check if this node is nested inside compound statements # that we shouldn't look for imports in - *ancestry, _ = self._current_loc + *ancestry, _ = self._current_node_ancestry svi = self._should_visit_imports should_profile = all( svi.get(cast(_CompoundNodeType, a_type), True) @@ -281,14 +455,19 @@ def _visit_import( ) if not should_profile: - self.generic_visit(node) + # No need for further descent, no other node of interest can + # reside import[-from] nodes return node - this_visit = cast(_Import, self.generic_visit(node)) - visited: list[_Import | ast.Expr] = [this_visit] - for target in get_import_targets(node): - name = target.name - if name in self._profiled_imports: + result: list[_Import | ast.Expr] = [node] + *_, index = self._current_loc + assert isinstance(index, int) + duplicate_checker = self._duplicate_checker + + for target in get_import_targets(index, node): + if not duplicate_checker.should_profile_import( + target, self._current_loc, + ): continue expr = _ast_create_node_from_import_target( target, profile_star_imports=self._profile_star_imports, @@ -296,9 +475,11 @@ def _visit_import( if expr is None: # Bookkeeping self._dropped_star_imports.add(target) else: - self._profiled_imports.add(name) - visited.append(expr) - return visited + duplicate_checker.record_profiled_import( + target, self._current_loc, + ) + result.append(expr) + return result def visit_Import( self, node: ast.Import, @@ -319,20 +500,7 @@ def visit_Import( a list containing the import node and the profiling node(s) """ - # Note: we don't actually care about the `ImportTarget.index` - # here; in fact, we're just reusing the name-resolution - # machinery in `ImportTarget` - get_targets = partial(ImportTarget._from_import_node, 0) - return self._visit_import(node, get_targets) - - def visit(self, node: ast.AST) -> ast.AST | list[ast.AST]: - # Bookkeeping - loc = self._current_loc - loc.append(type(node).__name__) - try: - return super().visit(node) - finally: - loc.pop() + return self._visit_import(node, ImportTarget._from_import_node) def visit_ImportFrom( self, node: ast.ImportFrom @@ -353,8 +521,70 @@ def visit_ImportFrom( a list containing the import node and the profiling node(s) """ - get_targets = partial(ImportTarget._from_import_from_node, 0) - return self._visit_import(node, get_targets) + return self._visit_import(node, ImportTarget._from_import_from_node) + + def visit(self, node: ast.AST) -> ast.AST | list[ast.AST]: + """ + :py:meth:`ast.NodeTransformer.visit` with extra bookkeeping. + """ + anc = self._current_node_ancestry + anc.append(type(node).__name__) + try: + return super().visit(node) + finally: + anc.pop() + + def generic_visit(self, node: ast.AST) -> ast.AST: + """ + :py:meth:`ast.NodeTransformer.generic_visit` with extra + bookkeeping. + """ + for field, value in ast.iter_fields(node): + if isinstance(value, ast.AST): + self._visit_generic_child(node, field, value) + elif isinstance(value, MutableSequence): # Compound node + if not all(isinstance(item, ast.AST) for item in value): + continue + self._visit_generic_children( + node, field, cast(MutableSequence[ast.AST], value), + ) + return node + + def _visit_generic_child( + self, node: ast.AST, field: str, child: ast.AST, + ) -> None: + self._current_loc.append(field) + try: + replacement: ast.AST | list[ast.AST] = self.visit(child) + if isinstance(replacement, ast.AST): + setattr(node, field, replacement) + else: + raise RuntimeError( + f'node = {node!r}: invalid field `.{field}` replacement ' + f'({child!r} -> {replacement!r})' + ) + finally: + self._current_loc.pop() + + def _visit_generic_children( + self, node: ast.AST, field: str, children: MutableSequence[ast.AST], + ) -> None: + self._current_loc.append(field) + try: + new_children: list[ast.AST] = [] + for i, item in enumerate(children): + self._current_loc.append(i) + try: + replacement = self.visit(item) + if isinstance(replacement, ast.AST): + new_children.append(replacement) + else: + new_children.extend(replacement) + finally: + self._current_loc.pop() + children[:] = new_children + finally: + self._current_loc.pop() @staticmethod def _get_profile_imports_in( diff --git a/line_profiler/autoprofile/ast_tree_profiler.py b/line_profiler/autoprofile/ast_tree_profiler.py index 750cd149..8fdf89f2 100644 --- a/line_profiler/autoprofile/ast_tree_profiler.py +++ b/line_profiler/autoprofile/ast_tree_profiler.py @@ -1,8 +1,9 @@ from __future__ import annotations import ast +import dataclasses import os -from collections.abc import Collection, MutableSequence, Sequence +from collections.abc import Collection, Mapping, MutableSequence, Sequence from typing import Any, cast from ._import_targets import ImportTarget @@ -122,8 +123,8 @@ def _get_script_ast_tree(script_file: str) -> ast.Module: def _profile_ast_tree( self, tree: ast.Module, - tree_imports_to_profile_dict: dict[ - tuple[str | int, ...], list[ImportTarget] + tree_imports_to_profile_dict: Mapping[ + tuple[str | int, ...], Sequence[ImportTarget] ], profile_full_script: bool = False, profile_imports: bool = False, @@ -146,14 +147,14 @@ def _profile_ast_tree( tree (_ast.Module): abstract syntax tree to be profiled. - tree_imports_to_profile_dict (dict[tuple[str | int, ...], \ -list[ImportTarget]]): + tree_imports_to_profile_dict \ +(Mapping[tuple[str | int, ...], Sequence[ImportTarget]]): dict of imports to profile key (tuple[str | int, ...]): Location of import in AST, e.g. ``('body', 0)`` for the case where it is the first statement in the :py:attr:`ast.Module.body` - value (list[ImportTarget]): + value (Sequence[ImportTarget]): list of import targets (see the documentation of :py:class:`line_profiler.autoprofile\ .profmod_extractor.ImportTarget`) @@ -180,12 +181,14 @@ def _profile_ast_tree( (_ast.Module): tree abstract syntax tree with profiling. """ - profiled_imports = [] + profiled_imports: dict[tuple[str | int, ...], list[ImportTarget]] = {} argsort_tree_indexes = sorted( - list(tree_imports_to_profile_dict), reverse=True + tree_imports_to_profile_dict, reverse=True, ) for tree_loc in argsort_tree_indexes: imports = tree_imports_to_profile_dict[tree_loc] + # Also handle bookkeeping tasks + updated_imports = profiled_imports[tree_loc] = [] *loc, tree_index = tree_loc assert isinstance(tree_index, int) body = cast(MutableSequence[ast.AST], self._descend(tree, loc)) @@ -196,9 +199,16 @@ def _profile_ast_tree( expr = _ast_create_node_from_import_target( imp, modnames_to_profile, profile_star_imports, ) - if expr is not None: - body.insert(tree_index + 1, expr) - profiled_imports.append(imp.name) + if expr is None: + continue + body.insert(tree_index + 1, expr) + # Bookkeeping: make sure that we're keeping track of + # the updated locations of the import targets + updated_imports[:] = [ + dataclasses.replace(imp, index=imp.index + 1) + for imp in updated_imports + ] + updated_imports.insert(0, imp) if profile_full_script: tree = self._ast_transformer_class_handler._transform( tree, self._script_file, diff --git a/tests/test_autoprofile.py b/tests/test_autoprofile.py index c6632f46..b38918ea 100644 --- a/tests/test_autoprofile.py +++ b/tests/test_autoprofile.py @@ -1906,7 +1906,6 @@ async def afunc(): ast.parse(test_case), case_fname, profile_imports=True, - profiled_imports=[], config=config, ) output = ast.unparse(module_ast) @@ -1926,18 +1925,30 @@ async def afunc(): # Nothing special happens when calling `first()` and `second()`, # there's only a single import target (`textwrap.indent`) inside # the function + ('first', 'profmod_extractor', ['indent']), ('first', 'ast_tree_profiler', ['indent']), ('first', 'ast_profile_transformer', ['indent']), + ('second', 'profmod_extractor', ['indent']), ('second', 'ast_tree_profiler', ['indent']), ('second', 'ast_profile_transformer', ['indent']), - # With `third()`, because `textwrap.dedent()` is also imported, - # it is also profiled when using `AstProfileTransformer` - ('third', 'ast_tree_profiler', ['indent']), + # With `third()`, because `textwrap.dedent()` is also imported: + ('third', 'profmod_extractor', ['indent']), + # - When using `AstTreeProfiler`, `ProfmodExtractor` first + # inserts a profiling node for `indent()`, then followed by + # another for `dedent()` inserted by `AstProfileTransformer`; + # since the profiling node for `dedent()` is created later, it + # is inserted bewteen the import statement and the profiling + # node for `indent()`, and is hence executed first + ('third', 'ast_tree_profiler', ['dedent', 'indent']), + # - When using `AstProfileTransformer`, profiling nodes are + # inserted for both `indent()` and `dedent()` in one go ('third', 'ast_profile_transformer', ['indent', 'dedent']), ]) def test_nested_imports_correct_deduplication_across_scopes( call: Literal['first', 'second', 'third'], - use_component: Literal['ast_tree_profiler', 'ast_profile_transformer'], + use_component: Literal[ + 'profmod_extractor', 'ast_tree_profiler', 'ast_profile_transformer', + ], expected_profiled_objects: Sequence[Literal['indent', 'dedent']], ) -> None: """ @@ -1991,21 +2002,27 @@ def third() -> str: print(_get_toml_import_discovery_section(), file=fobj) config = ConfigSource.from_config(cfg_fname) - if use_component == 'ast_tree_profiler': + if use_component == 'profmod_extractor': # Ditto comment in # `test_import_discovery_in_all_compound_statements()` mod_ast = AstTreeProfiler( mod_fname, ['textwrap.indent'], False, config=config, ).profile() - else: + elif use_component == 'ast_tree_profiler': + # Integration of both + mod_ast = AstTreeProfiler( + mod_fname, ['textwrap.indent', str(mod_fname)], True, + config=config, + ).profile() + else: # `ast_profile_transformer` mod_ast = AstProfileTransformer._transform( ast.parse(test_module), mod_fname, profile_imports=True, - profiled_imports=[], config=config, ) + # We need this to actually compile and exec the code mod_ast = ast.fix_missing_locations(mod_ast) print(ast.unparse(mod_ast)) @@ -2019,3 +2036,88 @@ def third() -> str: assert namespace[call]() == ' ' + call profiled_objects = [func.__name__ for func in mock_prof.profiled_objects] assert profiled_objects == list(expected_profiled_objects) + + +@pytest.mark.parametrize( + ('definitions', 'expected'), + [ + # The profiling statement is always inserted into the first + # function body where the import occurs... + (['foo'], {'indent'}), (['bar'], {'ind'}), + # ... but only the first + (['foo', 'bar'], {'indent'}), (['bar', 'foo'], {'ind'}), + ]) +def test_ast_profile_transformer_deprecated_profiled_imports( + definitions: Sequence[Literal['foo', 'bar']], + expected: Collection[Literal['indent', 'ind']], +) -> None: + """ + Test that the legacy invocation of + :py:class:`.AstProfileTransformer` with + ``profiled_imports: Collection[str]`` works "as expected": + + - :py:class:`DeprecationWarning` is issued, instructing users to + switch to the mapping form of the argument. + + - Deduplication of imports happens without regard of scopes. + + See also: + :py:func:\ +`test_nested_imports_correct_deduplication_across_scopes` + """ + defs = { + 'foo': """ + def foo() -> str: + from textwrap import indent + + return indent('foo', ' ') + """, + 'bar': """ + def bar() -> str: + from textwrap import indent as ind + + return ind('bar', ' ') + """, + } + test_module = '\n\n'.join( + ub.codeblock(defs[func]).strip('\n') for func in definitions + ) + + with contextlib.ExitStack() as stack: + tmp = stack.enter_context(tempfile.TemporaryDirectory()) + + mod_fname = os.path.join(tmp, 'test_module.py') + with open(mod_fname, 'w') as fobj: + print(test_module, file=fobj) + + cfg_fname = os.path.join(tmp, 'config.toml') + with open(cfg_fname, 'w') as fobj: + print(_get_toml_import_discovery_section(), file=fobj) + + stack.enter_context(pytest.warns( + DeprecationWarning, + match='.*'.join( + '{}{}{}'.format( + r'\b' if chunk[0].isalnum() else '', + chunk, + r'\b' if chunk[-1].isalnum() else '', + ) + for chunk in [ + 'profiled_imports=', r'Collection\[str\]', 'deprecated', + 'use', r'Mapping\[.+, .+\]', 'or', 'None', + ] + ) + )) + + config = ConfigSource.from_config(cfg_fname) + mod_ast = AstProfileTransformer._transform( + ast.parse(test_module), + mod_fname, + profile_imports=True, + profiled_imports=[], # This triggers legacy behavior + config=config, + ) + output = ast.unparse(mod_ast) + print(output) + + assert set(_grep_profiled_names(output)) == set(expected) From 6ec1ccd3bfe6dcf37a47c31e5f888de06c993120 Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Fri, 24 Jul 2026 03:08:59 +0200 Subject: [PATCH 08/12] Unit test for `add_star_import()` line_profiler/autoprofile/line_profiler_utils.py::add_star_import() Minor refactoring of internals: - `__all__` is assumed to be either nonexistent or a `Sequence[str]`, since any other values would've resulted in an error upon the preceding star-import anyway - Calls to `add_imported_function_or_module()` now made in import order of the names, where possible (i.e. there is an `__all__`) tests/test_autoprofile.py _RecordingProfiler.add_imported_function_or_module() Now returning 1 to be consistent with `line_profiler_utils.py::add_imported_function_or_module()` test_drop_and_warn_against_star_imports() Renamed from `test_handle_star_imports()` test_add_star_import() New unit test for `add_star_import()`, testing the retrieval and (optionally selective) profiling of the imported names --- .../autoprofile/line_profiler_utils.py | 30 ++--- tests/test_autoprofile.py | 104 ++++++++++++++++-- 2 files changed, 113 insertions(+), 21 deletions(-) diff --git a/line_profiler/autoprofile/line_profiler_utils.py b/line_profiler/autoprofile/line_profiler_utils.py index 90efac42..c3a74ef0 100644 --- a/line_profiler/autoprofile/line_profiler_utils.py +++ b/line_profiler/autoprofile/line_profiler_utils.py @@ -2,11 +2,11 @@ import inspect import operator -from collections.abc import Callable, Collection, MutableMapping +from collections.abc import Callable, Collection, Iterable, MutableMapping from functools import cached_property, partial, partialmethod from importlib import import_module from types import FunctionType, MethodType, ModuleType -from typing import TYPE_CHECKING, Any, Literal, cast, overload +from typing import TYPE_CHECKING, Any, Literal, overload from .profmod_extractor import _should_profile @@ -149,23 +149,27 @@ def add_star_import( """ # Dynamically inspect the module to see which names would've been # inserted by a star-import - module = import_module(import_from) # TODO + module = import_module(import_from) try: - all_names: set[str] | None = set(cast( - Collection[str], getattr(module, '__all__', None), - )) - except Exception: - # Could be for whatever reason, there's no guarantee that a - # module declares the `.__all__` CORRECTLY + all_names: list[str] | None = list(module.__all__) + except AttributeError: + # Note: we don't expect any other errors; the language standard + # dictates that `__all__` has to be `Sequence[str]`, and any + # other value would be an error upon star-import anyway all_names = None check_attr: Callable[[str], bool] + sort_items: Callable[[Iterable[tuple[str, Any]]], list[tuple[str, Any]]] if all_names is None: # Default behavior: take all public names check_attr = lambda attr: ( # noqa: E731 not attr.startswith('_') ) + sort_items = list else: # If we have a valid `.__all__`, take names therefrom - check_attr = partial(operator.contains, all_names) + check_attr = partial(operator.contains, set(all_names)) + sort_items = partial( + sorted, key=lambda kv: all_names.index(kv[0]), + ) imported_names: dict[str, Any] = { attr: value for attr, value in inspect.getmembers(module) @@ -173,7 +177,7 @@ def add_star_import( } # Decide on which of the names to pass to the profiler - # Note: :the name shoul've already been inserted into the namespace + # Note: the names should've already been inserted into the namespace # by the import statement itself, so this is just post-hoc # bookkeeping add: Callable[[Any], int] @@ -184,7 +188,7 @@ def add_star_import( add = partial(add_imported_function_or_module, self, **kwargs) count = 0 sentinel = object() - for name, value in imported_names.items(): + for name, value in sort_items(imported_names.items()): if not ( targets is None or _should_profile(targets, f'{import_from}.{name}') @@ -192,7 +196,7 @@ def add_star_import( # Check that the name should be profiled (if we have # constrained `targets`) continue - if namespace.get(name, sentinel) is not value: + if namespace.get(name, sentinel) is not value: # nocover # Check that the actual object in the namespace is # consistent with what is imported continue diff --git a/tests/test_autoprofile.py b/tests/test_autoprofile.py index b38918ea..546524ce 100644 --- a/tests/test_autoprofile.py +++ b/tests/test_autoprofile.py @@ -1,6 +1,7 @@ from __future__ import annotations import ast +import builtins import contextlib import os import re @@ -9,7 +10,8 @@ import sys import textwrap import tempfile -from collections.abc import Collection, Sequence +import uuid +from collections.abc import Collection, Generator, Sequence from typing import Any, Literal, get_args from warnings import catch_warnings, WarningMessage @@ -20,6 +22,7 @@ AstProfileTransformer, ) from line_profiler.autoprofile.ast_tree_profiler import AstTreeProfiler +from line_profiler.autoprofile.line_profiler_utils import add_star_import from line_profiler.autoprofile.profmod_extractor import ProfmodExtractor @@ -1437,8 +1440,9 @@ def __init__(self) -> None: def __call__(self, x: Any) -> Any: return x - def add_imported_function_or_module(self, obj) -> None: + def add_imported_function_or_module(self, obj) -> Literal[1]: self.profiled_objects.append(obj) + return 1 def test_multitarget_import_transformation_executes() -> None: @@ -1543,7 +1547,7 @@ def _grep_profiled_names(module_text: str) -> list[str]: # No whole-file rewriting, bu we explicitly ask to profile the # `spam.ham.*` import (which can't be done) (['spam.ham'], [], False, False, True)]) -def test_handle_star_imports( +def test_drop_and_warn_against_star_imports( prof_mod: list[str], expected_targets: Collection[Literal['bar', 'baz']], profile_imports: bool, @@ -1551,11 +1555,9 @@ def test_handle_star_imports( expect_warnings: bool, ) -> None: """ - Test that star-imports (``from ... import *``) don't cause - :py:meth:`AstTreeProfiler.profile` to choke, instead just issuing - warnings about ignoring them. - - TODO: actually handle star-imports + Test the default behavior of :py:meth:`AstTreeProfiler.profile`: + that star-imports (``from ... import *``) don't cause it to choke, + instead just issuing warnings about ignoring them. """ code = ub.codeblock( """ @@ -2121,3 +2123,89 @@ def bar() -> str: print(output) assert set(_grep_profiled_names(output)) == set(expected) + + +@pytest.mark.parametrize( + ('targets', 'dunder_all', 'expected_imports', 'expected_profiled'), + [ + # No `__all__` -> all the public names included + (None, None, {'indent', 'foo', 'bar'}, {'indent', 'foo', 'bar'}), + # `_baz` specified as a target, but is never imported to begin + # with + ({'__module__.bar', '__module__._baz'}, None, + {'indent', 'foo', 'bar'}, {'bar'}), + # With a valid `__all__`, only names inside will be imported + (None, ['foo', '_baz', '_dedent'], + {'foo', '_dedent', '_baz'}, {'foo', 'dedent', '_baz'}), + ({'__module__._baz', '__module__._dedent'}, ['foo', '_baz', '_dedent'], + {'foo', '_dedent', '_baz'}, {'dedent', '_baz'}), + ]) +def test_add_star_import( + targets: Collection[str] | None, + dunder_all: Sequence[str] | None, + expected_imports: Collection[str], + expected_profiled: Collection[str], +) -> None: + """ + Test that :py:func:`.add_star_import` works as expected, retriving + the correct names from the namespace and profiling them. + """ + def propose_module_names( + prefix: str = 'my_module' + ) -> Generator[str, None, None]: + while True: + random = str(uuid.uuid4()).replace('-', '_') + name = f'{prefix}_{random}' + if not name.isidentifier(): + continue + if name not in sys.modules: + yield name + + test_module = ub.codeblock(""" + from textwrap import indent, dedent as _dedent + + + def foo() -> None: + ... + + + def bar() -> None: + ... + + + def _baz() -> None: + ... + """).strip('\n') + if dunder_all is not None: + all_repr = repr(dunder_all) + assert ast.literal_eval(all_repr) == dunder_all + test_module = f'{test_module}\n\n__all__ = {all_repr}' + + module_name = next(propose_module_names()) + if targets is not None: + targets = [ + t.replace('__module__', module_name) for t in targets + ] + mock_prof = _RecordingProfiler() + with contextlib.ExitStack() as stack: + tmp = stack.enter_context(tempfile.TemporaryDirectory()) + mp = stack.enter_context(pytest.MonkeyPatch.context()) + mp.syspath_prepend(tmp) + + mod_fname = os.path.join(tmp, module_name + '.py') + with open(mod_fname, 'w') as fobj: + print(test_module, file=fobj) + + # Check that the correct names are imported by the star-import + namespace: dict[str, Any] = {'baz': None, '__builtins__': builtins} + preexisting = set(namespace) + exec(f'from {module_name} import *', namespace) + assert set(namespace) == set(expected_imports) | preexisting + + # Check that the same names are passed to the profiler by + # `add_star_import()` + add_star_import(mock_prof, module_name, targets, namespace) + profiled_objects = { + func.__name__ for func in mock_prof.profiled_objects + } + assert profiled_objects == set(expected_profiled) From 0b330b921c50367ed9be78d257b974c9341037c2 Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Fri, 24 Jul 2026 04:54:38 +0200 Subject: [PATCH 09/12] Manual overrides for `autoprofile.*` options line_profiler/autoprofile/ast_profile_transformer.py ::AstProfileTransformer ._get_profile_imports_in() New param `profile_nested_imports` allowing for overriding the `autoprofile.import_discovery` subtable in `config` ._transform() - New param `profile_nested_imports` (ditto above) - New param `profile_star_imports` for parsing the default value thereof to pass to the initializer from `config` line_profiler/autoprofile/ast_tree_profiler.py::AstTreeProfiler ._profile_ast_tree() New param `profile_star_imports` allowing for overriding the `autoprofile.prof_star_imports` value in `config` .profile() - New param `profile_nested_imports` (ditto above) - Param `profile_star_imports` now keyword-only line_profiler/autoprofile/autoprofile.py::run() New params `profile_nested_imports` and `profile_star_imports` (ditto above) line_profiler/autoprofile/profmod_extractor.py _ImportFinder._get_filter_args() ProfmodExtractor._ast_get_imports_from_tree() New param `find_nested_imports` allowing for overriding the `autoprofile.import_discovery` subtable in `config` ProfmodExtractor.extract_all() - New param `find_nested_imports` (ditto above) - Param `filter_star_imports` now keyword-only tests/test_autoprofile.py ::test_import_discovery_in_all_compound_statements() Updated parametrization and implementation to test both when the nested-import toggles are passed via the `config` file and the `profile_nested_imports` argument --- .../autoprofile/ast_profile_transformer.py | 40 ++++++++++++-- .../autoprofile/ast_tree_profiler.py | 37 +++++++++++-- line_profiler/autoprofile/autoprofile.py | 55 +++++++++++++------ .../autoprofile/profmod_extractor.py | 50 +++++++++++++---- tests/test_autoprofile.py | 31 ++++++++--- 5 files changed, 166 insertions(+), 47 deletions(-) diff --git a/line_profiler/autoprofile/ast_profile_transformer.py b/line_profiler/autoprofile/ast_profile_transformer.py index 6b8d773c..d48af83e 100644 --- a/line_profiler/autoprofile/ast_profile_transformer.py +++ b/line_profiler/autoprofile/ast_profile_transformer.py @@ -12,7 +12,10 @@ from .. import _diagnostics as diagnostics from ..toml_config import ConfigSource from ._import_targets import ImportTarget -from .profmod_extractor import _CompoundNodeType, _ImportFinder +from .profmod_extractor import ( + _CompoundNodeType, _CompoundStatement, _ImportFinder, + _should_profile_star_imports, +) _Import = TypeVar('_Import', ast.Import, ast.ImportFrom) @@ -340,7 +343,7 @@ def __init__( 'FunctionDef', 'AsyncFunctionDef', 'ClassDef', \ 'For', `AsyncFor`, `While`, 'If', 'match_case', \ 'With', 'AsyncWith'. 'Try', 'TryStar', 'ExceptHandler'], bool]): - for each of the compound-statemnt node type, whether to + for each of the compound-statement node type, whether to profile import statements residing therein. """ self._profile_imports = bool(profile_imports) @@ -589,11 +592,12 @@ def _visit_generic_children( @staticmethod def _get_profile_imports_in( config: ConfigSource | None = None, + profile_nested_imports: Collection[_CompoundStatement] | None = None, ) -> dict[_CompoundNodeType, bool]: if config is None: config = ConfigSource.from_default() return _ImportFinder.filter_node_types( - **_ImportFinder._get_filter_args(config), + **_ImportFinder._get_filter_args(config, profile_nested_imports), ) @classmethod @@ -603,19 +607,38 @@ def _transform( filename: PathLike[str] | str | None = None, *, config: ConfigSource | None = None, + profile_star_imports: bool | None = None, + profile_nested_imports: Collection[_CompoundStatement] | None = None, **kwargs, ) -> ast.Module: """ - Wrapper around ``.visit()`` with extra bookkeeping. + Wrapper around ``.visit()`` with extra bookkeeping and + convenience args. Args: node (ast.Module): AST module node + filename (PathLike[str] | str | None): Optional filename to be used in error/warning messages + config (ConfigSource | None): Optional :py:class:`.ConfigSource` to load options from, controlling whether an import should be profiled + + profile_star_imports (bool | None): + Whether to profile star-imports (``from ... import *``); + if :py:const:`None`, it is loaded from the ``config`` + (from ``autoprofile.prof_star_imports``) + + profile_nested_imports \ +(Collection[Literal['func_defs', 'class_defs', \ +'loops', 'conditionals', 'contexts', 'try_except']] | None): + Which of the compound-statement types to look for nested + imports in; + if :py:const:`None`, it is loaded from the ``config`` + (from ``autoprofile.import_discovery``) + **kwargs Passed to the initializer @@ -623,10 +646,15 @@ def _transform( node (ast.Module): Input module node """ + if profile_star_imports is None: + profile_star_imports = _should_profile_star_imports(config) kwargs.setdefault( - 'profile_imports_in', cls._get_profile_imports_in(config), + 'profile_imports_in', + cls._get_profile_imports_in(config, profile_nested_imports), + ) + transformer = cls( + profile_star_imports=profile_star_imports, **kwargs, ) - transformer = cls(**kwargs) dropped_star_imports = transformer._dropped_star_imports if filename is None: filename = '???' diff --git a/line_profiler/autoprofile/ast_tree_profiler.py b/line_profiler/autoprofile/ast_tree_profiler.py index 8fdf89f2..1cb552d9 100644 --- a/line_profiler/autoprofile/ast_tree_profiler.py +++ b/line_profiler/autoprofile/ast_tree_profiler.py @@ -14,7 +14,9 @@ # Keep import below for compatibility ast_create_profile_node, ) -from .profmod_extractor import ProfmodExtractor, _should_profile_star_imports +from .profmod_extractor import ( + _CompoundStatement, ProfmodExtractor, _should_profile_star_imports, +) __docstubs__ = """ from .ast_profile_transformer import AstProfileTransformer @@ -130,6 +132,7 @@ def _profile_ast_tree( profile_imports: bool = False, modnames_to_profile: Collection[str] = (), profile_star_imports: bool = False, + profile_nested_imports: Collection[_CompoundStatement] | None = None, ) -> ast.Module: """ Add profiling to an abstract syntax tree by adding nodes to the @@ -177,8 +180,16 @@ def _profile_ast_tree( ``profile_full_script`` and ``profile_imports`` are true) + profile_nested_imports \ +(Collection[Literal['func_defs', 'class_defs', \ +'loops', 'conditionals', 'contexts', 'try_except']] | None): + Which of the compound-statement types to look for nested + imports in; + if :py:const:`None`, it is loaded from the ``config`` + (from ``autoprofile.import_discovery``). + Returns: - (_ast.Module): tree + tree (_ast.Module): abstract syntax tree with profiling. """ profiled_imports: dict[tuple[str | int, ...], list[ImportTarget]] = {} @@ -215,12 +226,18 @@ def _profile_ast_tree( profile_imports=profile_imports, profiled_imports=profiled_imports, profile_star_imports=profile_star_imports, + profile_nested_imports=profile_nested_imports, config=self._config, ) ast.fix_missing_locations(tree) return tree - def profile(self, profile_star_imports: bool | None = None) -> ast.Module: + def profile( + self, + *, + profile_star_imports: bool | None = None, + profile_nested_imports: Collection[_CompoundStatement] | None = None, + ) -> ast.Module: """ Create an abstract syntax tree of a script and add profiling to it: @@ -245,10 +262,18 @@ def profile(self, profile_star_imports: bool | None = None) -> ast.Module: if True, add targets imported by ``from ... import *`` statements to the profiler; if :py:const:`None`, it is loaded from the ``config`` - (from `autoprofile.prof_star_imports`). + (from ``autoprofile.prof_star_imports``). + + profile_nested_imports \ +(Collection[Literal['func_defs', 'class_defs', \ +'loops', 'conditionals', 'contexts', 'try_except']] | None): + Which of the compound-statement types to look for nested + imports in; + if :py:const:`None`, it is loaded from the ``config`` + (from ``autoprofile.import_discovery``). Returns: - (_ast.Module): tree + tree (_ast.Module): abstract syntax tree with profiling. """ if profile_star_imports is None: @@ -278,6 +303,7 @@ def profile(self, profile_star_imports: bool | None = None) -> ast.Module: ) tree_imports_to_profile_dict = extractor.extract_all( filter_star_imports=filter_star_imports_in_extract_all, + find_nested_imports=profile_nested_imports, ) tree_profiled = self._profile_ast_tree( @@ -287,6 +313,7 @@ def profile(self, profile_star_imports: bool | None = None) -> ast.Module: profile_imports=self._profile_imports, modnames_to_profile=extractor._modnames_to_profile, profile_star_imports=profile_star_imports, + profile_nested_imports=profile_nested_imports, ) return tree_profiled diff --git a/line_profiler/autoprofile/autoprofile.py b/line_profiler/autoprofile/autoprofile.py index a53dce51..5b8454d8 100644 --- a/line_profiler/autoprofile/autoprofile.py +++ b/line_profiler/autoprofile/autoprofile.py @@ -46,16 +46,17 @@ def main(): """ from __future__ import annotations + import importlib.util import os import sys import types -from collections.abc import MutableMapping +from collections.abc import Collection, MutableMapping from typing import Any, cast from ..toml_config import ConfigSource from ..line_profiler_utils import restore -from .ast_tree_profiler import AstTreeProfiler +from .ast_tree_profiler import AstTreeProfiler, _CompoundStatement from .run_module import AstTreeModuleProfiler from .line_profiler_utils import ( add_imported_function_or_module, add_star_import, @@ -100,32 +101,50 @@ def run( as_module: bool = False, *, config: os.PathLike[str] | str | None = None, + profile_star_imports: bool | None = None, + profile_nested_imports: Collection[_CompoundStatement] | None = None, ) -> None: - """Automatically profile a script and run it. - - Profile functions, classes & modules specified in prof_mod without needing to add - @profile decorators. + """ + Automatically profile a script and run it, profiling functions, + classes & modules specified in ``prof_mod`` without needing to add + ``@profile`` decorators. Args: script_file (str): - path to script being profiled. + path to the script being profiled. ns (dict): - "locals" from kernprof scope. + local names to injected into the namespace where + ``script_file``'s code is executed. prof_mod (List[str]): - list of imports to profile in script. - passing the path to script will profile the whole script. - the objects can be specified using its dotted path or full path (if applicable). + list of imports to profile in ``script_file``; + passing the path ``script_file`` will profile the whole + script via AST rewriting; + the objects can be specified using its dotted path or + file-system path (if applicable). profile_imports (bool): - if True, when auto-profiling whole script, profile all imports aswell. + if :py:const:`True`, when rewriting the AST, profile all its + imports aswell. as_module (bool): - whether we're running script_file as a module + whether we're running ``script_file`` as a module. config (os.PathLike[str] | str | None): - optional path to load the session config from + optional path to load the session config from. + + profile_star_imports (bool | None): + whether to profile star-imports (``from ... import *``); + if :py:const:`None`, the value is taken from ``config``. + + profile_nested_imports \ +(Collection[Literal['func_defs', 'class_defs', \ +'loops', 'conditionals', 'contexts', 'try_except']] | None): + Which of the compound-statement types to look for nested + imports in; + if :py:const:`None`, it is loaded from the ``config`` (from + ``autoprofile.import_discovery``) """ Profiler: type[AstTreeModuleProfiler] | type[AstTreeProfiler] @@ -134,7 +153,8 @@ def run( module_name = modpath_to_modname(script_file) if not module_name: raise ModuleNotFoundError( - f'script_file = {script_file!r}: cannot find corresponding module' + f'script_file = {script_file!r}: ' + 'cannot find corresponding module' ) module_obj = types.ModuleType(module_name) @@ -151,7 +171,10 @@ def run( script_file, prof_mod, profile_imports, config=ConfigSource.from_config(config), ) - tree_profiled = profiler.profile() + tree_profiled = profiler.profile( + profile_star_imports=profile_star_imports, + profile_nested_imports=profile_nested_imports, + ) _extend_line_profiler_for_profiling_imports(ns[PROFILER_LOCALS_NAME]) code_obj = compile(tree_profiled, script_file, 'exec') diff --git a/line_profiler/autoprofile/profmod_extractor.py b/line_profiler/autoprofile/profmod_extractor.py index 3caab517..a8439e14 100644 --- a/line_profiler/autoprofile/profmod_extractor.py +++ b/line_profiler/autoprofile/profmod_extractor.py @@ -35,6 +35,10 @@ # `try-except` nodes 'Try', 'TryStar', 'ExceptHandler', ] +_CompoundStatement = Literal[ + 'func_defs', 'class_defs', + 'loops', 'conditionals', 'contexts', 'try_except', +] class _ImportFinder(ast.NodeVisitor): @@ -160,12 +164,23 @@ def filter_node_types( } @staticmethod - def _get_filter_args(config: ConfigSource) -> dict[str, bool]: - cfg = ( - config - .get_subconfig('autoprofile', 'import_discovery') - .conf_dict - ) + def _get_filter_args( + config: ConfigSource, + find_nested_imports: Collection[_CompoundStatement] | None = None, + ) -> dict[str, bool]: + if find_nested_imports is None: + cfg = cast( + dict[_CompoundStatement, bool], + config + .get_subconfig('autoprofile', 'import_discovery') + .conf_dict, + ) + else: + cfg = { + cast(_CompoundStatement, stmt): + stmt in find_nested_imports + for stmt in get_args(_CompoundStatement) + } return { 'collect_from_conditionals': cfg['conditionals'], 'collect_from_try_except': cfg['try_except'], @@ -359,11 +374,13 @@ def _get_modnames_to_profile_from_prof_mod( @staticmethod def _ast_get_imports_from_tree( - node: ast.AST, config: ConfigSource | None = None, + node: ast.AST, + config: ConfigSource | None = None, + find_nested_imports: Collection[_CompoundStatement] | None = None, ) -> dict[tuple[str | int, ...], list[ImportTarget]]: if config is None: config = ConfigSource.from_default() - kwargs = _ImportFinder._get_filter_args(config) + kwargs = _ImportFinder._get_filter_args(config, find_nested_imports) return _ImportFinder.find(node, **kwargs) @staticmethod @@ -412,7 +429,10 @@ def _find_modnames_in_tree_imports( return filtered_imports def extract_all( - self, filter_star_imports: bool | None = None, + self, + *, + filter_star_imports: bool | None = None, + find_nested_imports: Collection[_CompoundStatement] | None = None, ) -> dict[tuple[str | int, ...], list[ImportTarget]]: """ Map ``prof_mod`` to imports in an abstract syntax tree. @@ -425,7 +445,15 @@ def extract_all( If true, filter out star imports (``from import *``) with a warning; if :py:const:`None`, it is loaded from the ``config`` - (as the negation of `autoprofile.prof_star_imports`). + (as the negation of ``autoprofile.prof_star_imports``). + + find_nested_imports \ +(Collection[Literal['func_defs', 'class_defs', \ +'loops', 'conditionals', 'contexts', 'try_except']] | None): + Which of the compound-statement types to look for nested + imports in; + if :py:const:`None`, it is loaded from the ``config`` + (from ``autoprofile.import_discovery``). Returns: tree_imports_to_profile_dict \ @@ -460,7 +488,7 @@ def extract_all( self._config, ) import_targets = self._ast_get_imports_from_tree( - self._tree, self._config, + self._tree, self._config, find_nested_imports, ) raw: dict[tuple[str | int, ...], list[ImportTarget]] = { (*loc, index): filtered_imports diff --git a/tests/test_autoprofile.py b/tests/test_autoprofile.py index 546524ce..603660c9 100644 --- a/tests/test_autoprofile.py +++ b/tests/test_autoprofile.py @@ -1720,6 +1720,7 @@ def __len__(self) -> int: assert set(_grep_profiled_names(output_module)) == expected +@pytest.mark.parametrize('inject_options_with', ['config', 'args']) @pytest.mark.parametrize('use_component', ['ast_tree_profiler', 'ast_profile_transformer']) @pytest.mark.parametrize( @@ -1752,6 +1753,7 @@ def test_import_discovery_in_all_compound_statements( compound_statement: _CompoundStatement, options: set[_ImportDiscoveryOption], use_component: Literal['ast_tree_profiler', 'ast_profile_transformer'], + inject_options_with: Literal['config', 'args'], should_be_profiled: bool, ) -> None: """ @@ -1881,34 +1883,45 @@ async def afunc(): test_case = ub.codeblock(test_cases[compound_statement]).strip('\n') version_bound: tuple[int, ...] = version_bounds.get(compound_statement, ()) if sys.version_info < version_bound: - pytest.skip( - reason=f'cannot test {compound_statement} on {sys.version_info}', - ) + version = '.'.join(str(v) for v in sys.version_info[:3]) + pytest.skip(reason=f'cannot test {compound_statement} on {version}') with tempfile.TemporaryDirectory() as tmp: case_fname = os.path.join(tmp, 'test_case.py') with open(case_fname, 'w') as fobj: print(test_case, file=fobj) - cfg_fname = os.path.join(tmp, 'config.toml') - with open(cfg_fname, 'w') as fobj: - print(_get_toml_import_discovery_section(options), file=fobj) + if inject_options_with == 'config': + cfg_fname = os.path.join(tmp, 'config.toml') + with open(cfg_fname, 'w') as fobj: + print(_get_toml_import_discovery_section(options), file=fobj) + + config: ConfigSource | None + profile_nested_imports: Collection[_ImportDiscoveryOption] | None + + config = ConfigSource.from_config(cfg_fname) + profile_nested_imports = None + else: # Explicitly passed via args + config, profile_nested_imports = None, options - config = ConfigSource.from_config(cfg_fname) if use_component == 'ast_tree_profiler': atp = AstTreeProfiler( # `profile_imports=False` prevents # `AstProfileTransformer` from rewriting the imports, so # we're really testing `ProfmodExtractor` here - case_fname, list(all_names), False, config=config, + case_fname, list(all_names), False, + config=config, + ) + module_ast = atp.profile( + profile_nested_imports=profile_nested_imports, ) - module_ast = atp.profile() else: module_ast = AstProfileTransformer._transform( ast.parse(test_case), case_fname, profile_imports=True, config=config, + profile_nested_imports=profile_nested_imports, ) output = ast.unparse(module_ast) From 7dc2d5d1d86e30b2b902972b76a05d2361420367 Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Fri, 24 Jul 2026 10:41:51 +0200 Subject: [PATCH 10/12] `kernprof` options for the new configs kernprof.py __doc__ Updated with the new options --prof-imports Updated help text --[no-]prof-star-imports New flag for whether to profile star-imports, corresponding to `~.autoprofile.autoprofile.run(profile_star_imports=...)` --prof-nested-imports New flag for selecting constructs wherein to look for and profile imports, corresponding to `~.autoprofile.autoprofile.run(profile_nested_imports=...)` _pre_profile() Migrated normalization of `options.prof_mod` to `_parse_arguments()` _main_profile() Now calling `line_profiler.autoprofile.autoprofile.run()` with the appropriate values for the params `profile_star_imports` and `profile_nested_imports` line_profiler/rc/line_profiler.toml::[tool.line_profiler.kernprof] Added new key-value pairs `prof-star-imports` and `prof-nested-imports`, corresponding to the defaults for the eponymous `kernprof` CLI options --- kernprof.py | 93 +++++++++++++++++++++++++---- line_profiler/rc/line_profiler.toml | 10 ++++ 2 files changed, 93 insertions(+), 10 deletions(-) diff --git a/kernprof.py b/kernprof.py index 988ee7f3..24164f48 100755 --- a/kernprof.py +++ b/kernprof.py @@ -135,8 +135,24 @@ def main(): profiling (`-l`/`--line-by-line`). (Default: True) --prof-imports [Y[es] | N[o] | T[rue] | F[alse] | on | off | 1 | 0] If the script/module profiled is in `--prof-mod`, autoprofile - all its imports. Only works with line profiling (`-l`/`--line- + its imports regardless of whether the import targets are + themselves in `--prof-mod`; restrictions from `--prof-star- + imports` and `--prof-nested-imports` apply. Only works with + line profiling (`-l`/`--line-by-line`). (Default: False) + --prof-star-imports [Y[es] | N[o] | T[rue] | F[alse] | on | off | 1 | 0] + Dynamically analyze the contents of star-imports (`from ... + import *`) and profile the imported names as one would the + explicit imports. Only works with line profiling (`-l`/`--line- by-line`). (Default: False) + --prof-nested-imports {conditional | loops | contexts | try-except | func-defs | class-defs}[,...] + List of compound-statement constructs in which to look for and + profile nested import statements. They can be supplied both as + comma-separated items, or separately with multiple copies of + this flag. Only works with line profiling (`-l`/ + `--line-by-line`). (Default: ['conditionals', 'try_except', + 'contexts', 'class_defs']; pass an empty string to clear the + defaults (or any `--prof-nested-imports` target specified + earlier)) output options: -o, --outfile OUTFILE @@ -566,10 +582,38 @@ def _add_core_parser_arguments(parser): '--prof-imports', action='store_true', help='If the script/module profiled is in `--prof-mod`, ' - 'autoprofile all its imports. ' + 'autoprofile its imports regardless of whether the import targets ' + 'are themselves in `--prof-mod`; ' + 'restrictions from `--prof-star-imports` and `--prof-nested-imports` ' + 'apply. ' 'Only works with line profiling (`-l`/`--line-by-line`). ' f'(Default: {default.conf_dict["prof_imports"]})', ) + add_argument( + prof_opts, + '--prof-star-imports', + action='store_true', + help='Dynamically analyze the contents of star-imports ' + '(`from ... import *`) and profile the imported names as ' + 'one would the explicit imports. ' + 'Only works with line profiling (`-l`/`--line-by-line`). ' + f'(Default: {default.conf_dict["prof_star_imports"]})', + ) + add_argument( + prof_opts, + '--prof-nested-imports', + action='append', + metavar='{conditionals | loops | contexts ' + '| try-except | func-defs | class-defs}[,...]', + help='List of compound-statement constructs in which to look for ' + 'and profile nested import statements. ' + 'They can be supplied both as comma-separated items, ' + 'or separately with multiple copies of this flag. ' + 'Only works with line profiling (`-l`/`--line-by-line`). ' + f'(Default: {default.conf_dict["prof_nested_imports"]}; ' + 'pass an empty string to clear the defaults ' + '(or any `--prof-nested-imports` target specified earlier))', + ) out_opts = parser.add_argument_group('output options') if default.conf_dict['outfile']: def_outfile = repr(default.conf_dict['outfile']) @@ -754,17 +798,28 @@ def _parse_arguments( else: return - # Parse the provided config file (if any), and resolve the values - # of the un-specified options + # Parse the provided config file (if any), normalize the specified + # options, and resolve the values of the un-specified options try: del options.help except AttributeError: pass + normalizers = { + # Note: `prof_mod` entries can be filenames (which can contain + # commas), so check against existing filenames before splitting + # them + 'prof_mod': _normalize_profiling_targets, + 'prof_nested_imports': _normalize_prof_nested_imports, + } default = get_cli_config('kernprof', options.config) options.config = default.path for key, default in default.conf_dict.items(): - if getattr(options, key, None) is None: + value = getattr(options, key, None) + if value is None: # Not specified setattr(options, key, default) + elif key in normalizers: # Normalize + value = normalizers[key](value) + setattr(options, key, value) # Add in the pre-partitioned arguments cut off by `-m ` or # `-c