Correct typing

apis/python/src/tiledbsoma/_collection.py

@@ -33,13 +33,11 @@
 from typing_extensions import Self
 
 from . import _funcs, _tdb_handles
-from . import pytiledbsoma as clib
 from ._common_nd_array import NDArray
 from ._dataframe import DataFrame
 from ._dense_nd_array import DenseNDArray
 from ._exception import SOMAError, is_does_not_exist_error
 from ._sparse_nd_array import SparseNDArray
-from ._tdb_handles import DataFrameWrapper, DenseNDArrayWrapper, SparseNDArrayWrapper
 from ._tiledb_object import AnyTileDBObject, TileDBObject
 from ._types import OpenTimestamp
 from ._util import (
@@ -428,7 +426,7 @@ def __getitem__(self, key: str) -> CollectionElementType:
             raise KeyError(err_str) from None
         if entry.soma is None:
             from . import _factory  # Delayed binding to resolve circular import.
-            
+
             uri = entry.entry.uri
             mode = self.mode
             context = self.context

apis/python/src/tiledbsoma/_dataframe.py

@@ -21,7 +21,7 @@
 from ._constants import SOMA_JOINID
 from ._query_condition import QueryCondition
 from ._read_iters import TableReadIter
-from ._tdb_handles import DataFrameWrapper
+from ._tdb_handles import ArrayWrapper, DataFrameWrapper
 from ._tiledb_array import TileDBArray
 from ._types import NPFloating, NPInteger, OpenTimestamp, Slice, is_slice_of
 from .options import SOMATileDBContext
@@ -264,7 +264,7 @@ def count(self) -> int:
             Experimental.
         """
         self._check_open_read()
-        
+
         # if is it in read open mode, then it is a DataFrameWrapper
         return cast(DataFrameWrapper, self._handle).count
 
@@ -384,6 +384,8 @@ def write(
         """
         _util.check_type("values", values, (pa.Table,))
 
+        # the handle will be the ArrayWrapper in write-mode
+        handle = cast(ArrayWrapper, self._handle)
         dim_cols_map: Dict[str, pd.DataFrame] = {}
         attr_cols_map: Dict[str, pd.DataFrame] = {}
         dim_names_set = self.index_column_names
@@ -393,8 +395,8 @@ def write(
             col = values.column(name)
             n = len(col)
 
-            if self._handle.schema.has_attr(name):
-                attr = self._handle.schema.attr(name)
+            if handle.schema.has_attr(name):
+                attr = handle.schema.attr(name)
 
                 # Add the enumeration values to the TileDB Array from ArrowArray
                 if attr.enum_label is not None and col.num_chunks != 0:
@@ -404,7 +406,7 @@ def write(
                             f"{name} but saw {col.type}"
                         )
 
-                    enmr = self._handle.enum(attr.name)
+                    enmr = handle.enum(attr.name)
 
                     # get new enumeration values, maintain original ordering
                     update_vals = []
@@ -434,7 +436,7 @@ def write(
                             [chunk.dictionary_decode() for chunk in col.chunks]
                         )
                     else:
-                        attr = self._handle.schema.attr(name)
+                        attr = handle.schema.attr(name)
                         if attr.enum_label is not None:
                             # Normal case: writing categorical data to categorical schema.
                             cols_map[name] = col.chunk(0).indices.to_pandas()
@@ -449,7 +451,7 @@ def write(
 
             else:
                 if name not in dim_names_set:
-                    attr = self._handle.schema.attr(name)
+                    attr = handle.schema.attr(name)
                     if attr.enum_label is not None:
                         raise ValueError(
                             f"Categorical column {name} must be presented with categorical data"

apis/python/src/tiledbsoma/_factory.py

@@ -39,7 +39,6 @@
 )
 from ._exception import SOMAError
 from ._funcs import typeguard_ignore
-from ._tdb_handles import DataFrameWrapper, DenseNDArrayWrapper, SparseNDArrayWrapper
 from ._types import OpenTimestamp
 from .options import SOMATileDBContext
 from .options._soma_tiledb_context import _validate_soma_tiledb_context

apis/python/src/tiledbsoma/_read_iters.py

@@ -457,7 +457,6 @@ def _arrow_table_reader(sr: clib.SOMAArray) -> Iterator[pa.Table]:
     while tbl is not None:
         yield tbl
         tbl = sr.read_next()
-
 
 
 def _coords_strider(

apis/python/src/tiledbsoma/_sparse_nd_array.py

@@ -158,10 +158,10 @@ def read(
         """
         del batch_size, platform_config  # Currently unused.
         _util.check_unpartitioned(partitions)
-        
+
         sr = self._handle.reader
         sr.reset(result_order=_util.to_clib_result_order(result_order))
-        
+
         return SparseNDArrayRead(sr, self, coords)
 
     def write(

apis/python/src/tiledbsoma/_tdb_handles.py

@@ -12,7 +12,6 @@
 import enum
 from typing import (
     Any,
-    Callable,
     Dict,
     Generic,
     Iterator,
@@ -31,7 +30,6 @@
 import numpy as np
 import pyarrow as pa
 import tiledb
-from numpy.typing import DTypeLike
 from somacore import options
 from typing_extensions import Literal, Self
 
@@ -72,7 +70,7 @@ def _open_with_clib_wrapper(
     mode: options.OpenMode,
     context: SOMATileDBContext,
     timestamp: Optional[OpenTimestamp] = None,
-) -> "DataFrameWrapper":
+) -> Union["DataFrameWrapper", "DenseNDArrayWrapper", "SparseNDArrayWrapper"]:
     open_mode = clib.OpenMode.read if mode == "r" else clib.OpenMode.write
     config = {k: str(v) for k, v in context.tiledb_config.items()}
     timestamp_ms = context._open_timestamp_ms(timestamp)
@@ -263,10 +261,10 @@ def domain(self) -> Tuple[Tuple[object, object], ...]:
     @property
     def ndim(self) -> int:
         return int(self._handle.schema.domain.ndim)
-    
+
     @property
-    def shape(self):
-        return self._handle.schema.shape
+    def shape(self) -> Tuple[int, ...]:
+        return cast(Tuple[int, ...], self._handle.schema.shape)
 
     @property
     def attr_names(self) -> Tuple[str, ...]:
@@ -333,20 +331,22 @@ def _do_initial_reads(self, reader: _RawHdl_co) -> None:  # type: ignore[misc]
         """
         # non–attrs-managed field
         self.metadata = MetadataWrapper(self, dict(reader.meta))
-        
+
     @property
     def schema(self) -> pa.Schema:
         return self._handle.schema
-    
+
     @property
-    def reader(self):
+    def reader(
+        self,
+    ) -> Union[clib.SOMADataFrame, clib.SOMADenseNDArray, clib.SOMASparseNDArray]:
         if self._handle.mode != "r":
             raise SOMAError("Reader handle is not in read-mode")
         return self._handle
 
     @property
-    def meta(self) -> Dict[str, Any]:
-        return MetadataWrapper(self, dict(self._handle.meta)) 
+    def meta(self) -> "MetadataWrapper":
+        return MetadataWrapper(self, dict(self._handle.meta))
 
     @property
     def domain(self) -> Tuple[Tuple[Any, Any], ...]:
@@ -484,7 +484,7 @@ def _opener(
         mode: options.OpenMode,
         context: SOMATileDBContext,
         timestamp: int,
-    ) -> clib.SOMADenseNDArray:  
+    ) -> 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] = []

apis/python/src/tiledbsoma/_tiledb_object.py

@@ -13,7 +13,6 @@
 from typing_extensions import Self
 
 from . import _constants, _tdb_handles
-from . import pytiledbsoma as clib
 from ._exception import SOMAError
 from ._tdb_handles import DataFrameWrapper, DenseNDArrayWrapper, SparseNDArrayWrapper
 from ._types import OpenTimestamp
@@ -100,7 +99,12 @@ def open(
     def __init__(
         self,
         # TODO DataFrameWrapper should be _WrapperType_co
-        handle: Union[_WrapperType_co, _tdb_handles.DataFrameWrapper],
+        handle: Union[
+            _WrapperType_co,
+            _tdb_handles.DataFrameWrapper,
+            _tdb_handles.DenseNDArrayWrapper,
+            _tdb_handles.SparseNDArrayWrapper,
+        ],
         *,
         _dont_call_this_use_create_or_open_instead: str = "unset",
     ):
@@ -133,7 +137,10 @@ def __init__(
 
     _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],
+        Type[_tdb_handles.SparseNDArrayWrapper],
     ]
     """Class variable of the Wrapper class used to open this object type."""
 

apis/python/tests/test_sparse_nd_array.py

@@ -1573,7 +1573,9 @@ def test_blockwise_scipy_iter_eager(
 @pytest.mark.parametrize("result_order", ["auto", "row-major", "column-major"])
 @pytest.mark.parametrize("axis", [0, 1])
 @pytest.mark.parametrize("compress", [True, False])
-def test_blockwise_scipy_iter_result_order(a_random_sparse_nd_array: str, result_order, axis, compress) -> None:
+def test_blockwise_scipy_iter_result_order(
+    a_random_sparse_nd_array: str, result_order, axis, compress
+) -> None:
     """
     Confirm behavior with different result_order.
     """