PY-34617 Support version check

GitOrigin-RevId: 3318ff79cdcc5ba0ce5e4feb65abad5ad0f4acfa
2026-04-18 04:21:24 +07:00 · 2024-07-26 22:12:08 +02:00
parent 16a7fb4b3e
commit 93b9066edf
60 changed files with 1086 additions and 219 deletions
--- a/python/python-psi-impl/src/com/jetbrains/python/PyElementTypesFacadeImpl.kt
+++ b/python/python-psi-impl/src/com/jetbrains/python/PyElementTypesFacadeImpl.kt
@@ -4,6 +4,7 @@ package com.jetbrains.python
 import com.intellij.lang.ASTNode
 import com.intellij.psi.PsiElement
 import com.intellij.psi.stubs.IStubElementType
+import com.intellij.psi.tree.IElementType
 import com.intellij.psi.tree.IReparseableElementType
 import com.jetbrains.python.psi.impl.*
 import java.util.function.Function
@@ -51,6 +52,12 @@ class PyElementTypesFacadeImpl : PyElementTypesFacade() {
    get() = PyStubElementTypes.TYPE_PARAMETER_LIST
  override val typeAliasStatement: IStubElementType<*, *>
    get() = PyStubElementTypes.TYPE_ALIAS_STATEMENT
+  override val ifPartIf: IElementType
+    get() = PyStubElementTypes.IF_PART_IF
+  override val ifPartElif: IElementType
+    get() = PyStubElementTypes.IF_PART_ELIF
+  override val elsePart: IElementType
+    get() = PyStubElementTypes.ELSE_PART

  override val statementList: IReparseableElementType
    get() = PyStatementListElementType()
@@ -170,10 +177,6 @@ class PyElementTypesFacadeImpl : PyElementTypesFacade() {
    get() = F { node -> PySetCompExpressionImpl(node) }
  override val dictCompExpressionConstructor: F
    get() = F { node -> PyDictCompExpressionImpl(node) }
-  override val ifPartIfConstructor: F
-    get() = F { node -> PyIfPartIfImpl(node) }
-  override val ifPartElifConstructor: F
-    get() = F { node -> PyIfPartElifImpl(node) }
  override val forPartConstructor: F
    get() = F { node -> PyForPartImpl(node) }
  override val whilePartConstructor: F
@@ -182,8 +185,6 @@ class PyElementTypesFacadeImpl : PyElementTypesFacade() {
    get() = F { node -> PyTryPartImpl(node) }
  override val finallyPartConstructor: F
    get() = F { node -> PyFinallyPartImpl(node) }
-  override val elsePartConstructor: F
-    get() = F { node -> PyElsePartImpl(node) }
  override val fStringNodeConstructor: F
    get() = F { node -> PyFormattedStringElementImpl(node) }
  override val fStringFragmentConstructor: F
--- a/python/python-psi-impl/src/com/jetbrains/python/PyStubElementTypes.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/PyStubElementTypes.java
@@ -33,4 +33,8 @@ public interface PyStubElementTypes {
  PyStubElementType<PyTypeParameterStub, PyTypeParameter> TYPE_PARAMETER = new PyTypeParameterElementType();
  PyStubElementType<PyTypeParameterListStub, PyTypeParameterList> TYPE_PARAMETER_LIST = new PyTypeParameterListElementType();
  PyStubElementType<PyTypeAliasStatementStub, PyTypeAliasStatement> TYPE_ALIAS_STATEMENT = new PyTypeAliasStatementElementType();
+
+  PyStubElementType<PyIfPartIfStub, PyIfPartIf> IF_PART_IF = new PyIfPartIfElementType();
+  PyStubElementType<PyIfPartElifStub, PyIfPartElif> IF_PART_ELIF = new PyIfPartElifElementType();
+  PyStubElementType<PyElsePartStub, PyElsePart> ELSE_PART = new PyElsePartElementType();
 }
--- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/dataflow/scope/impl/ScopeImpl.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/dataflow/scope/impl/ScopeImpl.java
@@ -18,6 +18,8 @@ import com.jetbrains.python.codeInsight.dataflow.scope.ScopeVariable;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.impl.PyAugAssignmentStatementNavigator;
 import com.jetbrains.python.psi.impl.PyPsiUtils;
+import com.jetbrains.python.psi.impl.PyVersionAwareElementVisitor;
+import com.jetbrains.python.pyi.PyiUtil;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -214,7 +216,14 @@ public class ScopeImpl implements Scope {
    final Set<String> nonlocals = new HashSet<>();
    final Set<String> augAssignments = new HashSet<>();
    final List<PyTargetExpression> targetExpressions = new ArrayList<>();
-    myFlowOwner.acceptChildren(new PyRecursiveElementVisitor() {
+    final LanguageLevel languageLevel;
+    if (myFlowOwner instanceof PyFile || myFlowOwner instanceof PyClass) {
+      languageLevel = PyiUtil.getOriginalLanguageLevel(myFlowOwner);
+    }
+    else {
+      languageLevel = null;
+    }
+    myFlowOwner.acceptChildren(new PyVersionAwareElementVisitor(languageLevel) {
      @Override
      public void visitPyTargetExpression(@NotNull PyTargetExpression node) {
        targetExpressions.add(node);
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/PyFileElementType.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/PyFileElementType.java
@@ -60,7 +60,7 @@ public class PyFileElementType extends IStubFileElementType<PyFileStub> {
  @Override
  public int getStubVersion() {
    // Don't forget to update versions of indexes that use the updated stub-based elements
-    return 89;
+    return 90;
  }

  @Nullable
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyClassImpl.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyClassImpl.java
@@ -642,8 +642,9 @@ public class PyClassImpl extends PyBaseElementImpl<PyClassStub> implements PyCla
    final Property decoratedProperty = processDecoratedProperties(filter);
    if (decoratedProperty != null) return decoratedProperty;

-    if (getStub() != null) {
-      return processStubProperties(filter);
+    PyClassStub stub = getStub();
+    if (stub != null) {
+      return processStubProperties(stub, filter);
    }
    else {
      // name = property(...) assignments from PSI
@@ -721,10 +722,12 @@ public class PyClassImpl extends PyBaseElementImpl<PyClassStub> implements PyCla
  }

  @Nullable
-  private Property processStubProperties(@Nullable Processor<? super Property> filter) {
-    final PyClassStub stub = getStub();
-    if (stub != null) {
-      for (StubElement<?> subStub : stub.getChildrenStubs()) {
+  private Property processStubProperties(@NotNull PyClassStub stub, @Nullable Processor<? super Property> filter) {
+    class StubPropertiesProcessor implements Processor<StubElement<?>> {
+      private @Nullable Property myResult;
+
+      @Override
+      public boolean process(StubElement<?> subStub) {
        if (subStub.getStubType() == PyElementTypes.TARGET_EXPRESSION) {
          final PyTargetExpressionStub targetStub = (PyTargetExpressionStub)subStub;
          final PropertyStubStorage prop = targetStub.getCustomStub(PropertyStubStorage.class);
@@ -735,13 +738,19 @@ public class PyClassImpl extends PyBaseElementImpl<PyClassStub> implements PyCla
            final String doc = prop.getDoc();
            if (getter != NONE || setter != NONE || deleter != NONE) {
              final PropertyImpl property = new PropertyImpl(targetStub.getName(), getter, setter, deleter, doc, targetStub.getPsi());
-              if (filter == null || filter.process(property)) return property;
+              if (filter == null || filter.process(property)) {
+                myResult = property;
+                return false;
+              }
            }
          }
        }
+        return true;
      }
    }
-    return null;
+    StubPropertiesProcessor processor = new StubPropertiesProcessor();
+    PyPsiUtils.processChildrenStubs(stub, PyiUtil.getOriginalLanguageLevel(this), processor);
+    return processor.myResult;
  }

  @Nullable
@@ -1014,28 +1023,39 @@ public class PyClassImpl extends PyBaseElementImpl<PyClassStub> implements PyCla

  @Override
  public List<PyTargetExpression> getClassAttributes() {
+    LanguageLevel languageLevel = PyiUtil.getOriginalLanguageLevel(this);
+    final ArrayList<PyTargetExpression> result = new ArrayList<>();
    PyClassStub stub = getStub();
    if (stub != null) {
-      final PyTargetExpression[] children = stub.getChildrenByType(PyElementTypes.TARGET_EXPRESSION, PyTargetExpression.EMPTY_ARRAY);
-      return Arrays.asList(children);
+      PyPsiUtils.processChildrenStubs(stub, languageLevel, child -> {
+        if (child.getStubType() == PyElementTypes.TARGET_EXPRESSION) {
+          result.add((PyTargetExpression)child.getPsi());
+        }
+        return true;
+      });
    }
-    List<PyTargetExpression> result = new ArrayList<>();
-    for (PsiElement psiElement : getStatementList().getChildren()) {
-      if (psiElement instanceof PyAssignmentStatement assignmentStatement) {
-        final PyExpression[] targets = assignmentStatement.getTargets();
-        for (PyExpression target : targets) {
-          if (target instanceof PyTargetExpression) {
-            result.add((PyTargetExpression)target);
+    else {
+      getStatementList().acceptChildren(new PyVersionAwareTopLevelElementVisitor(languageLevel) {
+        @Override
+        protected void checkAddElement(PsiElement psiElement) {
+          if (psiElement instanceof PyAssignmentStatement assignmentStatement) {
+            final PyExpression[] targets = assignmentStatement.getTargets();
+            for (PyExpression target : targets) {
+              if (target instanceof PyTargetExpression) {
+                result.add((PyTargetExpression)target);
+              }
+            }
+          }
+          else if (psiElement instanceof PyTypeDeclarationStatement) {
+            final PyExpression target = ((PyTypeDeclarationStatement)psiElement).getTarget();
+            if (target instanceof PyTargetExpression) {
+              result.add((PyTargetExpression)target);
+            }
          }
        }
-      }
-      else if (psiElement instanceof PyTypeDeclarationStatement) {
-        final PyExpression target = ((PyTypeDeclarationStatement)psiElement).getTarget();
-        if (target instanceof PyTargetExpression) {
-          result.add((PyTargetExpression)target);
-        }
-      }
+      });
    }
+    result.trimToSize();
    return result;
  }

@@ -1233,11 +1253,9 @@ public class PyClassImpl extends PyBaseElementImpl<PyClassStub> implements PyCla
  public boolean processClassLevelDeclarations(@NotNull PsiScopeProcessor processor) {
    final PyClassStub stub = getStub();
    if (stub != null) {
-      final List<StubElement<?>> children = stub.getChildrenStubs();
-      for (StubElement<?> child : children) {
-        if (!processor.execute(child.getPsi(), ResolveState.initial())) {
-          return false;
-        }
+      LanguageLevel languageLevel = PyiUtil.getOriginalLanguageLevel(this);
+      if (!PyPsiUtils.processChildrenStubs(stub, languageLevel, child -> processor.execute(child.getPsi(), ResolveState.initial()))) {
+        return false;
      }
    }
    else {
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyConditionalStatementPartImpl.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyConditionalStatementPartImpl.java
@@ -2,14 +2,22 @@
 package com.jetbrains.python.psi.impl;

 import com.intellij.lang.ASTNode;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.stubs.IStubElementType;
+import com.intellij.psi.stubs.StubElement;
 import com.jetbrains.python.psi.PyConditionalStatementPart;
 import com.jetbrains.python.psi.PyElementVisitor;

-public abstract class PyConditionalStatementPartImpl extends PyElementImpl implements PyConditionalStatementPart {
+public abstract class PyConditionalStatementPartImpl<T extends StubElement<? extends PsiElement>> extends PyBaseElementImpl<T>
+  implements PyConditionalStatementPart {
  public PyConditionalStatementPartImpl(ASTNode astNode) {
    super(astNode);
  }

+  protected PyConditionalStatementPartImpl(T stub, IStubElementType<?, ?> elementType) {
+    super(stub, elementType);
+  }
+
  @Override
  protected void acceptPyVisitor(PyElementVisitor pyVisitor) {
    pyVisitor.visitPyConditionalStatementPart(this);
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyElsePartImpl.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyElsePartImpl.java
@@ -16,11 +16,16 @@
 package com.jetbrains.python.psi.impl;

 import com.intellij.lang.ASTNode;
+import com.jetbrains.python.PyStubElementTypes;
 import com.jetbrains.python.psi.PyElsePart;
+import com.jetbrains.python.psi.stubs.PyElsePartStub;

-public class PyElsePartImpl extends PyElementImpl implements PyElsePart {
-  
+public class PyElsePartImpl extends PyBaseElementImpl<PyElsePartStub> implements PyElsePart {
  public PyElsePartImpl(ASTNode astNode) {
    super(astNode);
  }
+
+  public PyElsePartImpl(PyElsePartStub stub) {
+    super(stub, PyStubElementTypes.ELSE_PART);
+  }
 }
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyFileImpl.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyFileImpl.java
@@ -70,7 +70,8 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression {
      myModificationStamp = modificationStamp;

      final StubElement<?> stub = getStub();
-      processDeclarations(PyPsiUtils.collectAllStubChildren(PyFileImpl.this, stub), element -> {
+      LanguageLevel languageLevel = PyiUtil.getOriginalLanguageLevel(PyFileImpl.this);
+      processDeclarations(PyFileImpl.this, stub, languageLevel, element -> {
        if (element instanceof PsiNamedElement namedElement &&
            !(element instanceof PyKeywordArgument) &&
            !(stub == null && element.getParent() instanceof PyImportElement)) {
@@ -101,13 +102,16 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression {
      Collections.reverse(myImportedNameDefiners);
    }

-    private static boolean processDeclarations(@NotNull List<PsiElement> elements, @NotNull Processor<? super PsiElement> processor) {
-      for (PsiElement child : elements) {
+    private static boolean processDeclarations(@NotNull PsiElement element,
+                                               @Nullable StubElement<?> stub,
+                                               @NotNull LanguageLevel languageLevel,
+                                               @NotNull Processor<? super PsiElement> processor) {
+      for (PsiElement child : collectAllChildren(element, stub, languageLevel)) {
        if (!processor.process(child)) {
          return false;
        }
        if (child instanceof PyExceptPart part) {
-          if (!processDeclarations(PyPsiUtils.collectAllStubChildren(part, part.getStub()), processor)) {
+          if (!processDeclarations(part, part.getStub(), languageLevel, processor)) {
            return false;
          }
        }
@@ -228,7 +232,7 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression {

  @Override
  public @NotNull List<PyTypeAliasStatement> getTypeAliasStatements() {
-    return PyPsiUtils.collectStubChildren(this, getGreenStub(), PyTypeAliasStatement.class);
+    return collectChildren(PyTypeAliasStatement.class);
  }

  @Override
@@ -370,18 +374,18 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression {
  @Override
  @NotNull
  public List<PyClass> getTopLevelClasses() {
-    return PyPsiUtils.collectStubChildren(this, getGreenStub(), PyClass.class);
+    return collectChildren(PyClass.class);
  }

  @NotNull
  @Override
  public List<PyFunction> getTopLevelFunctions() {
-    return PyPsiUtils.collectStubChildren(this, getGreenStub(), PyFunction.class);
+    return collectChildren(PyFunction.class);
  }

  @Override
  public List<PyTargetExpression> getTopLevelAttributes() {
-    return PyPsiUtils.collectStubChildren(this, getGreenStub(), PyTargetExpression.class);
+    return collectChildren(PyTargetExpression.class);
  }

  @Override
@@ -490,7 +494,7 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression {
  @NotNull
  public List<PyImportElement> getImportTargets() {
    final List<PyImportElement> ret = new ArrayList<>();
-    final List<PyImportStatement> imports = PyPsiUtils.collectStubChildren(this, getGreenStub(), PyImportStatement.class);
+    final List<PyImportStatement> imports = collectChildren(PyImportStatement.class);
    for (PyImportStatement one : imports) {
      ContainerUtil.addAll(ret, one.getImportElements());
    }
@@ -500,7 +504,7 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression {
  @Override
  @NotNull
  public List<PyFromImportStatement> getFromImports() {
-    return PyPsiUtils.collectStubChildren(this, getGreenStub(), PyFromImportStatement.class);
+    return collectChildren(PyFromImportStatement.class);
  }

  @Nullable
@@ -845,4 +849,68 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression {
      }
    };
  }
+
+  @NotNull
+  private <T extends PyElement> List<T> collectChildren(Class<T> type) {
+    return collectChildren(this, getGreenStub(), PyiUtil.getOriginalLanguageLevel(this), type);
+  }
+
+  @NotNull
+  private static <T extends PyElement> List<T> collectChildren(@NotNull PsiElement element,
+                                                               @Nullable StubElement<?> stub,
+                                                               @NotNull LanguageLevel languageLevel,
+                                                               @NotNull Class<T> type) {
+    final List<T> result = new ArrayList<>();
+    if (stub != null) {
+      PyPsiUtils.processChildrenStubs(stub, languageLevel, child -> {
+        PsiElement childPsi = child.getPsi();
+        if (type.isInstance(childPsi)) {
+          result.add(type.cast(childPsi));
+        }
+        return true;
+      });
+    }
+    else {
+      element.acceptChildren(new PyVersionAwareTopLevelElementVisitor(languageLevel) {
+        @Override
+        protected void checkAddElement(PsiElement node) {
+          if (type.isInstance(node)) {
+            result.add(type.cast(node));
+          }
+        }
+
+        @Override
+        public void visitPyStatement(@NotNull PyStatement node) {
+          if (PyStatement.class.isAssignableFrom(type) && !(node instanceof PyCompoundStatement)) {
+            checkAddElement(node);
+            return;
+          }
+          super.visitPyStatement(node);
+        }
+      });
+    }
+    return result;
+  }
+
+  @NotNull
+  private static List<PsiElement> collectAllChildren(@NotNull PsiElement element,
+                                                     @Nullable StubElement<?> stub,
+                                                     @NotNull LanguageLevel languageLevel) {
+    List<PsiElement> result = new ArrayList<>();
+    if (stub != null) {
+      PyPsiUtils.processChildrenStubs(stub, languageLevel, child -> {
+        result.add(child.getPsi());
+        return true;
+      });
+    }
+    else {
+      element.acceptChildren(new PyVersionAwareTopLevelElementVisitor(languageLevel) {
+        @Override
+        protected void checkAddElement(PsiElement node) {
+          result.add(node);
+        }
+      });
+    }
+    return result;
+  }
 }
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyFunctionImpl.java
@@ -29,10 +29,7 @@ import com.jetbrains.python.documentation.docstrings.DocStringUtil;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.icons.PythonPsiApiIcons;
 import com.jetbrains.python.psi.resolve.QualifiedNameFinder;
-import com.jetbrains.python.psi.stubs.PyAnnotationOwnerStub;
-import com.jetbrains.python.psi.stubs.PyClassStub;
-import com.jetbrains.python.psi.stubs.PyFunctionStub;
-import com.jetbrains.python.psi.stubs.PyTargetExpressionStub;
+import com.jetbrains.python.psi.stubs.*;
 import com.jetbrains.python.psi.types.*;
 import com.jetbrains.python.sdk.PythonSdkUtil;
 import org.jetbrains.annotations.NotNull;
@@ -133,9 +130,9 @@ public class PyFunctionImpl extends PyBaseElementImpl<PyFunctionStub> implements
  public @Nullable PyClass getContainingClass() {
    final PyFunctionStub stub = getStub();
    if (stub != null) {
-      final StubElement parentStub = stub.getParentStub();
-      if (parentStub instanceof PyClassStub) {
-        return ((PyClassStub)parentStub).getPsi();
+      final StubElement<?> parentStub = PyPsiUtils.getParentStubSkippingVersionChecks(stub);
+      if (parentStub instanceof PyClassStub classStub) {
+        return classStub.getPsi();
      }

      return null;
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyIfPartElifImpl.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyIfPartElifImpl.java
@@ -2,13 +2,19 @@
 package com.jetbrains.python.psi.impl;

 import com.intellij.lang.ASTNode;
+import com.jetbrains.python.PyStubElementTypes;
 import com.jetbrains.python.psi.PyIfPartElif;
+import com.jetbrains.python.psi.stubs.PyIfPartElifStub;

 /**
 * PyIfPart that represents an 'elif' part.
 */
-public class PyIfPartElifImpl extends PyConditionalStatementPartImpl implements PyIfPartElif {
+public class PyIfPartElifImpl extends PyConditionalStatementPartImpl<PyIfPartElifStub> implements PyIfPartElif {
  public PyIfPartElifImpl(ASTNode astNode) {
    super(astNode);
  }
+
+  public PyIfPartElifImpl(PyIfPartElifStub stub) {
+    super(stub, PyStubElementTypes.IF_PART_ELIF);
+  }
 }
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyIfPartIfImpl.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyIfPartIfImpl.java
@@ -2,15 +2,19 @@
 package com.jetbrains.python.psi.impl;

 import com.intellij.lang.ASTNode;
+import com.jetbrains.python.PyStubElementTypes;
 import com.jetbrains.python.psi.PyIfPartIf;
+import com.jetbrains.python.psi.stubs.PyIfPartIfStub;

 /**
 * PyIfPart that represents an 'if' part.
 */
-public class PyIfPartIfImpl extends PyConditionalStatementPartImpl implements PyIfPartIf {
-
+public class PyIfPartIfImpl extends PyConditionalStatementPartImpl<PyIfPartIfStub> implements PyIfPartIf {
  public PyIfPartIfImpl(ASTNode astNode) {
    super(astNode);
  }

+  public PyIfPartIfImpl(PyIfPartIfStub stub) {
+    super(stub, PyStubElementTypes.IF_PART_IF);
+  }
 }
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyTargetExpressionImpl.java
@@ -590,7 +590,7 @@ public class PyTargetExpressionImpl extends PyBaseElementImpl<PyTargetExpression
  public PyClass getContainingClass() {
    final PyTargetExpressionStub stub = getStub();
    if (stub != null) {
-      final StubElement parentStub = stub.getParentStub();
+      final StubElement parentStub = PyPsiUtils.getParentStubSkippingVersionChecks(stub);
      if (parentStub instanceof PyClassStub) {
        return ((PyClassStub)parentStub).getPsi();
      }
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyWhilePartImpl.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/PyWhilePartImpl.java
@@ -16,9 +16,10 @@
 package com.jetbrains.python.psi.impl;

 import com.intellij.lang.ASTNode;
+import com.intellij.psi.stubs.StubElement;
 import com.jetbrains.python.psi.PyWhilePart;

-public class PyWhilePartImpl extends PyConditionalStatementPartImpl implements PyWhilePart {
+public class PyWhilePartImpl extends PyConditionalStatementPartImpl<StubElement<PyWhilePart>> implements PyWhilePart {
  public PyWhilePartImpl(ASTNode astNode) {
    super(astNode);
  }
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyClassElementType.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyClassElementType.java
@@ -161,7 +161,7 @@ public class PyClassElementType extends PyStubElementType<PyClassStub, PyClass>
    final List<String> slots = PyFileElementType.readNullableList(dataStream);

    final String docStringInStub = dataStream.readUTFFast();
-    final String docString = docStringInStub.length() > 0 ? docStringInStub : null;
+    final String docString = !docStringInStub.isEmpty() ? docStringInStub : null;

    final String deprecationMessage = dataStream.readNameString();

@@ -177,7 +177,7 @@ public class PyClassElementType extends PyStubElementType<PyClassStub, PyClass>
    if (name != null) {
      sink.occurrence(PyClassNameIndex.KEY, name);
      sink.occurrence(PyClassNameIndexInsensitive.KEY, StringUtil.toLowerCase(name));
-      if (stub.getParentStub() instanceof PyFileStub && PyUtil.getInitialUnderscores(name) == 0) {
+      if (PyPsiUtils.getParentStubSkippingVersionChecks(stub) instanceof PyFileStub && PyUtil.getInitialUnderscores(name) == 0) {
        sink.occurrence(PyExportedModuleAttributeIndex.KEY, name);
      }
    }
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyElsePartElementType.kt
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyElsePartElementType.kt
@@ -0,0 +1,41 @@
+package com.jetbrains.python.psi.impl.stubs
+
+import com.intellij.lang.ASTNode
+import com.intellij.psi.PsiElement
+import com.intellij.psi.stubs.StubElement
+import com.intellij.psi.stubs.StubInputStream
+import com.intellij.psi.stubs.StubOutputStream
+import com.jetbrains.python.psi.impl.PyVersionCheck
+import com.jetbrains.python.psi.PyElsePart
+import com.jetbrains.python.psi.PyIfStatement
+import com.jetbrains.python.psi.PyStubElementType
+import com.jetbrains.python.psi.impl.PyElsePartImpl
+import com.jetbrains.python.psi.stubs.PyElsePartStub
+
+class PyElsePartElementType : PyStubElementType<PyElsePartStub, PyElsePart>("ELSE_PART") {
+  override fun createPsi(stub: PyElsePartStub): PyElsePart {
+    return PyElsePartImpl(stub)
+  }
+
+  override fun createStub(psi: PyElsePart, parentStub: StubElement<out PsiElement>): PyElsePartStub {
+    return PyElsePartStubImpl(parentStub)
+  }
+
+  override fun serialize(stub: PyElsePartStub, dataStream: StubOutputStream) {
+  }
+
+  override fun deserialize(dataStream: StubInputStream, parentStub: StubElement<*>): PyElsePartStub {
+    return PyElsePartStubImpl(parentStub)
+  }
+
+  override fun createElement(node: ASTNode): PsiElement {
+    return PyElsePartImpl(node)
+  }
+
+  override fun shouldCreateStub(node: ASTNode): Boolean {
+    val ifStatement = node.treeParent?.psi as? PyIfStatement ?: return false
+    if (!isFileOrClassTopLevel(ifStatement)) return false
+    val ifParts = sequenceOf(ifStatement.ifPart) + ifStatement.elifParts.asSequence()
+    return ifParts.all { PyVersionCheck.fromCondition(it) != null }
+  }
+}
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyElsePartStubImpl.kt
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyElsePartStubImpl.kt
@@ -0,0 +1,9 @@
+package com.jetbrains.python.psi.impl.stubs
+
+import com.intellij.psi.stubs.StubBase
+import com.intellij.psi.stubs.StubElement
+import com.jetbrains.python.PyStubElementTypes
+import com.jetbrains.python.psi.PyElsePart
+import com.jetbrains.python.psi.stubs.PyElsePartStub
+
+class PyElsePartStubImpl(parent: StubElement<*>) : StubBase<PyElsePart>(parent, PyStubElementTypes.ELSE_PART), PyElsePartStub
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyFunctionElementType.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyFunctionElementType.java
@@ -19,7 +19,6 @@ import com.intellij.lang.ASTNode;
 import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.stubs.*;
-import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.PyStubElementTypes;
 import com.jetbrains.python.psi.PyFunction;
 import com.jetbrains.python.psi.PyStringLiteralExpression;
@@ -116,7 +115,7 @@ public class PyFunctionElementType extends PyStubElementType<PyFunctionStub, PyF
    final String name = stub.getName();
    if (name != null) {
      sink.occurrence(PyFunctionNameIndex.KEY, name);
-      if (stub.getParentStub() instanceof PyFileStub && PyUtil.getInitialUnderscores(name) == 0) {
+      if (PyPsiUtils.getParentStubSkippingVersionChecks(stub) instanceof PyFileStub && PyUtil.getInitialUnderscores(name) == 0) {
        sink.occurrence(PyExportedModuleAttributeIndex.KEY, name);
      }
    }
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyIfPartElifElementType.kt
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyIfPartElifElementType.kt
@@ -0,0 +1,44 @@
+package com.jetbrains.python.psi.impl.stubs
+
+import com.intellij.lang.ASTNode
+import com.intellij.psi.PsiElement
+import com.intellij.psi.stubs.StubElement
+import com.intellij.psi.stubs.StubInputStream
+import com.intellij.psi.stubs.StubOutputStream
+import com.intellij.util.takeWhileInclusive
+import com.jetbrains.python.psi.impl.PyVersionCheck
+import com.jetbrains.python.psi.PyIfPartElif
+import com.jetbrains.python.psi.PyIfStatement
+import com.jetbrains.python.psi.PyStubElementType
+import com.jetbrains.python.psi.impl.PyIfPartElifImpl
+import com.jetbrains.python.psi.stubs.PyIfPartElifStub
+
+class PyIfPartElifElementType : PyStubElementType<PyIfPartElifStub, PyIfPartElif>("IF_PART_ELIF") {
+  override fun createPsi(stub: PyIfPartElifStub): PyIfPartElif {
+    return PyIfPartElifImpl(stub)
+  }
+
+  override fun createStub(psi: PyIfPartElif, parentStub: StubElement<out PsiElement>): PyIfPartElifStub {
+    return PyIfPartElifStubImpl(parentStub, requireNotNull(PyVersionCheck.fromCondition (psi)))
+  }
+
+  override fun serialize(stub: PyIfPartElifStub, dataStream: StubOutputStream) {
+    serializeVersionCheck(stub.versionCheck, dataStream)
+  }
+
+  override fun deserialize(dataStream: StubInputStream, parentStub: StubElement<*>): PyIfPartElifStub {
+    return PyIfPartElifStubImpl(parentStub, deserializeVersionCheck(dataStream))
+  }
+
+  override fun createElement(node: ASTNode): PsiElement {
+    return PyIfPartElifImpl(node)
+  }
+
+  override fun shouldCreateStub(node: ASTNode): Boolean {
+    val ifStatement = node.treeParent?.psi as? PyIfStatement ?: return false
+    if (!isFileOrClassTopLevel(ifStatement)) return false
+    return (sequenceOf(ifStatement.ifPart) + ifStatement.elifParts.asSequence())
+      .takeWhileInclusive { it !== node.psi }
+      .all { PyVersionCheck.fromCondition(it) != null }
+  }
+}
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyIfPartElifStubImpl.kt
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyIfPartElifStubImpl.kt
@@ -0,0 +1,13 @@
+package com.jetbrains.python.psi.impl.stubs
+
+import com.intellij.psi.stubs.StubBase
+import com.intellij.psi.stubs.StubElement
+import com.jetbrains.python.PyStubElementTypes
+import com.jetbrains.python.psi.impl.PyVersionCheck
+import com.jetbrains.python.psi.PyIfPartElif
+import com.jetbrains.python.psi.stubs.PyIfPartElifStub
+
+class PyIfPartElifStubImpl(parent: StubElement<*>, private val versionCheck: PyVersionCheck)
+  : StubBase<PyIfPartElif>(parent, PyStubElementTypes.IF_PART_ELIF), PyIfPartElifStub {
+  override fun getVersionCheck(): PyVersionCheck = versionCheck
+}
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyIfPartIfElementType.kt
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyIfPartIfElementType.kt
@@ -0,0 +1,38 @@
+package com.jetbrains.python.psi.impl.stubs
+
+import com.intellij.lang.ASTNode
+import com.intellij.psi.PsiElement
+import com.intellij.psi.stubs.StubElement
+import com.intellij.psi.stubs.StubInputStream
+import com.intellij.psi.stubs.StubOutputStream
+import com.jetbrains.python.psi.impl.PyVersionCheck
+import com.jetbrains.python.psi.*
+import com.jetbrains.python.psi.impl.PyIfPartIfImpl
+import com.jetbrains.python.psi.stubs.PyIfPartIfStub
+
+class PyIfPartIfElementType : PyStubElementType<PyIfPartIfStub, PyIfPartIf>("IF_PART_IF") {
+  override fun createPsi(stub: PyIfPartIfStub): PyIfPartIf {
+    return PyIfPartIfImpl(stub)
+  }
+
+  override fun createStub(psi: PyIfPartIf, parentStub: StubElement<out PsiElement>): PyIfPartIfStub {
+    return PyIfPartIfStubImpl(parentStub, requireNotNull(PyVersionCheck.fromCondition (psi)))
+  }
+
+  override fun serialize(stub: PyIfPartIfStub, dataStream: StubOutputStream) {
+    serializeVersionCheck(stub.versionCheck, dataStream)
+  }
+
+  override fun deserialize(dataStream: StubInputStream, parentStub: StubElement<*>): PyIfPartIfStub {
+    return PyIfPartIfStubImpl(parentStub, deserializeVersionCheck(dataStream))
+  }
+
+  override fun createElement(node: ASTNode): PsiElement {
+    return PyIfPartIfImpl(node)
+  }
+
+  override fun shouldCreateStub(node: ASTNode): Boolean {
+    val psi = node.psi as PyIfPartIf
+    return isFileOrClassTopLevel(psi) && PyVersionCheck.fromCondition(psi) != null
+  }
+}
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyIfPartIfStubImpl.kt
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyIfPartIfStubImpl.kt
@@ -0,0 +1,13 @@
+package com.jetbrains.python.psi.impl.stubs
+
+import com.intellij.psi.stubs.StubBase
+import com.intellij.psi.stubs.StubElement
+import com.jetbrains.python.PyStubElementTypes
+import com.jetbrains.python.psi.impl.PyVersionCheck
+import com.jetbrains.python.psi.PyIfPartIf
+import com.jetbrains.python.psi.stubs.PyIfPartIfStub
+
+class PyIfPartIfStubImpl(parent: StubElement<*>, private val versionCheck: PyVersionCheck)
+  : StubBase<PyIfPartIf>(parent, PyStubElementTypes.IF_PART_IF), PyIfPartIfStub {
+  override fun getVersionCheck(): PyVersionCheck = versionCheck
+}
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyIfPartStubUtil.kt
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyIfPartStubUtil.kt
@@ -0,0 +1,29 @@
+package com.jetbrains.python.psi.impl.stubs
+
+import com.intellij.openapi.util.Version
+import com.intellij.psi.PsiElement
+import com.intellij.psi.stubs.StubInputStream
+import com.intellij.psi.stubs.StubOutputStream
+import com.intellij.psi.util.PsiTreeUtil
+import com.jetbrains.python.psi.impl.PyVersionCheck
+import com.jetbrains.python.psi.PyClass
+import com.jetbrains.python.psi.PyFile
+import com.jetbrains.python.psi.PyFunction
+
+internal fun isFileOrClassTopLevel(element: PsiElement): Boolean {
+  val parent = PsiTreeUtil.getParentOfType(element, PyFile::class.java, PyClass::class.java, PyFunction::class.java)
+  return parent is PyFile || parent is PyClass
+}
+
+internal fun serializeVersionCheck(versionCheck: PyVersionCheck, dataStream: StubOutputStream) {
+  dataStream.writeBoolean(versionCheck.isLessThan)
+  dataStream.writeVarInt(versionCheck.version.major)
+  dataStream.writeVarInt(versionCheck.version.minor)
+}
+
+internal fun deserializeVersionCheck(dataStream: StubInputStream): PyVersionCheck {
+  val isLessThan = dataStream.readBoolean()
+  val major = dataStream.readVarInt()
+  val minor = dataStream.readVarInt()
+  return PyVersionCheck(Version(major, minor, 0), isLessThan)
+}
--- a/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyTargetExpressionElementType.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/psi/impl/stubs/PyTargetExpressionElementType.java
@@ -15,6 +15,7 @@ import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.PythonDialectsTokenSetProvider;
 import com.jetbrains.python.documentation.docstrings.DocStringUtil;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.impl.PyPsiUtils;
 import com.jetbrains.python.psi.impl.PyTargetExpressionImpl;
 import com.jetbrains.python.psi.stubs.PyExportedModuleAttributeIndex;
 import com.jetbrains.python.psi.stubs.PyFileStub;
@@ -149,7 +150,7 @@ public class PyTargetExpressionElementType extends PyStubElementType<PyTargetExp
  public void indexStub(@NotNull PyTargetExpressionStub stub, @NotNull IndexSink sink) {
    String name = stub.getName();
    if (name != null && PyUtil.getInitialUnderscores(name) == 0) {
-      if (stub.getParentStub() instanceof PyFileStub) {
+      if (PyPsiUtils.getParentStubSkippingVersionChecks(stub) instanceof PyFileStub) {
        sink.occurrence(PyVariableNameIndex.KEY, name);
        sink.occurrence(PyExportedModuleAttributeIndex.KEY, name);
      }
--- a/python/python-psi-impl/src/com/jetbrains/python/pyi/PyiUtil.java
+++ b/python/python-psi-impl/src/com/jetbrains/python/pyi/PyiUtil.java
@@ -15,6 +15,7 @@
 */
 package com.jetbrains.python.pyi;

+import com.intellij.openapi.vfs.VirtualFile;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiFile;
 import com.intellij.psi.util.PsiTreeUtil;
@@ -105,14 +106,20 @@ public final class PyiUtil {
  @NotNull
  public static LanguageLevel getOriginalLanguageLevel(@NotNull PyElement element) {
    PsiFile containingFile = element.getContainingFile();
-    if (containingFile instanceof PyiFile) {
+    if (containingFile instanceof PyiFile pyiFile) {
      PsiElement impl = getOriginalElement(element);
      if (impl != null) {
        return LanguageLevel.forElement(impl);
      }
      else {
        // XXX: Relying on the fact .pyi files still have the language level key set by the pusher
-        return PythonLanguageLevelPusher.getLanguageLevelForVirtualFile(element.getProject(), containingFile.getVirtualFile());
+        VirtualFile vFile = containingFile.getVirtualFile();
+        if (vFile != null) {
+          return PythonLanguageLevelPusher.getLanguageLevelForVirtualFile(element.getProject(), vFile);
+        }
+        else {
+          return pyiFile.getLanguageLevel();
+        }
      }
    }
    return LanguageLevel.forElement(element);