enh(units): Load pint and pint-pandas only on demand (#57)

CHANGELOG.rst

@@ -3,6 +3,8 @@
 Changelog
 =========
 
+* Load pint and pint-pandas packages only on first use
+
 v0.5.1 (2024-05-20)
 ------------------------------------------------------------
 * Fix :func:`~selectors.ismatch` to not interpret brackets as regex symbols

src/pandas_indexing/units.py

@@ -56,27 +56,24 @@
 
 from pandas import DataFrame, Series
 
+from .core import assignlevel, uniquelevel
+from .types import Axis, Data
+from .utils import LazyLoader, doc
 
-INSTALL_PACKAGE_WARNING = "Please install {package} via conda or pip, or use the pandas-indexing[units] variant."
-
-try:
-    from pint_pandas import PintType
-
-    has_pint_pandas = True
-except ImportError:
-    has_pint_pandas = False
-
-try:
-    import pint
-
-    has_pint = True
-except ImportError:
-    has_pint = False
 
+INSTALL_PACKAGE_WARNING = (
+    "Please install {name} via conda or pip, or use the pandas-indexing[units] variant."
+)
 
-from .core import assignlevel, uniquelevel
-from .types import Axis, Data
-from .utils import doc
+pint_pandas = LazyLoader(
+    "pint_pandas",
+    globals(),
+    "pint_pandas",
+    INSTALL_PACKAGE_WARNING.format(name="pint_pandas"),
+)
+pint = LazyLoader(
+    "pint", globals(), "pint", INSTALL_PACKAGE_WARNING.format(name="pint")
+)
 
 
 @doc(
@@ -144,8 +141,6 @@ def quantify(
     dequantify
     convert_unit
     """
-    assert has_pint_pandas, INSTALL_PACKAGE_WARNING.format(package="pint-pandas")
-
     if unit is None:
         unit = uniquelevel(data, level, axis)
         if len(unit) != 1:
@@ -155,19 +150,17 @@ def quantify(
         unit = unit[0]
         data = data.droplevel(level, axis=axis)
 
-    return data.astype(PintType(unit), copy=copy)
+    return data.astype(pint_pandas.PintType(unit), copy=copy)
 
 
 def format_dtype(dtype):
-    if not isinstance(dtype, PintType):
+    if not isinstance(dtype, pint_pandas.PintType):
         return ""
 
     return ("{:" + dtype.ureg.default_format + "}").format(dtype.units)
 
 
 def dequantify(data: Data, level: str = "unit", axis: Axis = 0, copy: bool = False):
-    assert has_pint_pandas, INSTALL_PACKAGE_WARNING.format(package="pint-pandas")
-
     if isinstance(data, Series):
         unit = format_dtype(data.dtype)
         data = data.pint.magnitude
@@ -328,7 +321,6 @@ def get_openscm_registry(add_co2e: bool = True):
 def set_openscm_registry_as_default(add_co2e: bool = True):
     unit_registry = get_openscm_registry(add_co2e=add_co2e)
 
-    assert has_pint, INSTALL_PACKAGE_WARNING.format(package="pint")
     app_registry = pint.get_application_registry()
     if unit_registry is not app_registry:
         pint.set_application_registry(unit_registry)
@@ -337,7 +329,6 @@ def set_openscm_registry_as_default(add_co2e: bool = True):
 
 
 def is_unit(unit: str) -> bool:
-    assert has_pint, INSTALL_PACKAGE_WARNING.format(package="pint")
     ur = pint.get_application_registry()
     try:
         ur.Unit(unit)

src/pandas_indexing/utils.py

@@ -3,7 +3,9 @@
 Simple utility functions not of greater interest
 """
 
+import importlib
 import re
+from types import ModuleType
 from typing import Union
 
 from pandas import DataFrame, Index, Series
@@ -97,3 +99,51 @@ def quote_list(l):
 
 def s(l):
     return "s" if len(l) > 1 else ""
+
+
+class LazyLoader(ModuleType):
+    """Lazily import a module, mainly to avoid pulling in large dependencies.
+
+    `contrib`, and `ffmpeg` are examples of modules that are large and not always
+    needed, and this allows them to only be loaded when they are used.
+
+    Copied from tensorflow's agents.tf_agents.utils.lazy_loader by The TF-Agents Authors (2020), licensed under Apache 2.0
+    https://github.com/tensorflow/agents/blob/737d758452990dc3c81b8aeab1a6ae4f63afa12c/tf_agents/utils/lazy_loader.py#L28-L68
+    """
+
+    def __init__(self, local_name, parent_module_globals, name, importerrormsg=None):
+        self._local_name = local_name
+        self._parent_module_globals = parent_module_globals
+        self._importerrormsg = importerrormsg
+
+        super().__init__(name)
+
+    def _load(self):
+        """
+        Load the module and insert it into the parent's globals.
+        """
+        # Import the target module and insert it into the parent's namespace
+        try:
+            module = importlib.import_module(self.__name__)
+        except ImportError:
+            if self._importerrormsg is not None:
+                raise ImportError(self._importerrormsg) from None
+            else:
+                raise
+
+        self._parent_module_globals[self._local_name] = module
+
+        # Update this object's dict so that if someone keeps a reference to the
+        #   LazyLoader, lookups are efficient (__getattr__ is only called on lookups
+        #   that fail).
+        self.__dict__.update(module.__dict__)
+
+        return module
+
+    def __getattr__(self, item):
+        module = self._load()
+        return getattr(module, item)
+
+    def __dir__(self):
+        module = self._load()
+        return dir(module)

tests/test_units.py

@@ -5,17 +5,15 @@
 from pandas.testing import assert_frame_equal, assert_series_equal
 
 from pandas_indexing import assignlevel, convert_unit, set_openscm_registry_as_default
-from pandas_indexing.units import has_pint, has_pint_pandas, is_unit
+from pandas_indexing.units import is_unit
 
 
-has_openscm_units = bool(find_spec("openscm_units"))
-
-needs_pint = pytest.mark.skipif(not has_pint, reason="Needs pint package")
+needs_pint = pytest.mark.skipif(not find_spec("pint"), reason="Needs pint package")
 needs_openscm_units = pytest.mark.skipif(
-    not has_openscm_units, reason="Needs openscm-units package"
+    not find_spec("openscm_units"), reason="Needs openscm-units package"
 )
 needs_pint_pandas = pytest.mark.skipif(
-    not has_pint_pandas, reason="Needs pint-pandas package"
+    not find_spec("pint_pandas"), reason="Needs pint-pandas package"
 )