diff --git a/py/orbit/__init__.py b/py/orbit/__init__.py index fa5046e5..896f0802 100644 --- a/py/orbit/__init__.py +++ b/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", +] diff --git a/tests/py/orbit/test_import.py b/tests/py/orbit/test_import.py new file mode 100644 index 00000000..afc936bf --- /dev/null +++ b/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)