def test_star_args(self): # type: () -> None self.assert_type_comment( '(*str) -> Any', ([Argument(ClassType('str'), ARG_STAR)], AnyType())) self.assert_type_comment( '(int, *str) -> Any', ([class_arg('int'), Argument(ClassType('str'), ARG_STAR)], AnyType()))
def infer_annotation(type_comments): # type: (List[str]) -> Tuple[List[Argument], AbstractType] """Given some type comments, return a single inferred signature. Args: type_comments: Strings of form '(arg1, ... argN) -> ret' Returns: Tuple of (argument types and kinds, return type). """ assert type_comments args = {} # type: Dict[int, Set[Argument]] returns = set() for comment in type_comments: arg_types, return_type = parse_type_comment(comment) for i, arg_type in enumerate(arg_types): args.setdefault(i, set()).add(arg_type) returns.add(return_type) combined_args = [] for i in sorted(args): arg_infos = list(args[i]) kind = argument_kind(arg_infos) if kind is None: raise InferError('Ambiguous argument kinds:\n' + '\n'.join(type_comments)) types = [arg.type for arg in arg_infos] combined = combine_types(types) if kind != ARG_POS and (len(str(combined)) > 120 or isinstance(combined, UnionType)): # Avoid some noise. combined = AnyType() combined_args.append(Argument(combined, kind)) combined_return = combine_types(returns) return combined_args, combined_return
def test_union(self): # type: () -> None self.assert_type_comment('(Union[int, str]) -> Any', ([ Argument(UnionType([ClassType('int'), ClassType('str')]), ARG_POS) ], AnyType())) self.assert_type_comment('(Union[int]) -> Any', ([class_arg('int')], AnyType()))
def parse(self): # type: () -> Tuple[List[Argument], AbstractType] self.expect('(') arg_types = [] # type: List[Argument] stars_seen = set() # type: Set[str] while self.lookup() != ')': if self.lookup() == '*': self.expect('*') if self.lookup() == '*': if '**' in stars_seen: self.fail() self.expect('*') star_star = True else: if stars_seen: self.fail() star_star = False arg_type = self.parse_type() if star_star: arg_types.append(Argument(arg_type, ARG_STARSTAR)) stars_seen.add('**') else: arg_types.append(Argument(arg_type, ARG_STAR)) stars_seen.add('*') else: if stars_seen: self.fail() arg_type = self.parse_type() arg_types.append(Argument(arg_type, ARG_POS)) if self.lookup() == ',': self.expect(',') elif self.lookup() == ')': break self.expect(')') self.expect('->') ret_type = self.parse_type() if not isinstance(self.next(), End): self.fail() return arg_types, ret_type
def tuple_arg(items): # type: (List[AbstractType]) -> Argument return Argument(TupleType(items), ARG_POS)
def any_arg(): # type: () -> Argument return Argument(AnyType(), ARG_POS)
def class_arg(name, args=None): # type: (str, Optional[List[AbstractType]]) -> Argument return Argument(ClassType(name, args), ARG_POS)
def test_optional(self): # type: () -> None self.assert_type_comment('(Optional[int]) -> Any', ([ Argument(UnionType([ClassType('int'), ClassType('None')]), ARG_POS) ], AnyType()))