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

Support version checks for import statements.

GitOrigin-RevId: df52f60574962e1bc222121aadc082683de0a869

python/python-common-tests/com/jetbrains/python/PyCommonResolveTest.java

@@ -2094,6 +2094,49 @@ public abstract class PyCommonResolveTest extends PyCommonResolveTestCase {
     assertResolvedElement(LanguageLevel.PYTHON27, buz, e -> assertResolveResult(e, PyFunction.class, "buz", "mod.py"));
   }
 
+  // PY-34617
+  public void testImportUnderVersionCheckMultifile() {
+    myFixture.copyDirectoryToProject("resolve/ImportUnderVersionCheck", "");
+    String plainImport = """
+      from mod import *
+      math
+       <ref>""";
+    assertResolvedElement(LanguageLevel.PYTHON35, plainImport, e -> assertResolveResult(e, PyFile.class, "math.pyi", null));
+    assertResolvedElement(LanguageLevel.PYTHON34, plainImport, TestCase::assertNull);
+
+    String importAlias = """
+      from mod import *
+      cm
+       <ref>""";
+    assertResolvedElement(LanguageLevel.PYTHON35, importAlias, e -> assertResolveResult(e, PyFile.class, "cmath.pyi", null));
+    assertResolvedElement(LanguageLevel.PYTHON34, importAlias, TestCase::assertNull);
+  }
+
+  // PY-34617
+  public void testImportFromUnderVersionCheckMultifile() {
+    myFixture.copyDirectoryToProject("resolve/ImportUnderVersionCheck", "");
+    String plainImport = """
+      from mod import *
+      digits
+       <ref>""";
+    assertResolvedElement(LanguageLevel.PYTHON35, plainImport, e -> assertResolveResult(e, PyTargetExpression.class, "digits", null));
+    assertResolvedElement(LanguageLevel.PYTHON34, plainImport, TestCase::assertNull);
+
+    String importAlias = """
+      from mod import *
+      imported_name
+       <ref>""";
+    assertResolvedElement(LanguageLevel.PYTHON35, importAlias, e -> assertResolveResult(e, PyTargetExpression.class, "hexdigits", null));
+    assertResolvedElement(LanguageLevel.PYTHON34, importAlias, TestCase::assertNull);
+
+    String starImport = """
+      from mod import *
+      DivisionByZero
+       <ref>""";
+    assertResolvedElement(LanguageLevel.PYTHON35, starImport, e -> assertResolveResult(e, PyClass.class, "DivisionByZero", null));
+    assertResolvedElement(LanguageLevel.PYTHON34, starImport, TestCase::assertNull);
+  }
+
   private void assertResolvedElement(@NotNull LanguageLevel languageLevel, @NotNull String text, @NotNull Consumer<PsiElement> assertion) {
     runWithLanguageLevel(languageLevel, () -> {
       myFixture.configureByText(PythonFileType.INSTANCE, text);

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

@@ -7,7 +7,7 @@ import com.intellij.psi.util.QualifiedName;
 import org.jetbrains.annotations.Nullable;
 
 
-public interface PyFromImportStatementStub extends StubElement<PyFromImportStatement> {
+public interface PyFromImportStatementStub extends StubElement<PyFromImportStatement>, PyVersionSpecificStub {
   @Nullable
   QualifiedName getImportSourceQName();
 

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

@@ -5,5 +5,5 @@ import com.intellij.psi.stubs.StubElement;
 import com.jetbrains.python.psi.PyImportStatement;
 
 
-public interface PyImportStatementStub extends StubElement<PyImportStatement> {
+public interface PyImportStatementStub extends StubElement<PyImportStatement>, PyVersionSpecificStub {
 }

python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyFromImportStatementElementType.java

@@ -1,14 +1,15 @@
 // Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package com.jetbrains.python.psi.impl.stubs;
 
+import com.google.common.collect.RangeSet;
 import com.intellij.lang.ASTNode;
+import com.intellij.openapi.util.Version;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.stubs.IStubElementType;
 import com.intellij.psi.stubs.StubElement;
 import com.intellij.psi.stubs.StubInputStream;
 import com.intellij.psi.stubs.StubOutputStream;
 import com.intellij.psi.util.QualifiedName;
-import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.PyStubElementTypes;
 import com.jetbrains.python.psi.PyFromImportStatement;
 import com.jetbrains.python.psi.PyStubElementType;
@@ -43,8 +44,9 @@ public class PyFromImportStatementElementType extends PyStubElementType<PyFromIm
   @NotNull
   @Override
   public PyFromImportStatementStub createStub(@NotNull PyFromImportStatement psi, StubElement parentStub) {
+    final RangeSet<Version> versions = PyVersionSpecificStubBaseKt.evaluateVersionsForElement(psi);
     return new PyFromImportStatementStubImpl(psi.getImportSourceQName(), psi.isStarImport(), psi.getRelativeLevel(), parentStub,
-                                             getStubElementType());
+                                             getStubElementType(), versions);
   }
 
   @Override
@@ -53,6 +55,7 @@ public class PyFromImportStatementElementType extends PyStubElementType<PyFromIm
     QualifiedName.serialize(qName, dataStream);
     dataStream.writeBoolean(stub.isStarImport());
     dataStream.writeVarInt(stub.getRelativeLevel());
+    PyVersionSpecificStubBaseKt.serializeVersions(stub.getVersions(), dataStream);
   }
 
   @Override
@@ -61,7 +64,8 @@ public class PyFromImportStatementElementType extends PyStubElementType<PyFromIm
     QualifiedName qName = QualifiedName.deserialize(dataStream);
     boolean isStarImport = dataStream.readBoolean();
     int relativeLevel = dataStream.readVarInt();
-    return new PyFromImportStatementStubImpl(qName, isStarImport, relativeLevel, parentStub, getStubElementType());
+    RangeSet<Version> versions = PyVersionSpecificStubBaseKt.deserializeVersions(dataStream);
+    return new PyFromImportStatementStubImpl(qName, isStarImport, relativeLevel, parentStub, getStubElementType(), versions);
   }
 
   protected IStubElementType getStubElementType() {

python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyFromImportStatementStubImpl.java

@@ -1,22 +1,26 @@
 // Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package com.jetbrains.python.psi.impl.stubs;
 
+import com.google.common.collect.RangeSet;
+import com.intellij.openapi.util.Version;
 import com.intellij.psi.stubs.IStubElementType;
 import com.intellij.psi.stubs.StubBase;
 import com.intellij.psi.stubs.StubElement;
 import com.intellij.psi.util.QualifiedName;
 import com.jetbrains.python.psi.PyFromImportStatement;
+import com.jetbrains.python.psi.PyImportStatement;
 import com.jetbrains.python.psi.stubs.PyFromImportStatementStub;
+import org.jetbrains.annotations.NotNull;
 
 
-public class PyFromImportStatementStubImpl extends StubBase<PyFromImportStatement> implements PyFromImportStatementStub {
+public class PyFromImportStatementStubImpl extends PyVersionSpecificStubBase<PyFromImportStatement> implements PyFromImportStatementStub {
   private final QualifiedName myImportSourceQName;
   private final boolean myStarImport;
   private final int myRelativeLevel;
 
   public PyFromImportStatementStubImpl(QualifiedName importSourceQName, boolean isStarImport, int relativeLevel,
-                                       final StubElement parent, IStubElementType elementType) {
+                                       final StubElement parent, IStubElementType elementType, @NotNull RangeSet<Version> versions) {
-    super(parent, elementType);
+    super(parent, elementType, versions);
     myImportSourceQName = importSourceQName;
     myStarImport = isStarImport;
     myRelativeLevel = relativeLevel;

python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyImportStatementElementType.java

@@ -1,13 +1,14 @@
 // Copyright 2000-2018 JetBrains s.r.o. Use of this source code is governed by the Apache 2.0 license that can be found in the LICENSE file.
 package com.jetbrains.python.psi.impl.stubs;
 
+import com.google.common.collect.RangeSet;
 import com.intellij.lang.ASTNode;
+import com.intellij.openapi.util.Version;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.stubs.IStubElementType;
 import com.intellij.psi.stubs.StubElement;
 import com.intellij.psi.stubs.StubInputStream;
 import com.intellij.psi.stubs.StubOutputStream;
-import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.PyStubElementTypes;
 import com.jetbrains.python.psi.PyImportStatement;
 import com.jetbrains.python.psi.PyStubElementType;
@@ -42,17 +43,20 @@ public class PyImportStatementElementType extends PyStubElementType<PyImportStat
   @NotNull
   @Override
   public PyImportStatementStub createStub(@NotNull PyImportStatement psi, StubElement parentStub) {
-    return new PyImportStatementStubImpl(parentStub, getStubElementType());
+    final RangeSet<Version> versions = PyVersionSpecificStubBaseKt.evaluateVersionsForElement(psi);
+    return new PyImportStatementStubImpl(parentStub, getStubElementType(), versions);
   }
 
   @Override
   public void serialize(@NotNull PyImportStatementStub stub, @NotNull StubOutputStream dataStream) throws IOException {
+    PyVersionSpecificStubBaseKt.serializeVersions(stub.getVersions(), dataStream);
   }
 
   @Override
   @NotNull
   public PyImportStatementStub deserialize(@NotNull StubInputStream dataStream, StubElement parentStub) throws IOException {
-    return new PyImportStatementStubImpl(parentStub, getStubElementType());
+    RangeSet<Version> versions = PyVersionSpecificStubBaseKt.deserializeVersions(dataStream);
+    return new PyImportStatementStubImpl(parentStub, getStubElementType(), versions);
   }
 
   protected IStubElementType getStubElementType() {

python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyImportStatementStubImpl.java

@@ -15,16 +15,18 @@
  */
 package com.jetbrains.python.psi.impl.stubs;
 
+import com.google.common.collect.RangeSet;
+import com.intellij.openapi.util.Version;
 import com.intellij.psi.stubs.IStubElementType;
-import com.intellij.psi.stubs.StubBase;
 import com.intellij.psi.stubs.StubElement;
 import com.jetbrains.python.psi.PyImportStatement;
 import com.jetbrains.python.psi.stubs.PyImportStatementStub;
+import org.jetbrains.annotations.NotNull;
 
 
-public class PyImportStatementStubImpl extends StubBase<PyImportStatement> implements PyImportStatementStub {
+public class PyImportStatementStubImpl extends PyVersionSpecificStubBase<PyImportStatement> implements PyImportStatementStub {
-  public PyImportStatementStubImpl(StubElement parentStub, IStubElementType elementType) {
+  public PyImportStatementStubImpl(StubElement parentStub, IStubElementType elementType, @NotNull RangeSet<Version> versions) {
-    super(parentStub, elementType);
+    super(parentStub, elementType, versions);
   }
 
   @Override

python/testData/resolve/ImportUnderVersionCheck/mod.py

@@ -0,0 +1,8 @@
+import sys
+
+
+if sys.version_info >= (3,5):
+    import math
+    import cmath as cm
+    from string import digits, hexdigits as imported_name
+    from decimal import *

python/testData/stubs/StubStructureWithVersionChecks.py

@@ -22,4 +22,11 @@ if condition1:
                     i = 1
 
 if (sys.version_info > (2, 1) and ((sys.version_info <= (2, 2) or sys.version_info > (3, )))):
-    qux = ""
+    qux = ""
+
+if sys.version_info <= (2, 1):
+    import mod_aaa
+    import mod_bbb as bbb
+    from mod import ab
+    from mod import aba as abb
+    from mod import *

python/testSrc/com/jetbrains/python/PyStubsTest.java

@@ -154,7 +154,12 @@ public class PyStubsTest extends PyTestCase {
           ),
         element(PyTargetExpressionStub.class,
                 ImmutableRangeSet.unionOf(List.of(Range.openClosed(Version.parseVersion("2.1"), Version.parseVersion("2.2")),
-                                                  Range.greaterThan(Version.parseVersion("3.0")))))
+                                                  Range.greaterThan(Version.parseVersion("3.0"))))),
+        element(PyImportStatementStub.class, versionAtMost("2.1")),
+        element(PyImportStatementStub.class, versionAtMost("2.1")),
+        element(PyFromImportStatementStub.class, versionAtMost("2.1")),
+        element(PyFromImportStatementStub.class, versionAtMost("2.1")),
+        element(PyFromImportStatementStub.class, versionAtMost("2.1"))
       )
       .test(file.getStub());
   }
@@ -311,10 +316,12 @@ public class PyStubsTest extends PyTestCase {
   }
 
   public void testImportInExcept() {
-    final PyFileImpl file = (PyFileImpl) getTestFile();
+    runWithLanguageLevel(LanguageLevel.PYTHON26, () -> {
-    final PsiElement element = file.getElementNamed("tzinfo");
+      final PyFileImpl file = (PyFileImpl)getTestFile();
-    assertTrue(element != null ? element.toString() : "null", element instanceof PyClass);
+      final PsiElement element = file.getElementNamed("tzinfo");
-    assertNotParsed(file);
+      assertTrue(element != null ? element.toString() : "null", element instanceof PyClass);
+      assertNotParsed(file);
+    });
   }
 
 
@@ -1280,6 +1287,10 @@ public class PyStubsTest extends PyTestCase {
     return ImmutableRangeSet.of(Range.lessThan(Version.parseVersion(version)));
   }
 
+  private static @NotNull RangeSet<Version> versionAtMost(@NotNull String version) {
+    return ImmutableRangeSet.of(Range.atMost(Version.parseVersion(version)));
+  }
+
   private static @NotNull RangeSet<Version> versionRange(@NotNull String lowInclusive, @NotNull String highExclusive) {
     return ImmutableRangeSet.of(Range.closedOpen(Version.parseVersion(lowInclusive), Version.parseVersion(highExclusive)));
   }