mypy

Author: xadupre

onnx_diagnostic/torch_export_patches/onnx_export_serialization.py

@@ -140,9 +140,9 @@ def serialization_functions(
 ) -> Dict[type, Callable[[int], bool]]:
     """Returns the list of serialization functions."""
 
-    supported_classes = set()
-    classes = {}
-    all_functions = {}
+    supported_classes: Set[type] = set()
+    classes: Dict[type, Callable[[int], bool]] = {}
+    all_functions: Dict[type, Optional[str]] = {}
 
     if patch_transformers:
         from .serialization.transformers_impl import (

onnx_diagnostic/torch_export_patches/serialization/__init__.py

@@ -16,12 +16,12 @@ def make_serialization_function_for_dataclass(
     ``dataclasses.dataclass``.
     """
 
-    def flatten_cls(obj: cls) -> Tuple[List[Any], torch.utils._pytree.Context]:
+    def flatten_cls(obj: cls) -> Tuple[List[Any], torch.utils._pytree.Context]:  # type: ignore[valid-type]
         """Serializes a ``%s`` with python objects."""
         return list(obj.values()), list(obj.keys())
 
     def flatten_with_keys_cls(
-        obj: cls,
+        obj: cls,  # type: ignore[valid-type]
     ) -> Tuple[List[Tuple[torch.utils._pytree.KeyEntry, Any]], torch.utils._pytree.Context]:
         """Serializes a ``%s`` with python objects with keys."""
         values, context = list(obj.values()), list(obj.keys())
@@ -31,7 +31,7 @@ def flatten_with_keys_cls(
 
     def unflatten_cls(
         values: List[Any], context: torch.utils._pytree.Context, output_type=None
-    ) -> cls:
+    ) -> cls:  # type: ignore[valid-type]
         """Restores an instance of ``%s`` from python objects."""
         return cls(**dict(zip(context, values)))
 

onnx_diagnostic/torch_export_patches/serialization/diffusers_impl.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional
+from typing import Dict, Optional, Set
 
 try:
     from diffusers.models.unets.unet_2d_condition import UNet2DConditionOutput
@@ -14,15 +14,15 @@
 from . import make_serialization_function_for_dataclass
 
 
-def _make_wrong_registrations() -> Dict[str, Optional[str]]:
-    res = {}
+def _make_wrong_registrations() -> Dict[type, Optional[str]]:
+    res: Dict[type, Optional[str]] = {}
     for c in [UNet2DConditionOutput]:
         if c is not None:
             res[c] = None
     return res
 
 
-SUPPORTED_DATACLASSES = set()
+SUPPORTED_DATACLASSES: Set[type] = set()
 WRONG_REGISTRATIONS = _make_wrong_registrations()
 
 

onnx_diagnostic/torch_export_patches/serialization/transformers_impl.py

@@ -1,4 +1,4 @@
-from typing import Any, List, Tuple
+from typing import Any, List, Set, Tuple
 import torch
 import transformers
 from transformers.cache_utils import (
@@ -13,7 +13,7 @@
 from . import make_serialization_function_for_dataclass
 
 
-SUPPORTED_DATACLASSES = set()
+SUPPORTED_DATACLASSES: Set[type] = set()
 WRONG_REGISTRATIONS = {
     DynamicCache: "4.50",
     BaseModelOutput: None,