Fixed #495 Reduce Import Time (Lazy Imports)

environment.yml

@@ -27,3 +27,4 @@ dependencies:
   - jupyterlab-myst>=2.0.0
   - myst-nb>=1.0.0
   - polars>=1.14.0
+  - numba-cuda>=0.24.0

pyproject.toml

@@ -88,6 +88,7 @@ pip = "*"
 lxml = "*"
 fftw = "*"
 pyfftw = "*"
+numba-cuda = "*"
 # readthedocs
 sphinx = ">=3.5.3"
 pydata-sphinx-theme = "*"
@@ -112,5 +113,5 @@ tbb = ">=2019.5"
 tests = "./test.sh"
 coverage = "./test.sh coverage"
 docs = "cd docs && ./setup.sh"
-black = 'black --exclude=".*\.ipynb" --extend-exclude=".venv|.pixi" --diff ./'
+black = 'black --exclude=".*\.ipynb" --extend-exclude=".venv|.pixi" ./'
 isort = 'isort --profile black --skip .venv --skip .pixi ./'

stumpy/__init__.py

@@ -1,39 +1,63 @@
 import ast
+import importlib
 import os.path
 import pathlib
+import types
 from importlib.metadata import distribution
 from site import getsitepackages
 
 from numba import cuda
 
 from . import cache, config
-from .aamp import aamp  # noqa: F401
-from .aamp_mmotifs import aamp_mmotifs  # noqa: F401
-from .aamp_motifs import aamp_match, aamp_motifs  # noqa: F401
-from .aamp_ostinato import aamp_ostinato, aamp_ostinatoed  # noqa: F401
-from .aamp_stimp import aamp_stimp, aamp_stimped  # noqa: F401
-from .aampdist import aampdist, aampdisted  # noqa: F401
-from .aampdist_snippets import aampdist_snippets  # noqa: F401
-from .aamped import aamped  # noqa: F401
-from .aampi import aampi  # noqa: F401
-from .chains import allc, atsc  # noqa: F401
-from .core import mass  # noqa: F401
-from .floss import floss, fluss  # noqa: F401
-from .maamp import maamp, maamp_mdl, maamp_subspace  # noqa: F401
-from .maamped import maamped  # noqa: F401
-from .mmotifs import mmotifs  # noqa: F401
-from .motifs import match, motifs  # noqa: F401
-from .mpdist import mpdist, mpdisted  # noqa: F401
-from .mstump import mdl, mstump, subspace  # noqa: F401
-from .mstumped import mstumped  # noqa: F401
-from .ostinato import ostinato, ostinatoed  # noqa: F401
-from .scraamp import prescraamp, scraamp  # noqa: F401
-from .scrump import prescrump, scrump  # noqa: F401
-from .snippets import snippets  # noqa: F401
-from .stimp import stimp, stimped  # noqa: F401
-from .stump import stump  # noqa: F401
-from .stumped import stumped  # noqa: F401
-from .stumpi import stumpi  # noqa: F401
+
+# Define which functions belong to which module
+# Key: function name to expose at top level
+# Value: name of the module
+_lazy_imports = {
+    "aamp": "aamp",
+    "aamp_mmotifs": "aamp_mmotifs",
+    "aamp_match": "aamp_motifs",
+    "aamp_motifs": "aamp_motifs",
+    "aamp_ostinato": "aamp_ostinato",
+    "aamp_ostinatoed": "aamp_ostinato",
+    "aamp_stimp": "aamp_stimp",
+    "aamp_stimped": "aamp_stimp",
+    "aampdist": "aampdist",
+    "aampdisted": "aampdist",
+    "aampdist_snippets": "aampdist_snippets",
+    "aamped": "aamped",
+    "aampi": "aampi",
+    "allc": "chains",
+    "atsc": "chains",
+    "mass": "core",
+    "floss": "floss",
+    "fluss": "floss",
+    "maamp": "maamp",
+    "maamp_mdl": "maamp",
+    "maamp_subspace": "maamp",
+    "maamped": "maamped",
+    "mmotifs": "mmotifs",
+    "match": "motifs",
+    "motifs": "motifs",
+    "mpdist": "mpdist",
+    "mpdisted": "mpdist",
+    "mdl": "mstump",
+    "mstump": "mstump",
+    "subspace": "mstump",
+    "mstumped": "mstumped",
+    "ostinato": "ostinato",
+    "ostinatoed": "ostinato",
+    "prescraamp": "scraamp",
+    "scraamp": "scraamp",
+    "prescrump": "scrump",
+    "scrump": "scrump",
+    "snippets": "snippets",
+    "stimp": "stimp",
+    "stimped": "stimp",
+    "stump": "stump",
+    "stumped": "stumped",
+    "stumpi": "stumpi",
+}
 
 # Get the default fastmath flags for all njit functions
 # and update the _STUMPY_DEFAULTS dictionary
