PY-61639 Lift PyFile.getTopLevelAttributes(), PyFile.findTopLevelAttribute() to PyAstFile

GitOrigin-RevId: 373accde23bb2c55d1d23266da7874e34da96ea3
This commit is contained in:
Petr Golubev
2024-01-24 01:41:13 +01:00
committed by intellij-monorepo-bot
parent f818966298
commit c90db007c0
5 changed files with 92 additions and 54 deletions

View File

@@ -17,10 +17,13 @@ package com.jetbrains.python.ast;
import com.intellij.psi.PsiElement;
import com.intellij.psi.PsiFile;
import com.intellij.psi.PsiNamedElement;
import com.jetbrains.python.ast.controlFlow.AstScopeOwner;
import com.jetbrains.python.psi.FutureFeature;
import com.jetbrains.python.psi.LanguageLevel;
import org.jetbrains.annotations.ApiStatus;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import java.util.ArrayList;
import java.util.List;
@@ -37,10 +40,26 @@ public interface PyAstFile extends PyAstElement, PsiFile, PyAstDocStringOwner, A
return stmts;
}
List<? extends PyAstTargetExpression> getTopLevelAttributes();
@Nullable
default PyAstTargetExpression findTopLevelAttribute(@NotNull String name) {
return findByName(name, getTopLevelAttributes());
}
LanguageLevel getLanguageLevel();
/**
* Return true if the file contains a 'from __future__ import ...' statement with given feature.
*/
boolean hasImportFromFuture(FutureFeature feature);
private static <T extends PsiNamedElement> T findByName(@NotNull String name, @NotNull List<T> namedElements) {
for (T namedElement : namedElements) {
if (name.equals(namedElement.getName())) {
return namedElement;
}
}
return null;
}
}

View File

