[Python] Add type checks to Arrow using pyarrow-stubs and initially ignoring failing checks

.github/workflows/python.yml

@@ -138,6 +138,11 @@ jobs:
         continue-on-error: true
         run: archery docker push ${{ matrix.image }}
 
+      - name: Type check with ty
+        run: |-
+            python -m pip install ty
+            pushd python; python -m ty check
+
   macos:
     name: ${{ matrix.architecture }} macOS ${{ matrix.macos-version }} Python 3
     runs-on: macos-${{ matrix.macos-version }}

dev/release/rat_exclude_files.txt

@@ -95,6 +95,7 @@ python/pyarrow/tests/__init__.py
 python/pyarrow/vendored/*
 python/pyarrow/src/arrow/python/vendored/*
 python/requirements*.txt
+python/stubs/*
 pax_global_header
 MANIFEST.in
 __init__.pxd

python/benchmarks/parquet.py

@@ -21,7 +21,7 @@
 try:
     import pyarrow.parquet as pq
 except ImportError:
-    pq = None
+    pass
 from pyarrow.tests.util import rands
 
 

python/examples/dataset/write_dataset_encrypted.py

@@ -67,16 +67,19 @@ def kms_factory(kms_connection_configuration):
 
 
 crypto_factory = pe.CryptoFactory(kms_factory)
-parquet_encryption_cfg = ds.ParquetEncryptionConfig(
+parquet_encryption_cfg = ds.ParquetEncryptionConfig( \
+    # type: ignore[possibly-unbound-attribute]
     crypto_factory, kms_connection_config, encryption_config)
-parquet_decryption_cfg = ds.ParquetDecryptionConfig(crypto_factory,
-                                                    kms_connection_config,
-                                                    decryption_config)
+parquet_decryption_cfg = ds.ParquetDecryptionConfig( \
+    # type: ignore[possibly-unbound-attribute]
+    crypto_factory, kms_connection_config, decryption_config)
 
 # set encryption config for parquet fragment scan options
-pq_scan_opts = ds.ParquetFragmentScanOptions()
+pq_scan_opts = ds.ParquetFragmentScanOptions() \
+    # type: ignore[possibly-unbound-attribute]
 pq_scan_opts.parquet_decryption_config = parquet_decryption_cfg
-pformat = pa.dataset.ParquetFileFormat(default_fragment_scan_options=pq_scan_opts)
+pformat = pa.dataset.ParquetFileFormat(default_fragment_scan_options=pq_scan_opts) \
+    # type: ignore[possibly-unbound-attribute]
 
 if os.path.exists('sample_dataset'):
     shutil.rmtree('sample_dataset')

python/examples/flight/client.py

@@ -70,7 +70,7 @@ def do_action(args, client, connection_args={}):
         print('Running action', args.action_type)
         for result in client.do_action(action):
             print("Got result", result.body.to_pybytes())
-    except pyarrow.lib.ArrowIOError as e:
+    except pyarrow.lib.ArrowIOError as e:  # type: ignore[unresolved-attribute]
         print("Error calling action:", e)
 
 

python/py.typed

@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.

python/pyarrow/__init__.py

@@ -37,29 +37,30 @@
 import warnings as _warnings
 
 try:
-    from ._generated_version import version as __version__
+    from ._generated_version import version as __version__ \
+        # type: ignore[unresolved-import]
 except ImportError:
     # Package is not installed, parse git tag at runtime
     try:
-        import setuptools_scm
+        import setuptools_scm  # type: ignore[unresolved-import]
         # Code duplicated from setup.py to avoid a dependency on each other
 
         def parse_git(root, **kwargs):
             """
             Parse function for setuptools_scm that ignores tags for non-C++
             subprojects, e.g. apache-arrow-js-XXX tags.
             """
-            from setuptools_scm.git import parse
+            from setuptools_scm.git import parse  # type: ignore[unresolved-import]
             kwargs['describe_command'] = \
                 "git describe --dirty --tags --long --match 'apache-arrow-[0-9]*.*'"
-            return parse(root, **kwargs)
+            return parse(root, **kwargs)  # type: ignore[missing-argument]
         __version__ = setuptools_scm.get_version('../',
                                                  parse=parse_git)
     except ImportError:
-        __version__ = None
+        __version__ = None  # type: ignore[invalid-assignment]
 
-import pyarrow.lib as _lib
-from pyarrow.lib import (BuildInfo, CppBuildInfo, RuntimeInfo, set_timezone_db_path,
+import pyarrow.lib as _lib  # type: ignore[unresolved_import]
+from pyarrow.lib import (BuildInfo, CppBuildInfo, RuntimeInfo, set_timezone_db_path,  # type: ignore[unresolved_import]
                          MonthDayNano, VersionInfo, build_info, cpp_build_info,
                          cpp_version, cpp_version_info, runtime_info,
                          cpu_count, set_cpu_count, enable_signal_handlers,
@@ -153,7 +154,7 @@ def print_entry(label, value):
         print(f"  {codec: <20}: {status: <8}")
 
 
-from pyarrow.lib import (null, bool_,
+from pyarrow.lib import (null, bool_,  # type: ignore[unresolved_import]
                          int8, int16, int32, int64,
                          uint8, uint16, uint32, uint64,
                          time32, time64, timestamp, date32, date64, duration,
@@ -237,13 +238,13 @@ def print_entry(label, value):
                          FixedShapeTensorScalar, JsonScalar, OpaqueScalar, UuidScalar)
 
 # Buffers, allocation
-from pyarrow.lib import (DeviceAllocationType, Device, MemoryManager,
+from pyarrow.lib import (DeviceAllocationType, Device, MemoryManager,  # type: ignore[unresolved_import]
                          default_cpu_memory_manager)
 
-from pyarrow.lib import (Buffer, ResizableBuffer, foreign_buffer, py_buffer,
+from pyarrow.lib import (Buffer, ResizableBuffer, foreign_buffer, py_buffer,  # type: ignore[unresolved_import]
                          Codec, compress, decompress, allocate_buffer)
 
-from pyarrow.lib import (MemoryPool, LoggingMemoryPool, ProxyMemoryPool,
+from pyarrow.lib import (MemoryPool, LoggingMemoryPool, ProxyMemoryPool,  # type: ignore[unresolved_import]
                          total_allocated_bytes, set_memory_pool,
                          default_memory_pool, system_memory_pool,
                          jemalloc_memory_pool, mimalloc_memory_pool,
@@ -252,7 +253,7 @@ def print_entry(label, value):
                          supported_memory_backends)
 
 # I/O
-from pyarrow.lib import (NativeFile, PythonFile,
+from pyarrow.lib import (NativeFile, PythonFile,  # type: ignore[unresolved_import]
                          BufferedInputStream, BufferedOutputStream, CacheOptions,
                          CompressedInputStream, CompressedOutputStream,
                          TransformInputStream, transcoding_input_stream,
@@ -263,12 +264,12 @@ def print_entry(label, value):
                          input_stream, output_stream,
                          have_libhdfs)
 
-from pyarrow.lib import (ChunkedArray, RecordBatch, Table, table,
+from pyarrow.lib import (ChunkedArray, RecordBatch, Table, table,  # type: ignore[unresolved_import]
                          concat_arrays, concat_tables, TableGroupBy,
                          RecordBatchReader, concat_batches)
 
 # Exceptions
-from pyarrow.lib import (ArrowCancelled,
+from pyarrow.lib import (ArrowCancelled,  # type: ignore[unresolved_import]
                          ArrowCapacityError,
                          ArrowException,
                          ArrowKeyError,

python/pyarrow/acero.py

@@ -22,11 +22,11 @@
 # distutils: language = c++
 # cython: language_level = 3
 
-from pyarrow.lib import Table, RecordBatch, array
+from pyarrow.lib import Table, RecordBatch, array  # type: ignore[unresolved_import]
 from pyarrow.compute import Expression, field
 
 try:
-    from pyarrow._acero import (  # noqa
+    from pyarrow._acero import (  # type: ignore[unresolved_import]  # noqa
         Declaration,
         ExecNodeOptions,
         TableSourceNodeOptions,
@@ -45,15 +45,15 @@
 
 try:
     import pyarrow.dataset as ds
-    from pyarrow._dataset import ScanNodeOptions
+    from pyarrow._dataset import ScanNodeOptions  # type: ignore[unresolved_import]
 except ImportError:
     class DatasetModuleStub:
         class Dataset:
             pass
 
         class InMemoryDataset:
             pass
-    ds = DatasetModuleStub
+    ds = DatasetModuleStub  # type: ignore[invalid-assignment]
 
 
 def _dataset_to_decl(dataset, use_threads=True, implicit_ordering=False):

python/pyarrow/benchmark.py

@@ -18,4 +18,4 @@
 # flake8: noqa
 
 
-from pyarrow.lib import benchmark_PandasObjectIsNull
+from pyarrow.lib import benchmark_PandasObjectIsNull  # type: ignore[unresolved_import]

python/pyarrow/cffi.py

@@ -17,7 +17,7 @@
 
 from __future__ import absolute_import
 
-import cffi
+from cffi import FFI  # type: ignore[unresolved-import]
 
 c_source = """
     struct ArrowSchema {
@@ -77,5 +77,5 @@
     """
 
 # TODO use out-of-line mode for faster import and avoid C parsing
-ffi = cffi.FFI()
+ffi = FFI()
 ffi.cdef(c_source)

python/pyarrow/compute.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from pyarrow._compute import (  # noqa
+from pyarrow._compute import (  # type: ignore[unresolved_import]  # noqa
     Function,
     FunctionOptions,
     FunctionRegistry,

python/pyarrow/conftest.py

@@ -15,13 +15,13 @@
 # specific language governing permissions and limitations
 # under the License.
 
-import pytest
+import pytest  # type: ignore[unresolved-import]
 
 import os
 import pyarrow as pa
 from pyarrow import Codec
 from pyarrow import fs
-from pyarrow.lib import is_threading_enabled
+from pyarrow.lib import is_threading_enabled  # type: ignore[unresolved_import]
 from pyarrow.tests.util import windows_has_tzdata
 import sys
 
@@ -114,19 +114,19 @@
     defaults['timezone_data'] = os.path.exists("/usr/share/zoneinfo")
 
 try:
-    import cython  # noqa
+    import cython  # type: ignore[unresolved-import]  # noqa
     defaults['cython'] = True
 except ImportError:
     pass
 
 try:
-    import fastparquet  # noqa
+    import fastparquet  # type: ignore[unresolved_import]  # noqa
     defaults['fastparquet'] = True
 except ImportError:
     pass
 
 try:
-    import pyarrow.gandiva  # noqa
+    import pyarrow.gandiva  # type: ignore[unresolved_import]  # noqa
     defaults['gandiva'] = True
 except ImportError:
     pass
@@ -156,13 +156,13 @@
     pass
 
 try:
-    import pandas  # noqa
+    import pandas  # type: ignore[unresolved-import]  # noqa
     defaults['pandas'] = True
 except ImportError:
     defaults['nopandas'] = True
 
 try:
-    import numpy  # noqa
+    import numpy  # type: ignore[unresolved-import]  # noqa
     defaults['numpy'] = True
 except ImportError:
     defaults['nonumpy'] = True
@@ -186,25 +186,25 @@
     pass
 
 try:
-    from pyarrow.fs import AzureFileSystem  # noqa
+    from pyarrow.fs import AzureFileSystem  # type: ignore[possibly-unbound-import]  # noqa
     defaults['azure'] = True
 except ImportError:
     pass
 
 try:
-    from pyarrow.fs import GcsFileSystem  # noqa
+    from pyarrow.fs import GcsFileSystem  # type: ignore[possibly-unbound-import]  # noqa
     defaults['gcs'] = True
 except ImportError:
     pass
 
 try:
-    from pyarrow.fs import S3FileSystem  # noqa
+    from pyarrow.fs import S3FileSystem  # type: ignore[possibly-unbound-import]  # noqa
     defaults['s3'] = True
 except ImportError:
     pass
 
 try:
-    from pyarrow.fs import HadoopFileSystem  # noqa
+    from pyarrow.fs import HadoopFileSystem  # type: ignore[possibly-unbound-import]  # noqa
     defaults['hdfs'] = True
 except ImportError:
     pass
@@ -250,7 +250,7 @@ def pytest_ignore_collect(collection_path, config):
 
         if 'pyarrow/fs' in str(collection_path):
             try:
-                from pyarrow.fs import S3FileSystem  # noqa
+                from pyarrow.fs import S3FileSystem  # type: ignore[possibly-unbound-import]  # noqa
                 return False
             except ImportError:
                 return True
@@ -336,7 +336,7 @@ def unary_agg_func_fixture():
     Register a unary aggregate function (mean)
     """
     from pyarrow import compute as pc
-    import numpy as np
+    import numpy as np  # type: ignore[unresolved-import]
 
     def func(ctx, x):
         return pa.scalar(np.nanmean(x))
@@ -362,7 +362,7 @@ def varargs_agg_func_fixture():
     Register a unary aggregate function
     """
     from pyarrow import compute as pc
-    import numpy as np
+    import numpy as np  # type: ignore[unresolved-import]
 
     def func(ctx, *args):
         sum = 0.0

python/pyarrow/csv.py

@@ -16,7 +16,7 @@
 # under the License.
 
 
-from pyarrow._csv import (  # noqa
+from pyarrow._csv import (  # type: ignore[unresolved_import]  # noqa
     ReadOptions, ParseOptions, ConvertOptions, ISO8601,
     open_csv, read_csv, CSVStreamingReader, write_csv,
     WriteOptions, CSVWriter, InvalidRow)

python/pyarrow/cuda.py

@@ -18,7 +18,7 @@
 # flake8: noqa
 
 
-from pyarrow._cuda import (Context, IpcMemHandle, CudaBuffer,
+from pyarrow._cuda import (Context, IpcMemHandle, CudaBuffer,  # type: ignore[unresolved_import]
                            HostBuffer, BufferReader, BufferWriter,
                            new_host_buffer,
                            serialize_record_batch, read_message,

python/pyarrow/dataset.py

@@ -21,7 +21,7 @@
 from pyarrow.util import _is_iterable, _stringify_path, _is_path_like
 
 try:
-    from pyarrow._dataset import (  # noqa
+    from pyarrow._dataset import (  # type: ignore[unresolved_import]  # noqa
         CsvFileFormat,
         CsvFragmentScanOptions,
         JsonFileFormat,
@@ -70,7 +70,7 @@
 )
 
 try:
-    from pyarrow._dataset_orc import OrcFileFormat
+    from pyarrow._dataset_orc import OrcFileFormat  # type: ignore[unresolved_import]
     _orc_available = True
 except ImportError:
     pass
@@ -82,7 +82,7 @@
 )
 
 try:
-    from pyarrow._dataset_parquet import (  # noqa
+    from pyarrow._dataset_parquet import (  # type: ignore[unresolved_import]  # noqa
         ParquetDatasetFactory,
         ParquetFactoryOptions,
         ParquetFileFormat,
@@ -98,7 +98,7 @@
 
 
 try:
-    from pyarrow._dataset_parquet_encryption import (  # noqa
+    from pyarrow._dataset_parquet_encryption import (  # type: ignore[unresolved_import]  # noqa
         ParquetDecryptionConfig,
         ParquetEncryptionConfig,
     )