PY-23067 Pycharm not picking function metadata from functools.wraps with methods

GitOrigin-RevId: d95c1a8f64a6e58d1a6c6c65866b6ab08aaf71b3
This commit is contained in:
Petr
2024-06-18 17:33:22 +02:00
committed by intellij-monorepo-bot
parent ca17e28672
commit 5d2d6cc722
18 changed files with 325 additions and 45 deletions

View File

@@ -525,9 +525,13 @@
<customClassStubType implementation="com.jetbrains.python.psi.impl.stubs.PyDataclassStubType"/>
<customDecoratorStubType implementation="com.jetbrains.python.psi.stubs.PyTestFixtureDecoratorStubType"/>
<customDecoratorStubType implementation="com.jetbrains.python.psi.stubs.PyFunctoolsWrapsDecoratorStubType"/>
<typeProvider implementation="com.jetbrains.python.psi.types.PyCollectionTypeByModificationsProvider" order="last"/>
<typeProvider implementation="com.jetbrains.python.codeInsight.decorator.PyDecoratedFunctionTypeProvider"/>
<typeProvider implementation="com.jetbrains.python.codeInsight.decorator.PyDecoratedFunctionTypeProvider"
id="pyDecoratedFunctionTypeProvider"/>
<typeProvider implementation="com.jetbrains.python.codeInsight.decorator.PyFunctoolsWrapsDecoratedFunctionTypeProvider"
order="before pyDecoratedFunctionTypeProvider"/>
<!-- typing -->

View File

@@ -0,0 +1,63 @@
package com.jetbrains.python.codeInsight.decorator
import com.intellij.openapi.util.Ref
import com.intellij.psi.PsiElement
import com.intellij.psi.util.QualifiedName
import com.intellij.util.containers.ContainerUtil
import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil
import com.jetbrains.python.psi.PyCallable
import com.jetbrains.python.psi.PyFunction
import com.jetbrains.python.psi.PyKnownDecoratorUtil
import com.jetbrains.python.psi.PyUtil
import com.jetbrains.python.psi.impl.StubAwareComputation
import com.jetbrains.python.psi.resolve.PyResolveContext
import com.jetbrains.python.psi.resolve.PyResolveUtil
import com.jetbrains.python.psi.stubs.PyFunctoolsWrapsDecoratorStub
import com.jetbrains.python.psi.types.PyType
import com.jetbrains.python.psi.types.PyTypeProviderBase
import com.jetbrains.python.psi.types.TypeEvalContext
/**
* Infer type for reference of a function decorated with 'functools.wraps'.
* Has to be used before {@link com.jetbrains.python.codeInsight.decorator.PyDecoratedFunctionTypeProvider}
*/
class PyFunctoolsWrapsDecoratedFunctionTypeProvider : PyTypeProviderBase() {
override fun getReferenceType(referenceTarget: PsiElement, context: TypeEvalContext, anchor: PsiElement?): Ref<PyType>? {
if (referenceTarget !is PyFunction) return null
val wrappedFunction = ContainerUtil.findInstance(resolveWrapped(referenceTarget, context), PyFunction::class.java) ?: return null
return Ref.create(context.getType(wrappedFunction))
}
override fun getCallableType(callable: PyCallable, context: TypeEvalContext): PyType? {
return Ref.deref(getReferenceType(callable, context, null))
}
private fun resolveWrapped(function: PyFunction, context: TypeEvalContext): List<PsiElement> {
val decorator = function.decoratorList?.decorators?.find {
val qName = it.qualifiedName
qName != null && PyKnownDecoratorUtil.asKnownDecorators(qName).contains(PyKnownDecoratorUtil.KnownDecorator.FUNCTOOLS_WRAPS)
} ?: return emptyList()
return StubAwareComputation.on(decorator)
.withCustomStub { it.getCustomStub(PyFunctoolsWrapsDecoratorStub::class.java) }
.overStub {
if (it == null) return@overStub emptyList<PsiElement>()
var scopeOwner = ScopeUtil.getScopeOwner(decorator)
val wrappedQName = QualifiedName.fromDottedString(it.wrapped)
val resolved = mutableListOf<PsiElement>()
while (scopeOwner != null) {
resolved.addAll(PyResolveUtil.resolveQualifiedNameInScope(wrappedQName, scopeOwner, context))
scopeOwner = ScopeUtil.getScopeOwner(scopeOwner)
}
resolved
}
.overAst {
val wrappedExpr = it.argumentList?.getValueExpressionForParam(PyKnownDecoratorUtil.FunctoolsWrapsParameters.WRAPPED)
if (wrappedExpr == null)
emptyList()
else
PyUtil.multiResolveTopPriority(wrappedExpr, PyResolveContext.defaultContext(context))
}
.withStubBuilder(PyFunctoolsWrapsDecoratorStub::create)
.compute(context)
}
}

