Improve typing and docstrings

.flake8

@@ -23,6 +23,8 @@ ignore =
     # asserts are ok in test.
     S101
     C901
+extend-select = B950
+extend-ignore = E501,E701,E704
 exclude =
     .tox,
     .git,
@@ -37,7 +39,7 @@ exclude =
     .eggs,
     data.
     src/ptwt/__init__.py
-max-line-length = 90
+max-line-length = 80
 max-complexity = 20
 import-order-style = pycharm
 application-import-names =

setup.cfg

@@ -51,7 +51,6 @@ install_requires =
     torch
     scipy>=1.10
     pooch
-    matplotlib
     numpy
     pytest
     nox

src/ptwt/_stationary_transform.py

@@ -1,6 +1,7 @@
 """This module implements stationary wavelet transforms."""
 
-from typing import List, Optional, Union
+from collections.abc import Sequence
+from typing import Optional, Union
 
 import pywt
 import torch
@@ -19,7 +20,7 @@ def _swt(
     wavelet: Union[Wavelet, str],
     level: Optional[int] = None,
     axis: Optional[int] = -1,
-) -> List[torch.Tensor]:
+) -> list[torch.Tensor]:
     """Compute a multilevel 1d stationary wavelet transform.
 
     Args:
@@ -28,7 +29,7 @@ def _swt(
         level (Optional[int], optional): The number of levels to compute
 
     Returns:
-        List[torch.Tensor]: Same as wavedec.
+        list[torch.Tensor]: Same as wavedec.
         Equivalent to pywt.swt with trim_approx=True.
 
     Raises:
@@ -107,14 +108,15 @@ def _conv_transpose_dedilate(
 
 
 def _iswt(
-    coeffs: List[torch.Tensor],
+    coeffs: Sequence[torch.Tensor],
     wavelet: Union[pywt.Wavelet, str],
     axis: Optional[int] = -1,
 ) -> torch.Tensor:
     """Inverts a 1d stationary wavelet transform.
 
     Args:
-        coeffs (List[torch.Tensor]): The coefficients as computed by the swt function.
+        coeffs (Sequence[torch.Tensor]): The coefficients as computed
+            by the swt function.
         wavelet (Union[pywt.Wavelet, str]): The wavelet used for the forward transform.
         axis (int, optional): The axis the forward trasform was computed over.
             Defaults to -1.

src/ptwt/_util.py

@@ -1,11 +1,13 @@
 """Utility methods to compute wavelet decompositions from a dataset."""
 
 import typing
-from typing import Any, Callable, List, Optional, Protocol, Sequence, Tuple, Union
+from collections.abc import Sequence
+from typing import Any, Callable, Optional, Protocol, Union, cast, overload
 
 import numpy as np
 import pywt
 import torch
+from typing_extensions import Unpack
 
 from .constants import OrthogonalizeMethod
 
@@ -20,7 +22,7 @@ class Wavelet(Protocol):
     rec_hi: Sequence[float]
     dec_len: int
     rec_len: int
