PY-61639 Move PythonFormattingModelBuilder to python.syntax.core

GitOrigin-RevId: bdcbd076cd3c1531204eff35515b44cd0edc9f68

python/python-ast/src/com/jetbrains/python/ast/PyAstFunction.java

@@ -2,10 +2,7 @@
 package com.jetbrains.python.ast;
 
 import com.intellij.lang.ASTNode;
-import com.intellij.psi.PsiElement;
-import com.intellij.psi.PsiNameIdentifierOwner;
-import com.intellij.psi.StubBasedPsiElement;
-import com.intellij.psi.TokenType;
+import com.intellij.psi.*;
 import com.intellij.psi.stubs.StubElement;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.ArrayFactory;
@@ -282,4 +279,21 @@ public interface PyAstFunction extends PsiNameIdentifierOwner, PyAstCompoundStat
   default void acceptPyVisitor(PyAstElementVisitor pyVisitor) {
     pyVisitor.visitPyFunction(this);
   }
+
+  @Override
+  default @Nullable PsiComment getTypeComment() {
+    final PsiComment inlineComment = PyUtilCore.getCommentOnHeaderLine(this);
+    if (inlineComment != null && PyUtilCore.getTypeCommentValue(inlineComment.getText()) != null) {
+      return inlineComment;
+    }
+
+    final PyAstStatementList statements = getStatementList();
+    if (statements.getStatements().length != 0) {
+      final PsiComment comment = ObjectUtils.tryCast(statements.getFirstChild(), PsiComment.class);
+      if (comment != null && PyUtilCore.getTypeCommentValue(comment.getText()) != null) {
+        return comment;
+      }
+    }
+    return null;
+  }
 }

python/python-ast/src/com/jetbrains/python/ast/PyAstNamedParameter.java

@@ -16,9 +16,7 @@
 package com.jetbrains.python.ast;
 
 import com.intellij.lang.ASTNode;
-import com.intellij.psi.PsiElement;
-import com.intellij.psi.PsiNameIdentifierOwner;
-import com.intellij.psi.PsiNamedElement;
+import com.intellij.psi.*;
 import com.intellij.psi.stubs.IStubElementType;
 import com.intellij.psi.stubs.StubElement;
 import com.jetbrains.python.PyElementTypes;
@@ -165,4 +163,20 @@ public interface PyAstNamedParameter extends PyAstParameter, PsiNamedElement, Ps
   default void acceptPyVisitor(PyAstElementVisitor pyVisitor) {
     pyVisitor.visitPyNamedParameter(this);
   }
+
+  @Nullable
+  @Override
+  default PsiComment getTypeComment() {
+    for (PsiElement next = getNextSibling(); next != null; next = next.getNextSibling()) {
+      if (next.textContains('\n')) break;
+      if (!(next instanceof PsiWhiteSpace)) {
+        if (",".equals(next.getText())) continue;
+        if (next instanceof PsiComment && PyUtilCore.getTypeCommentValue(next.getText()) != null) {
+          return (PsiComment)next;
+        }
+        break;
+      }
+    }
+    return null;
+  }
 }

python/python-ast/src/com/jetbrains/python/ast/PyAstTargetExpression.java

@@ -2,6 +2,7 @@
 package com.jetbrains.python.ast;
 
 import com.intellij.lang.ASTNode;
+import com.intellij.psi.PsiComment;
 import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiNameIdentifierOwner;
 import com.intellij.psi.PsiNamedElement;
@@ -12,6 +13,7 @@ import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.PythonDialectsTokenSetProvider;
 import com.jetbrains.python.ast.impl.PyPsiUtilsCore;
 import com.jetbrains.python.ast.docstring.DocStringUtilCore;
+import com.jetbrains.python.ast.impl.PyUtilCore;
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
@@ -205,4 +207,33 @@ public interface PyAstTargetExpression extends PyAstQualifiedExpression, PsiName
   default void acceptPyVisitor(PyAstElementVisitor pyVisitor) {
     pyVisitor.visitPyTargetExpression(this);
   }
