POC: [Python] try out stubgen-pyx for auto-generated typing stubs

[jorisvandenbossche]

I recently stumbled upon https://github.com/jon-edward/stubgen-pyx, and this might be an interesting option to explore for the pyarrow typing situation.

It is essentially an improved stubgen specifically for cython code (in the past, when we explored options here, the existing stubgens didn't work sufficiently well for cython to be considered). And also https://github.com/NVIDIA/cuda-python is (experimenting with) using this.

This PR, just for illustration purposes, includes the full generated output for the entire pyarrow package (with .pyi files located next to the cython file, but this could of course be put elsewhere as well): stubgen-pyx python/pyarrow/*.pyx.

The interesting part that I learned from looking at the cuda-python code, is that you can actually use python type annotations inside the cython code, and that way we could have inline annotations. Whatever information the tool cannot extract from the cython information, we can augment with adding python type annotations. So essentially that would give an option to have all the custom type annotations inline (which the current PRs for #32609, such as #48622, are putting in separate stub files), and have the stub files be fully generated (i.e. allow us to have only one source of truth / one place to edit).

To illustrate that part, I am going to open a separate PR with just the generated stub file for types.pxi (so it can be compared to #48622), and with a few customizations to improve the generated stubs -> #51681

[rok]

It would be really great to automatically extract stubs from cython!

Whatever information the tool cannot extract from the cython information, we can augment with adding python type annotations.

By that you mean we add more precise type information or do you mean we additionally annotate symbols that are dynamically created at import time (e.g. all kernels)?
If we want to annotate dynamically generated symbols directly we could do that in a separately versioned .pyi files and dynamically inject them with docstrings at wheel build time as we do now. But then, maintenance wise, did we gain enough?

[jorisvandenbossche]

Specifically for the example of dynamically create compute kernels: those are actually pure-python functions (generated in compute.py, not in _compute.pyx), and so I assume it should be possible to update that generation code to directly include type annotations as well (e.g., we now also generate a Signature to add to the generated function.

But in general, I was talking about the cython files. As an example, see the second commit in #51681, which adds some additional type annotations in the cython file (for things that the stubgen cannot infer from the cython type defs), and so are then included in the generated stub file.

[rok]

Specifically for the example of dynamically create compute kernels: those are actually pure-python functions (generated in compute.py, not in _compute.pyx), and so I assume it should be possible to update that generation code to directly include type annotations as well (e.g., we now also generate a Signature to add to the generated function.

I suppose we could generate stubs for compute kernels with stubgen. But we couldn't annotate the source as in cython and then we lose the nice quality of being able to gradually annotate the source and export the improved annotations to stubs. I suppose we better keep the pyarrow-stubs approach for dynamically generated functions. What do you think?

But in general, I was talking about the cython files. As an example, see the second commit in #51681, which adds some additional type annotations in the cython file (for things that the stubgen cannot infer from the cython type defs), and so are then included in the generated stub file.

Right, for these it should be great.

Do these two styles (dynamically generated and cython) overlap in the namespace a lot? Would this make the annotation complex?

[rok]

To be more explicit, if we use auto-generated-stubs-from-cython we can proceed like so:

1. auto-generate the stubs for the cython and python surface with stubgen-pyx
2. cover the dynamically generated surface (kernels) with manual stubs
3. enrich auto-generated cython/ptyhon stubs

Both dynamic and cython/python stubs will be included into git without docstrings (we agreed on this somewhere in or in vicinity of #32609). We will want to use our existing docstring populating script or docstring-adder to inject docstrings into stubs at wheel build time.

What I'm a little worried about is potential namespaces (dynamic, cython and python) overlap which might require extra logic. If this logic gets too complex I would be against using the mix of auto-generated and manual stubs (and slightly prefer manual-only). This would be to keep maintenance cost of stubs as low as possible.

[rok]

Alternative process would be:

1. Cover all stubs with manual stubs
2. Enable conversion form cython/python into stubs with stubgen-pyx
3. Gradually enrich cython/python code with types from manual stubs (this will have the benefit of using internal types from the start) while removing manual stubs

[jorisvandenbossche]

I suppose we could generate stubs for compute kernels with stubgen

Maybe I am not familiar enough with how this all works, but why do we need stubs for python files? Why not add type annotations in compute.py itself? Or the problem for this file specifically is that it contains dynamically generated functions, and type checkers need to get that information statically?

Then we can still only provide a stub file for compute.py, and do inline annotations for the other python files?

Both dynamic and cython/python stubs will be included into git without docstrings (we agreed on this somewhere in or in vicinity of #32609). We will want to use our existing docstring populating script or docstring-adder to inject docstrings into stubs at wheel build time.

Why would we explicitly remove docstrings first (since the stubgen can include them) to then re-inject them with some build time magic? Is that for reducing the size of the git tracked sources? (as long as the files are auto-generated, this might not be too much of a problem?)

(I checked mentions of "docstring" in #32609, but those are all rather about wanting docstrings)

cover the dynamically generated surface (kernels) with manual stubs

I think we could relatively easily also generate those stubs? (we already generate the python code with signature on the fly, then we can also generate annotations on the fly and save that with some extra script to a file)

What I'm a little worried about is potential namespaces (dynamic, cython and python) overlap which might require extra logic. If this logic gets too complex I would be against using the mix of auto-generated and manual stubs (and slightly prefer manual-only). This would be to keep maintenance cost of stubs as low as possible.

I think the goal should be "no manual stubs" if we go stubgen(-pyx)

[jorisvandenbossche]

Alternative process would be:

1. Cover all stubs with manual stubs
2. Enable conversion form cython/python into stubs with stubgen-pyx
3. Gradually enrich cython/python code with types from manual stubs (this will have the benefit of using internal types from the start) while removing manual stubs

IMO the problem with this approach is that this first item is the bottleneck. The current manual stubs are super complex, and hard to all get in (the stalled PRs).
(unless you would go for starting with lighter manual stubs, and then also gradually enrich those, but lighter manual stubs are essentially the auto-generated ones, potentially)

My (not validated) proposal would be:

1. Start with adding the auto-generated stubs for cython and compute.py, and inline annotations for other python files.
2. Make minimal changes in the cython and python sources (eg add basic return types) to make the annotations minimally useful
3. Ensure we have good test coverage for the annoations (run type checkers on our tests, and add specific type-annotation tests)
4. Re-enable pyarrow as being py.typed, and try to have at least this in the next release
5. Gradually move richer annotations from the current manual stubs to the inline sources (so that the auto-generated annotations get updated as well)

As I mentioned in the pyarrow dev meeting last week, one aspect that might actually not work out with the outline above is the last step: I don't know if stubgen-pyx (well, actually cython) will support the more complex parts of the type annotations syntax in python functions in cython files. That is something that would have to be tested, and if that does not work, something to be discussed how important those additional features are (vs the maintenance benefit of being able to auto-generate stubs)

[rok]

I suppose we could generate stubs for compute kernels with stubgen

Maybe I am not familiar enough with how this all works, but why do we need stubs for python files? Why not add type annotations in compute.py itself? Or the problem for this file specifically is that it contains dynamically generated functions, and type checkers need to get that information statically?

Yes, kernels are never available statically. Can we technically put stub information into .py files? Seems unconventional.

Then we can still only provide a stub file for compute.py, and do inline annotations for the other python files?

Yes, and we have two separate ways and locations for storing annotations. I am not against this, but I'd note it complicates things.

Both dynamic and cython/python stubs will be included into git without docstrings (we agreed on this somewhere in or in vicinity of #32609). We will want to use our existing docstring populating script or docstring-adder to inject docstrings into stubs at wheel build time.

Why would we explicitly remove docstrings first (since the stubgen can include them) to then re-inject them with some build time magic? Is that for reducing the size of the git tracked sources? (as long as the files are auto-generated, this might not be too much of a problem?)

(I checked mentions of "docstring" in #32609, but those are all rather about wanting docstrings)

Found the thread #45919 (reply in thread)

cover the dynamically generated surface (kernels) with manual stubs

I think we could relatively easily also generate those stubs? (we already generate the python code with signature on the fly, then we can also generate annotations on the fly and save that with some extra script to a file)

Ok, but we can't have rich types. I don't think we want that.

What I'm a little worried about is potential namespaces (dynamic, cython and python) overlap which might require extra logic. If this logic gets too complex I would be against using the mix of auto-generated and manual stubs (and slightly prefer manual-only). This would be to keep maintenance cost of stubs as low as possible.

I think the goal should be "no manual stubs" if we go stubgen(-pyx)

If we can somehow annotate dynamically generated symbols then we can have "no manual stubs" otherwise we will always lose rich types.

IMO the problem with this approach is that this first item is the bottleneck. The current manual stubs are super complex, and hard to all get in (the stalled PRs).
(unless you would go for starting with lighter manual stubs, and then also gradually enrich those, but lighter manual stubs are essentially the auto-generated ones, potentially)

I think the stalled PRs is mostly on me not having full time availability for this.

As for the proposed types - do you believe they are overly complex for what's needed downstream? Let's try to set a target type complexity and then find a way to it. From memory types are pretty close to what you need if you want to e.g. infer what will be output type of a function given an input type. If this is something we ultimately want I don't think we can simplify much.
If we now throw the internal types away and start with simple types we'll have to recreate internal types later to get to desired functionality which I would like to avoid.

[rok]

@dangotbanned we could use your opinion here. If we go with auto-generated stubs (stubgen-pyx) plus inline annotations in Cython, we may not be able to express everything the current manual stubs do. From a downstream (narwhals) perspective, what level of precision would you eventually want: generics like Array[Int64Scalar] / ChunkedArray[T], overloads where the return type depends on the arguments (e.g. pa.array, list_), Literal narrowing? Which are must-haves and which are nice-to-haves?

[jorisvandenbossche]

Yes, kernels are never available statically. Can we technically put stub information into .py files? Seems unconventional.

Well, that's called type annotations? ;) I mean, whether a function is defined statically or defined by code, doesn't really matter for python (of course it matters to static analyzers). Just as we currently add a name, module, signature, etc to the generated function, we can also add a signature (in the compute.py source code).
That won't be useful for static analyzers (just as those currently also cannot even see that those functions exist), but generating a stub file from the dynamic python functions should be possible?

I think we could relatively easily also generate those [compute.py] stubs? (we already generate the python code with signature on the fly, then we can also generate annotations on the fly and save that with some extra script to a file)

Ok, but we can't have rich types. I don't think we want that.

Again as a non-typing expert, can you give a bit more details why that is not possible? Those rich types cannot be used as python variables? (i.e. you cannot write them as wrapper.__annotations__ = {"arg": ...}?)

Then we can still only provide a stub file for compute.py, and do inline annotations for the other python files?

Yes, and we have two separate ways and locations for storing annotations. I am not against this, but I'd note it complicates things.

AFAIK that is quite standard practice for libraries with compiled modules? Stub files for the compiled modules, and inline annotations for the python modules? (at least that is what pandas does, but I haven't been involved myself in setting that up) I do see that eg numpy also uses stub files for the python files.

If we would use inline annotations for python files in general, then we would have two "ways" anyhow. In that case compute.py would just be an exception on the rule (having a stub instead of inline annotations), rather than adding a separate way to store annotations.

Found the thread #45919 (reply in thread)

That thread is from a context discussing the manual stubs, though, were including docstrings means they have to be kept up to date manually. If we auto-generate stubs, I think that is a different situation.

Reading that thread, I like Marco's comment at #45919 (reply in thread), though, quoting: "start simple". Whether we go with auto-generated stubs (as discussed here) or with the manual stubs approach, I would highly recommend starting simple (and the current PRs are far from that).

I think the stalled PRs is mostly on me not having full time availability for this.

That is certainly a large part of it (and not only your time, but time of anyone who would want to move this forward). But those types being quite complicated also does not help. PRs adding simple type annotations could be reviewed / merged a lot more easily.

As for the proposed types - do you believe they are overly complex for what's needed downstream?

I don't know exactly what is needed downstream, so I cannot comment on "overly complex for what's needed", I can only say that I as a maintainer find them very complex.
And I also know that for a lot of downstream usage, already having basic types would be a gigantic step forward from having no types. For example, to enable a large part of tab completion / method calls evaluation in IDEs, just having basic return types would (AFAIU) already do a lot (e.g. the fact that an IDE can know that the result of pa.table() is a pa.Table and which methods are available for such a table, and for each of those methods what the accepted arguments are and the return type, etc)

Let's try to set a target type complexity and then find a way to it. From memory types are pretty close to what you need if you want to e.g. infer what will be output type of a function given an input type. If this is something we ultimately want I don't think we can simplify much.
If we now throw the internal types away and start with simple types we'll have to recreate internal types later to get to desired functionality which I would like to avoid.

I still think that even if we want to have such powerful type annotations such that analyzers can infer pyarrow types eventually, we should start with simple types (to ensure we provide most of the value on a shorter term). That might indeed mean multiple iterations on the types, and potentially changing approach along the way. But I think the end result (and the value of the intermediate steps) will be better.

[dangotbanned]

Thanks for the ping @rok

Which are must-haves and which are nice-to-haves?

This is a great question.
I can see there are already some opinions above, but would like to suggest a different starting point.

Questions

If you use these generated stubs and run a type checker ¹ with a strict config on the pyarrow test suite:

1. How many (type checking) errors do they report from correct use?
2. How many (runtime) errors do the stubs help catch?

These are the easiest things to measure, given that you have the cases already written.

However, you might find that the tests don't reflect the way users actually write code using pyarrow.
Getting the answer to that would look something like (https://github.com/numpy/numpy/tree/db6d7c2df8ecf7e1507209055037327512892b86/numpy/typing/tests/data)

Opinion

I think if generated stubs create an equal or better user-experience than manually written ones - then they are a great choice.
However, I expect that "simple" or otherwise generated types won't capture the real complexity of your API and mean lots of # type: ignores on correct code

Footnotes

1. Preferably multiple, because they will behave differently from what I'm seeing in this PR ↩

[dangotbanned]

@jorisvandenbossche do you know why most functions/methods do not have a return type?

I feel like maybe there's a configuration issue.

These should be:

    def __str__(self) -> str: ...
    def __repr__(self) -> str: ...
    def to_table(self, use_threads: bool=True) -> somewhere.Table:

Probably priority number 1 is fixing this

[jorisvandenbossche]

Because those functions do not have a return type in the source (no python type annotation return type and no cython return type; stubgen will not try to infer it from the body of the function)

But those can easily be added, see the second commit in #51681 as an illustration

[AlenkaF]

I have made some tests locally (stubgen-pyx 0.2.23, mypy 2.4.0 (stubgen), pyright 1.1.414 and Cython 3.2.4)
on

• lib.pyx/pyarrow.lib and
• _compute.pyx/pyarrow.compute.

This message is super long, sorry, so I added a short summary table first:

stubgen-pyx

• (+) If using inline annotations we could automatically regenerate stubs on code changes
• (+) stubgen-pyx keeps @overload in the generated stubs
• (−) Private classes are missing from the generated stub (e.g. _Weakrefable).
  The --include-private flag fixes that, but other internals then leak in (e.g. CFixedWidthType)
• (−) Missing return and argument types need to be added in Cython (inline annotations)
• (−) **kwargs annotation seems to be dropped from the stubs (didn't dive into this)

Note: output file has to match the module (lib.pyx -> pyarrow/lib.pyi)

stubgen --inspect-mode

• (+) Good stubs created for the compute module

But for Cython files:

• (−) Produces draft stubs that would have to be updated and maintained by hand
• (−) Adds an import _cython_3_2_4 to the stubs, which can't be resolved and changes with every Cython upgrade.
• (−) Module-level Cython functions (cython_function_or_method) are written as variables (e.g. null: _cython_3_2_4.cython_function_or_method).
• (−) Adds __init__ methods inherited at runtime, some with an erroneous @classmethod.

	stubgen-pyx	stubgen --inspect-mode
Properties	Kept as properties, no return type	Turned into writable attributes typed Incomplete
Overloads	Kept if added to the source	Fake overloads built from docstring signatures and examples
Docstrings	Can be added in full or skipped with --exclude-docstrings	Parsed for signatures; stubs become noisy and can contain invalid syntax (see below)
Dunders	Only the ones in the source	Also Cython internals (__reduce_cython__, __pyx_vtable__, ...)
Compute kernels	Not possible (kernels are generated at import time in compute.py); it does generates options classes which are needed!	All kernels generated with argument names and keyword-only markers; options parameters typed from their defaults (bool, int), but can cause false positives (e.g. q=0.5 -> float)

Example of invalid syntax from docstring parsing (stubgen --inspect-mode):

@classmethod
def __init__(cls, dense_union<a: fixed_size_binary[10] = ..., b: string = ...) -> Any: ...

Type-checking the generated lib stubs

Note: For inspect mode, the 10 lines with invalid syntax from docstring parsing had to be patched first, otherwise mypy stops on the syntax errors.

	stubgen-pyx	stubgen --inspect-mode
Total	54	364
Undefined names	42	178
Enum members annotated	—	73
__init__ from docstring parsing	—	51
Overloads from docstring parsing	—	57
Incompatible signatures in the PyArrow API (e.g. from_arrays)	6	4
Other (imports, ...)	6	1

Runing pyright with typeCheckingMode = "basic" (as currently configured in pyproject.toml) gives the same undefined-name errors as mypy but misses the incompatible signatures in the PyArrow API (see reportIncompatibleMethodOverride default). With standard mode they are reported together with one more: Method "dictionary_encode" overrides class "Array" in an incompatible manner.

Mix of stubgen --inspect-mode and stubgen-pyx for the compute module

I ran stubgen --inspect-mode on pyarrow.compute and stubgen-pyx on _compute.pyx, and the two complement each other. The inspect stub (compute.pyi) contains all the kernels, with correct signatures but missing types, and re-exports the options classes (e.g. ArraySortOptions, AssumeTimezoneOptions) from pyarrow._compute. Their definitions come from the stubgen-pyx _compute.pyi. The two are connected by a plain import, so no merging is needed.

compute.pyi itself has no errors in either mypy or pyright. pyright reports only on the file it is run on, while mypy also reports errors in the imported stubs: about 80 in _compute.pyi. Most of those come from the private classes. The rest:

• A C type alias leaks into the stub (CRegisterUdf, 6 errors).
• c_bool in a public signature, it should be bool.
• A bug in the code is also found: FunctionOptions.__eq__ raises TypeError for a different type instead of returning False (pc.ScalarAggregateOptions() == 5). CacheOptions.__eq__ has the same pattern.

Run type checkers on test_types.py

I have also used type checkers (mypy --check-untyped-defs and pyright on standard mode) with generated stubs and have run them on test_types.py. I found that:

• 1 of 19 TypeError from the test asserts has been found (19 out of 48 pytest.raises expect TypeError)
• 15 pyright/16 mypy TypeError's have been found if I used inline annotations in types.pxi similar to changes in POC: [Python] try out stubgen-pyx for auto-generated typing stubs - types.pxi #51681 (created them with an agent to have a quick test, added 16 signatures, 2 overloads and 3 return types).
• pyright found 40 false positives due to generated __init__.pyi not re-exporting the submodules (which is easily fixed)
• mypy found 6 false positives from the sparse_factories/dense_factories loops
• approx 10 mypy/ 2 pyright errors from optional imports (dateutil, etc.)

Run type checkers on test_compute.py

I have left Claude to make a change in _make_signature and _wrap_function that adds annotations to the sum kernel on runtime. The result of running type checkers on test_compute.py with and without added sum annotations on the mix of stubgen and stubgen-pyx generated stubs is as follows:

• generated stubs found 65 (mypy)/174 (pyright) errors,
• 3 real errors caught by both checkers, all structural: missing argument, unexpected keyword and missing indices,
• about 50 false positives (without injected type annotations)
  • Options parameters marked as required (20),
  • types guessed from default values (30)
• one rule in _make_signature removed 20 false positives from Options parameter, but at the cost of 1 real catch. Overloads could fix that but they can't be injected in the stubs this way. A separate script would be needed for that.
• Typing addition in _wrap_function works (adds stubs) and adds no false positives. Annotations must be strings, because stubgen's inspect mode drops runtime union objects.

We would also need to add annotations to option classes in _compute.pyx but I have not tested yet if Cython keeps annotations at runtime.

I hope this is helpful. I am just learning about typing so I have no experience on this field but feel the results are pretty good and would happily work on the generated simple type stubs first.

What do other think?

Note: I used Claude heavily in the last, test_compute.py section, as I felt it does not need a thorough look at right now but a fast test would do. On other fronts I mainly used Claude for guidance, error summary and notes review).

[dangotbanned]

#51680 (comment)

Thanks @AlenkaF. Two test modules is a a good start, but in reference to (#51680 (comment)) I was thinking to run type checkers on the entire tests directory?

I felt it does not need a thorough look at right now but a fast test would do.

If you want it to be fast, use ty and pyrefly.
If you want it to be thorough, use those + pyright, mypy too.

More recent versions of mypy will run much faster on a warm cache

[AlenkaF]

Thanks @dangotbanned!
I could use stubs from this PR (stubgen-pyx for Cython files) and use stubgen --inspect-mode additionaly for compute and then run the checkers on full test dir. Without any new inline annotations and changes to kernel signatures. Does that sound OK?

[rok]

For completeness I've done a similar exploration as Alenka yesterday. See rok#62.
I ran type checkers on the test suite and the result (IIRC) was somewhere between 1k-3k issues depending on the type checker. Which is comparable to what pyarrow-stubs initially showed.

I think it's valuable to explore autogen options now in case there is some upside.

However given I've already cleaned the test suite to make all type checkers happy I'm currently in favour of completing current (pyarrow-stubs port) PRs over auto generating.

[dangotbanned]

Thanks @dangotbanned! I could use stubs from this PR (stubgen-pyx for Cython files) and use stubgen --inspect-mode additionaly for compute and then run the checkers on full test dir. Without any new inline annotations and changes to kernel signatures. Does that sound OK?

Yep that sounds good.

I suppose if you get a good comparison between those - another thing you could try is comparing against @rok's most recent branch or (https://github.com/zen-xu/pyarrow-stubs)?
That would give you an idea of the gap between generated and (partially) handwritten stubs.

But most important is agreeing on something to test 🙂

[jorisvandenbossche]

However given I've already cleaned the test suite to make all type checkers happy I'm currently in favour of completing current (pyarrow-stubs port) PRs over auto generating.

That work to update the test suite can just be reused whichever type annotations we end up using (initially)?

[rok]

However given I've already cleaned the test suite to make all type checkers happy I'm currently in favour of completing current (pyarrow-stubs port) PRs over auto generating.

That work to update the test suite can just be reused whichever type annotations we end up using (initially)?

I suspect some of these corrections will depend on types we use. E.g. if you do a expect raise block on array[int].equals(array[float]) it will fail type checkers on generic types, but not on simple types.