Skip to content

Commit a192385

Browse files
committed
fixes
1 parent 2a31ec3 commit a192385

File tree

23 files changed

+213
-135
lines changed

23 files changed

+213
-135
lines changed

.github/workflows/python.yml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,7 @@ jobs:
140140

141141
- name: Type check with mypy and pyright
142142
run: |-
143-
python -m pip install mypy pyright griffe libcst pytest hypothesis fsspec scipy-stubs pandas-stubs types-python-dateutil types-psutil types-requests griffe libcst sphinx types-cffi
143+
python -m pip install mypy pyright ty griffe libcst pytest hypothesis fsspec scipy-stubs pandas-stubs types-python-dateutil types-psutil types-requests griffe libcst sphinx types-cffi
144144
pip install -i https://pypi.anaconda.org/scientific-python-nightly-wheels/simple pyarrow
145145
cd python
146146
mypy

python/pyarrow-stubs/pyarrow/_feather.pyi

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
# specific language governing permissions and limitations
1616
# under the License.
1717

18-
from typing import IO
18+
from typing import IO, Literal
1919
from collections.abc import Sequence
2020

2121
from _typeshed import StrPath
@@ -33,7 +33,7 @@ def write_feather(
3333
compression: str | None = None,
3434
compression_level: int | None = None,
3535
chunksize: int | None = None,
36-
version: int = 2,
36+
version: Literal[1, 2] = 2,
3737
): ...
3838

3939

python/pyarrow-stubs/pyarrow/_flight.pyi

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -106,7 +106,8 @@ class FlightWriteSizeExceededError(ArrowInvalid):
106106

107107
class Action(_Weakrefable):
108108

109-
def __init__(self, action_type: bytes | str, buf: Buffer | bytes | None) -> None: ...
109+
def __init__(
110+
self, action_type: bytes | str, buf: Buffer | bytes | None) -> None: ...
110111

111112
@property
112113
def type(self) -> str: ...
@@ -330,14 +331,11 @@ class _MetadataRecordBatchReader(_Weakrefable, _ReadPandasMixin):
330331

331332

332333
class MetadataRecordBatchReader(_MetadataRecordBatchReader):
333-
334334
@property
335335
def stats(self) -> ReadStats: ...
336-
...
337336

338337

339338
class FlightStreamReader(MetadataRecordBatchReader):
340-
341339
@property
342340
def stats(self) -> ReadStats: ...
343341

@@ -439,18 +437,21 @@ class FlightClient(_Weakrefable):
439437
) -> None: ...
440438

441439
def authenticate_basic_token(
442-
self, username: str | bytes, password: str | bytes, options: FlightCallOptions | None = None
440+
self, username: str | bytes, password: str | bytes,
441+
options: FlightCallOptions | None = None
443442
) -> tuple[str, str]: ...
444443

445444
def list_actions(self, options: FlightCallOptions |
446445
None = None) -> list[Action]: ...
447446

448447
def do_action(
449-
self, action: Action | tuple[bytes | str, bytes | str] | str, options: FlightCallOptions | None = None
448+
self, action: Action | tuple[bytes | str, bytes | str] | str,
449+
options: FlightCallOptions | None = None
450450
) -> Iterator[Result]: ...
451451

452452
def list_flights(
453-
self, criteria: str | bytes | None = None, options: FlightCallOptions | None = None
453+
self, criteria: str | bytes | None = None,
454+
options: FlightCallOptions | None = None
454455
) -> Generator[FlightInfo, None, None]: ...
455456

