ENH: handle nested and star imports - #440
Draft
TTsangSC wants to merge 12 commits into
Draft
Conversation
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
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
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
line_profiler/autoprofile/ast_profile_transformer.py
<General>
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`
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`
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)
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`
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
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
kernprof.py
__doc__
Updated with the new options
<CLI 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
line_profiler/autoprofile/_import_targets.py
_DROPPED_STAR_IMPORTS_MSG_TEMPLATE
New template for warning messages related to star-imports not
being profiled
ImportTarget._check_and_warn_dropped_imports()
Updated typing and default of param `category`
(`category: type[Warning] = UserWarning`
-> `category: type[Warning] | None = None`)
line_profiler/autoprofile/ast_profile_transformer.py
::AstProfileTransformer._transform()
- Updated warning message emitted (star-imports CAN be handled, we
just configured them not to be bydefault)
- Added new private argument so that we don't double-warn for the
same star-imports
line_profiler/autoprofile/ast_tree_profiler.py::AstTreeProfiler
._profile_ast_tree()
- Fixed bookkeeping of dropped imports to account for the
interpolated nodes
- Added code to issue a warning for dropped `--prof-mod` imports
(see `.profile()`)
- Now calling `AstProfileTransformer._transform()` with the
aforementioned private argument to ensure that we don't
double-warn when using both `--prof-mod` and full-module AST
rewrite with `--prof-imports`
.profile()
Now calling `ProfmodExtractor.extract_all()` with
`filter_star_imports=False` to let `._profile_ast_tree()` handle
the warnings
line_profiler/autoprofile/line_profiler_utils.py::add_star_import()
Updated because of the `_should_profile_regular_import()` rename
line_profiler/autoprofile/profmod_extractor.py
_should_profile_regular_import()
Renamed from `_should_profile()`
_should_profile_star_import()
New helper function analogous to the above but for star-imports
ProfmodExtractor._find_modnames_in_tree_imports()
Fixed behavior when a star-import is encountered: if e.g.
`'foo.bar.foobar'` is in `modnames_to_profile`, it counts as
a match against `ImportTarget.name == 'foo.bar.*'`, because
the `add_star_import()` call inserted after the
`from foo.bar import *` statement will decide whether
`foo.bar.foobar` was actually imported and proceed to profile it
if that is the case
ProfmodExtractor.extract_all()
Updated warning message emitted (star-imports CAN be handled, we
just configured them not to be bydefault)
tests/test_autoprofile.py
_propose_module_names()
Helper function refactored out from `test_add_star_import()`
test_autoprofile_star_imports()
New end-to-end test for testing `kernprof --prof-star-imports`
kernprof.py
__doc__
Updated sample help output
<Flag `--prof-nested-imports`>
- Revised help text
- Added processing for the special value 'all', which selects
all the construct types
tests/test_autoprofile.py
test_autoprofile_star_imports()
Removed redundant `MonkeyPatch.syspath_prepend()` call
test_autoprofile_nested_imports()
New end-to-end test for testing that
`kernprof --prof-nested-imports` behaves as expected
CHANGELOG.rst
Added entry
This was referenced Aug 14, 2026
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #440 +/- ##
==========================================
+ Coverage 84.95% 85.82% +0.87%
==========================================
Files 21 21
Lines 2412 2667 +255
Branches 376 419 +43
==========================================
+ Hits 2049 2289 +240
- Misses 263 276 +13
- Partials 100 102 +2
Continue to review full report in Codecov by Harness.
🚀 New features to boost your workflow:
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
(See earlier and more discussion at: #438, TTsangSC#6)
Motivation
The current
line_profiler.autoprofileimplementation doesn't allow for handling imports nested inside code blocks, unless--prof-mod=<executed/script/or/module> --prof-importsis used:However, the use of
--prof-importswould also mean that all imported names will be indiscriminately profiled; thus we either lose the fine-grained control that--prof-mod=some.specific.targetgives us, or the ability to capture nested imports like the ones found above.Plus, star-imports are still not handled (#434 mitigated the crash, but did not add actual profiling for them yet):
Changes (brief summary)
This draft PR remedies both by:
ProfmodExtractorto do an AST pass for imports, not just scanning overast.Module.body.ast_profile_transformer.py::ast_create_star_import_node()andline_profiler_utils.py::add_star_import()to handle the post-hoc profiling of star-imported names.kernprofCLI flags for controlling both:--prof-star-imports:Whether to handle star-imports (old behavior is false)
--prof-nested-imports=...:What kind(s) of code blocks
ProfmodExtractorshould descend into to scan for importsline_profiler.tomlconfig options:[tool.line_profiler.autoprofile]::prof_star_importsand::import_discovery[tool.line_profiler.kernprof]::prof-star-importsand::prof-nested-importsThe behavior of
--prof-importshas also become more consistent with that of--prof-modsince now both are reined in by--prof-nested-imports=.... But while this is IMO a positive change, it is still a semantic change nonetheless.See TTsangSC#6 for details.
Why this is still in draft
There's quite a lot of code in here, but as I was writing it it felt increasingly apparent to me that:
ProfmodExtractorandAstProfileTransformercrawl and transform (in the case ofProfmodExtractor, indirectly viaAstTreeProfiler._profile_ast_tree()) the AST only added (1) overhead and (2) unnecessary bookkeeping for consistency and prevention of duplication.Cutting out the middleman (
ProfmodExtractor) and lettingAstProfileTransformeralone handle the profiling of imports would fix both, but it is probably not trivial to do it in a backward-compatible way. Hence I figured I'd just put the code out here first, so that we sit on it for a while, and weigh whether to incorporate the changes as-is, or to wait until the next release where we can be more haphazard with changing (esp. removing) and refactoring repo components.