+
+  @Nullable
+  @Override
+  default PsiComment getTypeComment() {
+    PsiComment comment = null;
+    final PyAstAssignmentStatement assignment = PsiTreeUtil.getParentOfType(this, PyAstAssignmentStatement.class);
+    if (assignment != null) {
+      final PyAstExpression assignedValue = assignment.getAssignedValue();
+      if (assignedValue != null && !PsiTreeUtil.isAncestor(assignedValue, this, false)) {
+        comment = ObjectUtils.tryCast(PyPsiUtilsCore.getNextNonWhitespaceSiblingOnSameLine(assignedValue), PsiComment.class);
+      }
+    }
+    else {
+      PyAstStatementListContainer forOrWith = null;
+      final PyAstForPart forPart = PsiTreeUtil.getParentOfType(this, PyAstForPart.class);
+      if (forPart != null && PsiTreeUtil.isAncestor(forPart.getTarget(), this, false)) {
+        forOrWith = forPart;
+      }
+      final PyAstWithItem withPart = PsiTreeUtil.getParentOfType(this, PyAstWithItem.class);
+      if (withPart != null && PsiTreeUtil.isAncestor(withPart.getTarget(), this, false)) {
+        forOrWith = ObjectUtils.tryCast(withPart.getParent(), PyAstWithStatement.class);
+      }
+
+      if (forOrWith != null) {
+        comment = PyUtilCore.getCommentOnHeaderLine(forOrWith);
+      }
+    }
+    return comment != null && PyUtilCore.getTypeCommentValue(comment.getText()) != null ? comment : null;
+  }
 }

python/python-ast/src/com/jetbrains/python/ast/PyAstTypeCommentOwner.java

@@ -15,9 +15,24 @@
  */
 package com.jetbrains.python.ast;
 
+import com.intellij.psi.PsiComment;
 import com.intellij.psi.PsiElement;
 import org.jetbrains.annotations.ApiStatus;
+import org.jetbrains.annotations.Nullable;
 
 @ApiStatus.Experimental
 public interface PyAstTypeCommentOwner extends PsiElement {
+  /**
+   * Returns a special comment that follows element definition and starts with conventional "type:" prefix.
+   * It is supposed to contain type annotation in PEP 484 compatible format. For further details see sections
+   * <a href="https://www.python.org/dev/peps/pep-0484/#type-comments">Type Comments</a> and
+   * <a href="https://www.python.org/dev/peps/pep-0484/#suggested-syntax-for-python-2-7-and-straddling-code">Suggested syntax for Python 2.7 and straddling code</a> and
+   * in PEP 484.
+   * <p/>
+   * Use {@link #getTypeCommentAnnotation()} to get its content with the prefix stripped accessing either stubs or AST.
+   *
+   * @see #getTypeCommentAnnotation()
+   */
+  @Nullable
+  PsiComment getTypeComment();
 }

python/python-ast/src/com/jetbrains/python/ast/impl/PyPsiUtilsCore.java

@@ -117,6 +117,24 @@ public final class PyPsiUtilsCore {
     return result;
   }
 
