def parse_type(self):
# type: () -> AbstractType
t = self.next()
if not isinstance(t, DottedName):
self.fail()
if t.text == 'Any':
return AnyType()
elif t.text == 'Tuple':
self.expect('[')
args = self.parse_type_list()
self.expect(']')
return TupleType(args)
elif t.text == 'Union':
self.expect('[')
items = self.parse_type_list()
self.expect(']')
if len(items) == 1:
return items[0]
elif len(items) == 0:
self.fail()
else:
return UnionType(items)
else:
if self.lookup() == '[':
self.expect('[')
args = self.parse_type_list()
self.expect(']')
if t.text == 'Optional' and len(args) == 1:
return UnionType([args[0], ClassType('None')])
return ClassType(t.text, args)
else:
return ClassType(t.text)
def test_nested_tuples(self):
# type: () -> None
assert flatten_types(
[UnionType([UnionType([CT('int'), CT('str')]),
CT('X')])]) == [CT('int'),
CT('str'),
CT('X')]
def test_optional(Self):
# type: () -> None
assert str(UnionType([ClassType('str'),
ClassType('None')])) == 'Optional[str]'
assert str(UnionType([ClassType('None'),
ClassType('str')])) == 'Optional[str]'
assert str(
UnionType([ClassType('None'),
ClassType('str'),
ClassType('int')])) == 'Union[None, str, int]'
def test_infer_union_arg(self):
# type: () -> None
self.assert_infer(['(int) -> None',
'(str) -> None'],
([(UnionType([ClassType('int'),
ClassType('str')]), ARG_POS)],
ClassType('None')))
def test_merge_union_of_same_length_tuples(self):
# type: () -> None
assert merge_items([TupleType([CT('str')]),
TupleType([CT('int')])]) == [TupleType([UnionType([CT('str'),
CT('int')])])]
assert merge_items([TupleType([CT('str')]),
TupleType([CT('Text')])]) == [TupleType([CT('Text')])]
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 str(combined) == 'None':
# It's very rare for an argument to actually be typed `None`, more likely than
# not we simply don't have any data points for this argument.
combined = UnionType([ClassType('None'), AnyType()])
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_merge_unions(self):
# type: () -> None
self.assert_infer(['(Union[int, str]) -> None',
'(Union[str, None]) -> None'],
([(UnionType([ClassType('int'),
ClassType('str'),
ClassType('None')]), ARG_POS)],
ClassType('None')))
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 combine_types(types):
# type: (Iterable[AbstractType]) -> AbstractType
"""Given some types, return a combined and simplified type.
For example, if given 'int' and 'List[int]', return Union[int, List[int]]. If given
'int' and 'int', return just 'int'.
"""
items = simplify_types(types)
if len(items) == 1:
return items[0]
else:
return UnionType(items)
def test_simplify_potential_typed_dict(self):
# type: () -> None
# Fall back to Dict[x, Any] in case of a complex Dict type.
self.assert_infer(['(Dict[str, Union[int, str]]) -> Any'],
([(ClassType('Dict', [ClassType('str'), AnyType()]), ARG_POS)],
AnyType()))
self.assert_infer(['(Dict[Text, Union[int, str]]) -> Any'],
([(ClassType('Dict', [ClassType('Text'), AnyType()]), ARG_POS)],
AnyType()))
# Not a potential TypedDict so ordinary simplification applies.
self.assert_infer(['(Dict[str, Union[str, Text]]) -> Any'],
([(ClassType('Dict', [ClassType('str'), ClassType('Text')]), ARG_POS)],
AnyType()))
self.assert_infer(['(Dict[str, Union[int, None]]) -> Any'],
([(ClassType('Dict', [ClassType('str'),
UnionType([ClassType('int'),
ClassType('None')])]), ARG_POS)],
AnyType()))
def test_union_type_str(Self):
# type: () -> None
assert str(UnionType([ClassType('int'),
ClassType('str')])) == 'Union[int, str]'
def test_infer_union_return(self):
# type: () -> None
self.assert_infer(['() -> int',
'() -> str'],
([],
UnionType([ClassType('int'), ClassType('str')])))
def test_infer_ignore_mock_fallback_to_any(self):
# type: () -> None
self.assert_infer(['(mock.mock.Mock) -> str',
'(mock.mock.Mock) -> int'],
([(AnyType(), ARG_POS)],
UnionType([ClassType('str'), ClassType('int')])))
def test_optional(self):
# type: () -> None
self.assert_type_comment('(Optional[int]) -> Any', ([
Argument(UnionType([ClassType('int'),
ClassType('None')]), ARG_POS)
], AnyType()))
def test_infer_none_argument(self):
# type: () -> None
self.assert_infer(['(None) -> None'],
([(UnionType([ClassType('None'), AnyType()]), ARG_POS)],
ClassType('None')))