Make sure param hints are working for functions

This commit is contained in:
Dave Halter
2020-02-02 18:36:56 +01:00
parent 4c7179bc87
commit e3c4b5b77e
3 changed files with 44 additions and 1 deletions

View File

@@ -167,6 +167,10 @@ class GenericClass(ClassMixin, DefineGenericBase):
def get_type_hint(self, add_class_info=True):
n = self.py__name__()
# Not sure if this is the best way to do this, but all of these types
# are a bit special in that they have type aliases and other ways to
# become lower case. It's probably better to make them upper case,
# because that's what you can use in annotations.
n = dict(list="List", dict="Dict", set="Set", tuple="Tuple").get(n, n)
s = n + self._generics_manager.get_type_hint()
if add_class_info:

View File

@@ -86,6 +86,33 @@ class FunctionMixin(object):
def py__name__(self):
return self.name.string_name
def get_type_hint(self, add_class_info=True):
return_annotation = self.tree_node.annotation
if return_annotation is None:
def param_name_to_str(n):
s = n.string_name
annotation = n.infer().get_type_hint()
if annotation is not None:
s += ': ' + annotation
if n.default_node is not None:
s += '=' + n.default_node.get_code(include_prefix=False)
return s
function_execution = self.as_context()
result = function_execution.infer()
return_hint = result.get_type_hint()
body = self.py__name__() + '(%s)' % ', '.join([
param_name_to_str(n)
for n in function_execution.get_param_names()
])
if return_hint is None:
return body
else:
return_hint = return_annotation.get_code(include_prefix=False)
body = self.py__name__() + self.tree_node.children[2].get_code(include_prefix=False)
return body + ' -> ' + return_hint
def py__call__(self, arguments):
function_execution = self.as_context(arguments)
return function_execution.infer()
@@ -399,6 +426,9 @@ class OverloadedFunctionValue(FunctionMixin, ValueWrapper):
def get_signature_functions(self):
return self._overloaded_functions
def get_type_hint(self, add_class_info=True):
return 'Union[%s]' % ', '.join(f.get_type_hint() for f in self._overloaded_functions)
def _find_overload_functions(context, tree_node):
def _is_overload_decorated(funcdef):

View File

@@ -573,6 +573,7 @@ def test_definition_goto_follow_imports(Script):
('n: Type[List[int]]; n', 'Type[List[int]]'),
('n: Type[List]; n', 'Type[list]'),
('n: List; n', 'list'),
('n: List[int]; n', 'List[int]'),
('n: Iterable[int]; n', 'Iterable[int]'),
@@ -583,9 +584,17 @@ def test_definition_goto_follow_imports(Script):
('n = (1,); n', 'Tuple[int]'),
('n = {1: ""}; n', 'Dict[int, str]'),
('n = {1: "", 1.0: b""}; n', 'Dict[Union[float, int], Union[bytes, str]]'),
('n = next; n', 'Union[next(__i: Iterator[_T]) -> _T, '
'next(__i: Iterator[_T], default: _VT) -> Union[_T, _VT]]'),
('abs', 'abs(__n: SupportsAbs[_T]) -> _T'),
('def foo(x, y): return x if xxxx else y\nfoo(str(), 1)\nfoo',
'foo(x: str, y: int) -> Union[int, str]'),
('def foo(x, y = None): return x if xxxx else y\nfoo(str(), 1)\nfoo',
'foo(x: str, y: int=None) -> Union[int, str]'),
]
)
def test_get_type_hint(Script, code, expected, skip_python2):
def test_get_type_hint(Script, code, expected, skip_pre_python35):
code = 'from typing import *\n' + code
d, = Script(code).goto()
assert d.get_type_hint() == expected