Skip to content

Commit 17bc410

Browse files
replace typing module with native collection types (#120)
Fix #110 Replaced all deprecated typing module aliases across core.py, df.py, and reader.py with their native Python equivalents, as the project targets Python ≥ 3.10 which fully supports them. This includes switching things like typing.List, typing.Dict, typing.Tuple to their built-in forms, replacing typing.Optional[X] with the X | None union syntax, and moving Iterator, Callable, Mapping, and Hashable to collections.abc where they belong. typing.Any and typing.TYPE_CHECKING were kept from typing as they have no native equivalents. The changes were applied using pyupgrade --py310-plus to ensure consistency and correctness. also i ran the existing test suite after the changes to confirm nothing was broken. Co-authored-by: Ibraam-Ashraf <168275494+Ibraam-Ashraf@users.noreply.github.com>
1 parent 2cde4e5 commit 17bc410

File tree

3 files changed

+23
-21
lines changed

3 files changed

+23
-21
lines changed

xarray_sql/core.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,20 @@
11
import itertools
2-
import typing as t
2+
from collections.abc import Iterator
3+
from typing import Any
34

45
import numpy as np
56
import xarray as xr
67

7-
Row = t.List[t.Any]
8+
Row = list[Any]
89

910

1011
# deprecated
11-
def get_columns(ds: xr.Dataset) -> t.List[str]:
12+
def get_columns(ds: xr.Dataset) -> list[str]:
1213
return list(ds.sizes.keys()) + list(ds.data_vars.keys())
1314

1415

1516
# Deprecated
16-
def unravel(ds: xr.Dataset) -> t.Iterator[Row]:
17+
def unravel(ds: xr.Dataset) -> Iterator[Row]:
1718
dim_keys, dim_vals = zip(*ds.sizes.items())
1819

1920
for idx in itertools.product(*(range(d) for d in dim_vals)):

xarray_sql/df.py

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,20 @@
11
import itertools
2-
import typing as t
32
import warnings
3+
from collections.abc import Callable, Hashable, Iterator, Mapping
44

55
import numpy as np
66
import pandas as pd
77
import pyarrow as pa
88
import xarray as xr
99
from datafusion.context import ArrowStreamExportable
1010

11-
Block = t.Dict[t.Hashable, slice]
12-
Chunks = t.Optional[t.Dict[str, int]]
11+
Block = dict[Hashable, slice]
12+
Chunks = dict[str, int] | None
1313

1414

1515
# Borrowed from Xarray
1616
def _get_chunk_slicer(
17-
dim: t.Hashable, chunk_index: t.Mapping, chunk_bounds: t.Mapping
17+
dim: Hashable, chunk_index: Mapping, chunk_bounds: Mapping
1818
):
1919
if dim in chunk_index:
2020
which_chunk = chunk_index[dim]
@@ -25,7 +25,7 @@ def _get_chunk_slicer(
2525

2626

2727
# Adapted from Xarray `map_blocks` implementation.
28-
def block_slices(ds: xr.Dataset, chunks: Chunks = None) -> t.Iterator[Block]:
28+
def block_slices(ds: xr.Dataset, chunks: Chunks = None) -> Iterator[Block]:
2929
"""Compute block slices for a chunked Dataset."""
3030
if chunks is not None:
3131
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]:
5454
yield from blocks
5555

5656

57-
def explode(ds: xr.Dataset, chunks: Chunks = None) -> t.Iterator[xr.Dataset]:
57+
def explode(ds: xr.Dataset, chunks: Chunks = None) -> Iterator[xr.Dataset]:
5858
"""Explodes a dataset into its chunks."""
5959
yield from (ds.isel(b) for b in block_slices(ds, chunks=chunks))
6060

@@ -64,9 +64,9 @@ def _block_len(block: Block) -> int:
6464

6565

6666
def from_map_batched(
67-
func: t.Callable[..., pd.DataFrame],
67+
func: Callable[..., pd.DataFrame],
6868
*iterables,
69-
args: t.Optional[t.Tuple] = None,
69+
args: tuple | None = None,
7070
schema: pa.Schema = None,
7171
**kwargs,
7272
) -> pa.RecordBatchReader:
@@ -99,7 +99,7 @@ def map_batches():
9999

100100

101101
def from_map(
102-
func: t.Callable, *iterables, args: t.Optional[t.Tuple] = None, **kwargs
102+
func: Callable, *iterables, args: tuple | None = None, **kwargs
103103
) -> pa.Table:
104104
"""Create a PyArrow Table by mapping a function over iterables.
105105

xarray_sql/reader.py

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -10,14 +10,15 @@
1010

1111
from __future__ import annotations
1212

13-
import typing as t
13+
from collections.abc import Callable, Iterator
14+
from typing import TYPE_CHECKING
1415

1516
import pyarrow as pa
1617
import xarray as xr
1718

1819
from .df import Block, Chunks, block_slices, pivot, _parse_schema
1920

20-
if t.TYPE_CHECKING:
21+
if TYPE_CHECKING:
2122
from ._native import LazyArrowStreamTable
2223

2324

@@ -53,7 +54,7 @@ def __init__(
5354
ds: xr.Dataset,
5455
chunks: Chunks = None,
5556
*,
56-
_iteration_callback: t.Optional[t.Callable[[Block], None]] = None,
57+
_iteration_callback: Callable[[Block], None] | None = None,
5758
):
5859
"""Initialize the lazy reader.
5960
@@ -83,7 +84,7 @@ def schema(self) -> pa.Schema:
8384
"""The Arrow schema for this stream."""
8485
return self._schema
8586

86-
def _generate_batches(self) -> t.Iterator[pa.RecordBatch]:
87+
def _generate_batches(self) -> Iterator[pa.RecordBatch]:
8788
"""Generate RecordBatches lazily from xarray blocks.
8889
8990
This generator is only consumed when the Arrow stream's get_next
@@ -99,7 +100,7 @@ def _generate_batches(self) -> t.Iterator[pa.RecordBatch]:
99100
yield pa.RecordBatch.from_pandas(df, schema=self._schema)
100101

101102
def __arrow_c_stream__(
102-
self, requested_schema: t.Optional[object] = None
103+
self, requested_schema: object | None = None
103104
) -> object:
104105
"""Export as Arrow C Stream via PyCapsule.
105106
@@ -135,7 +136,7 @@ def __arrow_c_stream__(
135136
return reader.__arrow_c_stream__(requested_schema)
136137

137138
def __arrow_c_schema__(
138-
self, requested_schema: t.Optional[object] = None
139+
self, requested_schema: object | None = None
139140
) -> object:
140141
"""Export the schema as Arrow C Schema via PyCapsule.
141142
@@ -171,7 +172,7 @@ def read_xarray_table(
171172
ds: xr.Dataset,
172173
chunks: Chunks = None,
173174
*,
174-
_iteration_callback: t.Optional[t.Callable[[Block], None]] = None,
175+
_iteration_callback: Callable[[Block], None] | None = None,
175176
) -> "LazyArrowStreamTable":
176177
"""Create a lazy DataFusion table from an xarray Dataset.
177178
@@ -232,7 +233,7 @@ def read_xarray_table(
232233
# Each factory produces a RecordBatchReader for its specific chunk
233234
def make_partition_factory(
234235
block: Block,
235-
) -> t.Callable[[], pa.RecordBatchReader]:
236+
) -> Callable[[], pa.RecordBatchReader]:
236237
"""Create a factory function for a specific block/chunk."""
237238

238239
def make_stream() -> pa.RecordBatchReader:

0 commit comments

Comments
 (0)