-    filter_bank: Tuple[
+    filter_bank: tuple[
         Sequence[float], Sequence[float], Sequence[float], Sequence[float]
     ]
 
@@ -29,6 +31,15 @@ def __len__(self) -> int:
         return len(self.dec_lo)
 
 
+WaveletDetailTuple2d = tuple[torch.Tensor, torch.Tensor, torch.Tensor]
+WaveletDetailDict = dict[str, torch.Tensor]
+
+WaveletCoeffDetailTuple2d = tuple[
+    torch.Tensor, Unpack[tuple[WaveletDetailTuple2d, ...]]
+]
+WaveletCoeffDetailDict = tuple[torch.Tensor, Unpack[tuple[WaveletDetailDict, ...]]]
+
+
 def _as_wavelet(wavelet: Union[Wavelet, str]) -> Wavelet:
     """Ensure the input argument to be a pywt wavelet compatible object.
 
@@ -63,15 +74,15 @@ def _outer(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
     return a_mul * b_mul
 
 
-def _get_len(wavelet: Union[Tuple[torch.Tensor, ...], str, Wavelet]) -> int:
+def _get_len(wavelet: Union[tuple[torch.Tensor, ...], str, Wavelet]) -> int:
     """Get number of filter coefficients for various wavelet data types."""
     if isinstance(wavelet, tuple):
         return wavelet[0].shape[0]
     else:
         return len(_as_wavelet(wavelet))
 
 
-def _pad_symmetric_1d(signal: torch.Tensor, pad_list: Tuple[int, int]) -> torch.Tensor:
+def _pad_symmetric_1d(signal: torch.Tensor, pad_list: tuple[int, int]) -> torch.Tensor:
     padl, padr = pad_list
     dimlen = signal.shape[0]
     if padl > dimlen or padr > dimlen:
@@ -88,11 +99,11 @@ def _pad_symmetric_1d(signal: torch.Tensor, pad_list: Tuple[int, int]) -> torch.
             cat_list.insert(0, signal[:padl].flip(0))
         if padr > 0:
             cat_list.append(signal[-padr::].flip(0))
-        return torch.cat(cat_list, axis=0)  # type: ignore
+        return torch.cat(cat_list, dim=0)
 
 
 def _pad_symmetric(
-    signal: torch.Tensor, pad_lists: List[Tuple[int, int]]
+    signal: torch.Tensor, pad_lists: Sequence[tuple[int, int]]
 ) -> torch.Tensor:
     if len(signal.shape) < len(pad_lists):
         raise ValueError("not enough dimensions to pad.")
@@ -106,15 +117,15 @@ def _pad_symmetric(
     return signal
 
 
-def _fold_axes(data: torch.Tensor, keep_no: int) -> Tuple[torch.Tensor, List[int]]:
+def _fold_axes(data: torch.Tensor, keep_no: int) -> tuple[torch.Tensor, list[int]]:
     """Fold unchanged leading dimensions into a single batch dimension.
 
     Args:
-        data ( torch.Tensor): The input data array.
+        data (torch.Tensor): The input data array.
         keep_no (int): The number of dimensions to keep.
 
     Returns:
-        Tuple[ torch.Tensor, List[int]]:
+        tuple[torch.Tensor, list[int]]:
             The folded result array, and the shape of the original input.
     """
     dshape = list(data.shape)
@@ -124,7 +135,7 @@ def _fold_axes(data: torch.Tensor, keep_no: int) -> Tuple[torch.Tensor, List[int
     )
 
 
-def _unfold_axes(data: torch.Tensor, ds: List[int], keep_no: int) -> torch.Tensor:
+def _unfold_axes(data: torch.Tensor, ds: list[int], keep_no: int) -> torch.Tensor:
     """Unfold i.e. [batch*channel,height,widht] to [batch,channel,height,width]."""
     return torch.reshape(data, ds[:-keep_no] + list(data.shape[-keep_no:]))
 
@@ -137,49 +148,75 @@ def _check_if_tensor(array: Any) -> torch.Tensor:
     return array
 
 
-def _check_axes_argument(axes: List[int]) -> None:
+def _check_axes_argument(axes: Sequence[int]) -> None:
     if len(set(axes)) != len(axes):
         raise ValueError("Cant transform the same axis twice.")
 
 
 def _get_transpose_order(
-    axes: List[int], data_shape: List[int]
-) -> Tuple[List[int], List[int]]:
+    axes: Sequence[int], data_shape: Sequence[int]
+) -> tuple[list[int], list[int]]:
     axes = list(map(lambda a: a + len(data_shape) if a < 0 else a, axes))
     all_axes = list(range(len(data_shape)))
     remove_transformed = list(filter(lambda a: a not in axes, all_axes))
     return remove_transformed, axes
 
 
-def _swap_axes(data: torch.Tensor, axes: List[int]) -> torch.Tensor:
+def _swap_axes(data: torch.Tensor, axes: Sequence[int]) -> torch.Tensor:
     _check_axes_argument(axes)
     front, back = _get_transpose_order(axes, list(data.shape))
     return torch.permute(data, front + back)
 
 
-def _undo_swap_axes(data: torch.Tensor, axes: List[int]) -> torch.Tensor:
+def _undo_swap_axes(data: torch.Tensor, axes: Sequence[int]) -> torch.Tensor:
     _check_axes_argument(axes)
     front, back = _get_transpose_order(axes, list(data.shape))
     restore_sorted = torch.argsort(torch.tensor(front + back)).tolist()
     return torch.permute(data, restore_sorted)
 
 
+@overload
+def _map_result(
+    data: WaveletCoeffDetailTuple2d,
+    function: Callable[[torch.Tensor], torch.Tensor],
+) -> WaveletCoeffDetailTuple2d: ...
+
+
+@overload
 def _map_result(
-    data: List[Union[torch.Tensor, Any]],  # following jax tree_map typing can be Any
-    function: Callable[[Any], torch.Tensor],
-) -> List[Union[torch.Tensor, Any]]:
-    # Apply the given function to the input list of tensor and tuples.
-    result_lst: List[Union[torch.Tensor, Any]] = []
-    for element in data:
-        if isinstance(element, torch.Tensor):
-            result_lst.append(function(element))
-        elif isinstance(element, tuple):
+    data: WaveletCoeffDetailDict,
+    function: Callable[[torch.Tensor], torch.Tensor],
+) -> WaveletCoeffDetailDict: ...
+
+
+def _map_result(
+    data: Union[WaveletCoeffDetailTuple2d, WaveletCoeffDetailDict],
+    function: Callable[[torch.Tensor], torch.Tensor],
+) -> Union[WaveletCoeffDetailTuple2d, WaveletCoeffDetailDict]:
+    approx = function(data[0])
+    result_lst: list[
+        Union[
+            tuple[torch.Tensor, torch.Tensor, torch.Tensor],
+            dict[str, torch.Tensor],
+        ]
+    ] = []
+    for element in data[1:]:
+        if isinstance(element, tuple):
             result_lst.append(
-                (function(element[0]), function(element[1]), function(element[2]))
+                (
+                    function(element[0]),
+                    function(element[1]),
+                    function(element[2]),
+                )
             )
         elif isinstance(element, dict):
-            new_dict = {}
-            for key, value in element.items():
-                new_dict[key] = function(value)
+            new_dict = {key: function(value) for key, value in element.items()}
             result_lst.append(new_dict)
-    return result_lst
+        else:
+            raise AssertionError(f"Unexpected input type {type(element)}")
+
+    return_val = approx, *result_lst
+    return_val = cast(
+        Union[WaveletCoeffDetailTuple2d, WaveletCoeffDetailDict], return_val
+    )
+    return return_val

src/ptwt/continuous_transform.py

@@ -3,7 +3,7 @@
 This module is based on pywt's cwt implementation.
 """
 
-from typing import Any, Tuple, Union
+from typing import Any, Union
 
 import numpy as np
 import torch
@@ -27,7 +27,7 @@ def cwt(
     scales: Union[np.ndarray, torch.Tensor],  # type: ignore
     wavelet: Union[ContinuousWavelet, str],
     sampling_period: float = 1.0,
-) -> Tuple[torch.Tensor, np.ndarray]:  # type: ignore
+) -> tuple[torch.Tensor, np.ndarray]:  # type: ignore
     """Compute the single-dimensional continuous wavelet transform.
 
     This function is a PyTorch port of pywt.cwt as found at:
@@ -50,7 +50,7 @@ def cwt(
         ValueError: If a scale is too small for the input signal.
 
     Returns:
-        Tuple[torch.Tensor, np.ndarray]: The first tuple-element contains
+        tuple[torch.Tensor, np.ndarray]: The first tuple-element contains
             the transformation matrix of shape [scales, batch, time].
             The second element contains an array with frequency information.
 
@@ -267,7 +267,7 @@ def center(self) -> torch.Tensor:
 
     def wavefun(
         self, precision: int, dtype: torch.dtype = torch.float64
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         """Define a grid and evaluate the wavelet on it."""
         length = 2**precision
         # load the bounds from untyped pywt code.