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.

[pitrou]

It's both reducing the size and making maintenance (IMHO) easier.

If the overloads can be somehow generated by some clever use of type mappings, it can alleviate that concern.

[WillAyd]

I can speak a bit to the pandas perspective since that was brought up here. Pandas does have stubs with an entirely separate team that manages them (the owner of that repo is a pandas core team member). They aren't always up to date, but generally the team that manages them does a decent job of keeping them consistent with pandas releases. Typing issues are reported in the separate stubs repository and usually not noticed by pandas core teammembers.

The inline annotations are used for internal automated type checking validation, with the idea being that these would catch bugs. However, there are also plenty of ignores scattered throughout the code base, because pandas has worked on stubs for many years, dealt with different versions of mypy / pyright / etc..., and been contributed by people with different levels of interest when it comes to type maintenance. So its questionable if that has actually helped or been burdensome.

It's both reducing the size and making maintenance (IMHO) easier.

The importance of this can't be stressed enough. We have had issues in the past with not having a clear directive on how important types are, so some contributors come in leaving them on the backburner and then feel like they get stonewalled during review to fix types. I do think its unfortunate to leave features on the table that don't satisfy the type system, but a clearer directive from a project team could probably avoid that issue.

[MarcoGorelli]

Thanks all for the discussion!

If I may add a further perspective, as someone building on PyArrow:

• I find it extremely useful to know whether a method will return pa.Array or pa.ChunkedArray or np.ndarray or pd.Series or pd.Index. This has helped avoid bugs
• I've never found it useful to narrow those types down further to the data type level, if anything it's usually fairly noisy. As @pitrou said, you often don't know the types anyway and so end up either having them as Any or scattering # type: ignore / Incomplete around

I'd find it extremely useful to have simple type annotations from PyArrow, and if the myriad of overloads necessary to narrow down the types are left out due to size and maintenance concerns, I wouldn't personally see that as a loss

---

Furthermore, it's generally easier to start simple and increase complexity later than it is to go the other way round 😄

[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