Repository navigation
Start annotating ufuncs as np.ufunc instances - #349
Merged
Merged
Conversation
Ufuncs are not simple functions like the type stubs so far have said they are. They are instead instances of the `np.ufunc` class. For typing purposes it is useful to define generic typing-only subtypes for each combination of `nin` and `nout`, which then allow defining the `pyerfa` ufuncs by specifying the type parameters. Numpy does that for a few `nin` and `nout` values, but the implementation is private and `pyerfa` has ufuncs with `nin` and `nout` not covered by numpy, so it is necessary to implement the generic ufunc types ourselves. `(nin, nout)=(0, 1)` is the simplest case and the only ufunc (not gufunc) in `pyerfa` like that is `ufunc.zpv()`. It now has better annotations that (among other things) distinguish when does it return an `np.ndarray` and when does it return an array scalar. For now the annotations ignore the parameters of the ufunc interface except `out` because only `out` is used in the existing tests.
eerovaher
commented
Oct 5, 2026
Collaborator
Author
There was a problem hiding this comment.
This template is based on the implementation of the typing-only _PyFunc_Nin1_Nout1 in current numpy main.
eerovaher
commented
Oct 5, 2026
Comment on lines
+26
to
+27
| @overload | ||
| def __call__(self, /, out: NDArray[ReturnT] | tuple[NDArray[ReturnT]]) -> NDArray[ReturnT]: ... |
Collaborator
Author
There was a problem hiding this comment.
Note that this overload does not yet try to track the shapes of the arrays.
avalentino
approved these changes
Oct 5, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
I was planning to remove all explicit uses of
Anyfrom the generated type stubs before starting to work on the shapes of the ufunc outputs, but it has now become apparent that the type annotations have to be able to distinguish when does a ufunc return annp.ndarrayand when does it return an array scalar or else type checkers will report spurious errors intests/test_ufunc.py.The type stubs
erfa_generatorcreates in currentmaindeclare ufuncs to be functions, but at runtime they are actually instances of thenp.ufuncclass. For typing purposes it is useful to define generic typing-only subtypes for each combination ofninandnout, which then allow defining thepyerfaufuncs by specifying the type parameters. Numpy defines suitable generic types for a fewninandnoutvalues (in https://github.com/numpy/numpy/blob/v2.5.3/numpy/_typing/_ufunc.pyi), but the implementation is private andpyerfahas ufuncs withninandnoutnot covered by numpy, so it is necessary to implement the generic ufunc types ourselves.(nin, nout) = (0, 1)is the simplest case and the only ufunc (not gufunc) inpyerfalike that isufunc.zpv(). This PR therefore annotates onlyufunc.zpv()as anp.ufuncinstance, but the code implemented here can be reused for other(nin, nout)combinations (although not without any changes). For now the annotations ignore the parameters of the ufunc interface exceptoutbecause onlyoutis used in the existing tests.