annotation fix

Author: rok

python/pyarrow-stubs/pyarrow/_compute.pyi

@@ -391,7 +391,7 @@ class RoundTemporalOptions(FunctionOptions):
 
 
 class RoundToMultipleOptions(FunctionOptions):
-    def __init__(self, multiple: float = 1.0,
+    def __init__(self, multiple: int | float | lib.Scalar = 1.0,
                  round_mode: _RoundMode = "half_to_even") -> None: ...
 
 

python/pyarrow-stubs/pyarrow/_ipc.pyi

@@ -175,7 +175,7 @@ class _ReadPandasMixin:
     def read_pandas(self, **options) -> pd.DataFrame: ...
 
 
-class RecordBatchReader(_Weakrefable):
+class RecordBatchReader(_ReadPandasMixin, _Weakrefable):
     def __iter__(self) -> Self: ...
     def read_next_batch(self) -> RecordBatch: ...
 
@@ -191,7 +191,6 @@ class RecordBatchReader(_Weakrefable):
 
     def read_all(self) -> Table: ...
 
-    read_pandas = _ReadPandasMixin.read_pandas
     def close(self) -> None: ...
 
     def __enter__(self) -> Self: ...
@@ -237,7 +236,7 @@ class RecordBatchWithMetadata(NamedTuple):
     custom_metadata: KeyValueMetadata
 
 
-class _RecordBatchFileReader(_Weakrefable):
+class _RecordBatchFileReader(_ReadPandasMixin, _Weakrefable):
     @property
     def num_record_batches(self) -> int: ...
 
@@ -248,7 +247,6 @@ class _RecordBatchFileReader(_Weakrefable):
 
     def read_all(self) -> Table: ...
 
-    read_pandas = _ReadPandasMixin.read_pandas
     def __enter__(self) -> Self: ...
     def __exit__(self, exc_type, exc_val, exc_tb): ...
     @property

python/pyarrow-stubs/pyarrow/compute.pyi

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from collections.abc import Callable, Iterable, Sequence, Mapping
+from collections.abc import Callable, Hashable, Iterable, Sequence, Mapping
 from typing import Literal, TypeAlias, TypeVar, Any, ParamSpec
 
 import numpy as np
@@ -654,7 +654,7 @@ def round(
 def round_to_multiple(
     x: _NumericScalarT | _NumericArrayT | list | Expression,
     /,
-    multiple: int = 0,
+    multiple: int | float | NumericScalar = 1.0,
     round_mode: Literal[
         "down",
         "up",
@@ -1521,7 +1521,7 @@ def local_timestamp(
 def random(
     n: int,
     *,
-    initializer: Literal["system"] | int = "system",
+    initializer: Hashable = "system",
     options: RandomOptions | None = None,
     memory_pool: lib.MemoryPool | None = None,
 ) -> lib.DoubleArray: ...
@@ -1533,7 +1533,7 @@ def random(
 def cumulative_sum(
     values: _NumericArrayT | ArrayLike | Expression,
     /,
-    start: lib.Scalar | None = None,
+    start: int | float | lib.Scalar | None = None,
     *,
     skip_nulls: bool = False,
     options: CumulativeSumOptions | None = None,

python/pyarrow-stubs/pyarrow/lib.pyi

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from typing import NamedTuple
+from typing import Iterator, NamedTuple
 
 from .array import *  # noqa: F401, F403
 from .builder import *  # noqa: F401, F403
@@ -36,12 +36,15 @@ from .array import Array
 from ._types import DataType
 
 
-class MonthDayNano(NamedTuple):
-    days: int
+class MonthDayNano(tuple):
     months: int
+    days: int
     nanoseconds: int
 
-    def __init__(self, *args, **kwargs) -> None: ...  # type: ignore[misc]
+    def __new__(
+        cls,
+        sequence: tuple[int, int, int] | list[int] = ...,
+    ) -> MonthDayNano: ...
 
 
 def cpu_count() -> int: ...

python/pyarrow/parquet/core.py

@@ -1413,7 +1413,7 @@ def __init__(self, path_or_paths, filesystem=None, schema=None, *, filters=None,
             else:
                 single_file = path_or_paths
 
-        parquet_format = ds.ParquetFileFormat(**read_options)
+        parquet_format = ds.ParquetFileFormat(**read_options)  # type: ignore[invalid-argument-type]
 
         if single_file is not None:
             fragment = parquet_format.make_fragment(single_file, filesystem)

python/pyarrow/tests/parquet/test_dataset.py

@@ -971,7 +971,7 @@ def _test_write_to_dataset_with_partitions(base_path,
     input_df_cols = input_df.columns.tolist()
     assert partition_by == input_df_cols[-1 * len(partition_by):]
 
-    input_df = input_df[cols]  # type: ignore[assignment]
+    input_df = input_df.loc[:, cols]
     # Partitioned columns become 'categorical' dtypes
     for col in partition_by:
         output_df[col] = output_df[col].astype('category')
@@ -980,6 +980,7 @@ def _test_write_to_dataset_with_partitions(base_path,
         expected_date_type = schema.field('date').type.to_pandas_dtype()
         output_df["date"] = output_df["date"].astype(expected_date_type)
 
+    assert isinstance(input_df, pd.DataFrame)
     tm.assert_frame_equal(output_df, input_df)
 
 
@@ -1027,7 +1028,8 @@ def _test_write_to_dataset_no_partitions(base_path,
     ).read()
     input_df = input_table.to_pandas()
     input_df = input_df.drop_duplicates()
-    input_df = input_df[cols]  # type: ignore[assignment]
+    input_df = input_df[cols]
+    assert isinstance(input_df, pd.DataFrame)
     tm.assert_frame_equal(output_df, input_df)
 
 

python/pyarrow/tests/test_acero.py

@@ -16,6 +16,7 @@
 # under the License.
 
 import pytest
+from typing import Literal, cast
 
 import pyarrow as pa
 import pyarrow.compute as pc
@@ -268,7 +269,8 @@ def test_order_by():
     table = pa.table({'a': [1, 2, 3, 4], 'b': [1, 3, None, 2]})
     table_source = Declaration("table_source", TableSourceNodeOptions(table))
 
-    ord_opts = OrderByNodeOptions([("b", "ascending")])
+    sort_keys: list[tuple[str, Literal["ascending", "descending"]]] = [("b", "ascending")]
+    ord_opts = OrderByNodeOptions(sort_keys)
     decl = Declaration.from_sequence([table_source, Declaration("order_by", ord_opts)])
     result = decl.to_table()
     expected = pa.table({"a": [1, 4, 2, 3], "b": [1, 2, 3, None]})

python/pyarrow/tests/test_compute.py

@@ -1801,7 +1801,7 @@ def test_round_to_multiple():
 
     for multiple in [object, 99999999999999999999999]:
         with pytest.raises(TypeError, match="is not a valid multiple type"):
-            pc.round_to_multiple(values, multiple=multiple)
+            pc.round_to_multiple(values, multiple=multiple)  # type: ignore[arg-type]
 
 
 def test_round_binary():
@@ -3190,7 +3190,7 @@ def test_cumulative_sum(start, skip_nulls):
 
     for strt in ['a', pa.scalar('arrow'), 1.1]:
         with pytest.raises(pa.ArrowInvalid):
-            pc.cumulative_sum([1, 2, 3], start=strt)
+            pc.cumulative_sum([1, 2, 3], start=strt)  # type: ignore[arg-type]
 
 
 @pytest.mark.numpy

python/pyarrow/tests/test_fs.py

@@ -234,7 +234,7 @@ def gcsfs(request, gcs_server):
 
     yield dict(
         fs=fs,
-        pathfn=bucket.__add__,
+        pathfn=lambda p: bucket + p,
         allow_move_dir=False,
         allow_append_to_file=False,
     )
@@ -261,7 +261,7 @@ def s3fs(request, s3_server):
 
     yield dict(
         fs=fs,
-        pathfn=bucket.__add__,
+        pathfn=lambda p: bucket + p,
         allow_move_dir=False,
         allow_append_to_file=False,
     )
@@ -273,7 +273,7 @@ def subtree_s3fs(request, s3fs):
     prefix = 'pyarrow-filesystem/prefix/'
     return dict(
         fs=SubTreeFileSystem(prefix, s3fs['fs']),
-        pathfn=prefix.__add__,
+        pathfn=lambda p: prefix + p,
         allow_move_dir=False,
         allow_append_to_file=False,
     )
@@ -333,7 +333,7 @@ def azurefs(request, azure_server):
 
     yield dict(
         fs=fs,
-        pathfn=container.__add__,
+        pathfn=lambda p: container + p,
         allow_move_dir=True,
         allow_append_to_file=True,
     )
@@ -402,7 +402,7 @@ def py_fsspec_s3fs(request, s3_server):
 
     yield dict(
         fs=fs,
-        pathfn=bucket.__add__,
+        pathfn=lambda p: bucket + p,
         allow_move_dir=False,
         allow_append_to_file=True,
     )