diff --git a/python/ferro_ta/_utils.py b/python/ferro_ta/_utils.py index 8e24683..a981848 100644 --- a/python/ferro_ta/_utils.py +++ b/python/ferro_ta/_utils.py @@ -20,6 +20,26 @@ DEFAULT_OHLCV_COLUMNS = { } +@functools.lru_cache(maxsize=1) +def _optional_pandas_module(): + """Import pandas lazily once and cache absence for low-overhead hot paths.""" + try: + import pandas as pd + except ImportError: + return None + return pd + + +@functools.lru_cache(maxsize=1) +def _optional_polars_module(): + """Import polars lazily once and cache absence for low-overhead hot paths.""" + try: + import polars as pl + except ImportError: + return None + return pl + + def _to_f64(data: ArrayLike) -> np.ndarray: """Convert any array-like to a contiguous 1-D float64 NumPy array. @@ -159,9 +179,8 @@ def pandas_wrap(func): @functools.wraps(func) def wrapper(*args, **kwargs): - try: - import pandas as pd # local import — pandas is optional - except ImportError: + pd = _optional_pandas_module() + if pd is None: return func(*args, **kwargs) pd_index = None @@ -232,9 +251,8 @@ def polars_wrap(func): @functools.wraps(func) def wrapper(*args, **kwargs): - try: - import polars as pl # local import — polars is optional - except ImportError: + pl = _optional_polars_module() + if pl is None: return func(*args, **kwargs) pl_name: Optional[str] = None diff --git a/tests/unit/test_optional_dependency_wrappers.py b/tests/unit/test_optional_dependency_wrappers.py new file mode 100644 index 0000000..ed7b03f --- /dev/null +++ b/tests/unit/test_optional_dependency_wrappers.py @@ -0,0 +1,57 @@ +from __future__ import annotations + +from unittest.mock import patch + +import numpy as np + +from ferro_ta._utils import ( + _optional_pandas_module, + _optional_polars_module, + pandas_wrap, + polars_wrap, +) + + +def _missing_only(module_name: str): + real_import = __import__ + attempts: list[str] = [] + + def side_effect(name, globals=None, locals=None, fromlist=(), level=0): + if name == module_name: + attempts.append(name) + raise ImportError(f"{module_name} not installed") + return real_import(name, globals, locals, fromlist, level) + + return attempts, side_effect + + +def test_pandas_wrap_caches_missing_optional_import() -> None: + _optional_pandas_module.cache_clear() + wrapped = pandas_wrap(lambda arr: arr) + arr = np.array([1.0, 2.0, 3.0], dtype=np.float64) + attempts, side_effect = _missing_only("pandas") + + try: + with patch("builtins.__import__", side_effect=side_effect): + np.testing.assert_array_equal(wrapped(arr), arr) + np.testing.assert_array_equal(wrapped(arr), arr) + finally: + _optional_pandas_module.cache_clear() + + assert attempts == ["pandas"] + + +def test_polars_wrap_caches_missing_optional_import() -> None: + _optional_polars_module.cache_clear() + wrapped = polars_wrap(lambda arr: arr) + arr = np.array([1.0, 2.0, 3.0], dtype=np.float64) + attempts, side_effect = _missing_only("polars") + + try: + with patch("builtins.__import__", side_effect=side_effect): + np.testing.assert_array_equal(wrapped(arr), arr) + np.testing.assert_array_equal(wrapped(arr), arr) + finally: + _optional_polars_module.cache_clear() + + assert attempts == ["polars"]