diff --git a/jedi/plugins/stdlib.py b/jedi/plugins/stdlib.py index 448f2b50..8678051e 100644 --- a/jedi/plugins/stdlib.py +++ b/jedi/plugins/stdlib.py @@ -513,6 +513,49 @@ class PartialObject(object): ) +class PartialMethodObject(object): + def __init__(self, actual_value, arguments): + self._actual_value = actual_value + self._arguments = arguments + + def __getattr__(self, name): + return getattr(self._actual_value, name) + + def _get_function(self, unpacked_arguments): + key, lazy_value = next(unpacked_arguments, (None, None)) + if key is not None or lazy_value is None: + debug.warning("Partial should have a proper function %s", self._arguments) + return None + return lazy_value.infer() + + def get_signatures(self): + unpacked_arguments = self._arguments.unpack() + func = self._get_function(unpacked_arguments) + if func is None: + return [] + + arg_count = 1 + keys = set() + for key, _ in unpacked_arguments: + if key is None: + arg_count += 1 + else: + keys.add(key) + return [PartialSignature(s, arg_count, keys) for s in func.get_signatures()] + + def py__get__(self, instance, class_value): + return ValueSet([self]) + + def py__call__(self, arguments): + func = self._get_function(self._arguments.unpack()) + if func is None: + return NO_VALUES + + return func.execute( + MergedPartialArguments(self._arguments, arguments) + ) + + class PartialSignature(SignatureWrapper): def __init__(self, wrapped_signature, skipped_arg_count, skipped_arg_set): super(PartialSignature, self).__init__(wrapped_signature) @@ -547,6 +590,13 @@ def functools_partial(value, arguments, callback): ) +def functools_partialmethod(value, arguments, callback): + return ValueSet( + PartialMethodObject(instance, arguments) + for instance in value.py__call__(arguments) + ) + + @argument_clinic('first, /') def _return_first_param(firsts): return firsts @@ -744,6 +794,7 @@ _implemented = { }, 'functools': { 'partial': functools_partial, + 'partialmethod': functools_partialmethod, 'wraps': _functools_wraps, }, '_weakref': { diff --git a/test/completion/stdlib.py b/test/completion/stdlib.py index 636b66fb..4fa5a238 100644 --- a/test/completion/stdlib.py +++ b/test/completion/stdlib.py @@ -158,6 +158,23 @@ tup[0] #? float() tup[1] +class X: + def function(self, a, b): + return a, b + a = functools.partialmethod(function, 0) + kw = functools.partialmethod(function, b=1.0) + +#? int() +X().a('')[0] +#? str() +X().a('')[1] + +tup = X().kw(1) +#? int() +tup[0] +#? float() +tup[1] + def my_decorator(f): @functools.wraps(f) def wrapper(*args, **kwds): diff --git a/test/test_inference/test_signature.py b/test/test_inference/test_signature.py index 186b5770..967198b3 100644 --- a/test/test_inference/test_signature.py +++ b/test/test_inference/test_signature.py @@ -65,6 +65,19 @@ d = functools.partial() ''' +partialmethod_code = ''' +import functools + +class X: + def func(self, a, b, c): + pass + a = functools.partialmethod(func) + b = functools.partialmethod(func, 1) + c = functools.partialmethod(func, 1, c=2) + d = functools.partialmethod() +''' + + @pytest.mark.parametrize( 'code, expected', [ ('def f(a, * args, x): pass\n f(', 'f(a, *args, x)'), @@ -82,6 +95,11 @@ d = functools.partial() (partial_code + 'b(', 'func(b, c)'), (partial_code + 'c(', 'func(b)'), (partial_code + 'd(', None), + + (partialmethod_code + 'X().a(', 'func(a, b, c)'), + (partialmethod_code + 'X().b(', 'func(b, c)'), + (partialmethod_code + 'X().c(', 'func(b)'), + (partialmethod_code + 'X().d(', None), ] ) def test_tree_signature(Script, environment, code, expected):