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_drop_fn, 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__ and not _is_drop_fn(m) ) 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 _is_drop_fn(fn): # Dropping potentially unsupported return type annotation for jit._drop fn_def.returns = None fn_def.type_comment = None # 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: 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 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: # , , ... = torch.jit.frontend.(, ) 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])