Lazy load objects at top level

openff/units/__init__.py

@@ -1,12 +1,27 @@
+"""
+Registry of units provided by OpenFF Units.
+
+``unit`` may be used similarly to a module. It makes constants and units of
+measure available as attributes. Available units can be found in the
+:download:`constants <../../../openff/units/data/constants.txt>` and
+:download:`defaults <../../../openff/units/data/defaults.txt>` data files.
+"""
+
+import importlib
+from types import ModuleType
+from typing import TYPE_CHECKING
+
 from openff.units._version import get_versions
-from openff.units.openmm import ensure_quantity
-from openff.units.units import (  # type: ignore[attr-defined]
-    DEFAULT_UNIT_REGISTRY,
-    Measurement,
-    Quantity,
-    Unit,
-    UnitRegistry,
-)
+
+if TYPE_CHECKING:
+    # Type checkers can't see lazy-imported objects
+    from openff.units.openmm import ensure_quantity
+    from openff.units.units import DEFAULT_UNIT_REGISTRY, Measurement, Quantity, Unit, UnitRegistry
+
+
+versions = get_versions()
+__version__ = get_versions()["version"]
+__git_revision__ = versions["full-revisionid"]
 
 __all__ = [
     "unit",
@@ -16,18 +31,34 @@
     "ensure_quantity",
 ]
 
-unit: UnitRegistry = DEFAULT_UNIT_REGISTRY
-"""
-Registry of units provided by OpenFF Units.
+_objects: dict[str, str] = {
+    "ensure_quantity": "openff.units.openmm",
+    "Measurement": "openff.units.units",
+    "Unit": "openff.units.units",
+    "UnitRegistry": "openff.units.units",
+    "Quantity": "openff.units.units",
+    "DEFAULT_UNIT_REGISTRY": "openff.units.units",
+    "unit": "openff.units.units",
+}
 
-``unit`` may be used similarly to a module. It makes constants and units of
-measure available as attributes. Available units can be found in the
-:download:`constants <../../../openff/units/data/constants.txt>` and
-:download:`defaults <../../../openff/units/data/defaults.txt>` data files.
-"""
 
-# Handle versioneer
-versions = get_versions()
-__version__ = versions["version"]
-__git_revision__ = versions["full-revisionid"]
-del get_versions, versions
+def __getattr__(name) -> ModuleType:
+    """
+    Lazily import objects from submodules.
+
+    Taken from openff/toolkit/__init__.py
+    """
+    module = _objects.get(name)
+    if module is not None:
+        try:
+            return importlib.import_module(module).__dict__[name]
+        except ImportError as error:
+            raise ImportError from error
+
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__():
+    """Add _objects to dir()."""
+    keys = (*globals().keys(), *_objects.keys())
+    return sorted(keys)

openff/units/units.py

@@ -22,6 +22,7 @@
     "Quantity",
     "Measurement",
     "Unit",
+    "unit",
 ]
 
 
@@ -77,6 +78,8 @@ class UnitRegistry(pint.UnitRegistry):
 
 DEFAULT_UNIT_REGISTRY = UnitRegistry(get_defaults_path())
 
+unit = DEFAULT_UNIT_REGISTRY
+
 Unit: _Unit = DEFAULT_UNIT_REGISTRY.Unit
 Quantity: _Quantity = DEFAULT_UNIT_REGISTRY.Quantity
 Measurement: _Measurement = DEFAULT_UNIT_REGISTRY.Measurement