[Typing][C-129,C-130,C-131] Add type annotations for python/paddle/audio/{features/layers.py,functional/functional.py,functional/window.py} (#67079)

enkilee · web-flow · commit 47d8b8f6573c · 2024-08-10T12:46:40.000+08:00
diff --git a/python/paddle/audio/features/layers.py b/python/paddle/audio/features/layers.py
@@ -11,15 +11,36 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations
+
 from functools import partial
-from typing import Optional, Union
+from typing import TYPE_CHECKING, Literal
+
+from typing_extensions import TypeAlias
 
 import paddle
-from paddle import Tensor, nn
+from paddle import nn
 
 from ..functional import compute_fbank_matrix, create_dct, power_to_db
 from ..functional.window import get_window
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+_WindowLiteral: TypeAlias = Literal[
+    'hamming',
+    'hann',
+    'kaiser',
+    'gaussian',
+    'exponential',
+    'triang',
+    'bohman',
+    'blackman',
+    'cosine',
+    'tukey',
+    'taylor',
+]
+
 
 class Spectrogram(nn.Layer):
     """Compute spectrogram of given signals, typically audio waveforms.
@@ -57,15 +78,18 @@ class Spectrogram(nn.Layer):
             >>> feats = feature_extractor(waveform)
     """
 
+    power: float
+    fft_window: Tensor
+
     def __init__(
         self,
         n_fft: int = 512,
-        hop_length: Optional[int] = 512,
-        win_length: Optional[int] = None,
-        window: str = 'hann',
+        hop_length: int | None = 512,
+        win_length: int | None = None,
+        window: _WindowLiteral = 'hann',
         power: float = 1.0,
         center: bool = True,
-        pad_mode: str = 'reflect',
+        pad_mode: Literal['reflect'] = 'reflect',
         dtype: str = 'float32',
     ) -> None:
         super().__init__()
@@ -142,21 +166,28 @@ class MelSpectrogram(nn.Layer):
             >>> feats = feature_extractor(waveform)
     """
 
+    n_mels: int
+    f_min: float
+    f_max: float
+    htk: bool
+    norm: Literal['slaney'] | float
+    fbank_matrix: Tensor
+
     def __init__(
         self,
         sr: int = 22050,
         n_fft: int = 2048,
-        hop_length: Optional[int] = 512,
-        win_length: Optional[int] = None,
-        window: str = 'hann',
+        hop_length: int | None = 512,
+        win_length: int | None = None,
+        window: _WindowLiteral = 'hann',
         power: float = 2.0,
         center: bool = True,
-        pad_mode: str = 'reflect',
+        pad_mode: Literal['reflect'] = 'reflect',
         n_mels: int = 64,
         f_min: float = 50.0,
-        f_max: Optional[float] = None,
+        f_max: float | None = None,
         htk: bool = False,
-        norm: Union[str, float] = 'slaney',
+        norm: Literal['slaney'] | float = 'slaney',
         dtype: str = 'float32',
     ) -> None:
         super().__init__()
@@ -245,24 +276,28 @@ class LogMelSpectrogram(nn.Layer):
             >>> feats = feature_extractor(waveform)
     """
 
+    ref_value: float
+    amin: float
+    top_db: float | None
+
     def __init__(
         self,
         sr: int = 22050,
         n_fft: int = 512,
-        hop_length: Optional[int] = None,
-        win_length: Optional[int] = None,
-        window: str = 'hann',
+        hop_length: int | None = None,
+        win_length: int | None = None,
+        window: _WindowLiteral = 'hann',
         power: float = 2.0,
         center: bool = True,
-        pad_mode: str = 'reflect',
+        pad_mode: Literal['reflect'] = 'reflect',
         n_mels: int = 64,
         f_min: float = 50.0,
-        f_max: Optional[float] = None,
+        f_max: float | None = None,
         htk: bool = False,
-        norm: Union[str, float] = 'slaney',
+        norm: Literal['slaney'] | float = 'slaney',
         ref_value: float = 1.0,
         amin: float = 1e-10,
-        top_db: Optional[float] = None,
+        top_db: float | None = None,
         dtype: str = 'float32',
     ) -> None:
         super().__init__()
@@ -349,25 +384,27 @@ class MFCC(nn.Layer):
             >>> feats = feature_extractor(waveform)
     """
 
+    dct_matrix: Tensor
+
     def __init__(
         self,
         sr: int = 22050,
         n_mfcc: int = 40,
         n_fft: int = 512,
-        hop_length: Optional[int] = None,
-        win_length: Optional[int] = None,
-        window: str = 'hann',
+        hop_length: int | None = None,
+        win_length: int | None = None,
+        window: _WindowLiteral = 'hann',
         power: float = 2.0,
         center: bool = True,
-        pad_mode: str = 'reflect',
+        pad_mode: Literal['reflect'] = 'reflect',
         n_mels: int = 64,
         f_min: float = 50.0,
-        f_max: Optional[float] = None,
+        f_max: float | None = None,
         htk: bool = False,
-        norm: Union[str, float] = 'slaney',
+        norm: Literal['slaney'] | float = 'slaney',
         ref_value: float = 1.0,
         amin: float = 1e-10,
-        top_db: Optional[float] = None,
+        top_db: float | None = None,
         dtype: str = 'float32',
     ) -> None:
         super().__init__()
diff --git a/python/paddle/audio/functional/functional.py b/python/paddle/audio/functional/functional.py
@@ -12,18 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # Modified from librosa(https://github.com/librosa/librosa)
+from __future__ import annotations
+
 import math
-from typing import Optional, Union
+from typing import TYPE_CHECKING, Literal, TypeVar
 
 import paddle
 from paddle import Tensor
 from paddle.base.framework import Variable
 from paddle.pir import Value
 
+if TYPE_CHECKING:
+    _TensorOrFloat = TypeVar("_TensorOrFloat", Tensor, float)
+
 
-def hz_to_mel(
-    freq: Union[Tensor, Value, Variable, float], htk: bool = False
-) -> Union[Tensor, Value, Variable, float]:
+def hz_to_mel(freq: _TensorOrFloat, htk: bool = False) -> _TensorOrFloat:
     """Convert Hz to Mels.
 
     Args:
@@ -77,9 +80,7 @@ def hz_to_mel(
     return mels
 
 
-def mel_to_hz(
-    mel: Union[float, Tensor, Variable, Value], htk: bool = False
-) -> Union[float, Tensor, Variable, Value]:
+def mel_to_hz(mel: _TensorOrFloat, htk: bool = False) -> _TensorOrFloat:
     """Convert mel bin numbers to frequencies.
 
     Args:
@@ -128,7 +129,7 @@ def mel_frequencies(
     f_max: float = 11025.0,
     htk: bool = False,
     dtype: str = 'float32',
-) -> Union[Tensor, Variable, Value]:
+) -> Tensor:
     """Compute mel frequencies.
 
     Args:
@@ -190,9 +191,9 @@ def compute_fbank_matrix(
     n_fft: int,
     n_mels: int = 64,
     f_min: float = 0.0,
-    f_max: Optional[float] = None,
+    f_max: float | None = None,
     htk: bool = False,
-    norm: Union[str, float] = 'slaney',
+    norm: Literal['slaney'] | float = 'slaney',
     dtype: str = 'float32',
 ) -> Tensor:
     """Compute fbank matrix.
@@ -259,11 +260,11 @@ def compute_fbank_matrix(
 
 
 def power_to_db(
-    spect: Union[Tensor, Variable, Value],
+    spect: Tensor,
     ref_value: float = 1.0,
     amin: float = 1e-10,
-    top_db: Optional[float] = 80.0,
-) -> Union[Tensor, Variable, Value]:
+    top_db: float | None = 80.0,
+) -> Tensor:
     """Convert a power spectrogram (amplitude squared) to decibel (dB) units. The function computes the scaling `10 * log10(x / ref)` in a numerically stable way.
 
     Args:
@@ -305,7 +306,7 @@ def power_to_db(
 def create_dct(
     n_mfcc: int,
     n_mels: int,
-    norm: Optional[str] = 'ortho',
+    norm: Literal['ortho'] | None = 'ortho',
     dtype: str = 'float32',
 ) -> Tensor:
     """Create a discrete cosine transform(DCT) matrix.
diff --git a/python/paddle/audio/functional/window.py b/python/paddle/audio/functional/window.py
@@ -13,11 +13,16 @@
 from __future__ import annotations
 
 import math
+from typing import TYPE_CHECKING
 
 import numpy as np
 
 import paddle
-from paddle import Tensor
+
+if TYPE_CHECKING:
+    from paddle import Tensor
+
+    from ..features.layers import _WindowLiteral
 
 
 class WindowFunctionRegister:
@@ -334,7 +339,7 @@ def _cosine(M: int, sym: bool = True, dtype: str = 'float64') -> Tensor:
 
 
 def get_window(
-    window: str | tuple[str, float],
+    window: _WindowLiteral | tuple[_WindowLiteral, float],
     win_length: int,
     fftbins: bool = True,
     dtype: str = 'float64',
