Use lazy imports so tests pass without C-API

Author: adsharma

src_py/_backend.py

@@ -0,0 +1,27 @@
+from __future__ import annotations
+
+from importlib import import_module
+from typing import Any
+
+_CAPI_MODULE: Any | None = None
+_PYBIND_MODULE: Any | None = None
+_PYBIND_IMPORT_ATTEMPTED = False
+
+
+def get_capi_module() -> Any:
+    global _CAPI_MODULE
+    if _CAPI_MODULE is None:
+        _CAPI_MODULE = import_module("._lbug_capi", __package__)
+    return _CAPI_MODULE
+
+
+def get_pybind_module() -> Any | None:
+    global _PYBIND_MODULE, _PYBIND_IMPORT_ATTEMPTED
+    if _PYBIND_IMPORT_ATTEMPTED:
+        return _PYBIND_MODULE
+    _PYBIND_IMPORT_ATTEMPTED = True
+    try:
+        _PYBIND_MODULE = import_module("._lbug", __package__)
+    except ImportError:
+        _PYBIND_MODULE = None
+    return _PYBIND_MODULE

src_py/connection.py

@@ -6,17 +6,10 @@
 from typing import TYPE_CHECKING, Any
 from weakref import WeakSet
 
-from . import _lbug_capi as _lbug
+from ._backend import get_capi_module, get_pybind_module
 from .prepared_statement import PreparedStatement
 from .query_result import QueryResult
 
-try:
-    from . import _lbug as _lbug_pybind
-except (
-    ImportError
-):  # pragma: no cover - pybind module may be unavailable in some builds
-    _lbug_pybind = None
-
 if TYPE_CHECKING:
     import sys
     from collections.abc import Callable
@@ -73,12 +66,18 @@ def init_connection(self) -> None:
         self.database.init_database()
         if self._connection is None:
             backend_module = (
-                _lbug_pybind if self.database._use_pybind_backend else _lbug
+                get_pybind_module()
+                if self.database._use_pybind_backend
+                else get_capi_module()
+            )
+            self._connection = backend_module.Connection(
+                self.database._database, self.num_threads
             )
-            self._connection = backend_module.Connection(self.database._database, self.num_threads)  # type: ignore[union-attr]
 
     def _using_pybind_backend(self) -> bool:
-        return bool(self.database._use_pybind_backend and _lbug_pybind is not None)
+        return bool(
+            self.database._use_pybind_backend and get_pybind_module() is not None
+        )
 
     def set_max_threads_for_exec(self, num_threads: int) -> None:
         """
@@ -212,7 +211,7 @@ def _rewrite_local_scan_object(
     def _should_use_pybind_for_scan(
         self, query: str, parameters: dict[str, Any]
     ) -> bool:
-        if _lbug_pybind is None:
+        if get_pybind_module() is None:
             return False
         if not self._has_scan_pattern(query):
             return False
@@ -230,7 +229,8 @@ def _should_use_pybind_for_scan(
         return False
 
     def _get_pybind_connection(self) -> Any | None:
-        if _lbug_pybind is None:
+        pybind_module = get_pybind_module()
+        if pybind_module is None:
             return None
         if self._using_pybind_backend():
             return self._connection
@@ -239,7 +239,7 @@ def _get_pybind_connection(self) -> Any | None:
         if pybind_db is None:
             return None
         if self._py_connection is None:
-            self._py_connection = _lbug_pybind.Connection(pybind_db, self.num_threads)
+            self._py_connection = pybind_module.Connection(pybind_db, self.num_threads)
         return self._py_connection
 
     def _execute_with_pybind(

src_py/database.py

@@ -5,16 +5,9 @@
 from typing import TYPE_CHECKING, Any, ClassVar
 from weakref import WeakSet
 
-from . import _lbug_capi as _lbug
+from ._backend import get_capi_module, get_pybind_module
 from .types import Type
 
-try:
-    from . import _lbug as _lbug_pybind
-except (
-    ImportError
-):  # pragma: no cover - pybind module may be unavailable in some builds
-    _lbug_pybind = None
-
 if TYPE_CHECKING:
     import sys
     from types import TracebackType
@@ -157,12 +150,13 @@ def _resolve_backend_preference(cls, backend: str) -> str:
     def _should_use_pybind_backend(self) -> bool:
         if self.backend == "capi":
             return False
+        pybind_module = get_pybind_module()
         if self.backend == "pybind":
-            if _lbug_pybind is None:
+            if pybind_module is None:
                 msg = "Requested pybind backend, but ladybug._lbug is not available."
                 raise RuntimeError(msg)
             return True
-        return _lbug_pybind is not None
+        return pybind_module is not None
 
     def __enter__(self) -> Self:
         return self
@@ -185,7 +179,7 @@ def get_version() -> str:
         str
             The version of the database.
         """
-        return _lbug.Database.get_version()  # type: ignore[union-attr]
+        return get_capi_module().Database.get_version()
 
     @staticmethod
     def get_storage_version() -> int:
@@ -197,7 +191,7 @@ def get_storage_version() -> int:
         int
             The storage version of the database.
         """
-        return _lbug.Database.get_storage_version()  # type: ignore[union-attr]
+        return get_capi_module().Database.get_storage_version()
 
     def __getstate__(self) -> dict[str, Any]:
         state = {
@@ -217,7 +211,7 @@ def init_database(self) -> None:
             if self._use_pybind_backend:
                 self._database = self.init_pybind_database()
             else:
-                self._database = _lbug.Database(  # type: ignore[union-attr]
+                self._database = get_capi_module().Database(
                     self.database_path,
                     self.buffer_pool_size,
                     self.max_num_threads,
@@ -234,10 +228,11 @@ def init_database(self) -> None:
     def init_pybind_database(self) -> Any | None:
         """Initialize and return the optional pybind database backend."""
         self.check_for_database_close()
-        if _lbug_pybind is None:
+        pybind_module = get_pybind_module()
+        if pybind_module is None:
             return None
         if self._pybind_database is None:
-            self._pybind_database = _lbug_pybind.Database(
+            self._pybind_database = pybind_module.Database(
                 self.database_path,
                 self.buffer_pool_size,
                 self.max_num_threads,

src_py/query_result.py

@@ -18,8 +18,6 @@
     import pyarrow as pa
     import torch_geometric.data as geo
 
-    from . import _lbug_capi as _lbug
-
     if sys.version_info >= (3, 11):
         from typing import Self
     else:
@@ -29,7 +27,7 @@
 class QueryResult:
     """QueryResult stores the result of a query execution."""
 
-    def __init__(self, connection: _lbug.Connection, query_result: _lbug.QueryResult):  # type: ignore[name-defined]
+    def __init__(self, connection: Any, query_result: Any):
         """
         Parameters
         ----------