PY-59727 Infer types from usages only for collections assigned to variables

For everything else, e.g. d['key'] = {'foo': 42}, we cannot track subsequent
modifications of d['key'], such as d['key']['bar'] = 1, anyway. We need
a reference as a subscription operand, or a callee to make sure these mutating
operations are performed on the same value. It doesn't make sense to scan
through an entire containing scope looking for subscription expressions
and calls on such assignments targets.

(cherry picked from commit 8ebbcb73e43e5711cbe680da6936f3f1c4e54a55)

IJ-MR-105530

GitOrigin-RevId: 65ebcc9219fb05c4131a8598ab54a9e37c6489e7

python/python-psi-impl/src/com/jetbrains/python/psi/types/PyCollectionTypeUtil.kt

@@ -270,21 +270,7 @@ object PyCollectionTypeUtil {
     return visitor.elementTypes
   }
 
-  fun getTypedDictTypeByModificationsForDictConstructor(qualifiedName: String,
-                                                        element: PsiElement,
-                                                        context: TypeEvalContext): Pair<Boolean, PyTypedDictType?> {
-    if (qualifiedName == DICT_CONSTRUCTOR) {
-      val owner = ScopeUtil.getScopeOwner(element)
-      if (owner != null) {
-        val visitor = PyTypedDictTypeVisitor(element, context)
-        owner.accept(visitor)
-        return Pair(visitor.hasAllStrKeys, visitor.typedDictType)
-      }
-    }
-    return Pair(true, null)
-  }
-
-  fun getCollectionTypeByModifications(qualifiedName: String, element: PsiElement, context: TypeEvalContext): List<PyType?> {
+  fun getCollectionTypeByModifications(qualifiedName: String, element: PyTargetExpression, context: TypeEvalContext): List<PyType?> {
     val owner = ScopeUtil.getScopeOwner(element)
     if (owner != null) {
       val typeVisitor = getVisitorForQualifiedName(qualifiedName, element, context)
@@ -297,7 +283,7 @@ object PyCollectionTypeUtil {
   }
 
   private fun getVisitorForSequence(sequence: PySequenceExpression,
-                                    element: PsiElement,
+                                    element: PyTargetExpression,
                                     context: TypeEvalContext): PyCollectionTypeVisitor? {
     return when (sequence) {
       is PyListLiteralExpression -> PyListTypeVisitor(element, context)
@@ -307,7 +293,7 @@ object PyCollectionTypeUtil {
     }
   }
 
-  private fun getVisitorForQualifiedName(qualifiedName: String, element: PsiElement, context: TypeEvalContext): PyCollectionTypeVisitor? {
+  private fun getVisitorForQualifiedName(qualifiedName: String, element: PyTargetExpression, context: TypeEvalContext): PyCollectionTypeVisitor? {
     return when (qualifiedName) {
       LIST_CONSTRUCTOR, RANGE_CONSTRUCTOR -> PyListTypeVisitor(element, context)
       DICT_CONSTRUCTOR -> PyDictTypeVisitor(element, context)
@@ -316,10 +302,9 @@ object PyCollectionTypeUtil {
     }
   }
 
-  fun getTargetForValueInAssignment(value: PyExpression): PyExpression? {
-    val assignmentStatement = PsiTreeUtil.getParentOfType(value, PyAssignmentStatement::class.java, true, ScopeOwner::class.java)
-    assignmentStatement?.targetsToValuesMapping?.filter { it.second === value }?.forEach { return it.first }
-    return null
+  fun getTargetForValueInAssignment(value: PyExpression): PyTargetExpression? {
+    val assignment = PsiTreeUtil.getParentOfType(value, PyAssignmentStatement::class.java, true, ScopeOwner::class.java) ?: return null
+    return assignment.targetsToValuesMapping.firstOrNull { it.second === value }?.first as? PyTargetExpression
   }
 
   private fun getTypeForArgument(arguments: Array<PyExpression>, argumentIndex: Int, typeEvalContext: TypeEvalContext): PyType? {
@@ -434,7 +419,7 @@ object PyCollectionTypeUtil {
     return if (isModificationExist) Pair(keyTypes, valueTypes) else null
   }
 
-  private abstract class PyCollectionTypeVisitor(protected val myElement: PsiElement,
+  private abstract class PyCollectionTypeVisitor(protected val myElement: PyTargetExpression,
                                                  protected val myTypeEvalContext: TypeEvalContext) : PyRecursiveElementVisitor() {
     protected val scopeOwner: ScopeOwner? = ScopeUtil.getScopeOwner(myElement)
     protected open var isModificationExist = false
@@ -459,7 +444,7 @@ object PyCollectionTypeUtil {
     }
   }
 
-  private class PyListTypeVisitor(element: PsiElement,
+  private class PyListTypeVisitor(element: PyTargetExpression,
                                   typeEvalContext: TypeEvalContext) : PyCollectionTypeVisitor(element, typeEvalContext) {
     private val modificationMethods: Map<String, (Array<PyExpression>) -> List<PyType?>>
     private val valueTypes: MutableList<PyType?>
@@ -510,7 +495,7 @@ object PyCollectionTypeUtil {
     }
   }
 
-  private class PyDictTypeVisitor(element: PsiElement,
+  private class PyDictTypeVisitor(element: PyTargetExpression,
                                   typeEvalContext: TypeEvalContext) : PyCollectionTypeVisitor(element, typeEvalContext) {
     private val modificationMethods: Map<String, (Array<PyExpression>) -> List<PyType?>>
     private val keyTypes: MutableList<PyType?>
@@ -581,7 +566,7 @@ object PyCollectionTypeUtil {
     }
   }
 
-  private class PyTypedDictTypeVisitor(element: PsiElement,
+  private class PyTypedDictTypeVisitor(element: PyTargetExpression,
                                        typeEvalContext: TypeEvalContext) : PyCollectionTypeVisitor(element, typeEvalContext) {
     private val modificationMethods: Map<String, (Array<PyExpression>) -> List<PyType?>>
     var hasAllStrKeys = true
@@ -663,7 +648,7 @@ object PyCollectionTypeUtil {
     }
   }
 
-  private class PySetTypeVisitor(element: PsiElement,
+  private class PySetTypeVisitor(element: PyTargetExpression,
                                  typeEvalContext: TypeEvalContext) : PyCollectionTypeVisitor(element, typeEvalContext) {
     private val modificationMethods: Map<String, (Array<PyExpression>) -> List<PyType?>>
     private val valueTypes: MutableList<PyType?>