A new home for pyarrow-stubs?

[yangdanny97]

The maintainer of pyarrow-stubs (a set of Python type stubs for arrow) has expressed interest in donating the stubs since they no longer have time to maintain them: zen-xu/pyarrow-stubs#186

I'm opening a discussion here to see if there's any interest in merging the stubs into the main arrow project. They may remain as stubs, or if there is interest, they can also be merged into the Python bindings as inline types.

I'm happy to help set up the appropriate CI to support future maintenance and make sure the types are accurate/up-to-date going forward.

I see an existing issue for this topic (#30426), but not a lot of active discussion, so I'm hoping to hear what the maintainers think.

[AlenkaF]

Thank you for opening up this discussion!

There is a general consensus on the need for type annotations in PyArrow. In my opinion, the main challenge to progress is the lack of agreement on the approach to take.

To add some context, here are a few key comments from the ongoing discussion in #32609:

• The latest summary from @jorisvandenbossche: [Python] Type checking support #32609 (comment)
• Insights from @paleolimbot regarding the nanoarrow experience: [Python] Type checking support #32609 (comment)
• Advice from the pyarrow-stubs author: [Python] Type checking support #32609 (comment)

[raulcd]

Thanks for opening the discussion @yangdanny97 and thanks for the links @AlenkaF , really useful to get the history around it!
I've closed:

• [Python] Type stubs #30426

As a duplicate of:

• [Python] Type checking support #32609

And added some links to try and keep things organized on a single place.

[yangdanny97]

Thanks for the context!

If there isn't a strong consensus on inline types, I'd probably recommend the types be kept as stubs (with some tools like stubtest to make sure they're complete and in-sync). The reason is that if you change your mind, it's easier to move from stubs to inline types, whereas removing inline types will cause a lot of churn and merge conflicts.

Hand-written stubs will be much better quality than completely auto-generated stubs, and to the best of my knowledge all existing stub-generation tools do not produce stubs that are good enough to be used out-of-the-box without manual tweaks. Almost all the stubs in typeshed are maintained by hand, with a few exceptions like protobuf which has its own stub generation script.

I can see generated stubs being used as a temporary supplement to fill any gaps that pyarrow-stubs has (according to the author it's relatively complete), but after being generated the stubs will need to be fixed and maintained by hand.

[AlenkaF]

Just wanted to keep this discussion active—following up with my thoughts on a possible path forward. I believe a good approach would be to coordinate efforts with the pyarrow-stubs project and build on the work that has already been done there. That said, I’m a bit uncertain, since @paleolimbot and @jorisvandenbossche have both expressed support for a simpler, auto-generated solution—and they’re usually right. 😊

That said, PyArrow developers are currently managing several priorities, so this likely isn’t something we can kick off immediately.

Still, it would be valuable to start planning how we want to approach this. Based on the ongoing discussions, I’d like to highlight a couple of key points:

• As Joris suggested, we could start with minimal, lightweight stubs—something simpler than what currently exists in pyarrow-stubs—and gradually extend them as needed (comment).

  • If transferring ownership of the pyarrow-stubs project would still be relevant, bringing it under the Apache Arrow umbrella would involve the official Apache donation process, starting with a vote on the mailing list.

• It might also be worth creating a proof of concept for auto-generated stubs to compare approaches.

Lastly, I was recently pointed to a new type checker tool: https://github.com/astral-sh/ty. Does anyone have experience with it?

[yangdanny97]

Do you have an estimate for how much development resources/funding an upstreaming effort like this would require?

[AlenkaF]

Not really, no.

@paleolimbot — since you’ve worked on type checking in nanoarrow, could you give us a hand here?

We're wondering how much development effort might be needed for something similar in PyArrow, building on the existing pyarrow-stubs work — including both implementation and ongoing maintenance.

[paleolimbot]

My experience in nanoarrow is just with stub generation, mostly for the purposes of halfway decent autocomplete (we never quite got to setting up a type checker). That's still my main beef with writing pyarrow code (IDEs have no idea what most of the modules contain), although full on type checked everything would be great!

[rustyconover]

I'd like to help out with some pyarrow stub writing. I'd really like to be able to type check my PyArrow code with mypy.

[AlenkaF]

Just noting that I’d also like to be involved in this work.
I’d naively estimate that, with 2-3 people working part-time, this could be completed within one release cycle—or two at most, if additional review iterations are needed.

[rok]

As a sprint at EuroPython @paddyroddy and I created a draft for including @zen-xu's pyarrow-stubs with pyarrow (as proposal nr 2 in #32609 (comment)). We used ty and included a CI job that checks for annotation issues (following numpy's example).

If we decide to proceed with this we would of course need @zen-xu donating pyarrow-stubs first.

Things that come to mind after working on this:

• It would be nice to inline as much annotations as possible for maintainability.
• PR currently copies pyarrow-stubs verbatim, but we would probably want to remove docs from pyi files if possible.

[yangdanny97]

Given that arrow is a popular library, it's best to use multiple typecheckers in CI, to make sure the types work for users regardless of the tool they're using.

So that means for now starting with Mypy/pyright, and then adding ty/pyrefly whenever they're ready (Pyrefly and ty are both running in CI for pandas-stubs, so it may already be possible today).

Another good tool to use is stubtest, which will check the stubs against values at runtime to make sure that the stubs are complete & accurate. It has an allowlist functionality so that you can add it to CI without being fully passing, and then clean up the allowlist going forward.

[paddyroddy]

Does feel like we should use mypy_primer as well, which NumPy use in their type checking CI. Note, NumPy previously had a separate stubs library and moved them in.

[rok]

Thanks for the input on ty vs mypy vs pyrefly! Sounds like we want to aim for mypy first, but add others.

From your experience on the question @amoeba raises here:

Since inlining annotations is preferred, maybe we should have a section in the PyArrow docs for contributors who might want to help convert stubs to inlined annotations? Is the goal here to inline all type annotations? Should new contributions add stubs or add inline annotations?

Is inlining annotations generally preferred? Do you know of projects that inlined annotations and how was that organized?

[dangotbanned]

Is inlining annotations generally preferred? Do you know of projects that inlined annotations and how was that organized?

The typing spec has some helpful resources on why you might choose either

• https://typing.python.org/en/latest/guides/libraries.html
• https://typing.python.org/en/latest/guides/writing_stubs.html

[AlenkaF]

I am not against inline annotations—in fact, I’ll support whatever decision @rok and @paddyroddy make on this. That said, I wanted to link to a comment from March this year that might be relevant or helpful:

#45919 (reply in thread)

If there isn't a strong consensus on inline types, I'd probably recommend the types be kept as stubs (with some tools like stubtest to make sure they're complete and in-sync). The reason is that if you change your mind, it's easier to move from stubs to inline types, whereas removing inline types will cause a lot of churn and merge conflicts.

[rok]

Another update.

Following received feedback I've created a second draft. This one takes the approach of adding pyarrow-stubs, then checks annotations on a single test file (test_compute.py) using pyright. To make checks pass stubs are fixed in the PR. A question I have here is: do we work on a feature branch until we have full coverage (annotation checks pass on all files in arrow/python/pyarrow) or do we merge partial (when e.g. test_xyz.py passes checks)?

The idea would be that CI would report stub check failures and we can enforce annotations remain in sync.

Type coverage metric, stub linter and pre-commit hook were some ideas suggested so far that can come later.

I would propose we inline annotations where possible - @mpelko did a proof of concept here.

If there are no significant objections I'll open a PR against the main repo and start a ML discussion.

[rok]

Thanks for the discussion so far everyone.

I think we (can) agree on:

• Docs are ok as long as we can automate their maintenance.
• We want to start without overloads, with annotations that at least cover data structures but not necessarily the underlying data type (e.g. pa.Scalar[Any]).
• Types can be narrowed and overloaded in further work if there's need and consensus.

I've updated my draft PR with this in mind. I've also split the stubs so only about one third are left in the PR to make it easier to review.

If there's not strong opposition and we agree on principle I'll open a PR against the main branch and we can discuss implementation details there.

[dangotbanned]

Would there be a way to test these stubs out on narwhals?

I've had a look through (rok#45)
I do appreciate consensus here seems to be aiming for simplicity -but I'm feeling like we're gonna be fighting against [attr-defined]/[reportAttributeAccessIssue] warnings with zero overloads 🤔

Using the last doctest in pa.array as an example - the overloads in pyarrow-stubs will show index_type as being present only on DictionaryType.

The actual index type seems to be wrong 😅 - but we at least can know what attributes exist and that helps prevent hitting errors

[rok]

Would there be a way to test these stubs out on narwhals?

Maybe you could copy the stubs folder and point at them with $MYPYPATH.

I've had a look through (rok#45) I do appreciate consensus here seems to be aiming for simplicity -but I'm feeling like we're gonna be fighting against [attr-defined]/[reportAttributeAccessIssue] warnings with zero overloads 🤔

I'm not there yet, I hope this is solvable.

Using the last doctest in pa.array as an example - the overloads in pyarrow-stubs will show index_type as being present only on DictionaryType.

The actual index type seems to be wrong 😅 - but we at least can know what attributes exist and that helps prevent hitting errors

Nice!

[rok]

I've moved the PR to main to get more visibility and introduced the following changes:

• Removed docstrings and added a script that can import read docstrings dynamically and write them into stubfiles at wheel build time
• Stubfiles are now in arrow/python/pyarrow-stubs
• Overloads are still removed
• mypy and pyright annotation checks of pyarrow-stubs folder pass

I'd appreciate another round of feedback or if it's not necessary a round of reviews.

[rok]

Also, we're at ~13k LoC total. The ~1k LoC of comments I believe are mostly ellipses.

cloc --include-ext=pyi pyarrow-stubs 
      62 text files.
      61 unique files.                              
       1 file ignored.
-------------------------------------------------------------------------------
Language                     files          blank        comment           code
-------------------------------------------------------------------------------
Python                          61           2865           1090           8753
-------------------------------------------------------------------------------
SUM:                            61           2865           1090           8753
-------------------------------------------------------------------------------

[dangotbanned]

Easing maintenance burden

Note

TL;DR: If downstream users feel confident they can contribute typing-only changes, they'll be able to spread the burden 🙂

How the stubs/inline typing can be maintained in a sustainable way, seems to have been a concern of most voices here:

Show related comments

• A new home for pyarrow-stubs? #45919 (reply in thread)
• A new home for pyarrow-stubs? #45919 (reply in thread)
• A new home for pyarrow-stubs? #45919 (reply in thread)
• A new home for pyarrow-stubs? #45919 (reply in thread)
• A new home for pyarrow-stubs? #45919 (reply in thread)
• A new home for pyarrow-stubs? #45919 (reply in thread)
• A new home for pyarrow-stubs? #45919 (reply in thread)
• A new home for pyarrow-stubs? #45919 (reply in thread)

While there is Writing and Maintaining Stub Files to refer to for guidance, both that document and (I assume) the discussion so far have been asking/answering:

How can the current maintainers of pyarrow, maintain the stubs?

That is an important question, but I think making the process sustainable may benefit from answering this question as well:

How can the current maintainers of pyarrow, make it easy/easier for downstream users to contribute to the stubs?

I can only speak to my own experience, so as an example ...

Example story

I feel confident in identifying when there are issues in stubs, even if that requires reading a lil bit of cython
I also think I have enough of an understanding of the python typing spec to propose changes which utilize new features as they land.

I feel less confident that I'd pull off following docs/developers/python without issue - as someone with less experience on projects with a compile-step.
(although #36411 may shake things up here)