View File

@@ -18,16 +18,14 @@ import com.intellij.util.containers.FactoryMap;
import com.jetbrains.python.*;
import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtil;
import com.jetbrains.python.codeInsight.decorator.PyFunctoolsWrapsDecoratedFunctionTypeProvider;
import com.jetbrains.python.documentation.docstrings.DocStringUtil;
import com.jetbrains.python.psi.*;
import com.jetbrains.python.psi.impl.PyBuiltinCache;
import com.jetbrains.python.psi.resolve.PyResolveContext;
import com.jetbrains.python.psi.resolve.QualifiedNameFinder;
import com.jetbrains.python.psi.resolve.QualifiedResolveResult;
import com.jetbrains.python.psi.types.PyCallableParameter;
import com.jetbrains.python.psi.types.PyClassType;
import com.jetbrains.python.psi.types.PyType;
import com.jetbrains.python.psi.types.TypeEvalContext;
import com.jetbrains.python.psi.types.*;
import com.jetbrains.python.pyi.PyiUtil;
import com.jetbrains.python.toolbox.Maybe;
import one.util.streamex.StreamEx;
@@ -595,6 +593,16 @@ public class PyDocumentationBuilder {
return resolved;
}
}
// Return wrapped function for functools.wraps decorated function
if (myElement instanceof PyFunction function) {
PyType type = new PyFunctoolsWrapsDecoratedFunctionTypeProvider().getCallableType(function, myContext);
if (type instanceof PyCallableType callableType) {
PyCallable callable = callableType.getCallable();
if (callable != null) {
return callable;
}
}
}
return myElement;
}

View File

