Pass in as kw args; replace virtual method

single-cell-data · Feb 23, 2024 · d0a2f6c · d0a2f6c
1 parent 798dd0e
commit d0a2f6c
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 54 deletions.
diff --git a/apis/python/src/tiledbsoma/_tdb_handles.py b/apis/python/src/tiledbsoma/_tdb_handles.py
@@ -16,7 +16,6 @@
     Dict,
     Generic,
     Iterator,
-    List,
     Mapping,
     MutableMapping,
     Optional,
@@ -57,7 +56,7 @@ def open(
     # if there is not a valid SOMAObject at the given URI, this
     # returns None
     soma_object = clib.SOMAObject.open(
-        uri, open_mode, context.native_context, (0, timestamp_ms)
+        uri, open_mode, context.native_context, timestamp=(0, timestamp_ms)
     )
 
     # Avoid creating a TileDB-Py Ctx unless necessary
@@ -321,21 +320,31 @@ def _do_initial_reads(self, reader: tiledb.Group) -> None:
         }
 
 
-class SOMAArrayWrapper(Wrapper[clib.SOMAArray]):
+_ArrType = TypeVar("_ArrType", bound=clib.SOMAArray)
+
+
+class SOMAArrayWrapper(Wrapper[_ArrType]):
     """Base class for Pybind11 SOMAArrayWrapper handles."""
 
+    _WRAPPED_TYPE: Type[_ArrType]
+
     @classmethod
     def _opener(
         cls,
         uri: str,
         mode: options.OpenMode,
         context: SOMATileDBContext,
         timestamp: int,
-    ) -> clib.SOMAArray:
-        # Ensure SOMADataFrameWrapper, SOMASparseNDArrayWrapper, and 
-        # SOMADenseNDArrayWrapper have _opener implemented to open the correct
-        # clib object (clib.DataFrame.open, etc)
-        raise NotImplementedError
+    ) -> clib.SOMADenseNDArray:
+        open_mode = clib.OpenMode.read if mode == "r" else clib.OpenMode.write
+        return cls._WRAPPED_TYPE.open(
+            uri,
+            open_mode,
+            context=context.native_context,
+            column_names=[],
+            result_order=clib.ResultOrder.automatic,
+            timestamp=(0, timestamp),
+        )
 
     # Covariant types should normally not be in parameters, but this is for
     # internal use only so it's OK.
@@ -409,58 +418,20 @@ def enum(self, label: str) -> tiledb.Enumeration:
         raise NotImplementedError
 
 
-class DataFrameWrapper(SOMAArrayWrapper, Wrapper[clib.SOMADataFrame]):
+class DataFrameWrapper(SOMAArrayWrapper[clib.SOMADataFrame]):
     """Wrapper around a Pybind11 SOMADataFrame handle."""
 
-    @classmethod
-    def _opener(
-        cls,
-        uri: str,
-        mode: options.OpenMode,
-        context: SOMATileDBContext,
-        timestamp: int,
-    ) -> clib.SOMADataFrame:
-        open_mode = clib.OpenMode.read if mode == "r" else clib.OpenMode.write
-        config = {k: str(v) for k, v in context.tiledb_config.items()}
-        column_names: List[str] = []
-        result_order = clib.ResultOrder.automatic
-        return clib.SOMADataFrame.open(
-            uri,
-            open_mode,
-            config,
-            column_names,
-            result_order,
-            (0, timestamp),
-        )
+    _WRAPPED_TYPE = clib.SOMADataFrame
 
     @property
     def count(self) -> int:
         return int(self._handle.count)
 
 
-class DenseNDArrayWrapper(SOMAArrayWrapper, Wrapper[clib.SOMADenseNDArray]):
+class DenseNDArrayWrapper(SOMAArrayWrapper[clib.SOMADenseNDArray]):
     """Wrapper around a Pybind11 DenseNDArrayWrapper handle."""
 
-    @classmethod
-    def _opener(
-        cls,
-        uri: str,
-        mode: options.OpenMode,
-        context: SOMATileDBContext,
-        timestamp: int,
-    ) -> clib.SOMADenseNDArray:
-        open_mode = clib.OpenMode.read if mode == "r" else clib.OpenMode.write
-        config = {k: str(v) for k, v in context.tiledb_config.items()}
-        column_names: List[str] = []
-        result_order = clib.ResultOrder.automatic
-        return clib.SOMADenseNDArray.open(
-            uri,
-            open_mode,
-            config,
-            column_names,
-            result_order,
-            (0, timestamp),
-        )
+    _WRAPPED_TYPE = clib.SOMADenseNDArray
 
     @property
     def shape(self) -> Tuple[int, ...]:

diff --git a/apis/python/src/tiledbsoma/_tiledb_object.py b/apis/python/src/tiledbsoma/_tiledb_object.py
@@ -42,7 +42,9 @@ class TileDBObject(somacore.SOMAObject, Generic[_WrapperType_co]):
     """Class variable of the Wrapper class used to open this object type."""
     _wrapper_type: Type[_WrapperType_co]
     _reader_wrapper_type: Union[
-        Type[_WrapperType_co], Type[_tdb_handles.DataFrameWrapper]
+        Type[_WrapperType_co],
+        Type[_tdb_handles.DataFrameWrapper],
+        Type[_tdb_handles.DenseNDArrayWrapper],
     ]
 
     __slots__ = ("_close_stack", "_handle")

diff --git a/apis/python/src/tiledbsoma/soma_object.cc b/apis/python/src/tiledbsoma/soma_object.cc
@@ -52,10 +52,10 @@ void load_soma_object(py::module &m) {
 
     .def_static("open", [](std::string_view uri, 
                            OpenMode mode, 
-                           std::shared_ptr<SOMAContext> ctx, 
+                           std::shared_ptr<SOMAContext> context, 
                            std::optional<std::pair<uint64_t, uint64_t>> timestamp) -> py::object {
         try{
-            auto obj = SOMAObject::open(uri, mode, ctx, timestamp);
+            auto obj = SOMAObject::open(uri, mode, context, timestamp);
             if (obj->type() == "SOMADataFrame")
                 return py::cast(dynamic_cast<SOMADataFrame&>(*obj));
             else if (obj->type() == "SOMASparseNDArray")
@@ -72,7 +72,13 @@ void load_soma_object(py::module &m) {
         }catch(...){
             return py::none();
         }
-    })
+    },
+    "uri"_a,
+    "mode"_a,
+    "context"_a,
+    py::kw_only(),
+    "timestamp"_a = py::none())
+
     .def_property_readonly("type", &SOMAObject::type);
 
     };