Expose submodules from top-level orbit namespace

py/orbit/__init__.py

@@ -1,14 +1,3 @@
-## \namespace orbit
-## \brief The base module for all pyORBIT core python
-## modules.
-##
-## Some of the modules:
-## - lattice     - the base lattice implementation.
-## - parsers     - collection of parsers.
-## - teapot      - the TEAPOT-like ORBIT lattice elements.
-## - teapot_base - fundamental bricks of teapot elements.
-## - utils       - the collection of useful Python utilites.
-
 try:
     from importlib import metadata
 except ImportError:
@@ -20,3 +9,57 @@
 except ModuleNotFoundError:
     # package is not installed
     pass
+
+from . import aperture
+from . import bumps
+from . import bunch_generators
+from . import bunch_utils
+from . import collimation
+from . import core
+from . import diagnostics
+from . import errors
+from . import fieldtracker
+from . import foils
+from . import impedances
+from . import injection
+from . import kickernodes
+from . import lattice
+from . import matching
+from . import matrix_lattice
+from . import orbit_correction
+from . import parsers
+from . import py_linac
+from . import rf_cavities
+from . import space_charge
+from . import teapot
+from . import teapot_base
+from . import time_dep
+from . import utils
+
+__all__ = [
+    "aperture",
+    "bumps",
+    "bunch_generators",
+    "bunch_utils",
+    "collimation",
+    "core",
+    "diagnostics",
+    "errors",
+    "fieldtracker",
+    "foils",
+    "impedances",
+    "injection",
+    "kickernodes",
+    "lattice",
+    "matching",
+    "matrix_lattice",
+    "orbit_correction",
+    "parsers",
+    "py_linac",
+    "rf_cavities",
+    "space_charge",
+    "teapot",
+    "teapot_base",
+    "time_dep",
+    "utils",
+]

tests/py/orbit/test_import.py

@@ -0,0 +1,38 @@
+import orbit
+
+
+def test_import_orbit_exposes_all_submodules():
+    expected_submodules = {
+        "aperture",
+        "bumps",
+        "bunch_generators",
+        "bunch_utils",
+        "collimation",
+        "core",
+        "diagnostics",
+        "errors",
+        "fieldtracker",
+        "foils",
+        "impedances",
+        "injection",
+        "kickernodes",
+        "lattice",
+        "matching",
+        "matrix_lattice",
+        "orbit_correction",
+        "parsers",
+        "py_linac",
+        "rf_cavities",
+        "space_charge",
+        "teapot",
+        "teapot_base",
+        "time_dep",
+        "utils",
+    }
+    assert set(orbit.__all__) == expected_submodules
+    assert all(hasattr(orbit, name) for name in expected_submodules)
+
+
+def test_teapot_lattice_available_from_orbit_namespace():
+    lattice = orbit.teapot.TEAPOT_Lattice()
+    assert isinstance(lattice, orbit.teapot.TEAPOT_Lattice)