From d0a2f6c59949c88d2f37ab8da2d8cabfd9a012ef Mon Sep 17 00:00:00 2001 From: Vivian Nguyen Date: Fri, 23 Feb 2024 15:25:53 -0600 Subject: [PATCH] Pass in as kw args; replace virtual method --- apis/python/src/tiledbsoma/_tdb_handles.py | 71 ++++++-------------- apis/python/src/tiledbsoma/_tiledb_object.py | 4 +- apis/python/src/tiledbsoma/soma_object.cc | 12 +++- 3 files changed, 33 insertions(+), 54 deletions(-) diff --git a/apis/python/src/tiledbsoma/_tdb_handles.py b/apis/python/src/tiledbsoma/_tdb_handles.py index 0439b198a1..714a22485a 100644 --- 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,9 +320,14 @@ 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, @@ -331,11 +335,16 @@ def _opener( 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 index 288cd03bc3..42a7c78ba3 100644 --- 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 index 6e8784ff25..50f45faeab 100644 --- 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 ctx, + std::shared_ptr context, std::optional> 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(*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); };