PY-31442 Revive the idea of PyRichStringNodes -- LeafPsiElements with helper methods for strings

It will allow to uniformly work with formatted and literal string nodes in future,
and makes API prettier overall, since we don't need to use StringNodeInfo and
methods of PyStringLiteralUtil to get all this information about strings in Python.

python/psi-api/src/com/jetbrains/python/psi/PyLiteralStringNode.java

@@ -0,0 +1,7 @@
+package com.jetbrains.python.psi;
+
+/**
+ * @author Mikhail Golubev
+ */
+public interface PyLiteralStringNode extends PyRichStringNode {
+}

python/psi-api/src/com/jetbrains/python/psi/PyRichStringNode.java

@@ -0,0 +1,88 @@
+package com.jetbrains.python.psi;
+
+import com.intellij.openapi.util.TextRange;
+import com.intellij.psi.PsiElement;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.Set;
+
+/**
+ * @author Mikhail Golubev
+ */
+public interface PyRichStringNode extends PsiElement {
+
+  enum Modifier {
+    UNICODE,
+    BYTES,
+    RAW,
+    FORMATTED
+  }
+
+  /**
+   * @return string prefix, e.g. "UR", "b" etc.
+   */
+  @NotNull
+  String getPrefix();
+
+  int getPrefixLength();
+
+  /**
+   * @return the same text as {@code getNode().getText().substring(getPrefixLength())}
+   */
+  @NotNull
+  String getTextWithoutPrefix();
+
+  /**
+   * @return <em>relative</em> range of the content (excluding prefix and quotes)
+   * @see #getAbsoluteContentRange()
+   */
+  @NotNull
+  TextRange getContentRange();
+
+  /**
+   * @return <em>absolute</em> content range that accounts offset of the {@link #getNode() node} in the document
+   */
+  @NotNull
+  TextRange getAbsoluteContentRange();
+
+  /**
+   * @return content of the string node between quotes
+   */
+  @NotNull
+  String getContent();
+
+  @NotNull
+  String getQuote();
+
+  /**
+   * @return the first character of {@link #getQuote()}
+   */
+  char getQuoteChar();
+
+  boolean isTripleQuoted();
+
+  /**
+   * @return true if string literal ends with starting quote
+   */
+  boolean isTerminated();
+
+  @NotNull
+  Set<Modifier> getModifiers();
+
+  /**
+   * @return true if given string node contains "u" or "U" prefix
+   */
+  boolean isUnicode();
+
+  /**
+   * @return true if given string node contains "r" or "R" prefix
+   */
+  boolean isRaw();
+
+  /**
+   * @return true if given string node contains "b" or "B" prefix
+   */
+  boolean isBytes();
+
+  boolean isFormatted();
+}

python/psi-api/src/com/jetbrains/python/psi/PyStringLiteralExpression.java

@@ -27,6 +27,9 @@ public interface PyStringLiteralExpression extends PyLiteralExpression, StringLi
   @NotNull
   List<ASTNode> getStringNodes();
 
+  @NotNull
+  List<PyRichStringNode> getGluedStringNodes();
+
   int valueOffsetToTextOffset(int valueOffset);
 
   /**

python/src/META-INF/python-core-common.xml

@@ -27,6 +27,7 @@
     <lang.foldingBuilder language="Python" implementationClass="com.jetbrains.python.PythonFoldingBuilder"/>
     <lang.findUsagesProvider language="Python" implementationClass="com.jetbrains.python.findUsages.PythonFindUsagesProvider"/>
     <lang.formatter language="Python" implementationClass="com.jetbrains.python.formatter.PythonFormattingModelBuilder"/>
+    <lang.ast.factory language="Python" implementationClass="com.jetbrains.python.psi.impl.PythonASTFactory"/>
     <preFormatProcessor implementation="com.jetbrains.python.formatter.PyPreFormatProcessor"/>
     <postFormatProcessor implementation="com.jetbrains.python.formatter.PyTrailingBlankLinesPostFormatProcessor"/>
     <postFormatProcessor implementation="com.jetbrains.python.formatter.PyFromImportPostFormatProcessor"/>

python/src/com/jetbrains/python/psi/impl/PyLiteralStringNodeImpl.java

@@ -0,0 +1,134 @@
+package com.jetbrains.python.psi.impl;
+
+import com.intellij.openapi.util.Pair;
+import com.intellij.openapi.util.TextRange;
+import com.intellij.openapi.util.text.StringUtil;
+import com.intellij.psi.impl.source.tree.LeafPsiElement;
+import com.intellij.psi.tree.IElementType;
+import com.jetbrains.python.psi.PyLiteralStringNode;
+import com.jetbrains.python.psi.PyStringLiteralUtil;
+import org.jetbrains.annotations.NotNull;
+
+import java.util.EnumSet;
+import java.util.Set;
+
+/**
+ * @author Mikhail Golubev
+ */
+public class PyLiteralStringNodeImpl extends LeafPsiElement implements PyLiteralStringNode {
+  public PyLiteralStringNodeImpl(@NotNull IElementType type, CharSequence text) {
+    super(type, text);
+  }
+
+  @NotNull
+  @Override
+  public String getPrefix() {
+    return PyStringLiteralUtil.getPrefix(getText());
+  }
+
+  @Override
+  public int getPrefixLength() {
+    return PyStringLiteralUtil.getPrefixEndOffset(getText(), 0);
+  }
+
+  @NotNull
+  @Override
+  public String getContent() {
+    return getContentRange().substring(getText());
+  }
+
+  @NotNull
+  @Override
+  public String getTextWithoutPrefix() {
+    return getText().substring(getPrefixLength());
+  }
+
+  @NotNull
+  @Override
+  public TextRange getContentRange() {
+    final Pair<String, String> quotes = PyStringLiteralUtil.getQuotes(getText());
+    assert quotes != null;
+    final String prefixWithOpeningQuote = quotes.getFirst();
+    final String closingQuote = quotes.getSecond();
+    return TextRange.create(prefixWithOpeningQuote.length(), getTextLength() - closingQuote.length());
+  }
+
+  @NotNull
+  @Override
+  public TextRange getAbsoluteContentRange() {
+    return getContentRange().shiftRight(getStartOffset());
+  }
+
+  @Override
+  public char getQuoteChar() {
+    return getQuote().charAt(0);
+  }
+
+  @NotNull
+  @Override
+  public String getQuote() {
+    final Pair<String, String> quotes = PyStringLiteralUtil.getQuotes(getText());
+    assert quotes != null;
+    return quotes.getFirst();
+  }
+
+  @Override
+  public boolean isTripleQuoted() {
+    return getQuote().length() == 3;
+  }
+
+  @Override
+  public boolean isTerminated() {
+    final String unprefixed = getTextWithoutPrefix();
+    final String quote = getQuote();
+    return unprefixed.length() >= quote.length() * 2 && unprefixed.endsWith(quote);
+  }
+
+  @NotNull
+  @Override
+  public Set<Modifier> getModifiers() {
+    final EnumSet<Modifier> result = EnumSet.noneOf(Modifier.class);
+    if (isUnicode()) {
+      result.add(Modifier.UNICODE);
+    }
+    if (isBytes()) {
+      result.add(Modifier.BYTES);
+    }
+    if (isRaw()) {
+      result.add(Modifier.RAW);
+    }
+    if (isFormatted()) {
+      result.add(Modifier.FORMATTED);
+    }
+    return result;
+  }
+
+  /**
+   * @return true if given string node contains "u" or "U" prefix
+   */
+  @Override
+  public boolean isUnicode() {
+    return StringUtil.containsIgnoreCase(getPrefix(), "u");
+  }
+
+  /**
+   * @return true if given string node contains "r" or "R" prefix
+   */
+  @Override
+  public boolean isRaw() {
+    return StringUtil.containsIgnoreCase(getPrefix(), "r");
+  }
+
+  /**
+   * @return true if given string node contains "b" or "B" prefix
+   */
+  @Override
+  public boolean isBytes() {
+    return StringUtil.containsIgnoreCase(getPrefix(), "b");
+  }
+
+  @Override
+  public boolean isFormatted() {
+    return StringUtil.containsIgnoreCase(getPrefix(), "f");
+  }
+}