@@ -16,6 +16,7 @@ import org.jetbrains.annotations.ApiStatus;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;
@@ -64,6 +65,63 @@ public final class PyPsiUtilsCore {
return null;
}
@ApiStatus.Internal
@NotNull
public static <T extends PyAstElement> List<T> collectChildren(@NotNull PyAstFile pyFile, @NotNull Class<T> elementType) {
final List<T> result = new ArrayList<>();
pyFile.acceptChildren(new TopLevelVisitor() {
@Override
protected void checkAddElement(PsiElement node) {
if (elementType.isInstance(node)) {
result.add(elementType.cast(node));
}
}
@Override
public void visitPyStatement(@NotNull PyAstStatement node) {
if (PyAstStatement.class.isAssignableFrom(elementType) && !(node instanceof PyAstCompoundStatement)) {
checkAddElement(node);
return;
}
super.visitPyStatement(node);
}
});
return result;
}
@ApiStatus.Internal
@NotNull
public static List<PsiElement> collectAllChildren(PsiElement e) {
final List<PsiElement> result = new ArrayList<>();
e.acceptChildren(new TopLevelVisitor() {
@Override
protected void checkAddElement(PsiElement node) {
result.add(node);
}
});
return result;
}
private static abstract class TopLevelVisitor extends PyAstRecursiveElementVisitor {
@Override
public void visitPyElement(final @NotNull PyAstElement node) {
super.visitPyElement(node);
checkAddElement(node);
}
@Override
public void visitPyClass(final @NotNull PyAstClass node) {
checkAddElement(node); // do not recurse into functions
}
@Override
public void visitPyFunction(final @NotNull PyAstFunction node) {
checkAddElement(node); // do not recurse into classes
}
protected abstract void checkAddElement(PsiElement node);
}
@Nullable
public static String strValue(@Nullable PyAstExpression expression) {
return expression instanceof PyAstStringLiteralExpression ? ((PyAstStringLiteralExpression)expression).getStringValue() : null;

View File

@@ -39,6 +39,7 @@ public interface PyFile extends PyAstFile, PyElement, PsiFile, PyDocStringOwner,
@NotNull
List<PyFunction> getTopLevelFunctions();
@Override
List<PyTargetExpression> getTopLevelAttributes();
@Nullable
@@ -47,8 +48,11 @@ public interface PyFile extends PyAstFile, PyElement, PsiFile, PyDocStringOwner,
@Nullable
PyClass findTopLevelClass(@NonNls @NotNull String name);
@Override
@Nullable
PyTargetExpression findTopLevelAttribute(@NotNull String name);
default PyTargetExpression findTopLevelAttribute(@NotNull String name) {
return (PyTargetExpression)PyAstFile.super.findTopLevelAttribute(name);
}
@NotNull
List<PyTypeAliasStatement> getTypeAliasStatements();

View File

@@ -23,7 +23,10 @@ import com.jetbrains.python.psi.*;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import java.util.*;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
import java.util.stream.Stream;
public final class PyPsiUtils {
@@ -428,8 +431,8 @@ public final class PyPsiUtils {
static <T extends PyElement> List<T> collectStubChildren(@NotNull PyFile pyFile,
@Nullable StubElement<?> stub,
@NotNull Class<T> elementType) {
final List<T> result = new ArrayList<>();
if (stub != null) {
final List<T> result = new ArrayList<>();
@SuppressWarnings("rawtypes") final List<StubElement> children = stub.getChildrenStubs();
for (StubElement<?> child : children) {
PsiElement childPsi = child.getPsi();
@@ -437,46 +440,25 @@ public final class PyPsiUtils {
result.add(elementType.cast(childPsi));
}
}
return result;
}
else {
pyFile.acceptChildren(new TopLevelVisitor() {
@Override
protected void checkAddElement(PsiElement node) {
if (elementType.isInstance(node)) {
result.add(elementType.cast(node));
}
}
@Override
public void visitPyStatement(@NotNull PyStatement node) {
if (PyStatement.class.isAssignableFrom(elementType) && !(node instanceof PyCompoundStatement)) {
checkAddElement(node);
return;
}
super.visitPyStatement(node);
}
});
return PyPsiUtilsCore.collectChildren(pyFile, elementType);
}
return result;
}
static List<PsiElement> collectAllStubChildren(PsiElement e, StubElement stub) {
final List<PsiElement> result = new ArrayList<>();
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 {
e.acceptChildren(new TopLevelVisitor() {
@Override
protected void checkAddElement(PsiElement node) {
result.add(node);
}
});
return PyPsiUtilsCore.collectAllChildren(e);
}
return result;
}
public static int findArgumentIndex(PyCallExpression call, PsiElement argument) {
@@ -644,26 +626,6 @@ public final class PyPsiUtils {
.anyMatch(name -> name.matchesPrefix(sourceQName));
}
private static abstract class TopLevelVisitor extends PyRecursiveElementVisitor {
@Override
public void visitPyElement(final @NotNull PyElement node) {
super.visitPyElement(node);
checkAddElement(node);
}
@Override
public void visitPyClass(final @NotNull PyClass node) {
checkAddElement(node); // do not recurse into functions
}
@Override
public void visitPyFunction(final @NotNull PyFunction node) {
checkAddElement(node); // do not recurse into classes
}
protected abstract void checkAddElement(PsiElement node);
}
/**
* Returns text of the given PSI element. Unlike obvious {@link PsiElement#getText()} this method unescapes text of the element if latter
* belongs to injected code fragment using {@link InjectedLanguageManager#getUnescapedText(PsiElement)}.

View File

@@ -222,11 +222,6 @@ public class PyFileImpl extends PsiFileBase implements PyFile, PyExpression {
return findByName(name, getTopLevelClasses());
}
@Override
public PyTargetExpression findTopLevelAttribute(@NotNull String name) {
return findByName(name, getTopLevelAttributes());
}
@Override
public @NotNull List<PyTypeAliasStatement> getTypeAliasStatements() {
return PyPsiUtils.collectStubChildren(this, getGreenStub(), PyTypeAliasStatement.class);