| import ast |
| import enum |
| import inspect |
| import re |
| import builtins |
| import torch |
| import warnings |
| from .._jit_internal import List, Tuple, is_tuple, is_list, Dict, is_dict, Optional, \ |
| is_optional, _qualified_name, Any, Future, is_future, is_ignored_fn, Union, is_union |
| from .._jit_internal import BroadcastingList1, BroadcastingList2, BroadcastingList3 # type: ignore[attr-defined] |
| from ._state import _get_script_class |
| |
| from torch._C import TensorType, TupleType, FloatType, IntType, ComplexType, \ |
| ListType, StringType, DictType, BoolType, OptionalType, InterfaceType, AnyType, \ |
| NoneType, DeviceObjType, StreamObjType, FutureType, EnumType, UnionType, NumberType |
| |
| |
| from textwrap import dedent |
| from torch._sources import get_source_lines_and_file |
| from typing import Type |
| |
| if torch.distributed.rpc.is_available(): |
| from .._jit_internal import RRef, is_rref |
| from torch._C import RRefType |
| |
| from torch._ops import OpOverloadPacket |
| |
| class Module(object): |
| def __init__(self, name, members): |
| self.name = name |
| self.members = members |
| |
| def __getattr__(self, name): |
| try: |
| return self.members[name] |
| except KeyError: |
| raise RuntimeError(f"Module {self.name} has no member called {name}") from None |
| |
| |
| class EvalEnv(object): |
| env = { |
| 'torch': Module('torch', {'Tensor': torch.Tensor}), |
| 'Tensor': torch.Tensor, |
| 'typing': Module('typing', {'Tuple': Tuple}), |
| 'Tuple': Tuple, |
| 'List': List, |
| 'Dict': Dict, |
| 'Optional': Optional, |
| 'Union': Union, |
| 'Future': Future |
| } |
| |
| def __init__(self, rcb): |
| self.rcb = rcb |
| if torch.distributed.rpc.is_available(): |
| self.env['RRef'] = RRef |
| |
| def __getitem__(self, name): |
| if name in self.env: |
| return self.env[name] |
| if self.rcb is not None: |
| return self.rcb(name) |
| return getattr(builtins, name, None) |
| |
| def get_signature(fn, rcb, loc, is_method): |
| if isinstance(fn, OpOverloadPacket): |
| signature = try_real_annotations(fn.op, loc) |
| else: |
| signature = try_real_annotations(fn, loc) |
| if signature is not None and is_method: |
| # If this is a method, then the signature will include a type for |
| # `self`, but type comments do not contain a `self`. So strip it |
| # away here so everything is consistent (`inspect.ismethod` does |
| # not work here since `fn` is unbound at this point) |
| param_types, return_type = signature |
| param_types = param_types[1:] |
| signature = (param_types, return_type) |
| |
| if signature is None: |
| type_line, source = None, None |
| try: |
| source = dedent(''.join(get_source_lines_and_file(fn)[0])) |
| type_line = get_type_line(source) |
| except TypeError: |
| pass |
| # This might happen both because we failed to get the source of fn, or |
| # because it didn't have any annotations. |
| if type_line is not None: |
| signature = parse_type_line(type_line, rcb, loc) |
| |
| return signature |
| |
| |
| def is_function_or_method(the_callable): |
| # A stricter version of `inspect.isroutine` that does not pass for built-in |
| # functions |
| return inspect.isfunction(the_callable) or inspect.ismethod(the_callable) |
| |
| |
| def is_vararg(the_callable): |
| if not is_function_or_method(the_callable) and hasattr(the_callable, '__call__'): # noqa: B004 |
| # If `the_callable` is a class, de-sugar the call so we can still get |
| # the signature |
| the_callable = the_callable.__call__ |
| |
| if is_function_or_method(the_callable): |
| return inspect.getfullargspec(the_callable).varargs is not None |
| else: |
| return False |
| |
| |
| def get_param_names(fn, n_args): |
| if isinstance(fn, OpOverloadPacket): |
| fn = fn.op |
| |
| if not is_function_or_method(fn) and hasattr(fn, '__call__') and is_function_or_method(fn.__call__): # noqa: B004 |
| # De-sugar calls to classes |
| fn = fn.__call__ |
| |
| if is_function_or_method(fn): |
| if is_ignored_fn(fn): |
| fn = inspect.unwrap(fn) |
| return inspect.getfullargspec(fn).args |
| else: |
| # The `fn` was not a method or function (maybe a class with a __call__ |
| # method, so use a default param name list) |
| return [str(i) for i in range(n_args)] |
| |
| |
| def check_fn(fn, loc): |
| # Make sure the function definition is not a class instantiation |
| try: |
| source = dedent(''.join(get_source_lines_and_file(fn)[0])) |
| except (TypeError, IOError): |
| return |
| if source is None: |
| return |
| |
| py_ast = ast.parse(source) |
| if len(py_ast.body) == 1 and isinstance(py_ast.body[0], ast.ClassDef): |
| raise torch.jit.frontend.FrontendError( |
| loc, f"Cannot instantiate class '{py_ast.body[0].name}' in a script function") |
| if len(py_ast.body) != 1 or not isinstance(py_ast.body[0], ast.FunctionDef): |
| raise torch.jit.frontend.FrontendError(loc, "Expected a single top-level function") |
| |
| |
| def parse_type_line(type_line, rcb, loc): |
| """Parses a type annotation specified as a comment. |
| |
| Example inputs: |
| # type: (Tensor, torch.Tensor) -> Tuple[Tensor] |
| # type: (Tensor, Tuple[Tensor, Tensor]) -> Tensor |
| """ |
| arg_ann_str, ret_ann_str = split_type_line(type_line) |
| |
| try: |
| arg_ann = eval(arg_ann_str, {}, EvalEnv(rcb)) # type: ignore[arg-type] # noqa: P204 |
| except (NameError, SyntaxError) as e: |
| raise RuntimeError("Failed to parse the argument list of a type annotation") from e |
| |
| if not isinstance(arg_ann, tuple): |
| arg_ann = (arg_ann,) |
| |
| try: |
| ret_ann = eval(ret_ann_str, {}, EvalEnv(rcb)) # type: ignore[arg-type] # noqa: P204 |
| except (NameError, SyntaxError) as e: |
| raise RuntimeError("Failed to parse the return type of a type annotation") from e |
| |
| arg_types = [ann_to_type(ann, loc) for ann in arg_ann] |
| return arg_types, ann_to_type(ret_ann, loc) |
| |
| |
| def get_type_line(source): |
| """Tries to find the line containing a comment with the type annotation.""" |
| type_comment = '# type:' |
| |
| lines = source.split('\n') |
| lines = [(line_num, line) for line_num, line in enumerate(lines)] |
| type_lines = list(filter(lambda line: type_comment in line[1], lines)) |
| # `type: ignore` comments may be needed in JIT'ed functions for mypy, due |
| # to the hack in torch/_VF.py. |
| |
| # An ignore type comment can be of following format: |
| # 1) type: ignore |
| # 2) type: ignore[rule-code] |
| # This ignore statement must be at the end of the line |
| |
| # adding an extra backslash before the space, to avoid triggering |
| # one of the checks in .github/workflows/lint.yml |
| type_pattern = re.compile("# type:\\ ignore(\\[[a-zA-Z-]+\\])?$") |
| type_lines = list(filter(lambda line: not type_pattern.search(line[1]), |
| type_lines)) |
| |
| if len(type_lines) == 0: |
| # Catch common typo patterns like extra spaces, typo in 'ignore', etc. |
| wrong_type_pattern = re.compile("#[\t ]*type[\t ]*(?!: ignore(\\[.*\\])?$):") |
| wrong_type_lines = list(filter(lambda line: wrong_type_pattern.search(line[1]), lines)) |
| if len(wrong_type_lines) > 0: |
| raise RuntimeError("The annotation prefix in line " + str(wrong_type_lines[0][0]) |
| + " is probably invalid.\nIt must be '# type:'" |
| + "\nSee PEP 484 (https://www.python.org/dev/peps/pep-0484/#suggested-syntax-for-python-2-7-and-straddling-code)" # noqa: B950 |
| + "\nfor examples") |
| return None |
| elif len(type_lines) == 1: |
| # Only 1 type line, quit now |
| return type_lines[0][1].strip() |
| |
| # Parse split up argument types according to PEP 484 |
| # https://www.python.org/dev/peps/pep-0484/#suggested-syntax-for-python-2-7-and-straddling-code |
| return_line = None |
| parameter_type_lines = [] |
| for line_num, line in type_lines: |
| if '# type: (...) -> ' in line: |
| return_line = (line_num, line) |
| break |
| elif type_comment in line: |
| parameter_type_lines.append(line) |
| if return_line is None: |
| raise RuntimeError( |
| "Return type line '# type: (...) -> ...' not found on multiline " |
| "type annotation\nfor type lines:\n" + |
| '\n'.join([line[1] for line in type_lines]) + |
| "\n(See PEP 484 https://www.python.org/dev/peps/pep-0484/#suggested-syntax-for-python-2-7-and-straddling-code)") |
| |
| def get_parameter_type(line): |
| item_type = line[line.find(type_comment) + len(type_comment):] |
| return item_type.strip() |
| |
| types = map(get_parameter_type, parameter_type_lines) |
| parameter_types = ", ".join(types) |
| |
| return return_line[1].replace("...", parameter_types) |
| |
| |
| def split_type_line(type_line): |
| """Splits the comment with the type annotation into parts for argument and return types. |
| |
| For example, for an input of: |
| # type: (Tensor, torch.Tensor) -> Tuple[Tensor, Tensor] |
| |
| This function will return: |
| ("(Tensor, torch.Tensor)", "Tuple[Tensor, Tensor]") |
| |
| """ |
| start_offset = len('# type:') |
| try: |
| arrow_pos = type_line.index('->') |
| except ValueError: |
| raise RuntimeError("Syntax error in type annotation (cound't find `->`)") from None |
| return type_line[start_offset:arrow_pos].strip(), type_line[arrow_pos + 2:].strip() |
| |
| |
| def try_real_annotations(fn, loc): |
| """Tries to use the Py3.5+ annotation syntax to get the type.""" |
| try: |
| # Note: anything annotated as `Optional[T]` will automatically |
| # be returned as `Union[T, None]` per |
| # https://github.com/python/typing/blob/master/src/typing.py#L850 |
| sig = inspect.signature(fn) |
| except ValueError: |
| return None |
| |
| all_annots = [sig.return_annotation] + [p.annotation for p in sig.parameters.values()] |
| if all(ann is sig.empty for ann in all_annots): |
| return None |
| |
| arg_types = [ann_to_type(p.annotation, loc) |
| for p in sig.parameters.values()] |
| return_type = ann_to_type(sig.return_annotation, loc) |
| return arg_types, return_type |
| |
| |
| # Finds common type for enum values belonging to an Enum class. If not all |
| # values have the same type, AnyType is returned. |
| def get_enum_value_type(e: Type[enum.Enum], loc): |
| enum_values: List[enum.Enum] = list(e) |
| if not enum_values: |
| raise ValueError(f"No enum values defined for: '{e.__class__}'") |
| |
| types = {type(v.value) for v in enum_values} |
| ir_types = [try_ann_to_type(t, loc) for t in types] |
| |
| # If Enum values are of different types, an exception will be raised here. |
| # Even though Python supports this case, we chose to not implement it to |
| # avoid overcomplicate logic here for a rare use case. Please report a |
| # feature request if you find it necessary. |
| res = torch._C.unify_type_list(ir_types) |
| if not res: |
| return AnyType.get() |
| return res |
| |
| def is_tensor(ann): |
| if issubclass(ann, torch.Tensor): |
| return True |
| |
| if issubclass(ann, (torch.LongTensor, torch.DoubleTensor, torch.FloatTensor, |
| torch.IntTensor, torch.ShortTensor, torch.HalfTensor, |
| torch.CharTensor, torch.ByteTensor, torch.BoolTensor)): |
| warnings.warn("TorchScript will treat type annotations of Tensor " |
| "dtype-specific subtypes as if they are normal Tensors. " |
| "dtype constraints are not enforced in compilation either.") |
| return True |
| |
| return False |
| |
| |
| |
| def try_ann_to_type(ann, loc): |
| if ann is inspect.Signature.empty: |
| return TensorType.getInferred() |
| if ann is None: |
| return NoneType.get() |
| if inspect.isclass(ann) and is_tensor(ann): |
| return TensorType.get() |
| if is_tuple(ann): |
| # Special case for the empty Tuple type annotation `Tuple[()]` |
| if len(ann.__args__) == 1 and ann.__args__[0] == (): |
| return TupleType([]) |
| return TupleType([try_ann_to_type(a, loc) for a in ann.__args__]) |
| if is_list(ann): |
| elem_type = try_ann_to_type(ann.__args__[0], loc) |
| if elem_type: |
| return ListType(elem_type) |
| if is_dict(ann): |
| key = try_ann_to_type(ann.__args__[0], loc) |
| value = try_ann_to_type(ann.__args__[1], loc) |
| # Raise error if key or value is None |
| if key is None: |
| raise ValueError(f"Unknown type annotation: '{ann.__args__[0]}' at {loc.highlight()}") |
| if value is None: |
| raise ValueError(f"Unknown type annotation: '{ann.__args__[1]}' at {loc.highlight()}") |
| return DictType(key, value) |
| if is_optional(ann): |
| if issubclass(ann.__args__[1], type(None)): |
| contained = ann.__args__[0] |
| else: |
| contained = ann.__args__[1] |
| valid_type = try_ann_to_type(contained, loc) |
| msg = "Unsupported annotation {} could not be resolved because {} could not be resolved." |
| assert valid_type, msg.format(repr(ann), repr(contained)) |
| return OptionalType(valid_type) |
| if is_union(ann): |
| # TODO: this is hack to recognize NumberType |
| if set(ann.__args__) == set([int, float, complex]): |
| return NumberType.get() |
| inner: List = [] |
| # We need these extra checks because both `None` and invalid |
| # values will return `None` |
| # TODO: Determine if the other cases need to be fixed as well |
| for a in ann.__args__: |
| if a is None: |
| inner.append(NoneType.get()) |
| maybe_type = try_ann_to_type(a, loc) |
| msg = "Unsupported annotation {} could not be resolved because {} could not be resolved." |
| assert maybe_type, msg.format(repr(ann), repr(maybe_type)) |
| inner.append(maybe_type) |
| return UnionType(inner) # type: ignore[arg-type] |
| if torch.distributed.rpc.is_available() and is_rref(ann): |
| return RRefType(try_ann_to_type(ann.__args__[0], loc)) |
| if is_future(ann): |
| return FutureType(try_ann_to_type(ann.__args__[0], loc)) |
| if ann is float: |
| return FloatType.get() |
| if ann is complex: |
| return ComplexType.get() |
| if ann is int: |
| return IntType.get() |
| if ann is str: |
| return StringType.get() |
| if ann is bool: |
| return BoolType.get() |
| if ann is Any: |
| return AnyType.get() |
| if ann is type(None): |
| return NoneType.get() |
| if inspect.isclass(ann) and hasattr(ann, "__torch_script_interface__"): |
| return InterfaceType(ann.__torch_script_interface__) |
| if ann is torch.device: |
| return DeviceObjType.get() |
| if ann is torch.Stream: |
| return StreamObjType.get() |
| if ann is torch.dtype: |
| return IntType.get() # dtype not yet bound in as its own type |
| if inspect.isclass(ann) and issubclass(ann, enum.Enum): |
| if _get_script_class(ann) is None: |
| scripted_class = torch.jit._script._recursive_compile_class(ann, loc) |
| name = scripted_class.qualified_name() |
| else: |
| name = _qualified_name(ann) |
| return EnumType(name, get_enum_value_type(ann, loc), list(ann)) |
| if inspect.isclass(ann): |
| maybe_script_class = _get_script_class(ann) |
| if maybe_script_class is not None: |
| return maybe_script_class |
| if torch._jit_internal.can_compile_class(ann): |
| return torch.jit._script._recursive_compile_class(ann, loc) |
| |
| # Maybe resolve a NamedTuple to a Tuple Type |
| def fake_rcb(key): |
| return None |
| return torch._C._resolve_type_from_object(ann, loc, fake_rcb) |
| |
| |
| def ann_to_type(ann, loc): |
| the_type = try_ann_to_type(ann, loc) |
| if the_type is not None: |
| return the_type |
| raise ValueError(f"Unknown type annotation: '{ann}' at {loc.highlight()}") |
| |
| |
| __all__ = [ |
| 'Any', |
| 'List', |
| 'BroadcastingList1', |
| 'BroadcastingList2', |
| 'BroadcastingList3', |
| 'Tuple', |
| 'is_tuple', |
| 'is_list', |
| 'Dict', |
| 'is_dict', |
| 'is_optional', |
| 'is_union', |
| 'TensorType', |
| 'TupleType', |
| 'FloatType', |
| 'ComplexType', |
| 'IntType', |
| 'ListType', |
| 'StringType', |
| 'DictType', |
| 'AnyType', |
| 'Module', |
| # TODO: Consider not exporting these during wildcard import (reserve |
| # that for the types; for idiomatic typing code.) |
| 'get_signature', |
| 'check_fn', |
| 'get_param_names', |
| 'parse_type_line', |
| 'get_type_line', |
| 'split_type_line', |
| 'try_real_annotations', |
| 'try_ann_to_type', |
| 'ann_to_type', |
| ] |
