From cd34e009980c950ab23f43bb3b19b7e3379e8b3f Mon Sep 17 00:00:00 2001
From: Andrey Vlasovskikh <andrey.vlasovskikh@jetbrains.com>
Date: Mon, 21 Aug 2017 19:20:24 +0300
Subject: [PATCH] PY-8936 Infer return type for __get__ descriptor attributes

This fix is useful even if we cannot infer the type. Returning Any
instead of the descriptor instance suppresses the false positive error.
---
 .../src/com/jetbrains/python/PyNames.java     |  1 +
 .../psi/impl/PyReferenceExpressionImpl.java   | 24 ++++++++++++
 .../descriptorAttribute.py                    | 39 +++++++++++++++++++
 ...Py3UnresolvedReferencesInspectionTest.java |  5 +++
 4 files changed, 69 insertions(+)
 create mode 100644 python/testData/inspections/PyUnresolvedReferencesInspection3K/descriptorAttribute.py

diff --git a/python/psi-api/src/com/jetbrains/python/PyNames.java b/python/psi-api/src/com/jetbrains/python/PyNames.java
index b403b1f1f034..3fda5ece8f5a 100644
--- a/python/psi-api/src/com/jetbrains/python/PyNames.java
+++ b/python/psi-api/src/com/jetbrains/python/PyNames.java
@@ -85,6 +85,7 @@ public class PyNames {
   public static final String NEW = "__new__";
   public static final String GETATTR = "__getattr__";
   public static final String GETATTRIBUTE = "__getattribute__";
+  public static final String GET = "__get__";
   public static final String __CLASS__ = "__class__";
   public static final String DUNDER_METACLASS = "__metaclass__";
   public static final String METACLASS = "metaclass";
diff --git a/python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java b/python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java
index 350b15b8d603..bcef43656f31 100644
--- a/python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java
+++ b/python/src/com/jetbrains/python/psi/impl/PyReferenceExpressionImpl.java
@@ -48,6 +48,8 @@ import org.jetbrains.annotations.Nullable;
 
 import java.util.*;
 
+import static com.jetbrains.python.psi.PyUtil.as;
+
 /**
  * Implements reference expression PSI.
  *
@@ -252,6 +254,10 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere
       if (qualified && typeFromTargets instanceof PyNoneType) {
         return null;
       }
+      final Ref<PyType> descriptorType = getDescriptorType(typeFromTargets, context);
+      if (descriptorType != null) {
+        return descriptorType.get();
+      }
       return typeFromTargets;
     }
     finally {
@@ -259,6 +265,24 @@ public class PyReferenceExpressionImpl extends PyElementImpl implements PyRefere
     }
   }
 
+  @Nullable
+  private Ref<PyType> getDescriptorType(@Nullable PyType typeFromTargets, @NotNull TypeEvalContext context) {
+    if (!isQualified()) return null;
+    final PyClassLikeType targetType = as(typeFromTargets, PyClassLikeType.class);
+    if (targetType == null) return null;
+    final PyResolveContext resolveContext = PyResolveContext.noProperties().withTypeEvalContext(context);
+    final List<? extends RatedResolveResult> members = targetType.resolveMember(PyNames.GET, this, AccessDirection.READ,
+                                                                                resolveContext);
+    if (members == null || members.isEmpty()) return null;
+    final List<PyType> types = StreamEx.of(members)
+      .map((result) -> result.getElement())
+      .select(PyCallable.class)
+      .map((callable) -> context.getReturnType(callable))
+      .toList();
+    final PyType type = PyUnionType.union(types);
+    return Ref.create(type);
+  }
+
   @Nullable
   private Ref<PyType> getQualifiedReferenceType(@NotNull TypeEvalContext context) {
     if (!context.maySwitchToAST(this)) {
diff --git a/python/testData/inspections/PyUnresolvedReferencesInspection3K/descriptorAttribute.py b/python/testData/inspections/PyUnresolvedReferencesInspection3K/descriptorAttribute.py
new file mode 100644
index 000000000000..19df50cbd4d4
--- /dev/null
+++ b/python/testData/inspections/PyUnresolvedReferencesInspection3K/descriptorAttribute.py
@@ -0,0 +1,39 @@
+from typing import Any
+
+
+class StringDescriptor:
+    def __get__(self, instance, owner):
+        return 'foo'
+
+
+class AnyDescriptor:
+    def __get__(self, instance, owner) -> Any:
+        return 'bar'
+
+
+class ListDescriptor:
+    def __get__(self, instance: Any, owner: Any) -> list:
+        return 'baz'
+
+
+class C:
+    foo = StringDescriptor()
+    bar = AnyDescriptor()
+    baz = ListDescriptor()
+
+
+# Instance level
+c = C()
+c.foo.upper()
+c.foo.<warning descr="Unresolved attribute reference 'non_existent' for class 'str'">non_existent</warning>()
+c.bar.upper()
+c.bar.non_existent()
+c.baz.append()
+c.baz.<warning descr="Unresolved attribute reference 'non_existent' for class 'list'">non_existent</warning>()
+
+
+# Class level
+C.foo.upper()
+C.foo.<warning descr="Unresolved attribute reference 'non_existent' for class 'str'">non_existent</warning>()
+C.bar.upper()
+C.bar.non_existent()
diff --git a/python/testSrc/com/jetbrains/python/inspections/Py3UnresolvedReferencesInspectionTest.java b/python/testSrc/com/jetbrains/python/inspections/Py3UnresolvedReferencesInspectionTest.java
index eeece07ead9e..1e9e39dff1c1 100644
--- a/python/testSrc/com/jetbrains/python/inspections/Py3UnresolvedReferencesInspectionTest.java
+++ b/python/testSrc/com/jetbrains/python/inspections/Py3UnresolvedReferencesInspectionTest.java
@@ -220,4 +220,9 @@ public class Py3UnresolvedReferencesInspectionTest extends PyTestCase {
   public void testCallTypeGetAttributeAndSetAttrInInheritor() {
     doTest();
   }
+
+  // PY-8936
+  public void testDescriptorAttribute() {
+    doTest();
+  }
 }