Add overloads for 2d get_level

Author: felixblanke

src/ptwt/packets.py

@@ -6,7 +6,7 @@
 from collections.abc import Sequence
 from functools import partial
 from itertools import product
-from typing import TYPE_CHECKING, Callable, Literal, Optional, Union
+from typing import TYPE_CHECKING, Callable, Literal, Optional, Union, overload
 
 import numpy as np
 import pywt
@@ -530,6 +530,14 @@ def __getitem__(self, key: str) -> torch.Tensor:
             )
         return super().__getitem__(key)
 
+    @overload
+    @staticmethod
+    def get_level(level: int, order: Literal["freq"]) -> list[list[str]]: ...
+
+    @overload
+    @staticmethod
+    def get_level(level: int, order: Literal["natural"]) -> list[str]: ...
+
     @staticmethod
     def get_level(
         level: int, order: Literal["freq", "natural"] = "freq"