456457
def get_flight_info(
@@ -487,17 +488,21 @@ class FlightDataStream(_Weakrefable):
487488

488489

489490
class RecordBatchStream(FlightDataStream):
490-
491491
def __init__(self, data_source: RecordBatchReader | Table | None = None,
492492
options: IpcWriteOptions | None = None) -> None: ...
493493

494494

495495
class GeneratorStream(FlightDataStream):
496-
497496
def __init__(
498497
self,
499498
schema: Schema,
500-
generator: Iterable[FlightDataStream | Table | RecordBatch | RecordBatchReader | tuple[RecordBatch, bytes]],
499+
generator: Iterable[
500+
FlightDataStream
501+
| Table
502+
| RecordBatch
503+
| RecordBatchReader
504+
| tuple[RecordBatch, bytes]
505+
],
501506
options: IpcWriteOptions | None = None,
502507
) -> None: ...
503508

@@ -674,7 +679,8 @@ class FlightServerBase(_Weakrefable):
674679
def wait(self) -> None: ...
675680

676681
def __enter__(self) -> Self: ...
677-
def __exit__(self, exc_type: object, exc_value: object, traceback: object) -> None: ...
682+
def __exit__(
683+
self, exc_type: object, exc_value: object, traceback: object) -> None: ...
678684

679685

680686
def connect(

python/pyarrow-stubs/pyarrow/_fs.pyi

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@ import enum
2020
import sys
2121

2222
from abc import ABC, abstractmethod
23-
from typing import overload
2423
from _typeshed import StrPath
2524

2625
if sys.version_info >= (3, 11):
@@ -101,7 +100,9 @@ class FileSystem(_Weakrefable):
101100
@property
102101
def type_name(self) -> str: ...
103102

104-
def get_file_info(self, paths_or_selector: str | list[str] | FileSelector) -> list[FileInfo] | FileInfo: ...
103+
def get_file_info(
104+
self, paths_or_selector: str | list[str] | FileSelector
105+
) -> list[FileInfo] | FileInfo: ...
105106

106107
def create_dir(self, path: str, *, recursive: bool = True) -> None: ...
107108

python/pyarrow-stubs/pyarrow/_ipc.pyi

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -160,7 +160,8 @@ class _RecordBatchStreamWriter(_CRecordBatchWriter):
160160
@property
161161
def _use_legacy_format(self) -> bool: ...
162162
@property
163-
def _metadata_version(self) -> MetadataVersion: ... # noqa: Y011
163+
def _metadata_version(self) -> MetadataVersion: ...
164+
164165
def _open(
165166
self,
166167
sink,
@@ -265,6 +266,7 @@ class _RecordBatchFileReader(_Weakrefable):
265266
memory_pool: MemoryPool | None = None,
266267
) -> Self: ...
267268

269+
268270
def get_tensor_size(tensor: Tensor) -> int: ...
269271

270272

python/pyarrow-stubs/pyarrow/_orc.pyi

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
# specific language governing permissions and limitations
1616
# under the License.
1717

18-
from typing import IO, Literal
18+
from typing import IO, Literal, Any
1919

2020
from .lib import (
2121
Buffer,
@@ -64,10 +64,9 @@ class ORCWriter(_Weakrefable):
6464
file_version: str | None = None,
6565
batch_size: int | None = None,
6666
stripe_size: int | None = None,
67-
compression: Literal["UNCOMPRESSED", "ZLIB",
68-
"SNAPPY", "LZ4", "ZSTD"] | str | None = None,
67+
compression: Any = 'UNCOMPRESSED',
6968
compression_block_size: int | None = None,
70-
compression_strategy: Literal["COMPRESSION", "SPEED"] | str | None = None,
69+
compression_strategy: Any = 'SPEED',
7170
row_index_stride: int | None = None,
7271
padding_tolerance: float | None = None,
7372
dictionary_key_size_threshold: float | None = None,

python/pyarrow-stubs/pyarrow/_parquet.pyi

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -251,7 +251,9 @@ class SortingColumn:
251251
def from_ordering(
252252
cls,
253253
schema: Schema,
254-
sort_keys: Sequence[str] | Sequence[tuple[str, Order]] | Sequence[str | tuple[str, Order]],
254+
sort_keys: Sequence[str]
255+
| Sequence[tuple[str, Order]]
256+
| Sequence[str | tuple[str, Order]],
255257
null_placement: Literal["at_start", "at_end"] = "at_end",
256258
) -> tuple[SortingColumn, ...]: ...
257259

@@ -331,14 +333,14 @@ class FileMetaData(_Weakrefable):
331333

332334
class ParquetSchema(_Weakrefable):
333335
def __init__(self, container: FileMetaData) -> None: ...
334-
def __getitem__(self, i: int) -> "ColumnSchema": ...
336+
def __getitem__(self, i: int) -> ColumnSchema: ...
335337
def __hash__(self) -> int: ...
336338
def __len__(self) -> int: ...
337339
@property
338340
def names(self) -> list[str]: ...
339341
def to_arrow_schema(self) -> Schema: ...
340342
def equals(self, other: ParquetSchema) -> bool: ...
341-
def column(self, i: int) -> "ColumnSchema": ...
343+
def column(self, i: int) -> ColumnSchema: ...
342344

343345

344346
class ColumnSchema(_Weakrefable):
@@ -423,8 +425,10 @@ class ParquetReader(_Weakrefable):
423425
def read_all(
424426
self, column_indices: list[int] | None = None, use_threads: bool = True
425427
) -> Table: ...
428+
426429
def scan_contents(
427-
self, columns: Sequence[str] | Sequence[int] | None = None, batch_size: int = 65536
430+
self, columns: Sequence[str] | Sequence[int] | None = None,
431+
batch_size: int = 65536
428432
) -> int: ...
429433

430434
def column_name_idx(self, column_name: str) -> int: ...

python/pyarrow-stubs/pyarrow/_stubs_typing.pyi

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -45,8 +45,18 @@ Compression: TypeAlias = Literal[
4545
]
4646
NullEncoding: TypeAlias = Literal["mask", "encode"]
4747
NullSelectionBehavior: TypeAlias = Literal["drop", "emit_null"]
48-
Mask: TypeAlias = Sequence[bool | None] | NDArray[np.bool_] | BooleanArray | ChunkedArray[Any]
49-
Indices: TypeAlias = Sequence[int | None] | NDArray[np.integer[Any]] | IntegerArray | ChunkedArray[Any]
48+
Mask: TypeAlias = (
49+
Sequence[bool | None]
50+
| NDArray[np.bool_]
51+
| BooleanArray
52+
| ChunkedArray[Any]
53+
)
54+
Indices: TypeAlias = (
55+
Sequence[int | None]
56+
| NDArray[np.integer[Any]]
57+
| IntegerArray
58+
| ChunkedArray[Any]
59+
)
5060

5161
PyScalar: TypeAlias = (bool | int | float | Decimal | str | bytes |
5262
dt.date | dt.datetime | dt.time | dt.timedelta)

python/pyarrow-stubs/pyarrow/_types.pyi

Lines changed: 30 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -515,16 +515,20 @@ def unregister_extension_type(type_name: str) -> None: ...
515515

516516
class KeyValueMetadata(_Metadata, Mapping[bytes, bytes]):
517517

518-
def __init__(self, __arg0__: Mapping[str | bytes, str | bytes] | Iterable[tuple[str, str]] | KeyValueMetadata
519-
| None = None, **kwargs: str) -> None: ...
518+
def __init__(
519+
self, __arg0__: Mapping[str | bytes, str | bytes]
520+
| Iterable[tuple[str, str]]
521+
| KeyValueMetadata
522+
| None = None, **kwargs: str
523+
) -> None: ...
520524

521525
def equals(self, other: KeyValueMetadata) -> bool: ...
522526

523527
def __len__(self) -> int: ...
524528

525-
def __contains__(self, __key: object) -> bool: ... # type: ignore[override]
529+
def __contains__(self, /, __key: object) -> bool: ... # type: ignore[override]
526530

527-
def __getitem__(self, __key: Any) -> Any: ... # type: ignore[override]
531+
def __getitem__(self, /, __key: Any) -> Any: ... # type: ignore[override]
528532

529533
def __iter__(self) -> Iterator[bytes]: ...
530534

@@ -658,8 +662,11 @@ def unify_schemas(
658662
) -> Schema: ...
659663

660664

661-
def field(name: SupportArrowSchema | str | Any, type: _DataTypeT | str | None = None, nullable: bool = ...,
662-
metadata: dict[Any, Any] | None = None) -> Field[_DataTypeT] | Field[Any]: ...
665+
def field(
666+
name: SupportArrowSchema | str | Any, type: _DataTypeT | str | None = None,
667+
nullable: bool = ...,
668+
metadata: dict[Any, Any] | None = None
669+
) -> Field[_DataTypeT] | Field[Any]: ...
663670

664671

665672
def null() -> NullType: ...
@@ -692,7 +699,8 @@ def int64() -> Int64Type: ...
692699
def uint64() -> UInt64Type: ...
693700

694701

695-
def timestamp(unit: _Unit | str, tz: _Tz | None = None) -> TimestampType[_Unit, _Tz]: ...
702+
def timestamp(
703+
unit: _Unit | str, tz: _Tz | None = None) -> TimestampType[_Unit, _Tz]: ...
696704

697705

698706
def time32(unit: _Time32Unit | str) -> Time32Type[_Time32Unit]: ...
@@ -789,13 +797,18 @@ def map_(
789797

790798

791799
def dictionary(
792-
index_type: _IndexT | str, value_type: _BasicValueT | str, ordered: _Ordered | None = None
800+
index_type: _IndexT | str,
801+
value_type: _BasicValueT | str,
802+
ordered: _Ordered | None = None
793803
) -> DictionaryType[_IndexT, _BasicValueT, _Ordered]: ...
794804

795805

796806
def struct(
797-
fields: Iterable[Field[Any] | tuple[str, Field[Any] | None] | tuple[str, DataType | None]]
798-
| Mapping[str, Field[Any] | DataType | None],
807+
fields: Iterable[
808+
Field[Any]
809+
| tuple[str, Field[Any] | None]
810+
| tuple[str, DataType | None]
811+
] | Mapping[str, Field[Any] | DataType | None],
799812
) -> StructType: ...
800813

801814

@@ -809,9 +822,9 @@ def dense_union(
809822
) -> DenseUnionType: ...
810823

811824

812-
def union(child_fields: list[Field[Any]],
813-
mode: Literal["sparse"] | Literal["dense"] | int | str, # noqa: Y030
814-
type_codes: list[int] | None = None) -> SparseUnionType | DenseUnionType: ...
825+
def union(
826+
child_fields: list[Field[Any]], mode: Literal["sparse" | "dense"] | int | str,
827+
type_codes: list[int] | None = None) -> SparseUnionType | DenseUnionType: ...
815828

816829

817830
def run_end_encoded(
@@ -848,7 +861,10 @@ def schema(
848861
| Iterable[tuple[str, DataType | str | None]]
849862
| Mapping[Any, DataType | str | None]
850863
),
851-
metadata: Mapping[bytes, bytes] | Mapping[str, str] | Mapping[bytes, str] | Mapping[str, bytes] | None = None,
864+
metadata: Mapping[bytes, bytes]
865+
| Mapping[str, str]
866+
| Mapping[bytes, str]
867+
| Mapping[str, bytes] | None = None,
852868
) -> Schema: ...
853869

854870

python/pyarrow-stubs/pyarrow/array.pyi

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@
1717

1818
import sys
1919

20-
from collections.abc import Callable, Iterable, Iterator, Sequence
20+
from collections.abc import Iterable, Iterator, Sequence
2121

2222
if sys.version_info >= (3, 11):
2323
from typing import Self
@@ -29,13 +29,11 @@ from typing import (
2929
Generic,
3030
Literal,
3131
TypeVar,
32-
overload,
3332
)
3433

3534
import numpy as np
3635
import pandas as pd
3736

38-
from pandas.core.dtypes.base import ExtensionDtype
3937
from pyarrow._compute import CastOptions
4038
from pyarrow._stubs_typing import (
4139
ArrayLike,
@@ -819,7 +817,10 @@ class FixedShapeTensorArray(ExtensionArray[_ArrayT]):
819817
def to_tensor(self) -> Tensor: ...
820818

821819
@classmethod
822-
def from_numpy_ndarray(cls, obj: np.ndarray, dim_names: list[str] | tuple[str, ...] | None = None) -> Self: ...
820+
def from_numpy_ndarray(
821+
cls, obj: np.ndarray,
822+
dim_names: list[str] | tuple[str, ...] | None = None
823+
) -> Self: ...
823824

824825

825826
class OpaqueArray(ExtensionArray[_ArrayT]):

0 commit comments

Comments
 (0)