124 lines
3.1 KiB
Python
124 lines
3.1 KiB
Python
"""Track imports."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
from scooby.knowledge import get_standard_lib_modules
|
|
from scooby.report import Report
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import Mapping, Sequence
|
|
from types import ModuleType
|
|
|
|
|
|
TRACKING_SUPPORTED = False
|
|
SUPPORT_MESSAGE = (
|
|
'Tracking is not supported for this version of Python. Try using a modern version of Python.'
|
|
)
|
|
try:
|
|
import builtins
|
|
|
|
CLASSIC_IMPORT = builtins.__import__
|
|
TRACKING_SUPPORTED = True
|
|
except (ImportError, AttributeError):
|
|
pass
|
|
|
|
# The variable we track all imports in
|
|
TRACKED_IMPORTS: list[str | ModuleType] = ['scooby']
|
|
|
|
MODULES_TO_IGNORE = {
|
|
'pyMKL',
|
|
'mkl',
|
|
'vtkmodules',
|
|
'mpl_toolkits',
|
|
}
|
|
|
|
|
|
STDLIB_PKGS: set[str] = set()
|
|
|
|
|
|
def _criterion(name: str) -> bool:
|
|
return (
|
|
len(name) > 0
|
|
and name not in STDLIB_PKGS
|
|
and not name.startswith('_')
|
|
and name not in MODULES_TO_IGNORE
|
|
)
|
|
|
|
|
|
if TRACKING_SUPPORTED:
|
|
|
|
def scooby_import(
|
|
name: str,
|
|
globals: Mapping[str, object] | None = None, # noqa: A002
|
|
locals: Mapping[str, object] | None = None, # noqa: A002
|
|
fromlist: Sequence[str] = (),
|
|
level: int = 0,
|
|
) -> ModuleType:
|
|
"""Override of the import method to track package names."""
|
|
m = CLASSIC_IMPORT(name, globals=globals, locals=locals, fromlist=fromlist, level=level)
|
|
name = name.split('.')[0]
|
|
if level == 0 and _criterion(name):
|
|
TRACKED_IMPORTS.append(name)
|
|
return m
|
|
|
|
|
|
def track_imports() -> None:
|
|
"""Track all imported modules for the remainder of this session."""
|
|
if not TRACKING_SUPPORTED:
|
|
raise RuntimeError(SUPPORT_MESSAGE)
|
|
global STDLIB_PKGS
|
|
STDLIB_PKGS = get_standard_lib_modules()
|
|
builtins.__import__ = scooby_import
|
|
|
|
|
|
def untrack_imports() -> None:
|
|
"""Stop tracking imports and return to the builtin import method.
|
|
|
|
This will also clear the tracked imports.
|
|
"""
|
|
if not TRACKING_SUPPORTED:
|
|
raise RuntimeError(SUPPORT_MESSAGE)
|
|
builtins.__import__ = CLASSIC_IMPORT
|
|
TRACKED_IMPORTS.clear()
|
|
TRACKED_IMPORTS.append('scooby')
|
|
|
|
|
|
class TrackedReport(Report):
|
|
"""A class to inspect the active environment and generate a report.
|
|
|
|
Generates a report based on all imported modules. Simply pass the
|
|
``globals()`` dictionary.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
additional: list[str | ModuleType] | None = None,
|
|
ncol: int = 3,
|
|
text_width: int = 80,
|
|
sort: bool = False,
|
|
) -> None:
|
|
"""Initialize."""
|
|
if not TRACKING_SUPPORTED:
|
|
raise RuntimeError(SUPPORT_MESSAGE)
|
|
if len(TRACKED_IMPORTS) < 2:
|
|
msg = (
|
|
'There are no tracked imports, please use '
|
|
'`scooby.track_imports()` before running your '
|
|
'code.'
|
|
)
|
|
raise RuntimeError(
|
|
msg,
|
|
)
|
|
|
|
Report.__init__(
|
|
self,
|
|
additional=additional,
|
|
core=TRACKED_IMPORTS,
|
|
ncol=ncol,
|
|
text_width=text_width,
|
|
sort=sort,
|
|
optional=[],
|
|
)
|