A new home for pyarrow-stubs? #45919

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.

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.

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

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.

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!

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.

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.

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.

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.

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?

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

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.

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.

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

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!

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.

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
-------------------------------------------------------------------------------