From 6787719c2839f98a32b97b7eab564da35d154b2b Mon Sep 17 00:00:00 2001
From: Peter Law <PeterJCLaw@gmail.com>
Date: Sat, 24 Jul 2021 16:16:45 +0100
Subject: [PATCH] Ensure *args, **kwargs lookthrough works at module scope too

This means that passthrough signatures will be found for top level
functions, which is useful both where they're wrappered by
`functools.wraps` or not.

Fixes https://github.com/davidhalter/jedi/issues/1791.
---
 jedi/inference/star_args.py           | 10 +++++++++-
 test/test_inference/test_signature.py | 20 ++++++++++++++++++++
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/jedi/inference/star_args.py b/jedi/inference/star_args.py
index 521cd4eb..0c3a9e53 100644
--- a/jedi/inference/star_args.py
+++ b/jedi/inference/star_args.py
@@ -22,7 +22,15 @@ def _iter_nodes_for_param(param_name):
     from jedi.inference.arguments import TreeArguments
 
     execution_context = param_name.parent_context
-    function_node = execution_context.tree_node
+    # Walk up the parso tree to get the FunctionNode we want. We use the parso
+    # tree rather than going via the execution context so that we're agnostic of
+    # the specific scope we're evaluating within (i.e: module or function,
+    # etc.).
+    # - .tree_name is a Name
+    # - .parent is a Param
+    # - .parent is a PythonNode(parameters)
+    # - .parent is the FunctionNode we want.
+    function_node = param_name.tree_name.parent.parent.parent
     module_node = function_node.get_root_node()
     start = function_node.children[-1].start_pos
     end = function_node.children[-1].end_pos
diff --git a/test/test_inference/test_signature.py b/test/test_inference/test_signature.py
index beb405ca..9ee8a55c 100644
--- a/test/test_inference/test_signature.py
+++ b/test/test_inference/test_signature.py
@@ -292,6 +292,26 @@ def test_pow_signature(Script, environment):
                 return wrapper
 
             x(f)('''), 'f()'],
+        [dedent('''
+            # identifier:C
+            import functools
+            def f(x: int, y: float):
+                pass
+
+            @functools.wraps(f)
+            def wrapper(*args, **kwargs):
+                return f(*args, **kwargs)
+
+            wrapper('''), 'f(x: int, y: float)'],
+        [dedent('''
+            # identifier:D
+            def f(x: int, y: float):
+                pass
+
+            def wrapper(*args, **kwargs):
+                return f(*args, **kwargs)
+
+            wrapper('''), 'wrapper(x: int, y: float)'],
     ]
 )
 def test_wraps_signature(Script, code, signature):