Merge pull request #81 from v0lta/typing-everywhere

Type annotate the tests

Author: v0lta

noxfile.py

@@ -49,6 +49,7 @@ def mypy(session):
         "--no-warn-return-any",
         "--explicit-package-bases",
         "src",
+        "tests",
     )
 
 

src/ptwt/_util.py

@@ -7,7 +7,7 @@
 import pywt
 import torch
 
-from ptwt.constants import OrthogonalizeMethod
+from .constants import OrthogonalizeMethod
 
 
 class Wavelet(Protocol):

src/ptwt/continuous_transform.py

@@ -291,10 +291,10 @@ def __call__(self, grid_values: torch.Tensor) -> torch.Tensor:
         shannon = (
             torch.sqrt(self.bandwidth)
             * (
-                torch.sin(torch.pi * self.bandwidth * grid_values)  # type: ignore
+                torch.sin(torch.pi * self.bandwidth * grid_values)
                 / (torch.pi * self.bandwidth * grid_values)
             )
-            * torch.exp(1j * 2 * torch.pi * self.center * grid_values)  # type: ignore
+            * torch.exp(1j * 2 * torch.pi * self.center * grid_values)
         )
         return shannon
 
@@ -306,8 +306,8 @@ def __call__(self, grid_values: torch.Tensor) -> torch.Tensor:
         """Return numerical values for the wavelet on a grid."""
         morlet = (
             1.0
-            / torch.sqrt(torch.pi * self.bandwidth)  # type: ignore
+            / torch.sqrt(torch.pi * self.bandwidth)
             * torch.exp(-(grid_values**2) / self.bandwidth)
-            * torch.exp(1j * 2 * torch.pi * self.center * grid_values)  # type: ignore
+            * torch.exp(1j * 2 * torch.pi * self.center * grid_values)
         )
         return morlet

src/ptwt/matmul_transform.py

@@ -157,7 +157,11 @@ def orthogonalize(
     raise ValueError(f"Invalid orthogonalization method: {method}")
 
 
-class MatrixWavedec(object):
+class BaseMatrixWaveDec:
+    """A base class for matrix wavedec."""
+
+
+class MatrixWavedec(BaseMatrixWaveDec):
     """Compute the sparse matrix fast wavelet transform.
 
     Intermediate scale results must be divisible

src/ptwt/matmul_transform_2.py

@@ -30,7 +30,12 @@
     _preprocess_tensor_dec2d,
     _waverec2d_fold_channels_2d_list,
 )
-from .matmul_transform import construct_boundary_a, construct_boundary_s, orthogonalize
+from .matmul_transform import (
+    BaseMatrixWaveDec,
+    construct_boundary_a,
+    construct_boundary_s,
+    orthogonalize,
+)
 from .sparse_math import (
     batch_mm,
     cat_sparse_identity_matrix,
@@ -217,7 +222,7 @@ def _matrix_pad_2(height: int, width: int) -> Tuple[int, int, Tuple[bool, bool]]
     return height, width, pad_tuple
 
 
-class MatrixWavedec2(object):
+class MatrixWavedec2(BaseMatrixWaveDec):
     """Experimental sparse matrix 2d wavelet transform.
 
         For a completely pad-free transform,

src/ptwt/matmul_transform_3.py

@@ -58,7 +58,7 @@ def __init__(
         wavelet: Union[Wavelet, str],
         level: Optional[int] = None,
         axes: Tuple[int, int, int] = (-3, -2, -1),
-        boundary: Optional[str] = "qr",
+        boundary: OrthogonalizeMethod = "qr",
     ):
         """Create a *separable* three-dimensional fast boundary wavelet transform.
 
@@ -69,7 +69,7 @@ def __init__(
             wavelet (Union[Wavelet, str]): The wavelet to use.
             level (Optional[int]): The desired decomposition level.
                 Defaults to None.
-            boundary (Optional[str]): The matrix orthogonalization method.
+            boundary: The matrix orthogonalization method.
                 Defaults to "qr".
 
         Raises:

src/ptwt/packets.py

@@ -106,7 +106,7 @@ def __init__(
             if len(data.shape) == 1:
                 # add a batch dimension.
                 data = data.unsqueeze(0)
-            self.transform(data, maxlevel)  # type: ignore
+            self.transform(data, maxlevel)
         else:
             self.data = {}
 

src/ptwt/separable_conv_transform.py

@@ -167,9 +167,9 @@ def _separable_conv_waverecn(
 
     approx: torch.Tensor = coeffs[0]
     for level_dict in coeffs[1:]:
-        keys = list(level_dict.keys())
-        level_dict["a" * max(map(len, keys))] = approx
-        approx = _separable_conv_idwtn(level_dict, wavelet)
+        keys = list(level_dict.keys())  # type: ignore
+        level_dict["a" * max(map(len, keys))] = approx  # type: ignore
+        approx = _separable_conv_idwtn(level_dict, wavelet)  # type: ignore
     return approx
 
 

src/ptwt/sparse_math.py

@@ -6,7 +6,7 @@
 
 import torch
 
-from ptwt.constants import PaddingMode
+from .constants import PaddingMode
 
 
 def _dense_kron(

tests/_mackey_glass.py

@@ -59,7 +59,7 @@ def _mackey(
     return x[:, discard:]
 
 
-class MackeyGenerator(object):
+class MackeyGenerator:
     """Generates lorenz attractor data in 1 or 3d on the GPU."""
 
     def __init__(