@@ -60,7 +60,7 @@ public class PyFileElementType extends IStubFileElementType<PyFileStub> {
@Override
public int getStubVersion() {
// Don't forget to update versions of indexes that use the updated stub-based elements
return 91;
return 92;
}
@Nullable

View File

@@ -5,6 +5,7 @@ import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiFile;
import com.intellij.psi.util.QualifiedName;
import com.intellij.util.containers.ContainerUtil;
import com.jetbrains.python.FunctionParameter;
import com.jetbrains.python.PyNames;
import com.jetbrains.python.codeInsight.controlflow.ScopeOwner;
import com.jetbrains.python.psi.resolve.PyResolveContext;
@@ -12,7 +13,9 @@ import com.jetbrains.python.psi.resolve.PyResolveUtil;
import com.jetbrains.python.psi.types.TypeEvalContext;
import com.jetbrains.python.pyi.PyiFile;
import one.util.streamex.StreamEx;
import org.jetbrains.annotations.ApiStatus;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import java.util.*;
@@ -178,12 +181,18 @@ public final class PyKnownDecoratorUtil {
.toImmutableList();
}
else {
// The method might have been called during building of PSI stub indexes. Thus, we can't leave this file's boundaries.
// TODO Use proper local resolve to imported names here
return Collections.unmodifiableList(BY_SHORT_NAME.getOrDefault(qualifiedName.getLastComponent(), Collections.emptyList()));
return asKnownDecorators(qualifiedName);
}
}
@ApiStatus.Internal
@NotNull
public static List<KnownDecorator> asKnownDecorators(@NotNull QualifiedName qualifiedName) {
// The method might have been called during building of PSI stub indexes. Thus, we can't leave this file's boundaries.
// TODO Use proper local resolve to imported names here
return Collections.unmodifiableList(BY_SHORT_NAME.getOrDefault(qualifiedName.getLastComponent(), Collections.emptyList()));
}
/**
* Check that given element has any non-standard (read "unreliable") decorators.
*
@@ -270,4 +279,26 @@ public final class PyKnownDecoratorUtil {
? decorators.isEmpty()
: decoratorList.getDecorators().length == StreamEx.of(decorators).groupingBy(KnownDecorator::getShortName).size();
}
public enum FunctoolsWrapsParameters implements FunctionParameter {
WRAPPED(0, "wrapped");
private final int myPosition;
private final String myName;
FunctoolsWrapsParameters(int position, @NotNull String name) {
myPosition = position;
myName = name;
}
@Override
public int getPosition() {
return myPosition;
}
@Override
public @NotNull String getName() {
return myName;
}
}
}

View File

@@ -0,0 +1,38 @@
package com.jetbrains.python.psi.stubs
import com.intellij.psi.stubs.StubInputStream
import com.intellij.psi.stubs.StubOutputStream
import com.jetbrains.python.psi.PyDecorator
import com.jetbrains.python.psi.PyKnownDecoratorUtil
import com.jetbrains.python.psi.PyReferenceExpression
import com.jetbrains.python.psi.impl.stubs.PyCustomDecoratorStub
import com.jetbrains.python.psi.impl.stubs.PyCustomDecoratorStubType
class PyFunctoolsWrapsDecoratorStubType : PyCustomDecoratorStubType<PyFunctoolsWrapsDecoratorStub> {
override fun createStub(psi: PyDecorator): PyFunctoolsWrapsDecoratorStub? {
return PyFunctoolsWrapsDecoratorStub.create(psi)
}
override fun deserializeStub(stream: StubInputStream): PyFunctoolsWrapsDecoratorStub? {
val name = stream.readNameString() ?: return null
return PyFunctoolsWrapsDecoratorStub(name)
}
}
class PyFunctoolsWrapsDecoratorStub(val wrapped: String) : PyCustomDecoratorStub {
override fun getTypeClass(): Class<out PyCustomDecoratorStubType<*>> = PyFunctoolsWrapsDecoratorStubType::class.java
override fun serialize(stream: StubOutputStream) {
stream.writeName(wrapped)
}
companion object {
fun create(psi: PyDecorator): PyFunctoolsWrapsDecoratorStub? {
val qName = psi.qualifiedName ?: return null
if (!PyKnownDecoratorUtil.asKnownDecorators(qName).contains(PyKnownDecoratorUtil.KnownDecorator.FUNCTOOLS_WRAPS)) return null
val wrappedExpr = psi.argumentList?.getValueExpressionForParam(PyKnownDecoratorUtil.FunctoolsWrapsParameters.WRAPPED) as? PyReferenceExpression
val wrappedExprQName = wrappedExpr?.asQualifiedName() ?: return null
return PyFunctoolsWrapsDecoratorStub(wrappedExprQName.toString())
}
}
}

View File

@@ -0,0 +1,5 @@
from m import Router
r = Router()
r.route("", 13)
r.route(""<warning descr="Parameter 'i' unfilled">)</warning>

View File

@@ -0,0 +1,16 @@
import functools
class MyClass:
def foo(self, s: str, i: int):
pass
class Route:
@functools.wraps(MyClass.foo)
def __init__(self):
pass
class Router:
@functools.wraps(wrapped=Route.__init__)
def route(self, s: str):
pass

View File

@@ -0,0 +1,5 @@
from m import Router
router = Router()
router.route(-2)
router.route(<warning descr="Expected type 'int', got 'str' instead">""</warning>)

View File

@@ -0,0 +1,18 @@
import functools
class MyClass:
def foo(self, i: int):
pass
class Route:
@functools.wraps(MyClass.foo)
def __init__(self):
pass
class Router:
@functools.wraps(wrapped=Route.__init__)
def route(self, s: str):
pass

View File

@@ -1,21 +0,0 @@
from functools import wraps
import inspect
class Route:
def __init__(self, input_a: int, input_b: float):
...
class Router:
def __init__(self):
self.routes = []
@wraps(Route.__init__)
def route(self, *args, **kwargs):
route = Route(*args, **kwargs)
self.routes.append(route)
r = Router()
r.route(<arg1>)

View File

@@ -0,0 +1,22 @@
import functools
class MyClass:
def foo(self, s: str, b: bool):
pass
class Route:
@functools.wraps(MyClass.foo)
def __init__(self, a: int, b: float, c: object):
pass
class Router:
@functools.wraps(wrapped=Route.__init__)
def route(self, *args, **kwargs):
pass
r = Router()
r.route(<arg1>"", <arg2>True)

View File

@@ -0,0 +1,3 @@
<html><body><div class="bottom"><icon src="AllIcons.Nodes.Class"/>&nbsp;<code><a href="psi_element://#typename#FunctoolsWraps.Cls">FunctoolsWraps.Cls</a></code></div><div class="definition"><pre><span style="color:#000080;font-weight:bold;">def </span><span style="color:#000000;">foo</span><span style="">(</span><span style="color:#94558d;">self</span><span style="">,</span>
<span style="color:#000000;">s</span><span style="">: </span><span style="color:#000000;"><span style="color:#000080;"><a href="psi_element://#typename#str">str</a></span></span><span style="">,</span>
<span style="color:#000000;">b</span><span style="">: </span><span style="color:#000000;"><span style="color:#000080;"><a href="psi_element://#typename#bool">bool</a></span></span><span style="">)</span> -&gt; <span style="color:#000000;"><span style="color:#000080;font-weight:bold;">None</span></span></pre></div><div class="content">Unittest placeholder</div><table class="sections"><tr><td class="section" valign="top">Params:</td><td valign="top"><p><code>s</code> &ndash; str</p><p><code>b</code> &ndash; bool</p></td></tr><tr><td class="section" valign="top">Returns:</td><td valign="top">None</td></tr></table></body></html>

View File

@@ -0,0 +1,24 @@
import functools
class Cls:
def foo(self, s: str, b: bool):
"""
Doc text
:param s: str
:param b: bool
:return: None
"""
pass
class Route:
@functools.wraps(Cls.foo)
def __init__(self):
pass
class Router:
@functools.wraps(wrapped=Route.__init__)
def route(self, s: str):
pass
r = Router()
r.<the_ref>route(13)

View File

@@ -69,7 +69,7 @@ public class Py3QuickDocTest extends LightMarkedTestCase {
assertNotNull(stringValue);
PsiElement referenceElement = marks.get("<the_ref>").getParent(); // ident -> expr
final PyDocStringOwner docOwner = (PyDocStringOwner)((PyReferenceExpression)referenceElement).getReference().resolve();
final PyDocStringOwner docOwner = (PyDocStringOwner)referenceElement.getReference().resolve();
assertNotNull(docOwner);
assertEquals(docElement, docOwner.getDocStringExpression());
@@ -91,7 +91,7 @@ public class Py3QuickDocTest extends LightMarkedTestCase {
private void checkHover() {
Map<String, PsiElement> marks = loadTest();
final PsiElement originalElement = marks.get("<the_ref>");
final PsiElement docOwner = ((PyReferenceExpression)originalElement.getParent()).getReference().resolve();
final PsiElement docOwner = originalElement.getParent().getReference().resolve();
checkByHTML(myProvider.getQuickNavigateInfo(docOwner, originalElement));
}
@@ -165,18 +165,11 @@ public class Py3QuickDocTest extends LightMarkedTestCase {
}
public void testPropNewSetter() {
Map<String, PsiElement> marks = loadTest();
PsiElement referenceElement = marks.get("<the_ref>");
final PyDocStringOwner docStringOwner = (PyDocStringOwner)referenceElement.getParent().getReference().resolve();
checkByHTML(myProvider.generateDoc(docStringOwner, referenceElement));
checkHTMLOnly();
}
public void testPropNewDeleter() {
Map<String, PsiElement> marks = loadTest();
PsiElement referenceElement = marks.get("<the_ref>");
final PyDocStringOwner docStringOwner =
(PyDocStringOwner)((PyReferenceExpression)(referenceElement.getParent())).getReference().resolve();
checkByHTML(myProvider.generateDoc(docStringOwner, referenceElement));
checkHTMLOnly();
}
public void testPropOldGetter() {
@@ -185,10 +178,7 @@ public class Py3QuickDocTest extends LightMarkedTestCase {
public void testPropOldSetter() {
Map<String, PsiElement> marks = loadTest();
PsiElement referenceElement = marks.get("<the_ref>");
final PyDocStringOwner docStringOwner = (PyDocStringOwner)referenceElement.getParent().getReference().resolve();
checkByHTML(myProvider.generateDoc(docStringOwner, referenceElement));
checkHTMLOnly();
}
public void testPropOldDeleter() {
@@ -871,6 +861,11 @@ public class Py3QuickDocTest extends LightMarkedTestCase {
checkHTMLOnly();
}
// PY-23067
public void testFunctoolsWraps() {
checkHTMLOnly();
}
@Override
protected String getTestDataPath() {
return super.getTestDataPath() + "/quickdoc/";

View File

@@ -1258,6 +1258,14 @@ public class PyParameterInfoTest extends LightMarkedTestCase {
feignCtrlP(marks.get("<arg3>").getTextOffset()).check("a: int, *, name: str = ..., year: int", new String[]{"year: int"});
}
// PY-23067
public void testFunctoolsWraps() {
final Map<String, PsiElement> marks = loadTest(2);
feignCtrlP(marks.get("<arg1>").getTextOffset()).check("self: MyClass, s: str, b: bool", new String[]{"s: str, "}, new String[]{"self: MyClass, "});
feignCtrlP(marks.get("<arg2>").getTextOffset()).check("self: MyClass, s: str, b: bool", new String[]{"b: bool"}, new String[]{"self: MyClass, "});
}
// PY-58497
public void testSimplePopupWithHintsOff() {
Map<String, PsiElement> marks = loadTest(5);

View File

@@ -311,4 +311,35 @@ public class Py3ArgumentListInspectionTest extends PyInspectionTestCase {
Derived2(0, <warning descr="Unexpected argument">0</warning>, b=0<warning descr="Parameter 'qq' unfilled">)</warning>
""");
}
// PY-23067
public void testFunctoolsWraps() {
doTestByText("""
import functools
class MyClass:
def foo(self, s: str, i: int):
pass
class Route:
@functools.wraps(MyClass.foo)
def __init__(self):
pass
class Router:
@functools.wraps(wrapped=Route.__init__)
def route(self, s: str):
pass
r = Router()
r.route("", 13)
r.route(""<warning descr="Parameter 'i' unfilled">)</warning>
r.route("", 13, <warning descr="Unexpected argument">1</warning>)
""");
}
// PY-23067
public void testFunctoolsWrapsMultiFile() {
doMultiFileTest();
}
}

View File

@@ -2144,4 +2144,34 @@ def foo(param: str | int) -> TypeGuard[str]:
PosArgsT = TypeVarTuple("PosArgsT")
""");
}
// PY-23067
public void testFunctoolsWraps() {
doTestByText("""
import functools
class MyClass:
def foo(self, i: int):
pass
class Route:
@functools.wraps(MyClass.foo)
def __init__(self):
pass
class Router:
@functools.wraps(wrapped=Route.__init__)
def route(self, s: str):
pass
router = Router()
router.route(-2)
router.route(<warning descr="Expected type 'int', got 'str' instead">""</warning>)
""");
}
// PY-23067
public void testFunctoolsWrapsMultiFile() {
doMultiFileTest();
}
}