From abae2f744d4843721cfde307e1af8f90ff7658b0 Mon Sep 17 00:00:00 2001
From: Semyon Proshev <Semyon.Proshev@jetbrains.com>
Date: Tue, 6 Nov 2018 15:57:32 +0300
Subject: [PATCH] Provide fields' types for typing.NamedTuple(...) inheritors
 (PY-32240)

---
 .../stdlib/PyNamedTupleTypeProvider.kt        | 19 +++++++++++++++++++
 .../com/jetbrains/python/PyTypeTest.java      | 11 +++++++++++
 2 files changed, 30 insertions(+)

diff --git a/python/src/com/jetbrains/python/codeInsight/stdlib/PyNamedTupleTypeProvider.kt b/python/src/com/jetbrains/python/codeInsight/stdlib/PyNamedTupleTypeProvider.kt
index ce2a00904f3a..0bdce4cfd9a6 100644
--- a/python/src/com/jetbrains/python/codeInsight/stdlib/PyNamedTupleTypeProvider.kt
+++ b/python/src/com/jetbrains/python/codeInsight/stdlib/PyNamedTupleTypeProvider.kt
@@ -34,6 +34,11 @@ class PyNamedTupleTypeProvider : PyTypeProviderBase() {
       return fieldTypeForNamedTuple
     }
 
+    val fieldTypeForTypingNTFunctionInheritor = getFieldTypeForTypingNTFunctionInheritor(referenceExpression, context)
+    if (fieldTypeForTypingNTFunctionInheritor != null) {
+      return fieldTypeForTypingNTFunctionInheritor
+    }
+
     val namedTupleTypeForCallee = getNamedTupleTypeForCallee(referenceExpression, context)
     if (namedTupleTypeForCallee != null) {
       return namedTupleTypeForCallee
@@ -90,6 +95,20 @@ class PyNamedTupleTypeProvider : PyTypeProviderBase() {
       return qualifierNTType.fields[referenceExpression.name]?.type
     }
 
+    private fun getFieldTypeForTypingNTFunctionInheritor(referenceExpression: PyReferenceExpression, context: TypeEvalContext): PyType? {
+      val qualifierType = referenceExpression.qualifier?.let { context.getType(it) } as? PyWithAncestors
+      if (qualifierType == null || qualifierType is PyNamedTupleType) return null
+
+      return PyUnionType.union(
+        qualifierType
+          .getAncestorTypes(context)
+          .filterIsInstance<PyNamedTupleType>()
+          .mapNotNull { it.fields[referenceExpression.name] }
+          .map { it.type }
+          .toList()
+      )
+    }
+
     private fun getNamedTupleTypeForCallee(referenceExpression: PyReferenceExpression, context: TypeEvalContext): PyNamedTupleType? {
       if (PyCallExpressionNavigator.getPyCallExpressionByCallee(referenceExpression) == null) return null
 
diff --git a/python/testSrc/com/jetbrains/python/PyTypeTest.java b/python/testSrc/com/jetbrains/python/PyTypeTest.java
index 489b1bf72881..4e70b0e2d825 100644
--- a/python/testSrc/com/jetbrains/python/PyTypeTest.java
+++ b/python/testSrc/com/jetbrains/python/PyTypeTest.java
@@ -2588,6 +2588,17 @@ public class PyTypeTest extends PyTestCase {
     );
   }
 
+  // PY-32240
+  public void testTypingNTFunctionInheritorField() {
+    doTest("str",
+           "from typing import NamedTuple\n" +
+           "\n" +
+           "class A(NamedTuple(\"NT\", [(\"user\", str)])):\n" +
+           "    pass\n" +
+           "    \n" +
+           "expr = A(undefined).user");
+  }
+
   // PY-4351
   public void testCollectionsNTInheritorField() {
     // Seems that this case won't be supported because