From 173416a8985f9748c8c63fb3dd422f28992f0a1c Mon Sep 17 00:00:00 2001 From: Eero Vaher Date: Thu, 8 Oct 2026 23:57:26 +0200 Subject: [PATCH] Hardcode `Ufunc_Nin0_Nout1` in ufunc stub template A good way to annotate ufuncs is to generate typing-only generic types for each `(nin,nout)` combination and to declare individual ufuncs as instances of those generic types with appropriate values for the type parameters. The ufuncs with `nout==1` have a usable `at()` method and their `__call__()` method signature is also somewhat different from ufuncs with `nout!=1`, so it will be useful to have `erfa_generator` create the generic ufunc subtypes from two separate templates. However, if a ufunc has `(nin,nout)==(0,1)` then its `__call__()` method is analogous to the other ufuncs with `nout==1`, but it nonetheless does not have a usable `at()` method. That means the corresponding generic ufunc subtype is the simplest to implement by writing it directly in the type stub template instead of having `erfa_generator` create it from a separate template. --- erfa/ufunc.pyi.templ | 41 +++++++++++++++++++++++++++++++++++-- erfa/ufunc_no_methods.templ | 37 --------------------------------- erfa_generator.py | 12 ----------- 3 files changed, 39 insertions(+), 51 deletions(-) delete mode 100644 erfa/ufunc_no_methods.templ diff --git a/erfa/ufunc.pyi.templ b/erfa/ufunc.pyi.templ index ffb9847..86e2931 100644 --- a/erfa/ufunc.pyi.templ +++ b/erfa/ufunc.pyi.templ @@ -79,12 +79,49 @@ type YMDFDType = CustomVoidDType[np.intc] dt_ymdf: Final[YMDFDType] +@type_check_only +class Ufunc_Nin0_Nout1[NTypes: Literal[0, 1], ReturnT: np.generic](np.ufunc): # type: ignore[misc] + @property + @override + def identity(self) -> None: ... + @property + @override + def nin(self) -> Literal[0]: ... + @property + @override + def nout(self) -> Literal[1]: ... + @property + @override + def nargs(self) -> Literal[1]: ... + @property + @override + def ntypes(self) -> NTypes: ... + @property + @override + def signature(self) -> None: ... + @override + @overload + def __call__(self, /, out: tuple[None] | None = None) -> ReturnT: ... + @overload + def __call__(self, /, out: EllipsisType) -> np.ndarray[tuple[()], np.dtype[ReturnT]]: ... + @overload + def __call__(self, /, out: NDArray[ReturnT] | tuple[NDArray[ReturnT]]) -> NDArray[ReturnT]: ... + @override + def at(self, a: Never, indices: Never, /) -> NoReturn: ... # type: ignore[override] + @override + def accumulate(self, array: Never, /) -> NoReturn: ... # type: ignore[override] + @override + def reduce(self, array: Never, /) -> NoReturn: ... # type: ignore[override] + @override + def reduceat(self, array: Never, /, indices: Never) -> NoReturn: ... # type: ignore[override] + @override + def outer(self, A: Never, B: Never, /) -> NoReturn: ... # type: ignore[override] + + def get_leap_seconds() -> np.ndarray[tuple[int], np.dtype[LeapSecondDType]]: ... def set_leap_seconds(leap_seconds: ArrayLike) -> None: ... -$generic_ufunc_types - $funcs diff --git a/erfa/ufunc_no_methods.templ b/erfa/ufunc_no_methods.templ deleted file mode 100644 index e1f6b0b..0000000 --- a/erfa/ufunc_no_methods.templ +++ /dev/null @@ -1,37 +0,0 @@ -@type_check_only -class Ufunc_Nin${nin}_Nout$nout[NTypes: Literal[0, 1], ReturnT: np.generic](np.ufunc): # type: ignore[misc] - @property - @override - def identity(self) -> None: ... - @property - @override - def nin(self) -> Literal[$nin]: ... - @property - @override - def nout(self) -> Literal[$nout]: ... - @property - @override - def nargs(self) -> Literal[$nargs]: ... - @property - @override - def ntypes(self) -> NTypes: ... - @property - @override - def signature(self) -> None: ... - @override - @overload - def __call__(self, /, out: tuple[None] | None = None) -> ReturnT: ... - @overload - def __call__(self, /, out: EllipsisType) -> np.ndarray[tuple[()], np.dtype[ReturnT]]: ... - @overload - def __call__(self, /, out: NDArray[ReturnT] | tuple[NDArray[ReturnT]]) -> NDArray[ReturnT]: ... - @override - def at(self, a: Never, indices: Never, /) -> NoReturn: ... # type: ignore[override] - @override - def accumulate(self, array: Never, /) -> NoReturn: ... # type: ignore[override] - @override - def reduce(self, array: Never, /) -> NoReturn: ... # type: ignore[override] - @override - def reduceat(self, array: Never, /, indices: Never) -> NoReturn: ... # type: ignore[override] - @override - def outer(self, A: Never, B: Never, /) -> NoReturn: ... # type: ignore[override] diff --git a/erfa_generator.py b/erfa_generator.py index 32385ff..890e648 100644 --- a/erfa_generator.py +++ b/erfa_generator.py @@ -907,20 +907,8 @@ def main(srcdir: Path, templateloc: Path) -> None: ]), ) - ufunc_no_methods_template = Template( - (templateloc / "ufunc_no_methods.templ").read_text() - ) - generic_ufunc_types = [] - for nin, nout in {(f.nin, f.nout) for f in funcs if f.signature == "NULL"}: - if nin <= 0 and nout <= 1: - generic_ufunc_types.append( - ufunc_no_methods_template.substitute( - nin=nin, nout=nout, nargs=nin + nout - ) - ) _render_template( templateloc / "ufunc.pyi.templ", - generic_ufunc_types="\n".join(generic_ufunc_types), funcs="\n\n\n".join(func.ufunc_signature for func in funcs), )