[IR] Improve from_proto/to_proto typing with overloads

onnxscript/ir/serde.py

@@ -14,6 +14,7 @@
 from __future__ import annotations
 
 import functools
+import typing
 
 __all__ = [
     # Tensors
@@ -59,7 +60,6 @@
 import collections
 import logging
 import os
-import typing
 from typing import Any, Callable, List, Mapping, Sequence
 
 import numpy as np
@@ -121,16 +121,33 @@
     return array[::2] + 1j * array[1::2]
 
 
+@typing.overload
+def from_proto(proto: onnx.ModelProto) -> _core.Model: ...
+@typing.overload
+def from_proto(proto: onnx.GraphProto) -> _core.Graph: ...
+@typing.overload
+def from_proto(proto: onnx.NodeProto) -> _core.Node: ...
+@typing.overload
+def from_proto(proto: onnx.TensorProto) -> _core.Tensor: ...
+@typing.overload
+def from_proto(proto: onnx.AttributeProto) -> _core.Attr: ...
+@typing.overload
+def from_proto(proto: onnx.ValueInfoProto) -> _core.Value: ...
+@typing.overload
+def from_proto(proto: onnx.TypeProto) -> _core.TypeAndShape: ...
+@typing.overload
+def from_proto(proto: onnx.FunctionProto) -> _core.Function: ...
+@typing.overload
 def from_proto(
-    proto: onnx.ModelProto
-    | onnx.GraphProto
-    | onnx.NodeProto
-    | onnx.TensorProto
-    | onnx.AttributeProto
-    | onnx.ValueInfoProto
-    | onnx.TypeProto
-    | onnx.FunctionProto,
-) -> Any:
+    proto: onnx.TensorShapeProto.Dimension,
+) -> tuple[int | _core.SymbolicDim, str | None]: ...
+@typing.overload
+def from_proto(proto: Sequence[onnx.OperatorSetIdProto]) -> dict[str, int]: ...
+@typing.overload
+def from_proto(proto: Sequence[onnx.StringStringEntryProto]) -> dict[str, str]: ...
+
+
+def from_proto(proto: object) -> object:
     """Deserialize an ONNX proto message to an IR object."""
     if isinstance(proto, onnx.ModelProto):
         return deserialize_model(proto)
@@ -151,6 +168,16 @@
         )
     if isinstance(proto, onnx.FunctionProto):
         return deserialize_function(proto)
+    if isinstance(proto, onnx.TensorShapeProto.Dimension):
+        return deserialize_dimension(proto)
+    if isinstance(proto, Sequence) and all(
+        isinstance(p, onnx.OperatorSetIdProto) for p in proto
+    ):
+        return deserialize_opset_import(proto)
+    if isinstance(proto, Sequence) and all(
+        isinstance(p, onnx.StringStringEntryProto) for p in proto
+    ):
+        return deserialize_metadata_props(proto)
     raise NotImplementedError(
         f"Deserialization of {type(proto)} in from_proto is not implemented. "
         "Use a specific ir.serde.deserialize* function instead."