PY-34617 Take into account sys.version_info checks when analyzing Python files

Support and, or, <=, > operators in version checks.

GitOrigin-RevId: 5006e88b0f7935d0bf0841dfd5fad5c371e8ff12

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

@@ -1,6 +1,7 @@
 // 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.jetbrains.python.psi.LanguageLevel
 import com.jetbrains.python.psi.PyIfStatement
 import com.jetbrains.python.psi.PyRecursiveElementVisitor
@@ -10,20 +11,22 @@ 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() {
+open class PyVersionAwareElementVisitor(languageLevel: LanguageLevel?) : PyRecursiveElementVisitor() {
+  private val version = languageLevel?.let { Version(it.majorVersion, it.minorVersion, 0) }
+
   override fun visitPyIfStatement(node: PyIfStatement) {
-    if (languageLevel == null) {
+    if (version == null) {
       super.visitPyIfStatement(node)
       return
     }
-    val ifParts = sequenceOf(node.getIfPart()) + node.elifParts.asSequence()
+    val ifParts = sequenceOf(node.ifPart) + node.elifParts.asSequence()
     for (ifPart in ifParts) {
-      val versionCheck = PyVersionCheck.fromCondition(ifPart)
-      if (versionCheck == null) {
+      val versions = ifPart.condition?.let(PyVersionCheck::convertToVersionRanges)
+      if (versions == null) {
         super.visitPyIfStatement(node)
         return
       }
-      if (versionCheck.matches(languageLevel)) {
+      if (versions.contains(version)) {
         ifPart.statementList.accept(this)
         return
       }

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

@@ -1,85 +1,88 @@
 // 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.google.common.collect.ImmutableRangeSet
+import com.google.common.collect.Range
 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 == languageLevel.isLessThan(version)
-  }
-
-  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
+object PyVersionCheck {
+  /**
+   * @return Version ranges if {@code expression} is a version check, {@code null} otherwise
+   *
+   * @see <a href="https://peps.python.org/pep-0484/#version-and-platform-checking">Version and Platform Checks</a>
+   */
+  @JvmStatic
+  fun convertToVersionRanges(expression: PyAstExpression): ImmutableRangeSet<Version>? {
+    val binaryExpr = PyPsiUtilsCore.flattenParens(expression) as? PyAstBinaryExpression ?: return null
+    when (val operator = binaryExpr.operator) {
+      PyTokenTypes.AND_KEYWORD, PyTokenTypes.OR_KEYWORD -> {
+        val rhs = binaryExpr.rightExpression ?: return null
+        val ranges1 = convertToVersionRanges(binaryExpr.leftExpression) ?: return null
+        val ranges2 = convertToVersionRanges(rhs) ?: return null
+        return if (operator === PyTokenTypes.AND_KEYWORD)
+          ranges1.intersection(ranges2)
+        else
+          ranges1.union(ranges2)
       }
 
-      val major = evaluateNumber(elements[0])
-      if (major == null) {
-        return null
+      PyTokenTypes.LT, PyTokenTypes.GT, PyTokenTypes.LE, PyTokenTypes.GE -> {
+        val refExpr = PyPsiUtilsCore.flattenParens(binaryExpr.leftExpression) as? PyAstReferenceExpression ?: return null
+        if (SYS_VERSION_INFO_QUALIFIED_NAME != refExpr.asQualifiedName()) return null
+
+        val tuple = PyPsiUtilsCore.flattenParens(binaryExpr.rightExpression) as? PyAstTupleExpression<*> ?: return null
+        val version = evaluateVersion(tuple) ?: return null
+
+        val range = when (operator) {
+          PyTokenTypes.LT -> Range.lessThan(version)
+          PyTokenTypes.GT -> Range.greaterThan(version)
+          PyTokenTypes.LE -> Range.atMost(version)
+          PyTokenTypes.GE -> Range.atLeast(version)
+          else -> throw IllegalStateException()
+        }
+        return ImmutableRangeSet.of(range)
       }
 
-      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
-      val intValue = value.toInt()
-      return if (BigInteger.valueOf(intValue.toLong()) == value) intValue else null
+      else -> return null
     }
   }
-}
 
-fun LanguageLevel.isLessThan(version: Version): Boolean {
-  return version.compareTo(majorVersion, minorVersion) > 0
-}
+  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
+    val intValue = value.toInt()
+    return if (BigInteger.valueOf(intValue.toLong()) == value) intValue else null
+  }
+}

python/python-psi-api/src/com/jetbrains/python/psi/stubs/PyVersionSpecificStub.kt

@@ -1,10 +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.stubs
 
+import com.google.common.collect.RangeSet
 import com.intellij.openapi.util.Version
 
 interface PyVersionSpecificStub {
-  val versionRange: PyVersionRange
+  val versions: RangeSet<Version>
 }
-
-data class PyVersionRange(val lowInclusive: Version?, val highExclusive: Version?)