+  /**
+   * Returns the first non-whitespace sibling following the given element but within its line boundaries.
+   */
+  @Nullable
+  public static PsiElement getNextNonWhitespaceSiblingOnSameLine(@NotNull PsiElement element) {
+    PsiElement cur = element.getNextSibling();
+    while (cur != null) {
+      if (!(cur instanceof PsiWhiteSpace)) {
+        return cur;
+      }
+      else if (cur.textContains('\n')) {
+        break;
+      }
+      cur = cur.getNextSibling();
+    }
+    return null;
+  }
+
   private static abstract class TopLevelVisitor extends PyAstRecursiveElementVisitor {
     @Override
     public void visitPyElement(final @NotNull PyAstElement node) {

python/python-ast/src/com/jetbrains/python/ast/impl/PyUtilCore.java

@@ -1,8 +1,10 @@
 package com.jetbrains.python.ast.impl;
 
 import com.intellij.openapi.editor.Document;
-import com.intellij.psi.PsiDocumentManager;
-import com.intellij.psi.PsiElement;
+import com.intellij.openapi.util.TextRange;
+import com.intellij.openapi.util.text.StringUtil;
+import com.intellij.psi.*;
+import com.intellij.psi.stubs.StubElement;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.intellij.util.Consumer;
 import com.intellij.util.Function;
@@ -10,6 +12,7 @@ import com.intellij.util.ObjectUtils;
 import com.jetbrains.python.PyNames;
 import com.jetbrains.python.ast.*;
 import com.jetbrains.python.ast.controlFlow.AstScopeOwner;
+import com.jetbrains.python.codeInsight.dataflow.scope.ScopeUtilCore;
 import org.jetbrains.annotations.ApiStatus;
 import org.jetbrains.annotations.Contract;
 import org.jetbrains.annotations.NotNull;
@@ -17,6 +20,9 @@ import org.jetbrains.annotations.Nullable;
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Objects;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
 
 /**
  * Assorted utility methods for Python code insight.
@@ -28,6 +34,8 @@ import java.util.List;
 @ApiStatus.Experimental
 public final class PyUtilCore {
 
+  private static final Pattern TYPE_COMMENT_PATTERN = Pattern.compile("# *type: *([^#]+) *(#.*)?");
+
   private PyUtilCore() {
   }
 
@@ -80,6 +88,28 @@ public final class PyUtilCore {
     return unfoldParentheses(targets, new ArrayList<>(targets.length), false, true);
   }
 
+  public static boolean onSameLine(@NotNull PsiElement e1, @NotNull PsiElement e2) {
+    PsiFile firstFile = e1.getContainingFile();
+    PsiFile secondFile = e2.getContainingFile();
+    if (firstFile == null || secondFile == null) return false;
+    Document document = firstFile.getFileDocument();
+    if (document != secondFile.getFileDocument()) return false;
+    return document.getLineNumber(e1.getTextOffset()) == document.getLineNumber(e2.getTextOffset());
+  }
+
+  public static boolean isTopLevel(@NotNull PsiElement element) {
+    if (element instanceof StubBasedPsiElement) {
+      final StubElement stub = ((StubBasedPsiElement<?>)element).getStub();
+      if (stub != null) {
+        final StubElement parentStub = stub.getParentStub();
+        if (parentStub != null) {
+          return parentStub.getPsi() instanceof PsiFile;
+        }
+      }
+    }
+    return ScopeUtilCore.getScopeOwner(element) instanceof PsiFile;
+  }
+
   /**
    * Retrieves the document from {@link PsiDocumentManager} using the anchor PSI element and, if it's not null,
    * passes it to the consumer function.
@@ -115,6 +145,10 @@ public final class PyUtilCore {
     return null;
   }
 
+  public static boolean isSpecialName(@NotNull String name) {
+    return name.length() > 4 && name.startsWith("__") && name.endsWith("__");
+  }
+
   @Nullable
   public static PyAstLoopStatement getCorrespondingLoop(@NotNull PsiElement breakOrContinue) {
     return breakOrContinue instanceof PyAstContinueStatement || breakOrContinue instanceof PyAstBreakStatement
@@ -198,4 +232,67 @@ public final class PyUtilCore {
     }
     return null;
   }
+
+  public static boolean isAssignmentToModuleLevelDunderName(@Nullable PsiElement element) {
+    if (element instanceof PyAstAssignmentStatement statement && isTopLevel(element)) {
+      PyAstExpression[] targets = statement.getTargets();
+      if (targets.length == 1) {
+        String name = targets[0].getName();
+        return name != null && isSpecialName(name);
+      }
+    }
+    return false;
+  }
+
+  /**
+   * Returns the line comment that immediately precedes statement list of the given compound statement. Python parser ensures
+   * that it follows the statement header, i.e. it's directly after the colon, not on its own line.
+   */
+  @Nullable
+  public static PsiComment getCommentOnHeaderLine(@NotNull PyAstStatementListContainer container) {
+    return ObjectUtils.tryCast(getHeaderEndAnchor(container), PsiComment.class);
+  }
+
+  @NotNull
+  public static PsiElement getHeaderEndAnchor(@NotNull PyAstStatementListContainer container) {
+    final PyAstStatementList statementList = container.getStatementList();
+    return Objects.requireNonNull(PsiTreeUtil.skipWhitespacesBackward(statementList));
+  }
+
+  /**
+   * Checks that text of a comment starts with "# type:" prefix and returns trimmed type hint after it.
+   * The trailing part is supposed to contain type annotation in PEP 484 compatible format and an optional
+   * plain text comment separated from it with another "#".
+   * <p>
+   * For instance, for {@code # type: List[int]  # comment} it returns {@code List[int]}.
+   * <p>
+   * This method cannot return an empty string.
+   *
+   * @see #getTypeCommentValueRange(String)
+   */
+  @Nullable
+  public static String getTypeCommentValue(@NotNull String text) {
+    final Matcher m = TYPE_COMMENT_PATTERN.matcher(text);
+    if (m.matches()) {
+      return StringUtil.nullize(m.group(1).trim());
+    }
+    return null;
+  }
+
+  /**
+   * Returns the corresponding text range for a type hint as returned by {@link #getTypeCommentValue(String)}.
+   *
+   * @see #getTypeCommentValue(String)
+   */
+  @Nullable
+  public static TextRange getTypeCommentValueRange(@NotNull String text) {
+    final Matcher m = TYPE_COMMENT_PATTERN.matcher(text);
+    if (m.matches()) {
+      final String hint = getTypeCommentValue(text);
+      if (hint != null) {
+        return TextRange.from(m.start(1), hint.length());
+      }
+    }
+    return null;
+  }
 }

python/python-ast/src/com/jetbrains/python/codeInsight/dataflow/scope/ScopeUtilCore.java

@@ -0,0 +1,97 @@
+package com.jetbrains.python.codeInsight.dataflow.scope;
+
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.StubBasedPsiElement;
+import com.intellij.psi.stubs.StubElement;
+import com.intellij.psi.util.CachedValueProvider;
+import com.intellij.psi.util.CachedValuesManager;
+import com.intellij.psi.util.PsiModificationTracker;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.jetbrains.python.ast.*;
+import com.jetbrains.python.ast.controlFlow.AstScopeOwner;
+import org.jetbrains.annotations.ApiStatus;
+import org.jetbrains.annotations.Nullable;
+
+import static com.intellij.psi.util.PsiTreeUtil.getParentOfType;
+import static com.intellij.psi.util.PsiTreeUtil.isAncestor;
+
+@ApiStatus.Experimental
+public final class ScopeUtilCore {
+  /**
+   * Return the scope owner for the element.
+   *
+   * Scope owner is not always the first ScopeOwner parent of the element. Some elements are resolved in outer scopes.
+   *
+   * This method does not access AST if underlying PSI is stub based.
+   */
+  @Nullable
+  public static AstScopeOwner getScopeOwner(@Nullable final PsiElement element) {
+    if (element == null) {
+      return null;
+    }
+    if (element instanceof PyAstExpressionCodeFragment) {
+      final PsiElement context = element.getContext();
+      return context instanceof AstScopeOwner ? (AstScopeOwner)context : getScopeOwner(context);
+    }
+    if (element instanceof StubBasedPsiElement) {
+      final StubElement stub = ((StubBasedPsiElement<?>)element).getStub();
+      if (stub != null) {
+        StubElement parentStub = stub.getParentStub();
+        while (parentStub != null) {
+          final PsiElement parent = parentStub.getPsi();
+          if (parent instanceof AstScopeOwner) {
+            return (AstScopeOwner)parent;
+          }
+          parentStub = parentStub.getParentStub();
+        }
+        return null;
+      }
+    }
+    return CachedValuesManager.getCachedValue(element, () -> CachedValueProvider.Result
+      .create(calculateScopeOwnerByAST(element), PsiModificationTracker.MODIFICATION_COUNT));
+  }
+
+  @Nullable
+  private static AstScopeOwner calculateScopeOwnerByAST(@Nullable PsiElement element) {
+    final AstScopeOwner firstOwner = getParentOfType(element, AstScopeOwner.class);
+    if (firstOwner == null) {
+      return null;
+    }
+    final AstScopeOwner nextOwner = getParentOfType(firstOwner, AstScopeOwner.class);
+    // References in decorator expressions are resolved outside of the function (if the lambda is not inside the decorator)
+    final PyAstElement decoratorAncestor = getParentOfType(element, PyAstDecorator.class);
+    if (decoratorAncestor != null && !isAncestor(decoratorAncestor, firstOwner, true)) {
+      return nextOwner;
+    }
+    /*
+     * References in default values are resolved outside of the function (if the lambda is not inside the default value).
+     * Annotations of parameters are resolved outside of the function if the function doesn't have type parameters list
+     */
+    final PyAstNamedParameter parameterAncestor = getParentOfType(element, PyAstNamedParameter.class);
+    if (parameterAncestor != null && !isAncestor(parameterAncestor, firstOwner, true)) {
+      final PyAstExpression defaultValue = parameterAncestor.getDefaultValue();
+      final PyAstAnnotation annotation = parameterAncestor.getAnnotation();
+      if (firstOwner instanceof PyAstFunction function) {
+        PyAstTypeParameterList typeParameterList = function.getTypeParameterList();
+        if ((typeParameterList == null && isAncestor(annotation, element, false))
+            || (isAncestor(defaultValue, element, false))) {
+          return nextOwner;
+        }
+      }
+    }
+    // Superclasses are resolved outside of the class if the class doesn't have type parameters list
+    final PyAstClass containingClass = getParentOfType(element, PyAstClass.class);
+    if (containingClass != null && isAncestor(containingClass.getSuperClassExpressionList(), element, false) && containingClass.getTypeParameterList() == null) {
+      return nextOwner;
+    }
+    // Function return annotations and type comments are resolved outside of the function if the function doesn't have type parameters list
+    if (firstOwner instanceof PyAstFunction function) {
+      PyAstTypeParameterList typeParameterList = function.getTypeParameterList();
+      if ((typeParameterList == null && isAncestor(function.getAnnotation(), element, false)
+           || isAncestor(function.getTypeComment(), element, false))) {
+        return nextOwner;
+      }
+    }
+    return firstOwner;
+  }
+}