Repository navigation
Conversation
Co-authored-by: Dan Redding <125183946+dangotbanned@users.noreply.github.com>
Co-authored-by: Dan Redding <125183946+dangotbanned@users.noreply.github.com>
Co-authored-by: Dan Redding <125183946+dangotbanned@users.noreply.github.com>
Co-authored-by: Dan Redding <125183946+dangotbanned@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
Add type annotation stubs for array, table, tensor, builder, memory, device, config, and types modules. Includes type-ignore annotations in related tests.
Drop the TypeVar/Generic parameters from the DataType subclasses, Field, Array, ChunkedArray and the tabular classes. Where a property can only return one of a few types (dictionary index type, run end type, time units) it now returns a union or Literal instead. Most of these parameters weren't buying much yet: scalar.pyi and lib.pyi are still placeholders, so e.g. Array[Scalar[...]] collapsed to Any anyway, and several annotations were being flagged by mypy and pyright (TypeVars used only once, type arguments on subclasses that weren't actually generic). Plain classes are a lot easier to read and review, and we can bring parameters back later where they clearly pay off. This also removes the stub-only _BasicDataType and _HasPrecisionScale helpers, which have no runtime counterpart, plus the IntoArray and SingleOrList aliases.
Add pyarrow-stubs/tests with pandas-stubs style checks:
check(assert_type(lib.list_(lib.int32()), lib.ListType), lib.ListType)
assert_type is verified by mypy, pyright and ty, which already include
pyarrow-stubs/ in their runs, and check() verifies the runtime class when
the tests run under pytest, so the stubs can't silently drift from what
pyarrow returns. Usage that should be rejected goes under
`if TYPE_CHECKING_INVALID_USAGE:` with a `# type: ignore[...]`; unused
ignores are errors in these tests, so those fail if the usage starts
type checking.
To have something to test, lib.pyi now re-exports the stub modules that
are annotated so far, still falling back to Any for everything else.
The CI type annotation script runs the tests with pytest after the type
checkers.
|
|
table.pyi imports RecordBatchReader from ._ipc, which had no stub, so pyright and ty failed to resolve the import. Add a placeholder in the same style as io.pyi and scalar.pyi until _ipc gets fully annotated.
| pip install fsspec pandas-stubs scipy-stubs types-cffi types-psutil types-requests types-python-dateutil | ||
|
|
||
| # Install type checkers | ||
| pip install mypy pyright ty |
There was a problem hiding this comment.
scratch note: should we add "typing_extensions>=4.12.0" also here to be certain to have it? Although CI job log of Conda Python 3.12 Sphinx & Numpydoc already shows the package in the conda env where pyarrow-stubs/tests currently run OK; typing_extensions 4.16.0 pyhcf101f3_0 conda-forge
There was a problem hiding this comment.
I think mypy (at least) installs typing_extensions as transitive dependency. Not sure if we need to install typing_extensions separately (yet).
github-actions
Bot
added
awaiting change review
|
A note while I'm still reviewing Type checking without NumPy, pandas, SciPy
uv run --isolated --no-project --with mypy==2.4.0 --with pyright==1.1.414 --with ty==0.0.92 bash -c 'mypy --no-incremental; pyright; ty check' Resolved: Remaining: |
|
I am reviewing |
8e8a811 to
aac6c4b
Compare
github-actions
Bot
added
awaiting changes
github-actions
Bot
added
awaiting change review
|
@tadeja @rok I have checked
Note: I am not sure about @tadeja I also created |
jorisvandenbossche
left a comment
jorisvandenbossche
left a comment
There was a problem hiding this comment.
Great work!
I did a (partial) pass though, and added a bunch of comments (sorry for the overload!). Some themes emerging from the inline comments, partially repeating here:
-
How should the type checkers be run? Through pre-commit, through the
python_test_type_annotations.shscript? Manually callingmypyor one of the others (in your typical local dev env)?
An important aspect here is that (in my experience) we have to control the exact set of versions of all dependencies (type checkers + pyarrow (test) deps). This is going to get messy if we have three ways to run them.I suppose listing those manually in pre-commit will have to stay (that cannot reference an external file), but so at least let's add a comment there that this has to be kept in sync with the list of deps used in CI. And for the CI script, we could maybe have a dependency-group in pyproject.toml specifying those? (that makes it more easily reused)
-
Do we already want to run the type checkers on the pyarrow tests? (or a subset of the test files) The PR seems to not yet do this, but on the other hand has a bunch of (partly incorrect?) changes to the test files (I started adding comments on those changes, but then stopped doing that). I would make a choice one way or the other.
(if we are currently not actually running the type checkers against the tests code, I would leave out all changes to those files for a subsequent PR; the changes are not lost, they can be cherry-picked in a follow-up PR) -
But then related to actually running the type checkers on the tests:
-
Should we specify
--follow-untyped-importsfor mypy? That avoids errors for me when I run it on eg pyarrow/tests/test_array.py complaining about pandas (and others) not being typed (because pandas-stubs is not installed, but the main pandas package also has annotations that mypy can find).
It would also avoid that any annotation using a pandas object gets interpreted asAny(also for the stub tests). -
The type checkers actually hardly report any error for me locally, but I think that is because the stubs don't actually exposes the public API in
pyarrow-stubs/__init__.pyi(this might be an easyfrom pyarrow.lib import *to be added?).
-
-
This does not seem to work out of the box with a local development (editable) install, to have other tools (eg IDE) pick up the stubs in that case.
I tried with opening a new vscode window, select my dev env as python interpreter, and start a python file with some pyarrow code. It could infer types when hovering or doing tab completion (and this does work when selecting a python env that has pyarrow-stubs installed) -
Should we add stubtest to CI? (that seems important to ensure the code and stubs don't drift on subsequent changes)
-
It would be good to get some consistency in how certain input arguments are typed, eg for array-like input, I see all of
Iterable,Sequenceandlist, while I think those all three should be the same (often, we might also functions where we are in practice stricter) -
I left a few inline comments related to this (but there are many other places doing it as well), but we should discuss how to type input arguments requiring a type specific Array (like BooleanArray for a mask, or IntegerArray for indices). I think we should type those all as just Array, as otherwise users will get annoying type errors if their variables are inferred as just Array (which will often be the case)
| ty check --verbose --output-format concise | ||
|
|
||
| # The type annotation tests are type checked above; also run them to verify | ||
| # the annotated types against the installed pyarrow. |
There was a problem hiding this comment.
| # the annotated types against the installed pyarrow. | |
| # the expected annotated types against the runtime types. |
to make it clear that this pytest invocation does not actually check annotations (it only does because we repeat it twice, but so in itself just checks the runtime isinstance).
(I was confused about what this pytest would do, in addition to running the type checkers)
There was a problem hiding this comment.
Would it also help to mention the description of assert_type and check methods from the python/pyarrow-stubs/tests/_check.py docstring?
| Run ``pytest`` directly and pass it the directory: ``python -m pytest`` puts | ||
| ``arrow/python`` on ``sys.path``, so the source tree would be imported instead | ||
| of the installed PyArrow. |
There was a problem hiding this comment.
I don't know if this is hallucinated or because someone of you actually run into an issue with this, but using python -m pytest instead of pytest with the example above works just fine for me.
(in general this is just about running pytest, so that should generally work the same as running the normal pyarrow unit tests? And there we don't say anything like this: https://arrow.apache.org/docs/developers/python/development.html#unit-testing)
There was a problem hiding this comment.
python -m pytest pyarrow-stubs/tests from arrow/python works for me too (the tests pass).
| def test_array() -> None: | ||
| # pa.array() isn't annotated with a specific return type yet, so tell | ||
| # the type checker which class to expect. | ||
| arr = check(cast(lib.Int64Array, lib.array([1, 2, None])), lib.Int64Array) |
There was a problem hiding this comment.
Is there a reason this is all written with lib instead of the public import path? (to be honest, if using pa. for that, this is only about a single character being saved ;), but in general in tests we want to use public imports; that also makes things easier to copy or paste from interactive testing)
There was a problem hiding this comment.
I think the reason is the placeholder in the __init__.pyi. Until that file is also updated (following PRs) all pa.* will be resolved to Any (due to __getattr__).
There was a problem hiding this comment.
AFAIK there is no reason to not already expose everything currently typed in __init__.pyi, as that is compatible with keeping the __getattr__ as a catch-all for everything else. At least in normal python code, a module __getattr__ is only used if a name is not available in the module through normal means.
|
|
||
| def test_import_from_c_capsule() -> None: | ||
| capsules = lib.array([1, 2]).__arrow_c_array__() | ||
| imported = lib.BooleanArray._import_from_c_capsule(*capsules) |
There was a problem hiding this comment.
Is the BooleanArray (which contrasts with the actual type) intentional?
In general I think the usage pattern for calling _import_from_c_capsule it to always do that using the base class pa.Array
There was a problem hiding this comment.
Yeah, I wanted to show that irrespective of the array type the method is called on, object with the type from the PyCapsule will be constructed. If it is misleading, I will change to pa.Array.
| def from_pandas( | ||
| obj: pd.Series | np.ndarray | ArrayLike, | ||
| mask: Mask | None = None, | ||
| type: DataType | None = None, |
There was a problem hiding this comment.
| type: DataType | None = None, | |
| type: DataType | str | None = None, |
| def from_buffers( | ||
| type: DataType, | ||
| length: int, | ||
| buffers: Sequence[Buffer | None], |
There was a problem hiding this comment.
Strictly speaking, this does not need to be a Sequence (which has getitem), but Collection (or even Sized) is enough
(for children below, which has the same requirements, a Collection is used)
| buffers: Sequence[Buffer | None], | ||
| null_count: int = -1, | ||
| offset=0, | ||
| children: NullableCollection | None = None, |
There was a problem hiding this comment.
| children: NullableCollection | None = None, | |
| children: Collection | None = None, |
?
|
|
||
| from types import TracebackType | ||
|
|
||
| from typing_extensions import Self |
There was a problem hiding this comment.
| from typing_extensions import Self | |
| from typing import Self |
|
|
||
| class ChunkedArray(_PandasConvertible): | ||
| @property | ||
| def data(self) -> Self: ... |
There was a problem hiding this comment.
Sidenote, we should probably just remove this attribute, it has been deprecated for years
github-actions
Bot
added
awaiting changes
Alternatively, could we ensure those optional dependencies are installed whenever doing type checks? (which is already the case in practice?) That avoids having to add such ignores?
Would be good to add a test for that (so we don't accidentally change that to Capsule later) |
|
Yay, @jorisvandenbossche, thanks for the review! <3
Yeah, I currently run pre-commit and also all 3 type checkers separately in my local environment. The two runs differ due to possibility of missing deps. I agree we should try to make this more clear and unified, maybe using pre-commit only or providing a requirements file for type checking locally in the dev env? (connected to ignores question)
I think we would have to run type checking on the tests locally until we have added all the stubs. That is what we have planned after the reviews and before the merge. See: #55264 (comment)
Yes, I already did some testing, see: #55264 (comment)
Yes, but that would make sense only after all the parts of PyArrow are typed. We did a stubtest check for the part from this PR and already added changes. See #55264 (comment)
Yes, I think consistency would be good. But going through the code I see we have very diverse way of checking the types in PyArrow. For example, sometimes a parameter is typed in Cython which checks types at runtime so we have to be stricter there. Or we check specifically for list or tuple with
Yeah, I think that makes sense. If we agree on the first item (how to run stubs) I can push the change and also add some documentation already (or at least add it on my todo).
Yes! I have it on my todo list 👍 |
Co-authored-by: Joris Van den Bossche <jorisvandenbossche@gmail.com>
github-actions
Bot
added
awaiting change review
github-actions
Bot
added
awaiting changes
github-actions
Bot
added
awaiting change review
Rationale for this change
This is the next step in adding type annotations to PyArrow (#32609). It merges #49103 (internal type system stubs) and #49193 (core data structures) because the two depend on each other and can't be type checked separately.
What changes are included in this PR?
Are these changes tested?
Are there any user-facing changes?
Was AI used for this PR?
In accordance to the AI generation guidelines, please disclose below whether and how AI was used in this PR.
PR code and description written by:
Reviewed before submission by: