Rework arc atom lazy import to be simpler by using a constructor function

instead of a dict of imported classes.

Author: dihm

src/rydiqule/atom_utils.py

@@ -4,55 +4,44 @@
 
 from scipy.constants import epsilon_0, hbar, c
 import numpy as np
-import re 
+import re
 from .sensor_utils import expand_statespec
 
-from typing import Union, Optional, Literal, Tuple, List, Dict, Callable, NamedTuple, Any
-
-from .exceptions import RydiquleError
-
-__all__ = [
-    'QSpec',
-    'A_QState',
-    'QState',
-    'D1_excited',
-    'D1_states',
-    'D2_excited',
-    'D2_states',
-    'ground_state',
-    'expand_qnums',
-    'calc_eta',
-    'calc_kappa',
-    'validate_qnums',
-]
-__lazy_loads = [
-    'ATOMS'
-]
-
-def __dir__() -> List[str]:
-
-    return sorted(__all__ + __lazy_loads)
-
-
-def __getattr__(name: str) -> Any:
-
-    if name == 'ATOMS':
-        import arc.alkali_atom_data as arc_atoms
-
-        ATOMS = {
-            'H': arc_atoms.Hydrogen,
-            'Li6': arc_atoms.Lithium6, 'Li7': arc_atoms.Lithium7,
-            'Na': arc_atoms.Sodium,
-            'K39': arc_atoms.Potassium39, 'K40': arc_atoms.Potassium40, 'K41': arc_atoms.Potassium41,
-            'Rb85': arc_atoms.Rubidium85, 'Rb87': arc_atoms.Rubidium87,
-            'Cs': arc_atoms.Caesium
-        }
-        """
-        Alkali atoms defined by ARC that can be used with :class:`~.Cell`.
-        """
-        return ATOMS
-    else:
-        raise AttributeError(f'Module {__name__} does not have attribute {name}')
+from typing import Union, Optional, Literal, Tuple, List, Dict, Callable, NamedTuple, TYPE_CHECKING
+
+from .exceptions import RydiquleError, AtomError
+
+if TYPE_CHECKING:
+    import arc.alkali_atom_data
+
+ATOMS = {
+    'H': 'Hydrogen',
+    'Li6': 'Lithium6', 'Li7': 'Lithium7',
+    'Na': 'Sodium',
+    'K39': 'Potassium39', 'K40': 'Potassium40', 'K41': 'Potassium41',
+    'Rb85': 'Rubidium85', 'Rb87': 'Rubidium87',
+    'Cs': 'Caesium'
+}
+"""
+Alkali atoms defined by ARC that can be used with :class:`~.Cell`.
+"""
+
+def _load_arc_atom(atom_flag: str) -> 'arc.alkali_atom_data.Alkali_Atom':
+    """
+    Function that lazy loads ARC atoms from :external+arc:module:`~arc.alkali_atom_data`
+
+    Returns
+    -------
+    arc.alkali_atom_data.Alkali_Atom
+        ARC alkali atom class associated with the provided atom_flag.
+    """
+
+    import arc.alkali_atom_data as arc_atoms
+    
+    if atom_flag not in ATOMS.keys():
+        raise AtomError(f"Atom flag must be one of {ATOMS.keys()}")
+    
+    return getattr(arc_atoms, ATOMS[atom_flag])() # instantiate the class here
 
 
 ground_n = {

src/rydiqule/cell.py

@@ -14,7 +14,7 @@
 from .sensor import Sensor
 from .sensor_utils import scale_dipole
 from .sensor_utils import ScannableParameter, TimeFunc
-from .atom_utils import calc_kappa, calc_eta, expand_qnums, validate_qnums, A_QState, ground_state
+from .atom_utils import calc_kappa, calc_eta, expand_qnums, validate_qnums, A_QState, ground_state, _load_arc_atom
 from .arc_utils import RQ_AlkaliAtom
 from .exceptions import RydiquleError, AtomError, CouplingNotAllowedError
 from .exceptions import RydiquleWarning, debug_state
@@ -165,12 +165,9 @@ def __init__(self, atom_flag: AtomFlags, atomic_states: List[A_QState],
         (5, 0, 0.5, f=2.0, m_f=2.0)
 
         """
-        from .atom_utils import ATOMS
-        if atom_flag not in ATOMS.keys():
-            raise AtomError(f"Atom flag must be one of {ATOMS.keys()}")
 
         self.atom_flag = atom_flag
-        self.atom = RQ_AlkaliAtom(ATOMS[atom_flag]())
+        self.atom = RQ_AlkaliAtom(_load_arc_atom(atom_flag))
         self.I = self.atom.arc_atom.I
 
         #prepare states by expanding statespec into list and intialize graph