Merge pull request #118 from Avasam/patch-1

complete `scipy.fft.dct`, and add relevant type aliases

Author: jorenham

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

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: ...