diff --git a/xarray_sql/core.py b/xarray_sql/core.py index afbd2eca..11b909b7 100644 --- a/xarray_sql/core.py +++ b/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)): diff --git a/xarray_sql/df.py b/xarray_sql/df.py index c799626f..9215fe0d 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -1,6 +1,6 @@ import itertools -import typing as t import warnings +from collections.abc import Callable, Hashable, Iterator, Mapping import numpy as np import pandas as pd @@ -8,13 +8,13 @@ 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. diff --git a/xarray_sql/reader.py b/xarray_sql/reader.py index 3b57559d..1252a952 100644 --- a/xarray_sql/reader.py +++ b/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: