From c8002968ccee79ebf617641c4c32194ab3272299 Mon Sep 17 00:00:00 2001
From: Mikhail Golubev <mikhail.golubev@jetbrains.com>
Date: Thu, 6 Feb 2025 11:40:52 +0200
Subject: [PATCH] PY-78236 Migrate
 PyDocStringTypeProvider#setTypeVarScopeOwners to PyCloningTypeVisitor

It allowed to get rid of PyGenericType.setScopeOwner hack, modifying the type inplace.

GitOrigin-RevId: dd9389ae496a7e1afa3bc7101d907fcc3d170d5e
---
 .../docstrings/PyDocStringTypeProvider.java   | 39 +++++++++----------
 .../python/psi/types/PyGenericType.java       | 11 +-----
 2 files changed, 20 insertions(+), 30 deletions(-)

diff --git a/python/python-psi-impl/src/com/jetbrains/python/documentation/docstrings/PyDocStringTypeProvider.java b/python/python-psi-impl/src/com/jetbrains/python/documentation/docstrings/PyDocStringTypeProvider.java
index 74e6f8bd7897..ecc0c2121e0e 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/documentation/docstrings/PyDocStringTypeProvider.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/documentation/docstrings/PyDocStringTypeProvider.java
@@ -11,7 +11,6 @@ import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
 import java.util.Set;
-import java.util.function.Function;
 
 /**
  * @author Mikhail Golubev
@@ -69,7 +68,17 @@ public final class PyDocStringTypeProvider extends PyTypeProviderBase {
     if (type != null) {
       type.assertValid("from docstring");
     }
-    setTypeVarScopeOwners(type, callable, context);
+    if (callable instanceof PyFunction pyFunction) {
+      return Ref.create(PyCloningTypeVisitor.clone(type, new PyCloningTypeVisitor(context) {
+        @Override
+        public PyType visitPyTypeVarType(@NotNull PyTypeVarType typeVarType) {
+          if (typeVarType instanceof PyTypeVarTypeImpl impl) {
+            return impl.withScopeOwner(findScopeOwner(typeVarType, pyFunction, context));
+          }
+          return typeVarType;
+        }
+      }));
+    }
     return Ref.create(type);
   }
 
@@ -89,31 +98,21 @@ public final class PyDocStringTypeProvider extends PyTypeProviderBase {
     return null;
   }
 
-  // A hack to update scope owners of type parameters parsed out of docstrings
-  private void setTypeVarScopeOwners(@Nullable PyType type, @NotNull PyCallable callable, @NotNull TypeEvalContext context) {
-    if (!(callable instanceof PyFunction pyFunction)) return;
-    PyTypeChecker.Generics typeParameters = PyTypeChecker.collectGenerics(type, context);
-    if (typeParameters.isEmpty()) return;
-
-    PyClass pyClass = pyFunction.getContainingClass();
-
-    Function<PyTypeVarType, PyQualifiedNameOwner> findScopeOwner = typeVar -> pyFunction;
-    if (PyUtil.isInitOrNewMethod(callable)) {
-      findScopeOwner = typeVar -> pyClass;
+  private PyQualifiedNameOwner findScopeOwner(@NotNull PyTypeVarType typeVar,
+                                              @NotNull PyFunction function,
+                                              @NotNull TypeEvalContext context) {
+    PyClass pyClass = function.getContainingClass();
+    if (PyUtil.isInitOrNewMethod(function)) {
+      return pyClass;
     }
     else if (pyClass != null) {
       PyType classGenericType = getGenericType(pyClass, context);
       if (classGenericType != null) {
         PyTypeChecker.Generics classTypeParameters = PyTypeChecker.collectGenerics(classGenericType, context);
         Set<String> classTypeVarNames = ContainerUtil.map2Set(classTypeParameters.getTypeVars(), PyTypeVarType::getName);
-        findScopeOwner = typeVar -> classTypeVarNames.contains(typeVar.getName()) ? pyClass : pyFunction;
-      }
-    }
-
-    for (PyTypeParameterType typeParam : typeParameters.getAllTypeParameters()) {
-      if (typeParam instanceof PyTypeVarTypeImpl typeVar) {
-        typeVar.setScopeOwner(findScopeOwner.apply(typeVar));
+        return classTypeVarNames.contains(typeVar.getName()) ? pyClass : function;
       }
     }
+    return function;
   }
 }
diff --git a/python/python-psi-impl/src/com/jetbrains/python/psi/types/PyGenericType.java b/python/python-psi-impl/src/com/jetbrains/python/psi/types/PyGenericType.java
index 12caf35be504..583230a33623 100644
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/types/PyGenericType.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/types/PyGenericType.java
@@ -11,7 +11,6 @@ import com.jetbrains.python.psi.PyQualifiedNameOwner;
 import com.jetbrains.python.psi.PyTargetExpression;
 import com.jetbrains.python.psi.resolve.PyResolveContext;
 import com.jetbrains.python.psi.resolve.RatedResolveResult;
-import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -30,7 +29,7 @@ public class PyGenericType implements PyTypeVarType {
   private final @Nullable Ref<PyType> myDefaultType;
   private final boolean myIsDefinition;
   private final @Nullable PyQualifiedNameOwner myDeclarationElement;
-  private @Nullable PyQualifiedNameOwner myScopeOwner;
+  private final @Nullable PyQualifiedNameOwner myScopeOwner;
 
   public PyGenericType(@NotNull String name,
                        @NotNull List<@Nullable PyType> constraints,
@@ -166,14 +165,6 @@ public class PyGenericType implements PyTypeVarType {
     return new PyTypeVarTypeImpl(getName(), getConstraints(), getBound(), getDefaultType(), isDefinition(), declarationElement, getScopeOwner());
   }
 
-  @ApiStatus.Internal
-  public void setScopeOwner(@NotNull PyQualifiedNameOwner scopeOwner) {
-    if (myScopeOwner != null && myScopeOwner != scopeOwner) {
-      throw new IllegalStateException("Cannot override the existing scope owner");
-    }
-    myScopeOwner = scopeOwner;
-  }
-
   @Override
   public @NotNull PyGenericType toInstance() {
     return myIsDefinition ? new PyTypeVarTypeImpl(myName, myConstraints, myBound, myDefaultType, false, myDeclarationElement, myScopeOwner) : this;