Clean up onnxscript/_internal/converter.py (#2759)

Clean up onnxscript/_internal/converter.py and move sourceinfo to
internal

---------

Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>

Author: justinchuby

onnxscript/_internal/analysis.py

@@ -5,8 +5,7 @@
 import ast
 from typing import Any, Optional, Sequence, Set
 
-from onnxscript import sourceinfo
-from onnxscript._internal import ast_utils
+from onnxscript._internal import ast_utils, sourceinfo
 
 
 def _get_loop_var(for_stmt: ast.For, formatter: sourceinfo.Formatter) -> str:

onnxscript/_internal/analysis_test.py

@@ -7,8 +7,8 @@
 from typing import Any
 
 from onnxscript._internal import analysis, ast_utils
+from onnxscript._internal.sourceinfo import formatter
 from onnxscript.onnx_opset import opset15 as op
-from onnxscript.sourceinfo import formatter
 
 
 class AnalysisResultsVisitor(ast.NodeVisitor):

onnxscript/_internal/converter.py

@@ -7,26 +7,23 @@
 from typing import (
     TYPE_CHECKING,
     Any,
-    Dict,
-    List,
     NoReturn,
-    Optional,
     Sequence,
-    Tuple,
     Union,
 )
 
 import onnx
 import onnx_ir as ir
 
 import onnxscript
-from onnxscript import onnx_types, sourceinfo
+from onnxscript import onnx_types
 from onnxscript._internal import (
     analysis,
     ast_utils,
     autocast,
     irbuilder,
     param_manipulation,
+    sourceinfo,
     values,
 )
 from onnxscript._internal import (
@@ -171,10 +168,10 @@ class :class:`onnxscript.irbuilder.IRBuilder` is used
 
     def __init__(
         self,
-        opset: Optional[values.Opset] = None,
-        global_names: Optional[dict[str, Any]] = None,
-        source: Optional[str] = None,
-        default_opset: Optional[values.Opset] = None,
+        opset: values.Opset | None = None,
+        global_names: dict[str, Any] | None = None,
+        source: str | None = None,
+        default_opset: values.Opset | None = None,
     ):
         self.source = source
         if global_names is not None:
@@ -184,11 +181,11 @@ def __init__(
         self.default_opset_ = default_opset
 
         # States initialized by `_init_function_translation`
-        self._outer: List[irbuilder.IRFunction] = []
+        self._outer: list[irbuilder.IRFunction] = []
         self._current_fn: irbuilder.IRFunction = None
         self._nextvar: int = 0
         self._used_vars: set[str] = set()
-        self._locals: List[Dict[str, LocalSymValue]] = [{}]
+        self._locals: list[dict[str, LocalSymValue]] = [{}]
         self._analyzer: analysis.AstAnalyzer | None = None
         self._castable: set[str] = set()
 
@@ -225,7 +222,7 @@ def _set_default_opset(self, opset: values.Opset, node: ast.AST) -> None:
         else:
             self.default_opset_ = opset
 
-    def _find_onnx_opset(self, node: ast.AST) -> Optional[values.Opset]:
+    def _find_onnx_opset(self, node: ast.AST) -> values.Opset | None:
         """Find the (first) ONNX opset used in the function, if any."""
         # Search for a Call expression of form "op.OpName(...)"
         if isinstance(node, ast.Call):
@@ -245,13 +242,15 @@ def _find_onnx_opset(self, node: ast.AST) -> Optional[values.Opset]:
     def _init_function_translation(self) -> None:
         """Initialize self for translating a new (top-level) function."""
         self._outer = []
-        self._current_fn: Optional[irbuilder.IRFunction] = None
+        self._current_fn: irbuilder.IRFunction | None = None
         self._nextvar = 0
         self._used_vars = set()
-        self._locals: List[Dict[str, LocalSymValue]] = [{}]
+        self._locals: list[dict[str, LocalSymValue]] = [{}]
 
     def _source_of(self, node: ast.AST) -> sourceinfo.SourceInfo:
-        return sourceinfo.SourceInfo(node, self.source, self._current_fn.name)
+        return sourceinfo.SourceInfo(
+            node, code=self.source, function_name=self._current_fn.name
+        )
 
     def _message(self, node: ast.AST, error_msg: str) -> str:
         """Constructs an error _message containing source information about an ast node."""
@@ -287,7 +286,7 @@ def _exit_scope(self) -> irbuilder.IRFunction:
         self._locals.pop(0)
         return graph
 
-    def _current_scope(self) -> Dict[str, LocalSymValue]:
+    def _current_scope(self) -> dict[str, LocalSymValue]:
         return self._locals[0]
 
     def _bind(self, name: str, val: LocalSymValue) -> None:
@@ -337,7 +336,7 @@ def tensor_name_generator() -> str:
         return ir.from_proto(proto)
 
     def _to_onnx_attr_ref(
-        self, val: values.AttrRef, info: Optional[sourceinfo.SourceInfo]
+        self, val: values.AttrRef, info: sourceinfo.SourceInfo | None
     ) -> ir.Attr:
         attrtype = val.value.type
         attrname = None
@@ -357,8 +356,8 @@ def _to_onnx_attr_ref(
     def _to_onnx_var(
         self,
         val: values.SymbolValue | PyValue,
-        target: Optional[PreferredName] = None,
-        info: Optional[sourceinfo.SourceInfo] = None,
+        target: PreferredName | None = None,
+        info: sourceinfo.SourceInfo | None = None,
     ) -> ir.Value:
         if isinstance(val, values.AttrRef):
             # promote attribute to value
@@ -397,8 +396,8 @@ def emit(
         self,
         outputs: Sequence[str],
         callee: values.Op | str,
-        inputs: Sequence[Optional[ir.Value]],
-        attrs: Optional[Sequence[ir.Attr]] = None,
+        inputs: Sequence[ir.Value | None],
+        attrs: Sequence[ir.Attr] | None = None,
     ) -> Sequence[ir.Value] | ir.Value:
         if not isinstance(callee, values.Op):
             callee = values.Op(self.default_opset, callee)
@@ -416,6 +415,7 @@ def emit(
         if not isinstance(callee, values.Op):
             raise TypeError(f"Unexpected type {type(callee)} for callee.")
         node.meta.setdefault("callee", callee)
+        assert self._current_fn is not None
         self._current_fn.append_node(node)
 
         return output_values if len(output_values) > 1 else output_values[0]
@@ -429,7 +429,7 @@ def emit1(self, *args, **kwargs) -> ir.Value:
     def _emit_const(
         self,
         pyvalue: PyValue,
-        suggested_name: Optional[PreferredName],
+        suggested_name: PreferredName | None,
         info: sourceinfo.SourceInfo,
     ) -> ir.Value:
         if suggested_name is None:
@@ -491,10 +491,8 @@ def _eval_constant_expr(self, expr: ast.expr) -> PyValue:
         function.)
         """
         # TODO: assert (self._is_constant_expr(expr))
-        # TODO: Refine types
         locals: dict[Any, Any] = {}
-        expression = ast.Expression(expr)
-        cpl = compile(expression, filename="<ast>", mode="eval")
+        cpl = compile(ast.Expression(expr), filename="<ast>", mode="eval")
         try:
             return eval(cpl, self.globals, locals)  # pylint: disable=eval-used
         except NameError as e:
@@ -506,7 +504,7 @@ def _eval_constant_expr(self, expr: ast.expr) -> PyValue:
                 )
             ) from e
 
-    def _get_type_annotation(self, annotation: ast.expr) -> Optional[ta.TypeAnnotationValue]:
+    def _get_type_annotation(self, annotation: ast.expr) -> ta.TypeAnnotationValue | None:
         typeinfo = self._eval_constant_expr(annotation)
         if not ta.is_valid_type(typeinfo):
             self.warn(
@@ -520,8 +518,8 @@ def _translate_attr(
         self,
         attr_name: str,
         expr: ast.AST,
-        attr_meta: Optional[onnx.defs.OpSchema.Attribute] = None,
-    ) -> Optional[ir.Attr]:
+        attr_meta: onnx.defs.OpSchema.Attribute | None = None,
+    ) -> ir.Attr | None:
         """Translate an attribute-value specification of the form `attr_name=<expr>`
         in a call to an op. expr is an AST. The following cases are supported:
         * Expr evaluates to a script-time constant (a python-value) that can be mapped
@@ -597,9 +595,7 @@ def _translate_docstring(self, node: ast.Expr) -> None:
             self.fail(node, "Docstring expression must be a constant.")
         self._current_fn.doc_string = node.value.value
 
-    def _translate_expr(
-        self, node: ast.AST, target: Optional[PreferredName] = None
-    ) -> ir.Value:
+    def _translate_expr(self, node: ast.AST, target: PreferredName | None = None) -> ir.Value:
         """Expression-translation generates "IR statements/nodes" that compute the value of
         the expression into a target-variable, and returns the variable that is
         assigned this value.
@@ -630,7 +626,7 @@ def _translate_expr(
         result = self.generate_unique_name(target)
         return self.emit1([result], callee, args, attrs)
 
-    def _translate_opt_expr(self, node: ast.expr) -> Optional[ir.Value]:
+    def _translate_opt_expr(self, node: ast.expr) -> ir.Value | None:
         """Translation of an expression where "None" is permitted (eg., for an optional argument).
         None is represented as a Constant in Python 3.9+.
         """
@@ -639,7 +635,7 @@ def _translate_opt_expr(self, node: ast.expr) -> Optional[ir.Value]:
         return self._translate_expr(node)
 
     def _translate_subscript_expr(
-        self, node: ast.Subscript, target: Optional[PreferredName]
+        self, node: ast.Subscript, target: PreferredName | None
     ) -> ir.Value:
         """List of supported syntaxes is below.
         `A` is a tensor or an expression equivalent to a tensor.
@@ -689,7 +685,7 @@ def _translate_subscript_expr(
         # TODO: Do this at a graph-scope level.
         cached_int_consts: dict[int, ir.Value] = {}
 
-        def const_1d(value, name: Optional[str] = None) -> ir.Value:
+        def const_1d(value, name: str | None = None) -> ir.Value:
             nonlocal cached_int_consts
             if value not in cached_int_consts:
                 cached_int_consts[value] = self._emit_const([value], name, info)
@@ -703,8 +699,8 @@ def one_1d() -> ir.Value:
         minint = -(1 << 63)
 
         def translate_slice_component(
-            node_arg, default_value: Optional[int] = None
-        ) -> tuple[ir.Value, Optional[int]]:
+            node_arg, default_value: int | None = None
+        ) -> tuple[ir.Value, int | None]:
             """Translate optional start/stop/step component of a Slice expression."""
             if node_arg is None:
                 if default_value is None:
@@ -758,9 +754,9 @@ def translate_slice(slice_expr: ast.Slice) -> tuple[ir.Value, ir.Value, ir.Value
 
         # As the first step, we partition the index elements into four kinds: Slice (eg., 1:5:2),
         # known-to-be-scalar (eg., 2), other-tensor (eg., I), skip/no-op (that is, just ":")
-        sliced_indices: List[Tuple[int, ast.expr]] = []
-        scalar_indices: List[Tuple[int, ast.expr]] = []
-        non_scalar_indices: List[Tuple[int, ast.expr]] = []
+        sliced_indices: list[tuple[int, ast.expr]] = []
+        scalar_indices: list[tuple[int, ast.expr]] = []
+        non_scalar_indices: list[tuple[int, ast.expr]] = []
         for axis, elt in enumerate(indices):
             if isinstance(elt, ast.Slice):
                 # Add to sliced_indices, unless it is "::", which is a no-op.
@@ -872,7 +868,7 @@ def translate_slice(slice_expr: ast.Slice) -> tuple[ir.Value, ir.Value, ir.Value
 
     def _translate_call_expr(
         self, node: ast.Call
-    ) -> tuple[values.Op, list[Optional[ir.Value]], list[ir.Attr]]:
+    ) -> tuple[values.Op, list[ir.Value | None], list[ir.Attr]]:
         """Translates a call-expression."""
         callee = self._translate_callee_expr(node.func)
         param_schemas = callee.param_schemas()
@@ -932,14 +928,7 @@ def _translate_unary_op_expr(self, node):
             # should intercept this call and replace node
             # by node.operand.
             # This mechanism does not handle somthing like `(-(-5))`.
-            if hasattr(node.operand, "value"):
-                # python 3.8+
-                val = node.operand.value
-            else:
-                raise TypeError(
-                    f"Unable to guess constant value from type {type(node.operand)!r} "
-                    f"and attributes {dir(node.operand)!r}."
-                )
+            val = node.operand.value
             if op == ast.USub:
                 cst = ast.Constant(-val, lineno=node.lineno, col_offset=node.col_offset)
                 return self._translate_expr(cst)
@@ -1042,7 +1031,7 @@ def _translate_stmt(self, node: ast.stmt, index_of_stmt=None) -> None:
             return None
         raise ValueError(self._message(node, f"Unsupported statement type '{type(node)!r}'."))
 
-    def _translate_assign_stmt(self, stmt: Union[ast.Assign, ast.AnnAssign]) -> None:
+    def _translate_assign_stmt(self, stmt: ast.Assign | ast.AnnAssign) -> None:
         def assign(lhs: ast.AST, rhs: ast.AST) -> None:
             if isinstance(lhs, ast.Name):
                 # Assignments of the form "x = SomeExpression"
@@ -1202,7 +1191,7 @@ def rename(x):
                 values.SymbolValue(y, self._source_of(stmt)),
             )
 
-    def _translate_loop_stmt(self, loop_stmt: Union[ast.For, ast.While]) -> None:
+    def _translate_loop_stmt(self, loop_stmt: ast.For | ast.While) -> None:
         # loop-variable
         if isinstance(loop_stmt, ast.For):
             if not isinstance(loop_stmt.target, ast.Name):

onnxscript/_internal/converter_test.py

@@ -77,10 +77,10 @@ def validate_save(
                 model = f.to_model_proto(io_types=FLOAT)
                 if save_text:
                     with (TEST_OUTPUT_DIR / f"{f.name}.txt").open("w", encoding="utf-8") as fi:
-                        fi.write(onnx.helper.printable_graph(model.graph))
+                        fi.write(onnx.printer.to_text(model.graph))
                         for fct in model.functions:
                             fi.write("\n-------------------------\n")
-                            fi.write(onnx.helper.printable_graph(fct))
+                            fi.write(onnx.printer.to_text(fct))
                 if check_ort and (skip_check_ort is None or f.name not in skip_check_ort):
                     try:
                         create_cpu_inference_session(model.SerializeToString())
@@ -92,10 +92,10 @@ def validate_save(
                     model = onnx.shape_inference.infer_shapes(model)
                 if save_text:
                     with open(os.path.join(TEST_OUTPUT_DIR, f"{f.name}.shape.txt"), "w") as fi:
-                        fi.write(onnx.helper.printable_graph(model.graph))
+                        fi.write(onnx.printer.to_text(model.graph))
                         for fct in model.functions:
-                            f.write("\n-------------------------\n")
-                            f.write(onnx.helper.printable_graph(fct))
+                            fi.write("\n-------------------------\n")
+                            fi.write(onnx.printer.to_text(fct))
                 try:
                     onnx.checker.check_model(model)
                 except onnx.checker.ValidationError as e:

onnxscript/sourceinfo.py onnxscript/_internal/sourceinfo.pyonnxscript/sourceinfo.py renamed to onnxscript/_internal/sourceinfo.py

@@ -1,14 +1,12 @@
-# -------------------------------------------------------------------------
-# Copyright (c) Microsoft Corporation. All rights reserved.
+# Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
-# -------------------------------------------------------------------------
 
 """Source code information used for diagnostic messages."""
 
 from __future__ import annotations
 
 import ast
-from typing import Callable, Optional
+from typing import Callable
 
 
 class SourceInfo:
@@ -17,8 +15,9 @@ class SourceInfo:
     def __init__(
         self,
         ast_node: ast.AST,
-        code: Optional[str] = None,
-        function_name: Optional[str] = None,
+        *,
+        code: str | None = None,
+        function_name: str | None = None,
     ):
         self.ast_node = ast_node
         self.code = code
@@ -52,8 +51,8 @@ def __str__(self) -> str:
 Formatter = Callable[[ast.AST, str], str]
 
 
-def formatter(source_code: Optional[str]) -> Formatter:
+def formatter(source_code: str | None) -> Formatter:
     def format(node: ast.AST, message: str) -> str:
-        return SourceInfo(node, source_code).msg(message)
+        return SourceInfo(node, code=source_code).msg(message)
 
     return format

onnxscript/_internal/values.py

@@ -27,8 +27,7 @@
 import onnx_ir as ir
 from typing_extensions import ParamSpec
 
-from onnxscript import sourceinfo
-from onnxscript._internal import ast_utils, deprecation, irbuilder, type_annotation
+from onnxscript._internal import ast_utils, deprecation, irbuilder, sourceinfo, type_annotation
 from onnxscript._internal import converter as converter_module
 from onnxscript.ir import _schemas
 from onnxscript.onnx_types import ONNXType