@@ -61,14 +85,18 @@ def _get_fastmath_value(module_name, func_name):  # pragma: no cover
     config._STUMPY_DEFAULTS[key] = _get_fastmath_value(module_name, func_name)
 
 if cuda.is_available():
-    from .gpu_aamp import gpu_aamp  # noqa: F401
-    from .gpu_aamp_ostinato import gpu_aamp_ostinato  # noqa: F401
-    from .gpu_aamp_stimp import gpu_aamp_stimp  # noqa: F401
-    from .gpu_aampdist import gpu_aampdist  # noqa: F401
-    from .gpu_mpdist import gpu_mpdist  # noqa: F401
-    from .gpu_ostinato import gpu_ostinato  # noqa: F401
-    from .gpu_stimp import gpu_stimp  # noqa: F401
-    from .gpu_stump import gpu_stump  # noqa: F401
+    _lazy_imports.update(
+        {
+            "gpu_aamp": "gpu_aamp",
+            "gpu_aamp_ostinato": "gpu_aamp_ostinato",
+            "gpu_aamp_stimp": "gpu_aamp_stimp",
+            "gpu_aampdist": "gpu_aampdist",
+            "gpu_mpdist": "gpu_mpdist",
+            "gpu_ostinato": "gpu_ostinato",
+            "gpu_stimp": "gpu_stimp",
+            "gpu_stump": "gpu_stump",
+        }
+    )
 else:  # pragma: no cover
     from . import core
     from .core import _gpu_aamp_driver_not_found as gpu_aamp  # noqa: F401
@@ -220,3 +248,69 @@ def _get_fastmath_value(module_name, func_name):  # pragma: no cover
     __version__ = "Please install this project with setup.py"
 else:  # pragma: no cover
     __version__ = _dist.version
+
+
+# PEP 562: module-level __getattr__ for lazy imports
+def __getattr__(name):  # pragma: no cover
+    if name in _lazy_imports:
+        mod_name = _lazy_imports[name]
+        module = importlib.import_module(f"{__package__}.{mod_name}")
+        # Retrieve the attribute from the loaded module and cache it
+        attr = getattr(module, name)
+        globals()[name] = attr
+        return attr
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+# Ensure that if a module was imported during package import
+# (causing the package attribute to point to the module object), we
+# replace that entry with the actual attribute (e.g., function) so that
+# users get the expected callable at `stumpy.aamp` rather than the module.
+for _name, _sub in _lazy_imports.items():  # pragma: no cover
+    val = globals().get(_name)
+    if isinstance(val, types.ModuleType):
+        try:
+            replacement = getattr(val, _name)
+        except AttributeError:
+            # Nothing to do if the module doesn't define the attribute
+            continue
+        globals()[_name] = replacement
+
+
+# Eagerly import exports that would otherwise collide with
+# same-named modules. This keeps lazy imports for most names but
+# ensures that when a top-level exported name exactly matches its
+# module (e.g., `stump` -> `stump.py`), the exported attribute is
+# available immediately so REPL completers prefer the callable/class
+# instead of the module.
+for _name, _sub in _lazy_imports.items():  # pragma: no cover
+    try:
+        if _name == _sub:
+            filepath = pathlib.Path(__file__).parent / f"{_sub}.py"
+            if filepath.exists():
+                module = importlib.import_module(f"{__package__}.{_sub}")
+                try:
+                    globals()[_name] = getattr(module, _name)
+                except AttributeError:
+                    # If the module doesn't define the attribute, keep it lazy
+                    pass
+    except Exception:
+        # Be conservative: don't let eager-import attempts raise during package import
+        pass
+
+
+def __dir__():  # pragma: no cover
+    # Expose lazy names in dir() for discoverability
+    # Also include __all__ so tools that consult it will see the intended
+    # top-level exports (this helps some REPL completers prefer the
+    # callable/class exports over same-named modules).
+    all_names = list(globals().keys()) + list(_lazy_imports.keys())
+    all_names += list(globals().get("__all__", []))
+    return sorted(all_names)
+
+
+# Make the lazy-exported names explicit for tools that respect __all__.
+# This helps REPL tab-completion prefer functions/classes over modules
+# when names collide (e.g., `stumpy.stump` should point to the function
+# rather than the module during completion).
+__all__ = sorted(list(_lazy_imports.keys()))

