Improve typing for boundaries (#126)

Author: cthoyt

src/ptwt/_util.py

@@ -26,8 +26,27 @@
     WaveletDetailTuple2d,
 )
 
-
-def _translate_boundary_strings(pywt_mode: BoundaryMode) -> str:
+#: All the PyTorch boundary modes for :func:`torch.nn.functional.pad`
+PyTorchBoundaryMode = Literal["replicate", "constant", "reflect", "circular"]
+
+#: All the PyTorch boundary modes for :func:`torch.nn.functional.pad`
+#: plus `symmetric` for the custom ptwt boundary
+ExtendedPyTorchBoundaryMode = PyTorchBoundaryMode | Literal["symmetric"]
+
+translation_dict: dict[BoundaryMode, ExtendedPyTorchBoundaryMode] = {
+    "constant": "replicate",
+    "zero": "constant",
+    "reflect": "reflect",
+    "periodic": "circular",
+    # pytorch does not support symmetric mode,
+    # we have our own implementation.
+    "symmetric": "symmetric",
+}
+
+
+def _translate_boundary_strings(
+    pywt_mode: BoundaryMode | None,
+) -> ExtendedPyTorchBoundaryMode:
     """Translate pywt mode strings to PyTorch mode strings.
 
     We support ``constant``, ``zero``, ``reflect``,
@@ -38,15 +57,8 @@ def _translate_boundary_strings(pywt_mode: BoundaryMode) -> str:
     Raises:
         ValueError: If the padding mode is not supported.
     """
-    translation_dict = {
-        "constant": "replicate",
-        "zero": "constant",
-        "reflect": "reflect",
-        "periodic": "circular",
-        # pytorch does not support symmetric mode,
-        # we have our own implementation.
-        "symmetric": "symmetric",
-    }
+    if pywt_mode is None:
+        return translation_dict["reflect"]
     if pywt_mode in translation_dict:
         return translation_dict[pywt_mode]
     else:

src/ptwt/conv_transform.py

@@ -55,9 +55,6 @@ def _fwt_pad(
     Returns:
         A PyTorch tensor with the padded input data
     """
-    # convert pywt to pytorch convention.
-    if mode is None:
-        mode = "reflect"
     pytorch_mode = _translate_boundary_strings(mode)
 
     if padding is None:

src/ptwt/conv_transform_2.py

@@ -84,8 +84,6 @@ def _fwt_pad2(
         The padded output tensor.
 
     """
-    if mode is None:
-        mode = "reflect"
     pytorch_mode = _translate_boundary_strings(mode)
 
     if padding is None: