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()))