Make export function preserve types (#666)

#665 Co-authored-by: Mike Lin <[email protected]>
chanzuckerberg · Dec 16, 2023 · b82e9da · b82e9da
1 parent 3f684ea
commit b82e9da
Showing 1 changed file with 4 additions and 5 deletions.
diff --git a/WDL/_util.py b/WDL/_util.py
@@ -32,11 +32,13 @@
 from types import FrameType
 from pythonjsonlogger import jsonlogger
 
+T = TypeVar("T")
+
 __all__: List[str] = []
 
 
-def export(obj) -> str:  # pyre-ignore
-    __all__.append(obj.__name__)
+def export(obj: T) -> T:
+    __all__.append(obj.__name__)  # pyre-ignore
     return obj
 
 
@@ -67,9 +69,6 @@ def strip_leading_whitespace(txt: str) -> Tuple[int, str]:
     return (to_strip, "\n".join(lines))
 
 
-T = TypeVar("T")
-
-
 @export
 class AdjM(Generic[T]):
     # A sparse adjacency matrix for topological sorting