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

@@ -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: str | int | bytes = "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

@@ -41,7 +41,14 @@ class MonthDayNano(NamedTuple):
     months: int
     nanoseconds: int
 
-    def __init__(self, *args, **kwargs) -> None: ...  # type: ignore[misc]
+    def __new__(
+        cls,
+        sequence: tuple[int, int, int] | list[int] = (),
+    ) -> MonthDayNano: ...
+    def __init__(
+        self,
+        sequence: tuple[int, int, int] | list[int] = (),
+    ) -> None: ...
 
 
 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

@@ -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,
     )
@@ -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,
     )