fix: annotate script()

To allow mypy to analyze typing for annotated functions. Otherwise it complains that "Untyped decorator makes function "ones_like" untyped  [misc]"

ghstack-source-id: 02f281c
Pull Request resolved: #245

Author: justinchuby

onnxscript/main.py

@@ -9,6 +9,7 @@
 import inspect
 import sys
 import textwrap
+from typing import Any, Callable, Optional
 
 import onnx.helper
 
@@ -52,7 +53,9 @@ def script_check(f: ast.FunctionDef, opset, global_names, source, default_opset=
     return convert.top_level_stmt(f)
 
 
-def script(opset=None, default_opset=None, **kwargs):
+def script(
+    opset: Optional[values.Opset] = None, default_opset: Optional[Any] = None, **kwargs: Any
+) -> Callable[[Callable[..., Any]], onnxscript.OnnxFunction]:
     """Main decorator. Declares a function as an onnx function.
 
     Args: