Lazy loader

.pre-commit-config.yaml

@@ -70,14 +70,15 @@ repos:
     hooks:
       - id: mypy
         files: src
+        exclude: ^src/galax/__init__\.py$
         additional_dependencies:
           - pytest
 
   - repo: https://github.com/codespell-project/codespell
     rev: "v2.2.6"
     hooks:
       - id: codespell
-        exclude: ^(?!.*\/notebooks\/)(?!.AUTHORS\.rst$).+$
+        exclude: ^(notebooks|AUTHORS\.rst)$
 
   - repo: https://github.com/abravalheri/validate-pyproject
     rev: v0.16

pyproject.toml

@@ -35,6 +35,7 @@
     "equinox != 0.11.3",
     "jax",
     "jaxlib",
+    "lazy_loader",
     "typing_extensions",
   ]
 
@@ -123,7 +124,11 @@
   warn_redundant_casts = true
   warn_unused_configs = true
   warn_unreachable = true
-  exclude = '''(^|/)tests/|(^/)docs/|(^/)conftest\.py'''
+  exclude = [
+    '(^/)docs/',  # docs
+    '(^|/)tests/',  # tests
+    '^conftest\.py$',  # nox test configuration
+  ]
 
   enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"]
 
@@ -196,6 +201,7 @@
     ]
     "noxfile.py" = ["ERA001", "T20"]
     "__init__.py" = ["F403"]
+    "__init__.pyi" = ["F401"]
 
 
   [tool.ruff.lint.isort]

src/galax/__init__.py

@@ -1,36 +1,23 @@
 """Copyright (c) 2023 galax maintainers. All rights reserved."""
 # ruff:noqa: F401
 
-__all__ = [
-    "__version__",
-    # modules
-    "units",
-    "potential",
-    "integrate",
-    "dynamics",
-    "utils",
-    "typing",
-]
-
 import os
 
 from jax import config
 from jaxtyping import install_import_hook
-
-from ._version import version as __version__
+from lazy_loader import attach_stub as _attach_stub
 
 config.update("jax_enable_x64", True)  # noqa: FBT003
 
-TYPECHECKER: str | None
-if os.environ.get("GALDYNAMIX_ENABLE_RUNTIME_TYPECHECKS", "1") == "1":
-    TYPECHECKER = "beartype.beartype"
-else:
-    TYPECHECKER = None
+_RUNTIME_TYPECHECKER = (
+    "beartype.beartype"
+    if (os.environ.get("GALAX_ENABLE_RUNTIME_TYPECHECKS", "1") == "1")
+    else None
+)
 
-with install_import_hook("galax", TYPECHECKER):
-    from galax import dynamics, integrate, potential, typing, units, utils
+with install_import_hook("galax", _RUNTIME_TYPECHECKER):
+    __getattr__, __dir__, __all__ = _attach_stub(__name__, __file__)
 
 
-# Clean up the namespace
-for name in set(dir()) - set(__all__):
-    del globals()[name]
+# Install the runtime typechecker
+install_import_hook("galax", _RUNTIME_TYPECHECKER)

src/galax/__init__.pyi

@@ -0,0 +1,24 @@
+__all__ = [
+    "__version__",
+    "__version_tuple__",
+    # modules
+    "dynamics",
+    "integrate",
+    "potential",
+    "typing",
+    "units",
+    "utils",
+]
+
+from . import (
+    dynamics as dynamics,
+    integrate as integrate,
+    potential as potential,
+    typing as typing,
+    units as units,
+    utils as utils,
+)
+from ._version import (  # type: ignore[attr-defined]
+    __version__ as __version__,
+    __version_tuple__ as __version_tuple__,
+)

tests/smoke/test_package.py

@@ -3,5 +3,21 @@
 import galax as pkg
 
 
-def test_version():
+def test_version() -> None:
     assert version("galax") == pkg.__version__
+
+
+def test_all() -> None:
+    """Test the `galax` package contents."""
+    # Test detailed contents (not order)
+    assert set(pkg.__all__) == {
+        "__version__",
+        "__version_tuple__",
+        # modules
+        "dynamics",
+        "integrate",
+        "potential",
+        "typing",
+        "units",
+        "utils",
+    }