init

2026-02-12 23:22:11 +08:00
parent 7b09eb3d89
commit 89660bba4e
5988 changed files with 2517516 additions and 0 deletions
@@ -0,0 +1,123 @@
+"""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=[],
+        )