PY-78236 Introduce PyTypeVisitor API

GitOrigin-RevId: b6aa236b364bc55840be0a148205abba9c22167b

python/python-psi-api/src/com/jetbrains/python/psi/types/PyCallableType.java

@@ -60,4 +60,9 @@ public interface PyCallableType extends PyType {
   default int getImplicitOffset() {
     return 0;
   }
+
+  @Override
+  default <T> T acceptTypeVisitor(@NotNull PyTypeVisitor<T> visitor) {
+    return visitor.visitPyCallableType(this);
+  }
 }

python/python-psi-api/src/com/jetbrains/python/psi/types/PyClassLikeType.java

@@ -48,4 +48,9 @@ public interface PyClassLikeType extends PyCallableType, PyWithAncestors, PyInst
 
   @Nullable
   PyClassLikeType getMetaClassType(@NotNull TypeEvalContext context, boolean inherited);
+
+  @Override
+  default <T> T acceptTypeVisitor(@NotNull PyTypeVisitor<T> visitor) {
+    return visitor.visitPyClassLikeType(this);
+  }
 }

python/python-psi-api/src/com/jetbrains/python/psi/types/PyClassType.java

@@ -20,4 +20,9 @@ public interface PyClassType extends PyClassLikeType, UserDataHolder {
   default boolean isAttributeWritable(@NotNull String name, @NotNull TypeEvalContext context) {
     return true;
   }
+
+  @Override
+  default <T> T acceptTypeVisitor(@NotNull PyTypeVisitor<T> visitor) {
+    return visitor.visitPyClassType(this);
+  }
 }

python/python-psi-api/src/com/jetbrains/python/psi/types/PyType.java

@@ -9,6 +9,7 @@ import com.jetbrains.python.psi.PyExpression;
 import com.jetbrains.python.psi.PyQualifiedNameOwner;
 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;
 
@@ -72,4 +73,13 @@ public interface PyType {
   boolean isBuiltin();
 
   void assertValid(String message);
+
+  /**
+   * For nullable {@code PyType} instance use {@link PyTypeVisitor#visitPyType(PyType, PyTypeVisitor)}
+   * to visit {@code null} values with {@link PyTypeVisitor#visitUnknownType()}.
+   */
+  @ApiStatus.Experimental
+  default <T> T acceptTypeVisitor(@NotNull PyTypeVisitor<T> visitor) {
+    return visitor.visitPyType(this);
+  }
 }

python/python-psi-api/src/com/jetbrains/python/psi/types/PyTypeVisitor.java

@@ -0,0 +1,76 @@
+// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
+package com.jetbrains.python.psi.types;
+
+import org.jetbrains.annotations.ApiStatus;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+
+/**
+ * Similarly to {@link com.intellij.psi.PsiElementVisitor}, implements double dispatching for the {@link PyType} hierarchy.
+ * <p>
+ * Because the "unknown" type is historically represented as {@code null} in the type system, {@code PyTypeVisitor.visitPyType(type, visitor)}
+ * should be used instead of direct {@code type.acceptTypeVisitor(visitor)} to properly account for possible {@code null} values.
+ * <p>
+ * This class gives access only to the types declared in the <tt>intellij.python.psi</tt> module. 
+ * Most actual implementations should extend {@link PyTypeVisitorExt}.
+ * <p>
+ * There are helper {@link PyRecursiveTypeVisitor} and {@link PyCloningTypeVisitor} for recursive type
+ * traversal and deep cloning of a type respectively.
+ *
+ * @see PyRecursiveTypeVisitor
+ * @see PyCloningTypeVisitor
+ * @see PyType#acceptTypeVisitor(PyTypeVisitor)
+ * @see #visitPyType(PyType, PyTypeVisitor) 
+ * @see #visitUnknownType()
+ */
+@ApiStatus.Experimental
+public abstract class PyTypeVisitor<T> {
+  public static <T> T visitPyType(@Nullable PyType type, @NotNull PyTypeVisitor<T> visitor) {
+    return type == null ? visitor.visitUnknownType() : type.acceptTypeVisitor(visitor);
+  }
+
+  public T visitPyType(@NotNull PyType type) {
+    return null;
+  }
+
+  public T visitPyClassType(@NotNull PyClassType classType) {
+    return visitPyClassLikeType(classType);
+  }
+
+  public T visitPyClassLikeType(@NotNull PyClassLikeType classLikeType) {
+    return visitPyCallableType(classLikeType);
+  }
+
+  public T visitPyFunctionType(@NotNull PyFunctionType functionType) {
+    return visitPyCallableType(functionType);
+  }
+
+  public T visitPyCallableType(@NotNull PyCallableType callableType) {
+    return visitPyType(callableType);
+  }
+
+  public T visitPyTypeVarType(@NotNull PyTypeVarType typeVarType) {
+    return visitPyTypeParameterType(typeVarType);
+  }
+
+  public T visitPyTypeVarTupleType(@NotNull PyTypeVarTupleType typeVarTupleType) {
+    return visitPyTypeParameterType(typeVarTupleType);
+  }
+
+  public T visitPyTypeParameterType(@NotNull PyTypeParameterType typeParameterType) {
+    return visitPyType(typeParameterType);
+  }
+
+  public T visitPyUnpackedTupleType(@NotNull PyUnpackedTupleType unpackedTupleType) {
+    return visitPyType(unpackedTupleType);
+  }
+
+  public T visitPyCallableParameterListType(@NotNull PyCallableParameterListType callableParameterListType) {
+    return visitPyType(callableParameterListType);
+  }
+
+  public T visitUnknownType() {
+    return null;
+  }
+}