diff --git a/scipy-stubs/_typing.pyi b/scipy-stubs/_typing.pyi index 3d6dba15..44592345 100644 --- a/scipy-stubs/_typing.pyi +++ b/scipy-stubs/_typing.pyi @@ -93,6 +93,8 @@ CorrelateMode: TypeAlias = Literal["valid", "same", "full"] # scipy literals NanPolicy: TypeAlias = Literal["raise", "propagate", "omit"] Alternative: TypeAlias = Literal["two-sided", "less", "greater"] +DCTType: TypeAlias = Literal[1, 2, 3, 4] +NormalizationMode: TypeAlias = Literal["backward", "ortho", "forward"] # used in `scipy.linalg.blas` and `scipy.linalg.lapack` @type_check_only diff --git a/scipy-stubs/fft/_realtransforms.pyi b/scipy-stubs/fft/_realtransforms.pyi index 4d9cdd89..70c6386b 100644 --- a/scipy-stubs/fft/_realtransforms.pyi +++ b/scipy-stubs/fft/_realtransforms.pyi @@ -1,83 +1,88 @@ -from scipy._typing import Untyped +import numpy.typing as npt +from numpy._typing import _ArrayLikeNumber_co +from scipy._typing import DCTType, NormalizationMode, Untyped def dctn( x: Untyped, - type: int = 2, + type: DCTType = 2, s: Untyped | None = None, axes: Untyped | None = None, - norm: Untyped | None = None, + norm: NormalizationMode | None = None, overwrite_x: bool = False, - workers: Untyped | None = None, + workers: int | None = None, *, - orthogonalize: Untyped | None = None, + orthogonalize: bool | None = None, ) -> Untyped: ... def idctn( x: Untyped, - type: int = 2, + type: DCTType = 2, s: Untyped | None = None, axes: Untyped | None = None, - norm: Untyped | None = None, + norm: NormalizationMode | None = None, overwrite_x: bool = False, - workers: Untyped | None = None, - orthogonalize: Untyped | None = None, + workers: int | None = None, + orthogonalize: bool | None = None, ) -> Untyped: ... def dstn( x: Untyped, - type: int = 2, + type: DCTType = 2, s: Untyped | None = None, axes: Untyped | None = None, - norm: Untyped | None = None, + norm: NormalizationMode | None = None, overwrite_x: bool = False, - workers: Untyped | None = None, - orthogonalize: Untyped | None = None, + workers: int | None = None, + orthogonalize: bool | None = None, ) -> Untyped: ... def idstn( x: Untyped, - type: int = 2, + type: DCTType = 2, s: Untyped | None = None, axes: Untyped | None = None, - norm: Untyped | None = None, + norm: NormalizationMode | None = None, overwrite_x: bool = False, - workers: Untyped | None = None, - orthogonalize: Untyped | None = None, + workers: int | None = None, + orthogonalize: bool | None = None, ) -> Untyped: ... + +# We could use overloads based on the type of x to get more accurate return type +# see https://github.com/jorenham/scipy-stubs/pull/118#discussion_r1807957439 def dct( - x: Untyped, - type: int = 2, - n: Untyped | None = None, + x: _ArrayLikeNumber_co, + type: DCTType = 2, + n: int | None = None, axis: int = -1, - norm: Untyped | None = None, + norm: NormalizationMode | None = None, overwrite_x: bool = False, - workers: Untyped | None = None, - orthogonalize: Untyped | None = None, -) -> Untyped: ... + workers: int | None = None, + orthogonalize: bool | None = None, +) -> npt.NDArray[Untyped]: ... def idct( x: Untyped, - type: int = 2, - n: Untyped | None = None, + type: DCTType = 2, + n: int | None = None, axis: int = -1, - norm: Untyped | None = None, + norm: NormalizationMode | None = None, overwrite_x: bool = False, - workers: Untyped | None = None, - orthogonalize: Untyped | None = None, + workers: int | None = None, + orthogonalize: bool | None = None, ) -> Untyped: ... def dst( x: Untyped, - type: int = 2, - n: Untyped | None = None, + type: DCTType = 2, + n: int | None = None, axis: int = -1, - norm: Untyped | None = None, + norm: NormalizationMode | None = None, overwrite_x: bool = False, - workers: Untyped | None = None, - orthogonalize: Untyped | None = None, + workers: int | None = None, + orthogonalize: bool | None = None, ) -> Untyped: ... def idst( x: Untyped, - type: int = 2, - n: Untyped | None = None, + type: DCTType = 2, + n: int | None = None, axis: int = -1, - norm: Untyped | None = None, + norm: NormalizationMode | None = None, overwrite_x: bool = False, - workers: Untyped | None = None, - orthogonalize: Untyped | None = None, + workers: int | None = None, + orthogonalize: bool | None = None, ) -> Untyped: ...