stumpy/gpu_aamp_ostinato.py

@@ -2,8 +2,9 @@
 # Copyright 2019 TD Ameritrade. Released under the terms of the 3-Clause BSD license.
 # STUMPY is a trademark of TD Ameritrade IP Company, Inc. All rights reserved.
 
-from . import core, gpu_aamp
+from . import core
 from .aamp_ostinato import _aamp_ostinato, _get_aamp_central_motif
+from .gpu_aamp import gpu_aamp
 
 
 def gpu_aamp_ostinato(Ts, m, device_id=0, p=2.0):

stumpy/gpu_aamp_stimp.py

@@ -2,8 +2,8 @@
 # Copyright 2019 TD Ameritrade. Released under the terms of the 3-Clause BSD license.
 # STUMPY is a trademark of TD Ameritrade IP Company, Inc. All rights reserved.
 
-from . import gpu_aamp
 from .aamp_stimp import _aamp_stimp
+from .gpu_aamp import gpu_aamp
 
 
 class gpu_aamp_stimp(_aamp_stimp):

stumpy/gpu_aampdist.py

@@ -4,8 +4,8 @@
 
 import functools
 
-from . import gpu_aamp
 from .core import _mpdist
+from .gpu_aamp import gpu_aamp
 
 
 def gpu_aampdist(T_A, T_B, m, percentage=0.05, k=None, device_id=0, p=2.0):

stumpy/stimp.py

@@ -4,8 +4,9 @@
 
 import numpy as np
 
-from . import core, scrump
+from . import core
 from .aamp_stimp import aamp_stimp, aamp_stimped
+from .scrump import scrump
 from .stump import stump
 from .stumped import stumped
 

stumpy/stomp.py

@@ -6,7 +6,8 @@
 
 import numpy as np
 
-from . import config, core, stamp
+from . import config, core
+from .stamp import _mass_PI
 
 
 def _stomp(T_A, m, T_B=None, ignore_trivial=True):
@@ -105,7 +106,7 @@ def _stomp(T_A, m, T_B=None, ignore_trivial=True):
         IR = -1
     else:
         if ignore_trivial:
-            P, I = stamp._mass_PI(
+            P, I = _mass_PI(
                 T_A[:m],
                 T_B,
                 M_T,
@@ -115,7 +116,7 @@ def _stomp(T_A, m, T_B=None, ignore_trivial=True):
                 T_subseq_isconstant=T_subseq_isconstant,
                 Q_subseq_isconstant=Q_subseq_isconstant[[0]],
             )
-            PR, IR = stamp._mass_PI(
+            PR, IR = _mass_PI(
                 T_A[:m],
                 T_B,
                 M_T,
@@ -127,7 +128,7 @@ def _stomp(T_A, m, T_B=None, ignore_trivial=True):
                 Q_subseq_isconstant=Q_subseq_isconstant[[0]],
             )
         else:
-            P, I = stamp._mass_PI(
+            P, I = _mass_PI(
                 T_A[:m],
                 T_B,
                 M_T,

stumpy/stumpi.py

@@ -4,8 +4,9 @@
 
 import numpy as np
 
-from . import config, core, stump
+from . import config, core
 from .aampi import aampi
+from .stump import stump
 
 
 @core.non_normalized(

tests/test_aamp.py

@@ -4,7 +4,8 @@
 import pandas as pd
 import pytest
 
-from stumpy import aamp, config
+from stumpy import config
+from stumpy.aamp import aamp
 
 test_data = [
     (

tests/test_aamp_mmotifs.py

@@ -3,7 +3,8 @@
 import numpy.testing as npt
 import pytest
 
-from stumpy import aamp_mmotifs, config
+from stumpy import config
+from stumpy.aamp_mmotifs import aamp_mmotifs
 
 test_data = [
     np.array(

tests/test_aamp_motifs.py

@@ -3,7 +3,8 @@
 import numpy.testing as npt
 import pytest
 
-from stumpy import aamp_match, aamp_motifs, core
+from stumpy import core
+from stumpy.aamp_motifs import aamp_match, aamp_motifs
 
 test_data = [
     (