From c733d1898f40a380c68ca22df4247fb0f4880914 Mon Sep 17 00:00:00 2001
From: Tatiana Ber <tatiana.ber@jetbrains.com>
Date: Thu, 12 Mar 2026 11:05:06 +0100
Subject: [PATCH] refactor python: replace utility helpers with kotlin
 equivalents in `PyTypingTypeProvider`

(cherry picked from commit a1d8d52a6bd5bdbcecd8698b7bed96260b11b3f0)

GitOrigin-RevId: c52d3c037743383fa9e5faabec3733695a78544d
---
 .../typing/PyTypingTypeProvider.kt            | 27 +++++++------------
 1 file changed, 10 insertions(+), 17 deletions(-)

diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.kt
index b802c0aba09f..86dbb6247dee 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.kt
+++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/typing/PyTypingTypeProvider.kt
@@ -22,9 +22,7 @@ import com.intellij.psi.util.CachedValuesManager
 import com.intellij.psi.util.PsiModificationTracker
 import com.intellij.psi.util.PsiTreeUtil
 import com.intellij.psi.util.QualifiedName
-import com.intellij.util.ArrayUtil
 import com.intellij.util.Function
-import com.intellij.util.containers.ContainerUtil
 import com.intellij.util.containers.Stack
 import com.jetbrains.python.PyCustomType
 import com.jetbrains.python.PyNames
@@ -1536,8 +1534,7 @@ class PyTypingTypeProvider : PyTypeProviderWithCustomContext<Context?>() {
         val containingClass = typeHintContext as? PyClass ?: PsiTreeUtil.getStubOrPsiParentOfType(typeHintContext, PyClass::class.java)
         if (containingClass == null) return null
 
-        val scopeClassType = PyUtil.`as`(containingClass.getType(context.typeContext), PyClassType::class.java)
-        if (scopeClassType == null) return null
+        val scopeClassType = containingClass.getType(context.typeContext) as? PyClassType ?: return null
 
         return Ref(PySelfType(scopeClassType).toInstance())
       }
@@ -2258,7 +2255,7 @@ class PyTypingTypeProvider : PyTypeProviderWithCustomContext<Context?>() {
             Ref()
         }
 
-        val declarationElement = PyUtil.`as`(element, PyQualifiedNameOwner::class.java)
+        val declarationElement = element as? PyQualifiedNameOwner
 
         when (element.kind) {
           PyAstTypeParameter.Kind.TypeVar -> {
@@ -2758,7 +2755,7 @@ class PyTypingTypeProvider : PyTypeProviderWithCustomContext<Context?>() {
       context: TypeEvalContext,
     ): MutableList<PsiElement?> {
       val qualifiedName = expression.asQualifiedName()
-      val pyFile = PyUtil.`as`(FileContextUtil.getContextFile(expression), PyFile::class.java)
+      val pyFile = FileContextUtil.getContextFile(expression) as? PyFile
 
       val anchor = expression.containingFile.context
       val scopeOwner: ScopeOwner?
@@ -2875,19 +2872,15 @@ class PyTypingTypeProvider : PyTypeProviderWithCustomContext<Context?>() {
 
     @JvmStatic
     fun coroutineOrGeneratorElementType(coroutineOrGeneratorType: PyType?): Ref<PyType?>? {
-      val genericType = PyUtil.`as`(coroutineOrGeneratorType, PyCollectionType::class.java)
-      val classType = PyUtil.`as`(coroutineOrGeneratorType, PyClassType::class.java)
+      if (coroutineOrGeneratorType !is PyCollectionType) return null
+      val qName = coroutineOrGeneratorType.classQName
 
-      if (genericType != null && classType != null) {
-        val qName = classType.classQName
+      if (AWAITABLE == qName) {
+        return Ref(coroutineOrGeneratorType.elementTypes.getOrNull(0))
+      }
 
-        if (AWAITABLE == qName) {
-          return Ref(ContainerUtil.getOrElse<PyType?>(genericType.elementTypes, 0, null))
-        }
-
-        if (ArrayUtil.contains(qName, COROUTINE, GENERATOR)) {
-          return Ref(ContainerUtil.getOrElse<PyType?>(genericType.elementTypes, 2, null))
-        }
+      if (qName in arrayOf(COROUTINE, PyNames.TYPES_COROUTINE_TYPE, GENERATOR)) {
+        return Ref(coroutineOrGeneratorType.elementTypes.getOrNull(2))
       }
 
       return null