refactor: replace try_import_xarray with the generic try_import

Drop the specific try_import_xarray helper now that a generic try_import exists, and check for xarray and xarray-sql directly at the call sites in the xarray source and the UI upload handlers.

Author: ghostiee-11

lumen/ai/ui.py

@@ -37,7 +37,7 @@
 from ..sources import Source
 from ..sources.duckdb import DuckDBSource
 from ..sources.xarray_sql import XArraySQLSource
-from ..util import log, normalize_table_name, try_import_xarray
+from ..util import log, normalize_table_name, try_import
 from .agents import (
     AnalysisAgent, BaseCodeAgent, ChatAgent, DocumentListAgent,
     DocumentSummarizerAgent, SourceAgent, SQLAgent, TableListAgent,
@@ -537,7 +537,7 @@ def _handle_xarray_upload(context, file_obj, alias, filename):
             _temp_files.append(tmp.name)
             return XArraySQLSource(uri=tmp.name, name=alias)
 
-        if try_import_xarray() is None:
+        if try_import("xarray") is None or try_import("xarray_sql") is None:
             return {}
 
         _temp_files: list[str] = []

lumen/sources/xarray_sql.py

@@ -14,7 +14,7 @@
 import param
 
 from ..transforms.sql import SQLFilter
-from ..util import try_import_xarray
+from ..util import try_import
 from .base import BaseSQLSource, cached
 
 
@@ -93,7 +93,7 @@ class XArraySQLSource(BaseSQLSource):
         The SQL expression template for table queries.""")
 
     def __init__(self, _dataset=None, _ctx=None, **params):
-        if try_import_xarray() is None:
+        if try_import("xarray") is None or try_import("xarray_sql") is None:
             raise ImportError(
                 "xarray and xarray-sql are required for XArraySQLSource. "
                 "Install them with: pip install lumen[xarray]"
@@ -161,7 +161,7 @@ def _load_dataset(cls, dataset, uri, engine, chunks, open_kwargs, variables):
                 kw["engine"] = resolved_engine
             if chunks is not None:
                 kw["chunks"] = chunks
-            xr = try_import_xarray()
+            xr = try_import("xarray")
             ds = xr.open_dataset(uri, **kw)
         else:
             raise ValueError("Either 'uri' or '_dataset' must be provided.")

lumen/tests/ai/test_ui.py

@@ -1165,7 +1165,7 @@ def test_resolve_data_netcdf_file(self):
 
     def test_resolve_data_xarray_import_error(self):
         """Test graceful error when xarray-sql not installed."""
-        with patch('lumen.sources.xarray_sql.try_import_xarray', return_value=None):
+        with patch('lumen.sources.xarray_sql.try_import', return_value=None):
             with pytest.raises(ImportError, match="xarray"):
                 UI._resolve_data('data.nc')
 
@@ -1244,7 +1244,7 @@ def test_upload_handler_creates_source_from_nc(self):
 
     def test_upload_handler_returns_empty_without_xarray(self):
         """Upload handlers should be empty if xarray-sql not installed."""
-        with patch('lumen.ai.ui.try_import_xarray', return_value=None):
+        with patch('lumen.ai.ui.try_import', return_value=None):
             handlers = UI._get_xarray_upload_handlers()
             assert handlers == {}
 

lumen/tests/sources/test_optional_imports.py

@@ -14,7 +14,7 @@
 
 import pytest
 
-from lumen.util import try_import_xarray
+from lumen.util import try_import
 
 # Each entry: (module_path, class_name, guard_package, pip_extra)
 OPTIONAL_SOURCES = [
@@ -88,28 +88,13 @@ def mock_import(name, *args, **kwargs):
         sys.modules.update(saved)
 
 
-def test_try_import_xarray_returns_module_when_installed():
-    """try_import_xarray returns the xarray module when xarray and xarray-sql are installed."""
-    xr = pytest.importorskip("xarray")
-    pytest.importorskip("xarray_sql")
-    assert try_import_xarray() is xr
+def test_try_import_returns_module_when_installed():
+    """try_import returns the imported module when it is available."""
+    import json
 
+    assert try_import("json") is json
 
-def test_try_import_xarray_none_when_missing():
-    """try_import_xarray returns None (not raises) when xarray-sql is absent."""
-    real_import = builtins.__import__
-
-    def mock_import(name, *args, **kwargs):
-        if name == "xarray_sql" or name.startswith("xarray_sql."):
-            raise ImportError("No module named 'xarray_sql'")
-        return real_import(name, *args, **kwargs)
 
-    saved = {
-        key: sys.modules.pop(key)
-        for key in [k for k in sys.modules if k == "xarray_sql" or k.startswith("xarray_sql.")]
-    }
-    try:
-        with patch("builtins.__import__", side_effect=mock_import):
-            assert try_import_xarray() is None
-    finally:
-        sys.modules.update(saved)
+def test_try_import_none_when_missing():
+    """try_import returns None (not raises) when the module is absent."""
+    assert try_import("a_module_that_definitely_does_not_exist_xyz") is None

lumen/util.py

@@ -505,13 +505,3 @@ def try_import(module: str):
         return importlib.import_module(module)
     except ImportError:
         return None
-
-
-def try_import_xarray():
-    """Import and return xarray, or None if xarray or xarray-sql is unavailable."""
-    try:
-        import xarray
-        import xarray_sql  # noqa
-        return xarray
-    except ImportError:
-        return None