✨ fft: sync (private) signatures of _realtransforms and _realtransforms_backend (#684)

Author: jorenham

scipy-stubs/fft/_realtransforms_backend.pyi

@@ -1,19 +1,110 @@
-from typing import TypeAlias
+from typing import TypeVar, overload
 
 import numpy as np
 import optype as op
 import optype.numpy as onp
+import optype.numpy.compat as npc
 
 from ._realtransforms import dct, dctn, dst, dstn, idct, idst
 from scipy._typing import AnyShape, DCTType, NormalizationMode
 
 __all__ = ["dct", "dctn", "dst", "dstn", "idct", "idctn", "idst", "idstn"]
 
-_RealND: TypeAlias = onp.ArrayND[np.float32 | np.float64 | np.longdouble]
+_ShapeT = TypeVar("_ShapeT", bound=tuple[int, ...])
+_DTypeT = TypeVar("_DTypeT", bound=np.dtype[np.float32 | np.float64 | np.longdouble | npc.complexfloating])
 
 # NOTE: Unlike the ones in `scipy.fft._realtransforms`, `orthogonalize` is keyword-only here.
+
+#
+@overload
 def idctn(
-    x: onp.ToComplexND,
+    x: onp.CanArrayND[npc.integer, _ShapeT],
+    type: DCTType = 2,
+    s: onp.ToInt | onp.ToIntND | None = None,
+    axes: AnyShape | None = None,
+    norm: NormalizationMode | None = None,
+    overwrite_x: op.CanBool = False,
+    workers: onp.ToInt | None = None,
+    *,
+    orthogonalize: op.CanBool | None = None,
+) -> onp.Array[_ShapeT, np.float64]: ...
+@overload
+def idctn(
+    x: onp.CanArrayND[np.float16, _ShapeT],
+    type: DCTType = 2,
+    s: onp.ToInt | onp.ToIntND | None = None,
+    axes: AnyShape | None = None,
+    norm: NormalizationMode | None = None,
+    overwrite_x: op.CanBool = False,
+    workers: onp.ToInt | None = None,
+    *,
+    orthogonalize: op.CanBool | None = None,
+) -> onp.Array[_ShapeT, np.float32]: ...
+@overload
+def idctn(
+    x: onp.ToJustFloat64_ND,
+    type: DCTType = 2,
+    s: onp.ToInt | onp.ToIntND | None = None,
+    axes: AnyShape | None = None,
+    norm: NormalizationMode | None = None,
+    overwrite_x: op.CanBool = False,
+    workers: onp.ToInt | None = None,
+    *,
+    orthogonalize: op.CanBool | None = None,
+) -> onp.ArrayND[np.float64]: ...
+@overload
+def idctn(
+    x: onp.ToFloatND,
+    type: DCTType = 2,
+    s: onp.ToInt | onp.ToIntND | None = None,
+    axes: AnyShape | None = None,
+    norm: NormalizationMode | None = None,
+    overwrite_x: op.CanBool = False,
+    workers: onp.ToInt | None = None,
+    *,
+    orthogonalize: op.CanBool | None = None,
+) -> onp.ArrayND[npc.floating]: ...
+
+#
+@overload
+def idstn(
+    x: onp.CanArrayND[npc.integer, _ShapeT],
+    type: DCTType = 2,
+    s: onp.ToInt | onp.ToIntND | None = None,
+    axes: AnyShape | None = None,
+    norm: NormalizationMode | None = None,
+    overwrite_x: op.CanBool = False,
+    workers: onp.ToInt | None = None,
+    *,
+    orthogonalize: op.CanBool | None = None,
+) -> onp.Array[_ShapeT, np.float64]: ...
+@overload
+def idstn(
+    x: onp.CanArrayND[np.float16, _ShapeT],
+    type: DCTType = 2,
+    s: onp.ToInt | onp.ToIntND | None = None,
+    axes: AnyShape | None = None,
+    norm: NormalizationMode | None = None,
+    overwrite_x: op.CanBool = False,
+    workers: onp.ToInt | None = None,
+    *,
+    orthogonalize: op.CanBool | None = None,
+) -> onp.Array[_ShapeT, np.float32]: ...
+@overload
+def idstn(
+    x: onp.CanArray[_ShapeT, _DTypeT],
+    type: DCTType = 2,
+    s: onp.ToInt | onp.ToIntND | None = None,
+    axes: AnyShape | None = None,
+    norm: NormalizationMode | None = None,
+    overwrite_x: op.CanBool = False,
+    workers: onp.ToInt | None = None,
+    *,
+    orthogonalize: op.CanBool | None = None,
+) -> np.ndarray[_ShapeT, _DTypeT]: ...
+@overload
+def idstn(
+    x: onp.ToJustFloat64_ND,
     type: DCTType = 2,
     s: onp.ToInt | onp.ToIntND | None = None,
     axes: AnyShape | None = None,
@@ -22,9 +113,10 @@ def idctn(
     workers: onp.ToInt | None = None,
     *,
     orthogonalize: op.CanBool | None = None,
-) -> _RealND: ...
+) -> onp.ArrayND[np.float64]: ...
+@overload
 def idstn(
-    x: onp.ToComplexND,
+    x: onp.ToFloatND,
     type: DCTType = 2,
     s: onp.ToInt | onp.ToIntND | None = None,
     axes: AnyShape | None = None,
@@ -33,4 +125,4 @@ def idstn(
     workers: onp.ToInt | None = None,
     *,
     orthogonalize: op.CanBool | None = None,
-) -> _RealND: ...
+) -> onp.ArrayND[npc.floating]: ...