ENH: ensure that dir(astropy) lists subpackages

mhvk · Jan 5, 2025 · Jan 2, 2025 · Jan 3, 2025 · Jan 3, 2025 · Jan 3, 2025
commit 06833cb1f562b6e34d5f8fa89f9c9bbb6543f369
diff --git a/astropy/__init__.py b/astropy/__init__.py
@@ -11,6 +11,54 @@
 
 from astropy.utils.system_info import system_info
 
+__all__ = [  # noqa: RUF100, RUF022
+    "__version__",
+    "__bibtex__",
+    # Subpackages (mostly lazy-loaded)
+    "config",
+    "constants",
+    "convolution",
+    "coordinates",
+    "cosmology",
+    "io",
+    "modeling",
+    "nddata",
+    "samp",
+    "stats",
+    "table",
+    "time",
+    "timeseries",
+    "uncertainty",
+    "units",
+    "utils",
+    "visualization",
+    "wcs",
+    # Functions
+    "test",
+    "log",
+    "find_api_page",
+    "online_help",
+    "online_docs_root",
+    "conf",
+    "physical_constants",
+    "astronomical_constants",
+    "system_info",
+]
+
+
+def __getattr__(attr):
+    if attr in __all__:
+        from importlib import import_module
+
+        return import_module("astropy." + attr)
+
+    raise AttributeError(f"module 'astropy' has no attribute {attr!r}")
+
+
+def __dir__():
+    return sorted(set(globals()).union(__all__))
+
+
 from .version import version as __version__
 
 # The location of the online documentation for astropy
@@ -198,30 +246,14 @@ def online_help(query):
     webbrowser.open(url)
 
 
-__dir_inc__ = [
-    "__version__",
-    "__githash__",
-    "__bibtex__",
-    "test",
-    "log",
-    "find_api_page",
-    "online_help",
-    "online_docs_root",
-    "conf",
-    "physical_constants",
-    "astronomical_constants",
-    "system_info",
-]
-
-
 from types import ModuleType as __module_type__
 
-# Clean up top-level namespace--delete everything that isn't in __dir_inc__
+# Clean up top-level namespace--delete everything that isn't in __all__
 # or is a magic attribute, and that isn't a submodule of this package
 for varname in dir():
     if not (
         (varname.startswith("__") and varname.endswith("__"))
-        or varname in __dir_inc__
+        or varname in __all__
         or (
             varname[0] != "_"
             and isinstance(locals()[varname], __module_type__)

diff --git a/astropy/tests/tests/test_imports.py b/astropy/tests/tests/test_imports.py
@@ -1,6 +1,14 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 
 import pkgutil
+import subprocess
+import sys
+from textwrap import dedent
+from types import ModuleType
+
+import pytest
+
+import astropy
 
 
 def test_imports():
@@ -32,3 +40,21 @@ def test_toplevel_namespace():
     assert "log" in d
     assert "test" in d
     assert "sys" not in d
+
+
+def test_toplevel_lazy_imports():
+    # Check that subpackages are loaded on demand.
+    cmd = dedent("""
+    import astropy, sys
+    assert 'astropy.units' not in sys.modules
+    astropy.units
+    assert 'astropy.units' in sys.modules
+    """)
+    cp = subprocess.check_call([sys.executable, "-c", cmd])
+    assert cp == 0
+
+
+def test_toplevel_attribute_error():
+    # Ensure that our __getattr__ does not leak an import error or so.
+    with pytest.raises(AttributeError, match="module 'astropy' has no"):
+        astropy.nonsense
diff --git a/docs/changes/17598.other.rst b/docs/changes/17598.other.rst
@@ -0,0 +1,4 @@
+After ``import astropy``, ``dir(astropy)`` will now list all subpackages,
+including those that have not yet been loaded. This also means tab
+completion will work as expected (e.g., ``from astropy.coo<TAB>`` will
+expand to ``from astropy.coordinates``).