| import torch |
| import sys |
| import ast |
| import dataclasses |
| import inspect |
| import string |
| import re |
| from collections import namedtuple |
| from textwrap import dedent |
| from typing import List, Tuple # noqa: F401 |
| from torch._C._jit_tree_views import ( |
| ClassDef, Ident, Stmt, Decl, Def, Var, |
| EmptyTypeAnnotation, Param, ExprStmt, Assign, |
| Delete, Return, Raise, Assert, AugAssign, While, |
| For, If, Pass, Break, Continue, Apply, Dots, Select, |
| TrueLiteral, FalseLiteral, NoneLiteral, Starred, |
| ListLiteral, TupleLiteral, DictLiteral, Const, |
| StringLiteral, ListComp, Attribute, BinOp, UnaryOp, |
| SliceExpr, Subscript, TernaryIf, With, WithItem, Property, |
| DictComp, |
| ) |
| from torch._sources import get_source_lines_and_file, parse_def, make_source_context |
| from torch._sources import ParsedDef as _ParsedDef |
| from torch.jit._dataclass_impls import DATACLASS_MAGIC_METHODS |
| from torch.jit._monkeytype_config import monkeytype_trace, get_qualified_name |
| from torch._jit_internal import should_drop, is_static_fn, FunctionModifiers # noqa: F401 |
| from torch import _jit_internal |
| import torch.jit.annotations |
| |
| _IS_ASTUNPARSE_INSTALLED = False |
| try: |
| import astunparse # type: ignore[import] |
| _IS_ASTUNPARSE_INSTALLED = True |
| except ImportError: |
| pass |
| |
| # Borrowed from cPython implementation |
| # https://github.com/python/cpython/blob/561612d8456cfab5672c9b445521113b847bd6b3/Lib/textwrap.py#L411# |
| |
| _reserved_prefix = '__jit' |
| _reserved_names = {'print'} |
| _identifier_chars = set(string.ascii_lowercase + string.ascii_uppercase + string.digits) |
| |
| |
| def is_reserved_name(name): |
| return name.startswith(_reserved_prefix) or name in _reserved_names |
| |
| |
| pretty_node_names = { |
| ast.FunctionDef: "function definitions", |
| ast.For: "for loops", |
| ast.Delete: "del statements", |
| ast.ClassDef: "class definitions", |
| ast.With: "with statements", |
| ast.Raise: "raise statements", |
| ast.Assert: "assertions", |
| ast.Import: "import statements", |
| ast.ImportFrom: "import statements", |
| ast.Global: "global variables", |
| ast.Break: "break statements", |
| ast.Continue: "continue statements", |
| } |
| |
| node_start_tokens = { |
| ast.FunctionDef: "def", |
| ast.For: "for", |
| ast.Delete: "del", |
| ast.ClassDef: "class", |
| ast.With: "with", |
| ast.Raise: "raise", |
| ast.Assert: "assert", |
| ast.Import: "import", |
| ast.ImportFrom: "from", |
| ast.Global: "global", |
| ast.Break: "break", |
| ast.Continue: "continue", |
| } |
| |
| pretty_node_names.update({ |
| ast.AsyncFunctionDef: "async function definitions", |
| ast.AsyncFor: "async for loops", |
| ast.AsyncWith: "async with statements", |
| ast.Try: "try blocks", |
| ast.Nonlocal: "nonlocal variables", |
| }) |
| |
| node_start_tokens.update({ |
| ast.AsyncFunctionDef: "async def", |
| ast.AsyncFor: "async for", |
| ast.AsyncWith: "async with", |
| ast.Try: "try", |
| ast.Nonlocal: "nonlocal", |
| }) |
| |
| if sys.version_info >= (3, 6): |
| pretty_node_names.update({ |
| ast.AnnAssign: "annotated assignments", |
| }) |
| # NB: no specific token for AnnAssign |
| |
| |
| class FrontendError(Exception): |
| def __init__(self, source_range, msg): |
| self.source_range = source_range |
| self.msg = msg |
| |
| # This has to be instantiated here so the ErrorReport is accurate to the |
| # call stack when the FrontendError was raised |
| self.error_report = torch._C.ErrorReport(self.source_range) |
| |
| def __str__(self): |
| return self.msg + self.error_report.what().lstrip() |
| |
| |
| class NotSupportedError(FrontendError): |
| pass |
| |
| |
| class UnsupportedNodeError(NotSupportedError): |
| def __init__(self, ctx, offending_node, reason=''): |
| # If we don't have a specific token, we default to length of 1 |
| node_type = type(offending_node) |
| range_len = len(node_start_tokens.get(node_type, ' ')) |
| source_range = ctx.make_range(offending_node.lineno, |
| offending_node.col_offset, |
| offending_node.col_offset + range_len) |
| feature_name = pretty_node_names.get(node_type, node_type.__name__) |
| msg = "{} {}aren't supported".format(feature_name, reason + ' ' if reason else '') |
| super(UnsupportedNodeError, self).__init__(source_range, msg) |
| |
| |
| class FrontendTypeError(FrontendError): |
| pass |
| |
| |
| def build_withitems(ctx, items): |
| items = [build_withitem(ctx, i) for i in items] |
| return list(items) |
| |
| |
| def build_stmts(ctx, stmts): |
| stmts = [build_stmt(ctx, s) for s in stmts] |
| return list(filter(None, stmts)) |
| |
| |
| def get_class_properties(cls, self_name): |
| """ |
| Get a list of Property objects representing the properties of a class. |
| |
| Args: |
| cls: The class to get properties of. |
| self_name: The name of the class that the properties should belong to. |
| Returns: |
| A list of Property objects corresponding to the properties of cls. Property |
| here refers to the subclass of TreeView. |
| """ |
| props = inspect.getmembers( |
| cls, predicate=lambda m: isinstance(m, property)) |
| # Any property that should not compiled must be in this list on the Module. |
| unused_properties = getattr(cls, "__jit_unused_properties__", []) |
| |
| # Create Property TreeView objects from inspected property objects. |
| properties = [] |
| for prop in props: |
| if prop[0] not in unused_properties and not should_drop(prop[1].fget): |
| getter = get_jit_def(prop[1].fget, f"__{prop[0]}_getter", self_name=self_name) |
| setter = get_jit_def(prop[1].fset, f"__{prop[0]}_setter", self_name=self_name) if prop[1].fset else None |
| properties.append(Property(getter.range(), Ident(getter.range(), prop[0]), getter, setter)) |
| |
| return properties |
| |
| |
| def get_class_assigns(ctx, cls_ast): |
| assigns = [] |
| |
| def maybe_build_assign(builder, entry): |
| nonlocal assigns |
| try: |
| assigns.append(builder(ctx, entry)) |
| except NotSupportedError: |
| pass |
| for entry in cls_ast.body: |
| if isinstance(entry, ast.Assign): |
| maybe_build_assign(StmtBuilder.build_Assign, entry) |
| elif isinstance(entry, ast.AnnAssign): |
| maybe_build_assign(StmtBuilder.build_AnnAssign, entry) |
| return assigns |
| |
| |
| def get_jit_class_def(cls, self_name): |
| # Get defs for each method within the current class independently |
| # TODO: proper overriding analysis when implementing class inheritance |
| methods = inspect.getmembers( |
| cls, |
| predicate=lambda m: (inspect.ismethod(m) or inspect.isfunction(m)) |
| and not is_static_fn(cls, m.__name__) |
| and m.__name__ in cls.__dict__ |
| ) |
| |
| def is_classmethod(fn): |
| return inspect.ismethod(fn) and getattr(fn, "__self__", None) == cls |
| |
| # Get and parse the source code for this class |
| sourcelines, file_lineno, filename = get_source_lines_and_file(cls, torch._C.ErrorReport.call_stack()) |
| source = ''.join(sourcelines) |
| |
| dedent_src = dedent(source) |
| py_ast = ast.parse(dedent_src) |
| |
| class_ast = py_ast.body[0] |
| assert isinstance(class_ast, ast.ClassDef) |
| |
| # Special case for dataclasses. In general we need access to the source code for |
| # an object in order to JIT compile it. But the dataclasses module dynamically synthesizes |
| # magic methods for classes, and we can't get the source code for these methods. As a |
| # workaround, we synthesize TorchScript-friendly implementations ourselves. |
| if dataclasses.is_dataclass(cls): |
| # Detect whether the user manually implemented any of the magic methods. If they did, |
| # we don't want to synthesize/override them. |
| overrides = { |
| method.name |
| for method in class_ast.body |
| if isinstance(method, ast.FunctionDef) and method.name in DATACLASS_MAGIC_METHODS |
| } |
| for i, (name, _) in enumerate(methods): |
| # Is this a magic method we can synthesize? |
| synthesizer_fn = DATACLASS_MAGIC_METHODS.get(name) |
| if synthesizer_fn and name not in overrides: |
| parsed_def = synthesizer_fn(cls) |
| methods[i] = name, parsed_def |
| func = getattr(cls, name) |
| _jit_internal.loader.cache(func, parsed_def.source) |
| |
| method_defs = [ |
| get_jit_def(obj, name, self_name=self_name, is_classmethod=is_classmethod(obj)) |
| for (name, obj) in methods |
| ] |
| properties = get_class_properties(cls, self_name) |
| |
| leading_whitespace_len = len(source.split('\n', 1)[0]) - len(dedent_src.split('\n', 1)[0]) |
| ctx = make_source_context(source, filename, file_lineno, leading_whitespace_len, False) |
| assigns = get_class_assigns(ctx, class_ast) |
| |
| return build_class_def(ctx, class_ast, method_defs, properties, self_name, assigns) |
| |
| |
| def get_jit_def(fn, def_name, self_name=None, is_classmethod=False): |
| """ |
| Build a JIT AST (TreeView) from the given function. |
| |
| Args: |
| fn: A function object to compile or a pre-parsed ParsedDef object |
| def_name: The name to give to the resulting AST object. This is not |
| always the same as `fn.__name__`, for example: |
| def _forward(self): |
| ... |
| forward = _forward |
| In this case, the `__name__` attribute of the function object is "_forward", |
| but we want the result AST to have the name "forward". |
| self_name: If this function is a method, what the type name of `self` is. |
| """ |
| parsed_def = parse_def(fn) if not isinstance(fn, _ParsedDef) else fn |
| type_line = torch.jit.annotations.get_type_line(parsed_def.source) |
| fn_def = parsed_def.ast.body[0] |
| |
| if is_classmethod: |
| arg_name = fn_def.args.args[0].arg |
| # Insert a statement that assigns the first argument to the class |
| assign_stmt = ast.parse(f"{arg_name} = {self_name}").body[0] |
| fn_def.body.insert(0, assign_stmt) |
| |
| # Swap out the function signature and body if it is unused |
| if should_drop(fn): |
| unused_fn_def = ast.parse("def unused_fn(self: Any):\n\traise RuntimeError(\"Cannot call @unused methods\")") |
| if len(unused_fn_def.body) != 1 or not isinstance(unused_fn_def.body[0], ast.FunctionDef): |
| raise RuntimeError(f"Expected a single top-level function: {parsed_def.filename}:{parsed_def.file_lineno}") |
| unused_def = unused_fn_def.body[0] |
| fn_def.body = unused_def.body |
| # kwarg/vararg not supported by `build_def` |
| fn_def.args.kwarg = fn_def.args.vararg = None |
| for arg in fn_def.args.args + fn_def.args.kwonlyargs: |
| # Replace potentially unsupported type annotations by "Any" |
| arg.annotation = unused_def.args.args[0].annotation |
| |
| # If MonkeyType is installed, get all the consolidated type traces |
| # for the arguments from type_trace_db |
| type_trace_db = torch.jit._script._get_type_trace_db() |
| pdt_arg_types = None |
| if monkeytype_trace and not isinstance(fn, _ParsedDef): |
| qualname = get_qualified_name(fn) |
| pdt_arg_types = type_trace_db.get_args_types(qualname) |
| |
| return build_def(parsed_def.ctx, fn_def, type_line, def_name, self_name=self_name, pdt_arg_types=pdt_arg_types) |
| |
| # TODO: more robust handling of recognizing ignore context manager |
| def is_torch_jit_ignore_context_manager(stmt): |
| # checks if the statement is torch.jit.ignore context manager |
| if isinstance(stmt.items[0].context_expr, ast.Call): |
| # extract torch part |
| function = stmt.items[0].context_expr.func |
| if isinstance(function, ast.Attribute): |
| attr_name = function.attr |
| attr_value = function.value |
| if attr_name == "_IgnoreContextManager" and isinstance(attr_value, ast.Attribute): |
| # there should be at most two nested attributes (e.g torch.jit._IgnoreContextManager) |
| if attr_value.attr == "jit" and isinstance(attr_value.value, ast.Name): |
| if attr_value.value.id == "torch": |
| return True |
| return False |
| |
| class Builder(object): |
| def __call__(self, ctx, node): |
| method = getattr(self, 'build_' + node.__class__.__name__, None) |
| if method is None: |
| raise UnsupportedNodeError(ctx, node) |
| return method(ctx, node) |
| |
| |
| def build_class_def(ctx, py_def, methods, properties, self_name, assigns): |
| r = ctx.make_range(py_def.lineno, py_def.col_offset, |
| py_def.col_offset + len("class")) |
| return ClassDef(Ident(r, self_name), [Stmt(method) for method in methods], properties, assigns) |
| |
| |
| def build_def(ctx, py_def, type_line, def_name, self_name=None, pdt_arg_types=None): |
| body = py_def.body |
| r = ctx.make_range(py_def.lineno, |
| py_def.col_offset, |
| py_def.col_offset + len("def")) |
| |
| param_list = build_param_list(ctx, py_def.args, self_name, pdt_arg_types) |
| return_type = None |
| if getattr(py_def, 'returns', None) is not None: |
| return_type = build_expr(ctx, py_def.returns) |
| |
| decl = Decl(r, param_list, return_type) |
| is_method = self_name is not None |
| if type_line is not None: |
| type_comment_decl = torch._C.parse_type_comment(type_line) |
| decl = torch._C.merge_type_from_type_comment(decl, type_comment_decl, is_method) |
| |
| return Def(Ident(r, def_name), |
| decl, |
| build_stmts(ctx, body)) |
| |
| |
| _vararg_kwarg_err = ("Compiled functions can't take variable number of arguments " |
| "or use keyword-only arguments with defaults") |
| |
| |
| def build_param_list(ctx, py_args, self_name, pdt_arg_types=None): |
| if py_args.kwarg is not None: |
| expr = py_args.kwarg |
| ctx_range = ctx.make_range(expr.lineno, expr.col_offset - 1, expr.col_offset + len(expr.arg)) |
| raise NotSupportedError(ctx_range, _vararg_kwarg_err) |
| if py_args.vararg is not None: |
| expr = py_args.vararg |
| ctx_range = ctx.make_range(expr.lineno, expr.col_offset - 1, expr.col_offset + len(expr.arg)) |
| raise NotSupportedError(ctx_range, _vararg_kwarg_err) |
| if len(py_args.kw_defaults) > 0: |
| # kw_defaults is a list of the values for the kwargs (which default to None), |
| # so they don't actually have line numbers. |
| for arg in py_args.kw_defaults: |
| if arg is not None: |
| ctx_range = build_expr(ctx, arg).range() |
| raise NotSupportedError(ctx_range, _vararg_kwarg_err) |
| |
| # List of Tuple of args and type as inferred by profile directed typing |
| arg_and_types = [(arg, pdt_arg_types[arg.arg] if pdt_arg_types and bool(pdt_arg_types[arg.arg]) else None) |
| for arg in py_args.args] |
| arg_and_types_kwonlyargs = [(arg, pdt_arg_types[arg.arg] if pdt_arg_types and bool(pdt_arg_types[arg.arg]) |
| else None) for arg in py_args.kwonlyargs] |
| |
| result = [build_param(ctx, arg, self_name, kwarg_only=False, pdt_arg_type=arg_type) |
| for arg, arg_type in arg_and_types] |
| result += [build_param(ctx, arg, self_name, kwarg_only=True, pdt_arg_type=arg_type) |
| for arg, arg_type in arg_and_types_kwonlyargs] |
| return result |
| |
| |
| def build_param(ctx, py_arg, self_name, kwarg_only, pdt_arg_type=None): |
| # NB: In Python3 py_arg is a pair of (str arg, expr? annotation) |
| name = py_arg.arg |
| r = ctx.make_range(py_arg.lineno, py_arg.col_offset, py_arg.col_offset + len(name)) |
| if getattr(py_arg, 'annotation', None) is not None: |
| annotation_expr = build_expr(ctx, py_arg.annotation) |
| elif pdt_arg_type: |
| annotation_expr = Var(Ident(r, pdt_arg_type)) |
| elif self_name is not None and name == 'self': |
| annotation_expr = Var(Ident(r, self_name)) |
| else: |
| annotation_expr = EmptyTypeAnnotation(r) |
| return Param(annotation_expr, Ident(r, name), kwarg_only) |
| |
| def build_ignore_context_manager(ctx, stmt): |
| InputType = namedtuple('InputType', ['name', 'ann']) |
| OutputType = namedtuple('OutputType', ['name', 'ann']) |
| |
| def process_ins_outs(args): |
| # parse the context manager to figure out inputs and outputs |
| # with their annotated types |
| # TODO: add input, output validator |
| inputs = [] |
| outputs = [] |
| for arg in args: |
| var_name = arg.arg |
| if sys.version_info < (3, 8): |
| # Starting python3.8 ast.Str is deprecated |
| var_ann = arg.value.s |
| else: |
| var_ann = arg.value.value |
| var_decl_type, var_ann = var_ann.split(":") |
| if var_decl_type == "inp": |
| inputs.append(InputType(var_name, var_ann)) |
| if var_decl_type == "out": |
| outputs.append(OutputType(var_name, var_ann)) |
| return inputs, outputs |
| |
| def create_unique_name_ext(ctx, stmt): |
| # extension will be based on the full path filename plus |
| # the line number of original context manager |
| fn = re.sub(r'[^a-zA-Z0-9_]', '_', ctx.filename) |
| return f"{fn}_{stmt.lineno}" |
| |
| def build_return_ann_stmt(outputs): |
| return_type_ann = "" |
| return_statement_str = "return " |
| if len(outputs) == 0: |
| return_type_ann += " -> None" |
| if len(outputs) == 1: |
| return_type_ann = " -> " + outputs[0].ann |
| return_statement_str += outputs[0].name |
| if len(outputs) > 1: |
| return_type_ann = " -> Tuple" |
| return_type_ann += "[" + ", ".join([var.ann for var in outputs]) + "]" |
| return_statement_str += ", ".join([var.name for var in outputs]) |
| return return_type_ann, return_statement_str |
| |
| def build_args(args): |
| return ", ".join([arg.name for arg in args]) |
| |
| inputs, outputs = process_ins_outs(stmt.items[0].context_expr.keywords) |
| |
| # build the replacement function str with given inputs and outputs |
| ignore_function_name = "func_ignore_" + create_unique_name_ext(ctx, stmt) |
| ignore_function_str = "\ndef " + ignore_function_name |
| ignore_function_str += "(" + ", ".join([var.name + " :" + var.ann for var in inputs]) + ")" |
| |
| return_ann, return_stmt = build_return_ann_stmt(outputs) |
| ignore_function_str += return_ann + ": pass" |
| |
| # first create the functionDef object from just declaration |
| ignore_function = ast.parse(ignore_function_str).body[0] |
| |
| # dump the body of context manager to dummy function |
| ignore_function.body = stmt.body # type: ignore[attr-defined] |
| |
| # insert return statement to the function |
| return_stmt = ast.parse(return_stmt).body[0] |
| ignore_function.body.append(return_stmt) # type: ignore[attr-defined] |
| |
| # registers the custom function in the global context |
| ignore_func_str = "@torch.jit.ignore\n" + astunparse.unparse(ignore_function) |
| ignore_func_str += "\nglobals()[\"{}\"] = {}".format(ignore_function_name, ignore_function_name) |
| exec(ignore_func_str) # noqa: P204 |
| |
| # build the statements as: |
| # <out_1>, <out_2>, ... = torch.jit.frontend.<func>(<in_1>, <in_2>) |
| assign_str_lhs = build_args(outputs) |
| # this function will be registered in torch.jit.frontend module by default |
| assign_str_rhs = "torch.jit.frontend.{}(".format(ignore_function_name) + build_args(inputs) + ")" |
| |
| if len(outputs) > 0: |
| assign_str = assign_str_lhs + " = " + assign_str_rhs |
| else: |
| assign_str = assign_str_rhs |
| assign_ast = ast.parse(assign_str).body[0] |
| return assign_ast |
| |
| def get_default_args(fn): |
| if fn is None: |
| return {} |
| |
| signature = inspect.signature(fn) |
| |
| return { |
| k: v.default |
| for k, v in signature.parameters.items() |
| if v.default is not inspect.Parameter.empty |
| } |
| |
| |
| def get_default_args_for_class(cls): |
| """ |
| Get default arguments for all methods in a class (except for static methods). |
| |
| Args: |
| cls: type - The class type to inspect for default arguments. |
| Returns: |
| A Dict[str, Dict[str, Any]] which maps each method name to a Dict[str, Any] |
| that maps each argument name to its default value. |
| """ |
| # Get methods (except static methods because those are compiled separately as |
| # if they were independent script functions). |
| methods = inspect.getmembers( |
| cls, |
| predicate=lambda m: (inspect.ismethod(m) or inspect.isfunction(m)) |
| and not is_static_fn(cls, m.__name__) |
| and m.__name__ in cls.__dict__ |
| ) |
| |
| # Get method defaults. Property defaults do not need to be considered |
| # because setters cannot be invoked without a value. |
| defaults = {method_name: get_default_args(method_impl) for method_name, method_impl in methods} |
| |
| return defaults |
| |
| |
| class WithItemBuilder(Builder): |
| @staticmethod |
| def build_withitem(ctx, item): |
| lineno = item.context_expr.lineno |
| start = item.context_expr.col_offset |
| end = start + len(pretty_node_names[ast.With]) |
| op_vars = item.optional_vars |
| r = ctx.make_range(lineno, start, end) |
| |
| return WithItem(r, build_expr(ctx, item.context_expr), build_expr(ctx, op_vars) if op_vars else None) |
| |
| |
| class StmtBuilder(Builder): |
| augassign_map = { |
| ast.Add: '+', |
| ast.Sub: '-', |
| ast.Mult: '*', |
| ast.Div: '/', |
| ast.Mod: '%', |
| ast.BitOr: '|', |
| ast.BitAnd: '&', |
| ast.BitXor: '^', |
| ast.LShift: '<<', |
| ast.RShift: '>>', |
| ast.Pow: '**', |
| } |
| |
| @staticmethod |
| def build_Expr(ctx, stmt): |
| value = stmt.value |
| if value.__class__.__name__ == 'Str': |
| # If a statement is a string literal expression, |
| # then it is a docstring. Just ignore it. |
| return None |
| else: |
| return ExprStmt(build_expr(ctx, value)) |
| |
| @staticmethod |
| def build_Assign(ctx, stmt): |
| rhs = build_expr(ctx, stmt.value) |
| lhs = [build_expr(ctx, x) for x in stmt.targets] |
| return Assign(lhs, rhs) |
| |
| @staticmethod |
| def build_AnnAssign(ctx, stmt): |
| if stmt.value is None: |
| raise UnsupportedNodeError(ctx, stmt, reason='without assigned value') |
| |
| # Disallow type annotations on instance attributes outside of __init__ |
| if type(stmt.target) == ast.Attribute and \ |
| stmt.target.value.id == "self" and ctx.funcname != "__init__": # type: ignore[attr-defined] |
| start = stmt.col_offset |
| end = start + len(f"self.{stmt.target.attr}") |
| if hasattr(stmt.annotation, 'id'): |
| end += len(f": {stmt.annotation.id}") |
| sr = ctx.make_range(stmt.lineno, start, end) |
| raise ValueError("Type annotations on instance attributes must be declared in " |
| f"__init__, not '{ctx.funcname}': {sr}") |
| |
| rhs = build_expr(ctx, stmt.value) |
| lhs = build_expr(ctx, stmt.target) |
| the_type = build_expr(ctx, stmt.annotation) |
| return Assign([lhs], rhs, the_type) |
| |
| @staticmethod |
| def build_Delete(ctx, stmt): |
| r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("del")) |
| |
| return Delete(r, [build_expr(ctx, target) for target in stmt.targets]) |
| |
| @staticmethod |
| def build_Return(ctx, stmt): |
| r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("return")) |
| return Return(r, None if stmt.value is None else build_expr(ctx, stmt.value)) |
| |
| @staticmethod |
| def build_Raise(ctx, stmt): |
| r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("raise")) |
| expr = build_expr(ctx, stmt.exc) |
| return Raise(r, expr) |
| |
| @staticmethod |
| def build_Assert(ctx, stmt): |
| r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("assert")) |
| test = build_expr(ctx, stmt.test) |
| msg = build_expr(ctx, stmt.msg) if stmt.msg is not None else None |
| return Assert(r, test, msg) |
| |
| @staticmethod |
| def build_AugAssign(ctx, stmt): |
| lhs = build_expr(ctx, stmt.target) |
| rhs = build_expr(ctx, stmt.value) |
| op = type(stmt.op) |
| if op in StmtBuilder.augassign_map: |
| op_token = StmtBuilder.augassign_map[op] |
| else: |
| raise NotSupportedError( |
| find_before(ctx, rhs.range().start, '=', offsets=(-1, 0)), |
| "unsupported kind of augmented assignment: " + op.__name__) |
| return AugAssign(lhs, op_token, rhs) |
| |
| @staticmethod |
| def build_While(ctx, stmt): |
| if stmt.orelse: |
| # TODO: try to recover the location of else:? Python doesn't give us useful |
| # annotations in this case |
| raise NotSupportedError(None, "else branches of while loops aren't supported") |
| r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("while")) |
| return While(r, build_expr(ctx, stmt.test), |
| build_stmts(ctx, stmt.body)) |
| |
| @staticmethod |
| def build_For(ctx, stmt): |
| r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("for")) |
| if stmt.orelse: |
| raise NotSupportedError(r, "else branches of for loops aren't supported") |
| |
| return For( |
| r, [build_expr(ctx, stmt.target)], |
| [build_expr(ctx, stmt.iter)], build_stmts(ctx, stmt.body)) |
| |
| @staticmethod |
| def build_If(ctx, stmt): |
| r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("if")) |
| return If(r, build_expr(ctx, stmt.test), |
| build_stmts(ctx, stmt.body), |
| build_stmts(ctx, stmt.orelse)) |
| |
| @staticmethod |
| def build_Print(ctx, stmt): |
| r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("print")) |
| if stmt.dest: |
| raise NotSupportedError(r, "print statements with non-default destinations aren't supported") |
| args = [build_expr(ctx, val) for val in stmt.values] |
| return ExprStmt(Apply(Var(Ident(r, "print")), args, [])) |
| |
| @staticmethod |
| def build_Pass(ctx, stmt): |
| r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("pass")) |
| return Pass(r) |
| |
| @staticmethod |
| def build_Break(ctx, stmt): |
| r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("break")) |
| return Break(r) |
| |
| @staticmethod |
| def build_Continue(ctx, stmt): |
| r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("continue")) |
| return Continue(r) |
| |
| @staticmethod |
| def build_With(ctx, stmt): |
| r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset + len("with")) |
| # Handle ignore context manager |
| if is_torch_jit_ignore_context_manager(stmt): |
| if not _IS_ASTUNPARSE_INSTALLED: |
| raise RuntimeError("torch.jit._IgnoreContextManager requires installing Python library `astunparse`,\ |
| please install it in your Python environment") |
| assign_ast = build_ignore_context_manager(ctx, stmt) |
| return build_stmt(ctx, assign_ast) |
| return With(r, build_withitems(ctx, stmt.items), build_stmts(ctx, stmt.body)) |
| |
| class ExprBuilder(Builder): |
| binop_map = { |
| ast.Add: '+', |
| ast.Sub: '-', |
| ast.Mult: '*', |
| ast.Div: '/', |
| ast.Pow: '**', |
| ast.Mod: '%', |
| ast.FloorDiv: '//', |
| ast.BitAnd: '&', |
| ast.BitXor: '^', |
| ast.BitOr: '|', |
| ast.LShift: '<<', |
| ast.RShift: '>>', |
| } |
| |
| binop_map[ast.MatMult] = '@' |
| |
| unop_map = { |
| ast.Not: 'not', |
| ast.USub: '-', |
| ast.Invert: '~', |
| } |
| |
| boolop_map = { |
| ast.And: 'and', |
| ast.Or: 'or', |
| } |
| |
| cmpop_map = { |
| ast.Eq: '==', |
| ast.NotEq: '!=', |
| ast.LtE: '<=', |
| ast.Lt: '<', |
| ast.GtE: '>=', |
| ast.Gt: '>', |
| ast.Is: 'is', |
| ast.IsNot: 'is not', |
| ast.In: 'in', |
| ast.NotIn: 'not in', |
| } |
| |
| @staticmethod |
| def build_Attribute(ctx, expr): |
| base = build_expr(ctx, expr.value) |
| # expr.attr is just a string, so it's not annotated in any way, so we have |
| # to build the range manually |
| source = ctx.source.encode('utf-8') |
| |
| def get_char(index): |
| return chr(source[index]) |
| |
| start_pos = base.range().end + 1 |
| while get_char(start_pos) in string.whitespace: # Skip whitespace |
| start_pos += 1 |
| end_pos = start_pos + len(expr.attr) |
| name_range = ctx.make_raw_range(start_pos, end_pos) |
| return Select(base, Ident(name_range, expr.attr)) |
| |
| @staticmethod |
| def build_Call(ctx, expr): |
| func = build_expr(ctx, expr.func) |
| args = [build_expr(ctx, py_arg) for py_arg in expr.args] |
| if hasattr(expr, 'starargs') and expr.starargs: |
| stararg_expr = build_expr(ctx, expr.starargs) |
| args += [Starred(stararg_expr.range(), stararg_expr)] |
| kwargs = [] |
| for kw in expr.keywords: |
| kw_expr = build_expr(ctx, kw.value) |
| # XXX: we could do a better job at figuring out the range for the name here |
| if not kw.arg: |
| raise NotSupportedError(kw_expr.range(), 'keyword-arg expansion is not supported') |
| kwargs.append(Attribute(Ident(kw_expr.range(), kw.arg), kw_expr)) |
| return Apply(func, args, kwargs) |
| |
| @staticmethod |
| def build_Ellipsis(ctx, expr): |
| r = ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + 3) # len("...") == 3 |
| return Dots(r) |
| |
| @staticmethod |
| def build_Name(ctx, expr): |
| r = ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + len(expr.id)) |
| if expr.id.startswith(_reserved_prefix): |
| raise NotSupportedError(r, "names of variables used in JIT-ed functions " |
| "can't start with " + _reserved_prefix) |
| if expr.id == "True": |
| return TrueLiteral(r) |
| elif expr.id == "False": |
| return FalseLiteral(r) |
| elif expr.id == "None": |
| return NoneLiteral(r) |
| elif expr.id == "Ellipsis": |
| return Dots(r) |
| return Var(Ident(r, expr.id)) |
| |
| @staticmethod |
| def build_NameConstant(ctx, expr): |
| r = ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + len(str(expr.value))) |
| if expr.value is True: |
| return TrueLiteral(r) |
| elif expr.value is False: |
| return FalseLiteral(r) |
| elif expr.value is None: |
| return NoneLiteral(r) |
| elif expr.value == Ellipsis: |
| return Dots(r) |
| else: |
| raise ValueError("Name constant value unsupported: " + str(expr.value)) |
| |
| @staticmethod |
| def build_BinOp(ctx, expr): |
| lhs = build_expr(ctx, expr.left) |
| rhs = build_expr(ctx, expr.right) |
| op = type(expr.op) |
| |
| if op == ast.Div and not ctx.uses_true_division: |
| err_range = ctx.make_raw_range(lhs.range().end, rhs.range().start) |
| raise FrontendError(err_range, 'Division of ints in TorchScript uses Python 3 true ' |
| 'division semantics. Please put `from __future__ ' |
| 'import division` at the top of your file') |
| op_token = ExprBuilder.binop_map.get(op) |
| if op_token is None: |
| err_range = ctx.make_raw_range(lhs.range().end, rhs.range().start) |
| raise NotSupportedError(err_range, "unsupported binary operator: " + op.__name__) |
| return BinOp(op_token, lhs, rhs) |
| |
| @staticmethod |
| def build_UnaryOp(ctx, expr): |
| sub_expr = build_expr(ctx, expr.operand) |
| op = type(expr.op) |
| op_token = ExprBuilder.unop_map.get(op) |
| if op_token is None: |
| raise NotSupportedError(expr.range(), "unsupported unary operator: " + op.__name__) |
| r = ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + len(op_token)) |
| return UnaryOp(r, op_token, sub_expr) |
| |
| @staticmethod |
| def build_BoolOp(ctx, expr): |
| if len(expr.values) < 2: |
| raise AssertionError("expected at least 2 values in BoolOp, but got " + str(len(expr.values))) |
| sub_exprs = [build_expr(ctx, sub_expr) for sub_expr in expr.values] |
| op = type(expr.op) |
| op_token = ExprBuilder.boolop_map.get(op) |
| if op_token is None: |
| err_range = ctx.make_raw_range(sub_exprs[0].range().end, sub_exprs[1].range().start) |
| raise NotSupportedError(err_range, "unsupported boolean operator: " + op.__name__) |
| lhs = sub_exprs[0] |
| for rhs in sub_exprs[1:]: |
| lhs = BinOp(op_token, lhs, rhs) |
| return lhs |
| |
| @staticmethod |
| def build_IfExp(ctx, expr): |
| return TernaryIf(build_expr(ctx, expr.test), |
| build_expr(ctx, expr.body), |
| build_expr(ctx, expr.orelse)) |
| |
| @staticmethod |
| def build_Compare(ctx, expr): |
| operands = [build_expr(ctx, e) for e in [expr.left] + list(expr.comparators)] |
| result = None |
| for lhs, op_, rhs in zip(operands, expr.ops, operands[1:]): |
| op = type(op_) |
| op_token = ExprBuilder.cmpop_map.get(op) |
| r = ctx.make_raw_range(lhs.range().end, rhs.range().start) |
| if op_token is None: |
| raise NotSupportedError(r, "unsupported comparison operator: " + op.__name__) |
| |
| if op == ast.NotIn: |
| # NB: `not in` is just `not( in )`, so we don't introduce new tree view |
| # but just make it a nested call in our tree view structure |
| in_expr = BinOp('in', lhs, rhs) |
| cmp_expr = UnaryOp(r, 'not', in_expr) |
| else: |
| cmp_expr = BinOp(op_token, lhs, rhs) |
| |
| if result is None: |
| result = cmp_expr |
| else: |
| result = BinOp('and', result, cmp_expr) |
| return result |
| |
| @staticmethod |
| def build_Subscript(ctx, expr): |
| def build_SliceExpr(ctx, base, slice_expr): |
| lower = build_expr(ctx, slice_expr.lower) if slice_expr.lower is not None else None |
| upper = build_expr(ctx, slice_expr.upper) if slice_expr.upper is not None else None |
| step = build_expr(ctx, slice_expr.step) if slice_expr.step is not None else None |
| return SliceExpr(base.range(), lower, upper, step) |
| |
| def build_Index(ctx, base, index_expr): |
| if isinstance(index_expr.value, ast.Tuple): |
| raise NotSupportedError(base.range(), |
| "slicing multiple dimensions with " |
| "tuples not supported yet") |
| return build_expr(ctx, index_expr.value) |
| |
| def build_ExtSlice(ctx, base, extslice): |
| sub_exprs = [] |
| for expr in extslice.dims: |
| sub_type = type(expr) |
| if sub_type is ast.Index: |
| sub_exprs.append(build_Index(ctx, base, expr)) |
| elif sub_type is ast.Slice: |
| sub_exprs.append(build_SliceExpr(ctx, base, expr)) |
| elif sub_type is ast.Ellipsis: |
| sub_exprs.append(Dots(base.range())) |
| else: |
| raise NotSupportedError(base.range(), |
| "slicing multiple dimensions with " |
| "{} not supported".format(sub_type)) |
| return sub_exprs |
| base = build_expr(ctx, expr.value) |
| sub_type = type(expr.slice) |
| if sub_type is ast.Index: |
| if isinstance(expr.slice.value, ast.Tuple): |
| # N-dimensional indexing using Tuple: x[(i, j, k)] is equivalent to x[i, j, k] |
| # XXX: Indexing using a list is **different**! It triggers advanced indexing. |
| indices = [build_expr(ctx, index_expr) for index_expr in expr.slice.value.elts] |
| if not indices: |
| # `col_offset` is an int, but `end_col_offset` is |
| # `Optional[int]`. The magic number is here to make |
| # sure we can parse `()` on any machine |
| r = ctx.make_range(expr.lineno, |
| expr.slice.value.col_offset, |
| expr.slice.value.col_offset + 2) |
| tup = TupleLiteral(r, []) |
| indices.append(tup) |
| return Subscript(base, indices) |
| else: |
| return Subscript(base, [build_expr(ctx, expr.slice.value)]) |
| elif sub_type is ast.Slice: |
| return Subscript(base, [build_SliceExpr(ctx, base, expr.slice)]) |
| elif sub_type is ast.ExtSlice: |
| return Subscript(base, build_ExtSlice(ctx, base, expr.slice)) |
| elif sys.version_info >= (3, 9): # In Python3.9 array indicies are not wrapped in ast.Index |
| if sub_type is ast.Tuple: |
| # N-dimensional indexing using Tuple: x[(i, j, k)] is equivalent to x[i, j, k] |
| indices = [] |
| for index_expr in expr.slice.elts: |
| if isinstance(index_expr, ast.Slice): |
| indices.append(build_SliceExpr(ctx, base, index_expr)) |
| else: |
| indices.append(build_expr(ctx, index_expr)) |
| # Special-case logic for `typing.Tuple[()]` |
| if not indices: |
| # See note above r.e. magic number |
| r = ctx.make_range(expr.lineno, |
| expr.slice.col_offset, |
| expr.slice.col_offset + 2) |
| tup = TupleLiteral(r, []) |
| indices.append(tup) |
| return Subscript(base, indices) |
| return Subscript(base, [build_expr(ctx, expr.slice)]) |
| else: # Ellipsis (can only happen in Python 2) |
| raise NotSupportedError(base.range(), "ellipsis is not supported") |
| |
| @staticmethod |
| def build_List(ctx, expr): |
| return ListLiteral(ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + 1), |
| [build_expr(ctx, e) for e in expr.elts]) |
| |
| @staticmethod |
| def build_Tuple(ctx, expr): |
| return TupleLiteral(ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + 1), |
| [build_expr(ctx, e) for e in expr.elts]) |
| |
| @staticmethod |
| def build_Dict(ctx, expr): |
| range = ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + 1) |
| if expr.keys and not expr.keys[0]: |
| raise NotSupportedError(range, "Dict expansion (e.g. `{**dict}`) is not supported") |
| return DictLiteral(range, [build_expr(ctx, e) for e in expr.keys], |
| [build_expr(ctx, e) for e in expr.values]) |
| |
| @staticmethod |
| def build_Num(ctx, expr): |
| value = str(expr.n) |
| r = ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + len(value)) |
| return Const(r, value) |
| |
| @staticmethod |
| def build_Constant(ctx, expr): |
| value = expr.value |
| if value is None or isinstance(value, bool): |
| # NB: this check has to happen before the int check because bool is |
| # a subclass of int |
| return ExprBuilder.build_NameConstant(ctx, expr) |
| if isinstance(value, (int, float, complex)): |
| return ExprBuilder.build_Num(ctx, expr) |
| elif isinstance(value, str): |
| return ExprBuilder.build_Str(ctx, expr) |
| elif isinstance(value, type(Ellipsis)): |
| return ExprBuilder.build_Ellipsis(ctx, expr) |
| else: |
| error_range = ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + len(str(value))) |
| raise FrontendError(error_range, "Unknown Constant expression type") |
| |
| @staticmethod |
| def build_Str(ctx, expr): |
| value = str(expr.s) |
| r = ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + len(value) + 1) |
| return StringLiteral(r, value) |
| |
| @staticmethod |
| def build_JoinedStr(ctx, expr): |
| s = '' |
| args = [] |
| for value in expr.values: |
| r = ctx.make_range(value.lineno, value.col_offset, value.col_offset + 1) |
| if isinstance(value, ast.FormattedValue): |
| if value.conversion != -1: |
| raise NotSupportedError(r, 'Don\'t support conversion in JoinedStr') |
| if value.format_spec is not None: |
| raise NotSupportedError(r, 'Don\'t support formatting in JoinedStr') |
| s += '{}' |
| args.append(build_expr(ctx, value.value)) |
| elif isinstance(value, ast.Str): |
| s += value.s |
| else: |
| raise NotSupportedError(r, 'Unsupported value in JoinedStr') |
| |
| r = ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + 1) |
| return Apply(Select(StringLiteral(r, s), Ident(r, 'format')), args, []) |
| |
| @staticmethod |
| def build_ListComp(ctx, stmt): |
| r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset) |
| if (len(stmt.generators) != 1): |
| raise NotSupportedError(r, "Only a single generator is currently supported") |
| |
| if (len(stmt.generators[0].ifs) != 0): |
| raise NotSupportedError(r, "Comprehension ifs are not supported yet") |
| |
| elt_expr = build_expr(ctx, stmt.elt) |
| target_expr = build_expr(ctx, stmt.generators[0].target) |
| iter_expr = build_expr(ctx, stmt.generators[0].iter) |
| |
| return ListComp(r, elt_expr, target_expr, iter_expr) |
| |
| @staticmethod |
| def build_GeneratorExp(ctx, stmt): |
| # Convert Generator expression to ListComp |
| return ExprBuilder.build_ListComp(ctx, stmt) |
| |
| @staticmethod |
| def build_DictComp(ctx, stmt): |
| r = ctx.make_range(stmt.lineno, stmt.col_offset, stmt.col_offset) |
| if (len(stmt.generators) != 1): |
| raise NotSupportedError(r, "Only a single generator is currently supported") |
| |
| if (len(stmt.generators[0].ifs) != 0): |
| raise NotSupportedError(r, "Comprehension ifs are not supported yet") |
| |
| key_expr = build_expr(ctx, stmt.key) |
| value_expr = build_expr(ctx, stmt.value) |
| target_expr = build_expr(ctx, stmt.generators[0].target) |
| iter_expr = build_expr(ctx, stmt.generators[0].iter) |
| |
| return DictComp(r, key_expr, value_expr, target_expr, iter_expr) |
| |
| @staticmethod |
| def build_Starred(ctx, expr): |
| r = ctx.make_range(expr.lineno, expr.col_offset, expr.col_offset + 1) |
| return Starred(r, build_expr(ctx, expr.value)) |
| |
| build_expr = ExprBuilder() |
| build_stmt = StmtBuilder() |
| build_withitem = WithItemBuilder() |
| |
| def find_before(ctx, pos, substr, offsets=(0, 0)): |
| new_pos = ctx.source[:pos].rindex(substr) |
| return ctx.make_raw_range(new_pos + offsets[0], new_pos + len(substr) + offsets[1]) |