python/src/com/jetbrains/python/psi/impl/PyStringLiteralExpressionImpl.java

@@ -37,6 +37,7 @@ import com.jetbrains.python.lexer.PythonHighlightingLexer;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
+import one.util.streamex.StreamEx;
 import org.intellij.lang.regexp.DefaultRegExpPropertiesProvider;
 import org.intellij.lang.regexp.RegExpLanguageHost;
 import org.intellij.lang.regexp.psi.*;
@@ -259,6 +260,14 @@ public class PyStringLiteralExpressionImpl extends PyElementImpl implements PySt
     return Arrays.asList(getNode().getChildren(PyTokenTypes.STRING_NODES));
   }
 
+  @NotNull
+  @Override
+  public List<PyRichStringNode> getGluedStringNodes() {
+    return StreamEx.of(getStringNodes())
+      .select(PyRichStringNode.class)
+      .toList();
+  }
+
   @NotNull
   @Override
   public String getStringValue() {

python/src/com/jetbrains/python/psi/impl/PythonASTFactory.java

@@ -0,0 +1,23 @@
+package com.jetbrains.python.psi.impl;
+
+import com.intellij.lang.ASTFactory;
+import com.intellij.psi.impl.source.tree.LeafElement;
+import com.intellij.psi.tree.IElementType;
+import com.jetbrains.python.PyTokenTypes;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+/**
+ * @author Mikhail Golubev
+ */
+public class PythonASTFactory extends ASTFactory {
+
+  @Nullable
+  @Override
+  public LeafElement createLeaf(@NotNull IElementType type, CharSequence text) {
+    if (PyTokenTypes.STRING_NODES.contains(type)) {
+      return new PyLiteralStringNodeImpl(type, text);
+    }
+    return super.createLeaf(type, text);
+  }
+}

python/testSrc/com/jetbrains/python/PyStringLiteralTest.java

@@ -154,8 +154,8 @@ public class PyStringLiteralTest extends PyTestCase {
   }
 
   private PyStringLiteralExpression createLiteralFromText(final String text) {
-    final PsiFile file = PsiFileFactory.getInstance(myFixture.getProject()).createFileFromText("test.py", "a = " + text);
-    final PyStringLiteralExpression expr = PsiTreeUtil.getParentOfType(file.findElementAt(5), PyStringLiteralExpression.class);
+    final PsiFile file = PsiFileFactory.getInstance(myFixture.getProject()).createFileFromText("test.py", PythonFileType.INSTANCE, "a = (" + text + ")");
+    final PyStringLiteralExpression expr = PsiTreeUtil.getParentOfType(file.findElementAt(6), PyStringLiteralExpression.class);
     assert expr != null;
     return expr;
   }
@@ -169,4 +169,9 @@ public class PyStringLiteralTest extends PyTestCase {
     }
     return characters;
   }
+
+  public void testRichStringNodes() {
+    final PyStringLiteralExpression string = createLiteralFromText("'foo' 'bar' 'baz'");
+    assertSize(3, string.getGluedStringNodes());
+  }
 }