replace typing module with native collection types

xarray_sql/core.py

@@ -1,19 +1,20 @@
 import itertools
-import typing as t
+from collections.abc import Iterator
+from typing import Any
 
 import numpy as np
 import xarray as xr
 
-Row = t.List[t.Any]
+Row = list[Any]
 
 
 # deprecated
-def get_columns(ds: xr.Dataset) -> t.List[str]:
+def get_columns(ds: xr.Dataset) -> list[str]:
   return list(ds.sizes.keys()) + list(ds.data_vars.keys())
 
 
 # Deprecated
-def unravel(ds: xr.Dataset) -> t.Iterator[Row]:
+def unravel(ds: xr.Dataset) -> Iterator[Row]:
   dim_keys, dim_vals = zip(*ds.sizes.items())
 
   for idx in itertools.product(*(range(d) for d in dim_vals)):

xarray_sql/df.py

@@ -1,20 +1,20 @@
 import itertools
-import typing as t
 import warnings
+from collections.abc import Callable, Hashable, Iterator, Mapping
 
 import numpy as np
 import pandas as pd
 import pyarrow as pa
 import xarray as xr
 from datafusion.context import ArrowStreamExportable
 
-Block = t.Dict[t.Hashable, slice]
-Chunks = t.Optional[t.Dict[str, int]]
+Block = dict[Hashable, slice]
+Chunks = dict[str, int] | None
 
 
 # Borrowed from Xarray
 def _get_chunk_slicer(
-    dim: t.Hashable, chunk_index: t.Mapping, chunk_bounds: t.Mapping
+    dim: Hashable, chunk_index: Mapping, chunk_bounds: Mapping
 ):
   if dim in chunk_index:
     which_chunk = chunk_index[dim]
@@ -25,7 +25,7 @@ def _get_chunk_slicer(
 
 
 # Adapted from Xarray `map_blocks` implementation.
-def block_slices(ds: xr.Dataset, chunks: Chunks = None) -> t.Iterator[Block]:
+def block_slices(ds: xr.Dataset, chunks: Chunks = None) -> Iterator[Block]:
   """Compute block slices for a chunked Dataset."""
   if chunks is not None:
     for_chunking = ds.copy(data=None, deep=False).chunk(chunks)
@@ -54,7 +54,7 @@ def block_slices(ds: xr.Dataset, chunks: Chunks = None) -> t.Iterator[Block]:
   yield from blocks
 
 
-def explode(ds: xr.Dataset, chunks: Chunks = None) -> t.Iterator[xr.Dataset]:
+def explode(ds: xr.Dataset, chunks: Chunks = None) -> Iterator[xr.Dataset]:
   """Explodes a dataset into its chunks."""
   yield from (ds.isel(b) for b in block_slices(ds, chunks=chunks))
 
@@ -64,9 +64,9 @@ def _block_len(block: Block) -> int:
 
 
 def from_map_batched(
-    func: t.Callable[..., pd.DataFrame],
+    func: Callable[..., pd.DataFrame],
     *iterables,
-    args: t.Optional[t.Tuple] = None,
+    args: tuple | None = None,
     schema: pa.Schema = None,
     **kwargs,
 ) -> pa.RecordBatchReader:
@@ -99,7 +99,7 @@ def map_batches():
 
 
 def from_map(
-    func: t.Callable, *iterables, args: t.Optional[t.Tuple] = None, **kwargs
+    func: Callable, *iterables, args: tuple | None = None, **kwargs
 ) -> pa.Table:
   """Create a PyArrow Table by mapping a function over iterables.
 

xarray_sql/reader.py

@@ -10,14 +10,15 @@
 
 from __future__ import annotations
 
-import typing as t
+from collections.abc import Callable, Iterator
+from typing import TYPE_CHECKING
 
 import pyarrow as pa
 import xarray as xr
 
 from .df import Block, Chunks, block_slices, pivot, _parse_schema
 
-if t.TYPE_CHECKING:
+if TYPE_CHECKING:
   from ._native import LazyArrowStreamTable
 
 
@@ -53,7 +54,7 @@ def __init__(
       ds: xr.Dataset,
       chunks: Chunks = None,
       *,
-      _iteration_callback: t.Optional[t.Callable[[Block], None]] = None,
+      _iteration_callback: Callable[[Block], None] | None = None,
   ):
     """Initialize the lazy reader.
 
@@ -83,7 +84,7 @@ def schema(self) -> pa.Schema:
     """The Arrow schema for this stream."""
     return self._schema
 
-  def _generate_batches(self) -> t.Iterator[pa.RecordBatch]:
+  def _generate_batches(self) -> Iterator[pa.RecordBatch]:
     """Generate RecordBatches lazily from xarray blocks.
 
     This generator is only consumed when the Arrow stream's get_next
@@ -99,7 +100,7 @@ def _generate_batches(self) -> t.Iterator[pa.RecordBatch]:
       yield pa.RecordBatch.from_pandas(df, schema=self._schema)
 
   def __arrow_c_stream__(
-      self, requested_schema: t.Optional[object] = None
+      self, requested_schema: object | None = None
   ) -> object:
     """Export as Arrow C Stream via PyCapsule.
 
@@ -135,7 +136,7 @@ def __arrow_c_stream__(
     return reader.__arrow_c_stream__(requested_schema)
 
   def __arrow_c_schema__(
-      self, requested_schema: t.Optional[object] = None
+      self, requested_schema: object | None = None
   ) -> object:
     """Export the schema as Arrow C Schema via PyCapsule.
 
@@ -171,7 +172,7 @@ def read_xarray_table(
     ds: xr.Dataset,
     chunks: Chunks = None,
     *,
-    _iteration_callback: t.Optional[t.Callable[[Block], None]] = None,
+    _iteration_callback: Callable[[Block], None] | None = None,
 ) -> "LazyArrowStreamTable":
   """Create a lazy DataFusion table from an xarray Dataset.
 
@@ -232,7 +233,7 @@ def read_xarray_table(
   # Each factory produces a RecordBatchReader for its specific chunk
   def make_partition_factory(
       block: Block,
-  ) -> t.Callable[[], pa.RecordBatchReader]:
+  ) -> Callable[[], pa.RecordBatchReader]:
     """Create a factory function for a specific block/chunk."""
 
     def make_stream() -> pa.RecordBatchReader: