Add type stubs for Cython classes in pywt._extensions

[NiklasMueller2001]

Currently, many public classes and functions are implemented in the Cython modules inside pywt._extensions (e.g. wavelist, Wavelet, ...). Their public API is currently documented in the pywavelet docs; however, when writing Python code that uses these functions or classes Python type checkers such as mypy or Pyright can not determine their signature and return types or even detect their existence.

Possible solution

To enable static type checking I propose to add .pyi interface files in the pywt._extensions module. I have already created a draft version that works in my locally built pywavelets installation. Please let me know if you think this would be a useful enhancement for the package.

# pywt/_extensions/_pywt.pyi

from typing import Literal, Optional

_WaveletFamily = Literal[
    "haar",
    "db",
    "sym",
    "coif",
    "bior",
    "rbio",
    "dmey",
    "gaus",
    "mexh",
    "morl",
    "cgau",
    "shan",
    "fbsp",
    "cmor",
]

_Kind = Literal["all", "continuous", "discrete"]

_Symmetry = Literal[
    "asymmetric",
    "near symmetric",
    "symmetric",
    "anti-symmetric",
    "unknown",
]

def wavelist(family: _WaveletFamily | None = ..., kind: _Kind = ...) -> list[str]: ...
def families(short: bool = ...) -> list[str]: ...

class Wavelet:
    def __init__(self, name: str, filter_bank=None): ...
    def __len__(self) -> int: ...
    @property
    def name(self) -> str: ...
    @property
    def dec_lo(self) -> list[float]: ...
    @property
    def dec_hi(self) -> list[float]: ...
    @property
    def rec_lo(self) -> list[float]: ...
    @property
    def rec_hi(self) -> list[float]: ...
    @property
    def rec_len(self) -> int: ...
    @property
    def dec_len(self) -> int: ...
    @property
    def family_number(self) -> int: ...
    @property
    def family_name(self) -> str: ...
    @property
    def short_family_name(self) -> str: ...
    @property
    def orthogonal(self) -> bool: ...
    @orthogonal.setter
    def orthogonal(self, value: bool) -> None: ...
    @property
    def biorthogonal(self) -> bool: ...
    @biorthogonal.setter
    def biorthogonal(self, value: bool) -> None: ...
    @property
    def symmetry(self) -> _Symmetry: ...
    @property
    def vanishing_moments_psi(self) -> int | None: ...
    @property
    def vanishing_moments_phi(self) -> int | None: ...
    @property
    def filter_bank(
        self,
    ) -> tuple[list[float], list[float], list[float], list[float]]: ...
    def get_filters_coeffs(
        self,
    ) -> tuple[list[float], list[float], list[float], list[float]]: ...
    @property
    def inverse_filter_bank(
        self,
    ) -> tuple[list[float], list[float], list[float], list[float]]: ...
    def get_reverse_filters_coeffs(
        self,
    ) -> tuple[list[float], list[float], list[float], list[float]]: ...

class ContinuousWavelet:
    @property
    def family_number(self) -> int: ...
    @property
    def family_name(self) -> str: ...
    @property
    def short_family_name(self) -> str: ...
    @property
    def orthogonal(self) -> bool: ...
    @orthogonal.setter
    def orthogonal(self, value: bool) -> None: ...
    @property
    def biorthogonal(self) -> bool: ...
    @biorthogonal.setter
    def biorthogonal(self, value: bool) -> None: ...
    @property
    def complex_cwt(self) -> bool: ...
    @complex_cwt.setter
    def complex_cwt(self, value: bool) -> None: ...
    @property
    def lower_bound(self) -> Optional[float]: ...
    @lower_bound.setter
    def lower_bound(self, value: float) -> None: ...
    @property
    def upper_bound(self) -> Optional[float]: ...
    @upper_bound.setter
    def upper_bound(self, value: float) -> None: ...
    @property
    def center_frequency(self) -> Optional[float]: ...
    @center_frequency.setter
    def center_frequency(self, value: float) -> None: ...
    @property
    def bandwidth_frequency(self) -> Optional[float]: ...
    @bandwidth_frequency.setter
    def bandwidth_frequency(self, value: float) -> None: ...
    @property
    def fbsp_order(self) -> Optional[int]: ...
    @fbsp_order.setter
    def fbsp_order(self, value: int) -> None: ...
    @property
    def symmetry(self) -> _Symmetry: ...

[rgommers]

Hi @NiklasMueller2001, thanks for suggesting this, that seems like a good thing to add indeed. This may be a bit of a project given the size of the API and the support for fused types (32/64-bit float and complex dtypes) in the core functions. It might fit in one large PR, or be broken into multiple PRs - whatever seems best to you.

Ideally it'd come with a CI job to verify that the added stubs are correct and complete. Given pyright is now unmaintained and mypy is slow and under-maintained, I think I'd prefer using pyrefly (it has a coverage report now: https://pyrefly.org/en/docs/report/).

[NiklasMueller2001]

I would preferably create one large PR and add commits incrementally to it. Just to clarify the scope of this project, when looking through the pywt._extensions module I identified the following functions/classes that could use some stub files:

_cwt.pyx:

• cwt_psi_single

_dwt.pyx:

• dwt_max_level
• dwt_coeff_len
• dwt_single
• idwt_single
• dwt_axis
• idwt_axis
• upcoef
• downcoef

_pywt.pyx:

• wavelist
• families
• DiscreteContinuousWavelet
• Wavelet
• ContinuousWavelet

_swt.pyx:

• swt_max_level
• swt
• swt_axis

[rgommers]

That plan sounds good to me. Those are the most important functions indeed. There is more in pure Python code though that is also untyped (everything in https://pywavelets.readthedocs.io/en/latest/ref/index.html). Perhaps that can be inferred easily though, e.g. with https://pyrefly.org/en/docs/stubgen/. It can also be left for later if the scope of that is too much.

[rgommers]

@jorenham pointed out to me that __all__ being dynamic prevents typecheckers from figuring out what is public, so changing that to a regular static list could already be quite helpful to see the API surface:

pywt/pywt/__init__.py

Line 26 in 45f51e3

__all__ = [s for s in dir() if not s.startswith('_')]

[NiklasMueller2001]

Interestingly, the API of public functions seem to be correctly propagated even with the current dynamic version with wildcard imports like this:

pywt/pywt/__init__.py

Lines 14 to 15 in 45f51e3

	from ._functions import *
	from ._multilevel import *

When import from pywt I can see the signature of public functions correctly:

import pywt

pywt.central_frequency  # type checking works for public function
pywt._integrate  # private functions are not known