PY-34617 Support version check

GitOrigin-RevId: 3318ff79cdcc5ba0ce5e4feb65abad5ad0f4acfa

python/python-psi-api/intellij.python.psi.iml

@@ -17,5 +17,6 @@
     <orderEntry type="module" module-name="intellij.python.parser" exported="" />
     <orderEntry type="module" module-name="intellij.python.ast" exported="" />
     <orderEntry type="module" module-name="intellij.python.syntax.core" />
+    <orderEntry type="library" name="kotlin-stdlib" level="project" />
   </component>
 </module>

python/python-psi-api/src/com/jetbrains/python/psi/PyElsePart.java

@@ -15,10 +15,12 @@
  */
 package com.jetbrains.python.psi;
 
+import com.intellij.psi.StubBasedPsiElement;
 import com.jetbrains.python.ast.PyAstElsePart;
+import com.jetbrains.python.psi.stubs.PyElsePartStub;
 
 /**
  * The 'else:' part of various compound statements.
  */
-public interface PyElsePart extends PyAstElsePart, PyStatementPart {
+public interface PyElsePart extends PyAstElsePart, PyStatementPart, StubBasedPsiElement<PyElsePartStub> {
 }

python/python-psi-api/src/com/jetbrains/python/psi/PyIfPartElif.java

@@ -1,7 +1,9 @@
 // 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;
 
+import com.intellij.psi.StubBasedPsiElement;
 import com.jetbrains.python.ast.PyAstIfPartElif;
+import com.jetbrains.python.psi.stubs.PyIfPartElifStub;
 
-public interface PyIfPartElif extends PyAstIfPartElif, PyIfPart {
+public interface PyIfPartElif extends PyAstIfPartElif, PyIfPart, StubBasedPsiElement<PyIfPartElifStub> {
 }

python/python-psi-api/src/com/jetbrains/python/psi/PyIfPartIf.java

@@ -1,7 +1,9 @@
 // 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;
 
+import com.intellij.psi.StubBasedPsiElement;
 import com.jetbrains.python.ast.PyAstIfPartIf;
+import com.jetbrains.python.psi.stubs.PyIfPartIfStub;
 
-public interface PyIfPartIf extends PyAstIfPartIf, PyIfPart {
+public interface PyIfPartIf extends PyAstIfPartIf, PyIfPart, StubBasedPsiElement<PyIfPartIfStub> {
 }

python/python-psi-api/src/com/jetbrains/python/psi/impl/PyPsiUtils.java

@@ -16,10 +16,15 @@ import com.intellij.psi.tree.TokenSet;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.psi.util.PsiUtilCore;
 import com.intellij.psi.util.QualifiedName;
+import com.intellij.util.Processor;
 import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.ast.impl.PyPsiUtilsCore;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.stubs.PyElsePartStub;
+import com.jetbrains.python.psi.stubs.PyIfPartElifStub;
+import com.jetbrains.python.psi.stubs.PyIfPartIfStub;
+import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 
@@ -408,40 +413,6 @@ public final class PyPsiUtils {
     return blocks;
   }
 
-  @NotNull
-  static <T extends PyElement> List<T> collectStubChildren(@NotNull PyFile pyFile,
-                                                           @Nullable StubElement<?> stub,
-                                                           @NotNull Class<T> elementType) {
-    if (stub != null) {
-      final List<T> result = new ArrayList<>();
-      final List<StubElement<?>> children = stub.getChildrenStubs();
-      for (StubElement<?> child : children) {
-        PsiElement childPsi = child.getPsi();
-        if (elementType.isInstance(childPsi)) {
-          result.add(elementType.cast(childPsi));
-        }
-      }
-      return result;
-    }
-    else {
-      return PyPsiUtilsCore.collectChildren(pyFile, elementType);
-    }
-  }
-
-  static List<PsiElement> collectAllStubChildren(PsiElement e, StubElement<?> stub) {
-    if (stub != null) {
-      final List<PsiElement> result = new ArrayList<>();
-      final List<StubElement<?>> children = stub.getChildrenStubs();
-      for (StubElement<?> child : children) {
-        result.add(child.getPsi());
-      }
-      return result;
-    }
-    else {
-      return PyPsiUtilsCore.collectAllChildren(e);
-    }
-  }
-
   public static int findArgumentIndex(PyCallExpression call, PsiElement argument) {
     final PyExpression[] args = call.getArguments();
     for (int i = 0; i < args.length; i++) {
@@ -504,10 +475,7 @@ public final class PyPsiUtils {
 
   @Nullable
   public static PyExpression flattenParens(@Nullable PyExpression expr) {
-    while (expr instanceof PyParenthesizedExpression) {
-      expr = ((PyParenthesizedExpression)expr).getContainedExpression();
-    }
-    return expr;
+    return (PyExpression)PyPsiUtilsCore.flattenParens(expr);
   }
 
   @Nullable
@@ -664,4 +632,59 @@ public final class PyPsiUtils {
     while (next != null);
     return null;
   }
+
+  /**
+   * @see <a href="https://typing.readthedocs.io/en/latest/source/stubs.html#version-and-platform-checks">Version and Platform Checks</a>
+   */
+  @ApiStatus.Internal
+  @Nullable
+  public static StubElement<?> getParentStubSkippingVersionChecks(@Nullable StubElement<?> stub) {
+    if (stub != null) {
+      for (StubElement<?> e = stub.getParentStub(); e != null; e = e.getParentStub()) {
+        boolean isVersionCheck = e instanceof PyIfPartIfStub || e instanceof PyIfPartElifStub || e instanceof PyElsePartStub;
+        if (!isVersionCheck) {
+          return e;
+        }
+      }
+    }
+    return null;
+  }
+
+  /**
+   * @see <a href="https://typing.readthedocs.io/en/latest/source/stubs.html#version-and-platform-checks">Version and Platform Checks</a>
+   */
+  @ApiStatus.Internal
+  public static boolean processChildrenStubs(@NotNull StubElement<?> stub,
+                                             @NotNull LanguageLevel languageLevel,
+                                             @NotNull Processor<? super StubElement<?>> processor) {
+    boolean recurse = false;
+    for (StubElement<?> child : stub.getChildrenStubs()) {
+      if (child instanceof PyIfPartIfStub ifStub) {
+        if (ifStub.getVersionCheck().matches(languageLevel)) {
+          recurse = false;
+          if (!processChildrenStubs(child, languageLevel, processor)) return false;
+        }
+        else {
+          recurse = true;
+        }
+      }
+      else if (child instanceof PyIfPartElifStub elifStub) {
+        if (recurse && elifStub.getVersionCheck().matches(languageLevel)) {
+          recurse = false;
+          if (!processChildrenStubs(child, languageLevel, processor)) return false;
+        }
+      }
+      else if (child instanceof PyElsePartStub) {
+        if (recurse) {
+          recurse = false;
+          if (!processChildrenStubs(child, languageLevel, processor)) return false;
+        }
+      }
+      else {
+        recurse = false;
+        if (!processor.process(child)) return false;
+      }
+    }
+    return true;
+  }
 }

python/python-psi-api/src/com/jetbrains/python/psi/impl/PyVersionAwareElementVisitor.kt

@@ -0,0 +1,33 @@
+// 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.impl
+
+import com.jetbrains.python.psi.LanguageLevel
+import com.jetbrains.python.psi.PyIfStatement
+import com.jetbrains.python.psi.PyRecursiveElementVisitor
+import org.jetbrains.annotations.ApiStatus
+
+/**
+ * @see [Version and Platform Checks](https://typing.readthedocs.io/en/latest/source/stubs.html.version-and-platform-checks)
+ */
+@ApiStatus.Internal
+open class PyVersionAwareElementVisitor(private val languageLevel: LanguageLevel?) : PyRecursiveElementVisitor() {
+  override fun visitPyIfStatement(node: PyIfStatement) {
+    if (languageLevel == null) {
+      super.visitPyIfStatement(node)
+      return
+    }
+    val ifParts = sequenceOf(node.getIfPart()) + node.elifParts.asSequence()
+    for (ifPart in ifParts) {
+      val versionCheck = PyVersionCheck.fromCondition(ifPart)
+      if (versionCheck == null) {
+        super.visitPyIfStatement(node)
+        return
+      }
+      if (versionCheck.matches(languageLevel)) {
+        ifPart.statementList.accept(this)
+        return
+      }
+    }
+    node.elsePart?.statementList?.accept(this)
+  }
+}

python/python-psi-api/src/com/jetbrains/python/psi/impl/PyVersionAwareTopLevelElementVisitor.kt

@@ -0,0 +1,27 @@
+// 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.impl
+
+import com.intellij.psi.PsiElement
+import com.jetbrains.python.psi.LanguageLevel
+import com.jetbrains.python.psi.PyClass
+import com.jetbrains.python.psi.PyElement
+import com.jetbrains.python.psi.PyFunction
+import org.jetbrains.annotations.ApiStatus
+
+@ApiStatus.Internal
+abstract class PyVersionAwareTopLevelElementVisitor(languageLevel: LanguageLevel) : PyVersionAwareElementVisitor(languageLevel) {
+  override fun visitPyElement(node: PyElement) {
+    super.visitPyElement(node)
+    checkAddElement(node)
+  }
+
+  override fun visitPyClass(node: PyClass) {
+    checkAddElement(node) // do not recurse into classes
+  }
+
+  override fun visitPyFunction(node: PyFunction) {
+    checkAddElement(node) // do not recurse into functions
+  }
+
+  protected abstract fun checkAddElement(node: PsiElement?)
+}

python/python-psi-api/src/com/jetbrains/python/psi/impl/PyVersionCheck.kt

@@ -0,0 +1,81 @@
+// 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.impl
+
+import com.intellij.openapi.util.Version
+import com.intellij.psi.util.QualifiedName
+import com.jetbrains.python.PyTokenTypes
+import com.jetbrains.python.ast.*
+import com.jetbrains.python.ast.impl.PyPsiUtilsCore
+import com.jetbrains.python.psi.LanguageLevel
+import org.jetbrains.annotations.ApiStatus
+import java.math.BigInteger
+
+@ApiStatus.Internal
+data class PyVersionCheck(val version: Version, val isLessThan: Boolean) {
+  fun matches(languageLevel: LanguageLevel): Boolean {
+    return isLessThan == version.compareTo(languageLevel.majorVersion, languageLevel.minorVersion) > 0
+  }
+
+  companion object {
+    /**
+     * Extracts the Python version comparison from {@code ifPart}'s condition if it's a version check as specified in
+     * <a href="https://typing.readthedocs.io/en/latest/source/stubs.html#version-and-platform-checks">Version and Platform Checks</a> E.g.
+     * <pre>{@code
+     * if sys.version_info >= (3,):
+     *   ...
+     * }</pre>
+     * @return A {@link VersionCheck} instance if {@code ifPart} is a (valid) version check, or {@code null} otherwise.
+     */
+    @JvmStatic
+    fun fromCondition(ifPart: PyAstIfPart): PyVersionCheck? {
+      val binaryExpr = PyPsiUtilsCore.flattenParens(ifPart.condition)
+      if (binaryExpr !is PyAstBinaryExpression) return null
+
+      val lhsRefExpr = PyPsiUtilsCore.flattenParens(binaryExpr.leftExpression)
+      if (lhsRefExpr !is PyAstReferenceExpression) return null
+      if (SYS_VERSION_INFO_QUALIFIED_NAME != lhsRefExpr.asQualifiedName()) return null
+
+      val versionTuple = PyPsiUtilsCore.flattenParens(binaryExpr.rightExpression)
+      if (versionTuple !is PyAstTupleExpression<*>) return null
+      val version = evaluateVersion(versionTuple)
+      if (version == null) return null
+
+      val operator = binaryExpr.getOperator()
+      if (operator !== PyTokenTypes.LT && operator !== PyTokenTypes.GE) return null
+      return PyVersionCheck(version, operator === PyTokenTypes.LT)
+    }
+
+    private val SYS_VERSION_INFO_QUALIFIED_NAME = QualifiedName.fromDottedString("sys.version_info")
+
+    private fun evaluateVersion(versionTuple: PyAstTupleExpression<*>): Version? {
+      val elements = versionTuple.elements
+      if (elements.size != 1 && elements.size != 2) {
+        return null
+      }
+
+      val major = evaluateNumber(elements[0])
+      if (major == null) {
+        return null
+      }
+
+      if (elements.size == 1) {
+        return Version(major, 0, 0)
+      }
+
+      val minor = evaluateNumber(elements[1])
+      if (minor == null) {
+        return null
+      }
+
+      return Version(major, minor, 0)
+    }
+
+    private fun evaluateNumber(expression: PyAstExpression?): Int? {
+      if (expression !is PyAstNumericLiteralExpression) return null
+      if (!expression.isIntegerLiteral) return null
+      val value = expression.bigIntegerValue ?: return null
+      val intValue = value.toInt()
+      return if (BigInteger.valueOf(intValue.toLong()) == value) intValue else null
+    }
+  }
+}

python/python-psi-api/src/com/jetbrains/python/psi/stubs/PyElsePartStub.java

@@ -0,0 +1,8 @@
+// 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.stubs;
+
+import com.intellij.psi.stubs.StubElement;
+import com.jetbrains.python.psi.PyElsePart;
+
+public interface PyElsePartStub extends StubElement<PyElsePart> {
+}

python/python-psi-api/src/com/jetbrains/python/psi/stubs/PyIfPartElifStub.java

@@ -0,0 +1,10 @@
+// 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.stubs;
+
+import com.intellij.psi.stubs.StubElement;
+import com.jetbrains.python.psi.impl.PyVersionCheck;
+import com.jetbrains.python.psi.PyIfPartElif;
+
+public interface PyIfPartElifStub extends StubElement<PyIfPartElif> {
+  PyVersionCheck getVersionCheck();
+}

python/python-psi-api/src/com/jetbrains/python/psi/stubs/PyIfPartIfStub.java

@@ -0,0 +1,10 @@
+// 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.stubs;
+
+import com.intellij.psi.stubs.StubElement;
+import com.jetbrains.python.psi.impl.PyVersionCheck;
+import com.jetbrains.python.psi.PyIfPartIf;
+
+public interface PyIfPartIfStub extends StubElement<PyIfPartIf> {
+  PyVersionCheck getVersionCheck();
+}