PY-79967 Support t-prefixed strings, Related PR: https://github.com/JetBrains/intellij-community/pull/302

* Add PYTHON314_PREFIXES to CompatibilityVisitor
* Patch typeshed, add t-string related stubs
* Infer `string.templatelib.Template` type for t-strings instead of plain `str`



(cherry picked from commit 0e913910ab9e0dca4052856b0585ce66265291c0)

GitOrigin-RevId: f03a68ae9aee458a74eccc19c38d8d3e4dc7491a

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

@@ -6,19 +6,26 @@ import com.intellij.lang.ASTNode;
 import com.intellij.navigation.ItemPresentation;
 import com.intellij.openapi.util.Pair;
 import com.intellij.openapi.util.TextRange;
+import com.intellij.openapi.util.text.StringUtil;
 import com.intellij.psi.ContributedReferenceHost;
+import com.intellij.psi.PsiElement;
 import com.intellij.psi.PsiLiteralValue;
 import com.intellij.psi.PsiReference;
 import com.intellij.psi.PsiReferenceService.Hints;
 import com.intellij.psi.impl.source.resolve.FileContextUtil;
 import com.intellij.psi.impl.source.resolve.reference.ReferenceProvidersRegistry;
 import com.intellij.psi.tree.IElementType;
+import com.intellij.psi.util.QualifiedName;
 import com.intellij.util.ArrayUtil;
 import com.intellij.util.containers.ContainerUtil;
 import com.jetbrains.python.PyElementTypes;
 import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.lexer.PythonHighlightingLexer;
 import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.resolve.PyResolveContext;
+import com.jetbrains.python.psi.resolve.PyResolveImportUtil;
+import com.jetbrains.python.psi.types.PyClassLikeType;
+import com.jetbrains.python.psi.types.PyClassType;
 import com.jetbrains.python.psi.types.PyType;
 import com.jetbrains.python.psi.types.TypeEvalContext;
 import one.util.streamex.StreamEx;
@@ -27,6 +34,7 @@ import org.jetbrains.annotations.Nullable;
 
 import javax.swing.*;
 import java.util.List;
+import java.util.Objects;
 
 public class PyStringLiteralExpressionImpl extends PyElementImpl implements PyStringLiteralExpression, PsiLiteralValue, ContributedReferenceHost {
 
@@ -112,8 +120,34 @@ public class PyStringLiteralExpressionImpl extends PyElementImpl implements PySt
     final ASTNode firstNode = ContainerUtil.getFirstItem(getStringNodes());
     if (firstNode != null) {
       if (firstNode.getElementType() == PyElementTypes.FSTRING_NODE) {
-        // f-strings can't have "b" prefix, so they are always unicode
-        return builtinCache.getUnicodeType(languageLevel);
+        String prefix = PyStringLiteralCoreUtil.getPrefix(firstNode.getText());
+        if (StringUtil.containsIgnoreCase(prefix, "t")) {
+          // For t-strings, return the string.templatelib.Template type
+          if (!context.maySwitchToAST(this)) {
+            return builtinCache.getStrType();
+          }
+          final List<PsiElement> templateModules = PyResolveImportUtil.resolveQualifiedName(
+            QualifiedName.fromDottedString("string.templatelib"), // find a stub file
+            PyResolveImportUtil.fromFoothold(this)
+          );
+          if (templateModules.isEmpty()) {
+            return builtinCache.getStrType();
+          }
+          final PsiElement templateModule = templateModules.get(0);
+          if (!(templateModule instanceof PyFile)) {
+            return builtinCache.getStrType();
+          }
+          final PyClass templateClassElement = PyUtil.as(((PyFile)templateModule).findExportedName("Template"), PyClass.class);
+
+          if (templateClassElement == null) {
+            return builtinCache.getStrType();
+          }
+          final PyClassType templateType = (PyClassType)templateClassElement.getType(context);
+
+          return templateType != null ? templateType.toInstance() : builtinCache.getStrType();
+        } else
+          // f-strings can't have "b" prefix, so they are always unicode
+          return builtinCache.getUnicodeType(languageLevel);
       }
       else if (firstNode.getElementType() == PyTokenTypes.DOCSTRING) {
         return builtinCache.getStrType();

python/python-psi-impl/src/com/jetbrains/python/validation/CompatibilityVisitor.java

@@ -49,6 +49,8 @@ public abstract class CompatibilityVisitor extends PyAnnotator {
 
   private static final @NotNull Set<String> PYTHON36_PREFIXES = Sets.newHashSet("R", "U", "B", "BR", "RB", "F", "FR", "RF");
 
+  private static final @NotNull Set<String> PYTHON314_PREFIXES = Sets.newHashSet("R", "U", "B", "BR", "RB", "F", "FR", "RF", "T", "TR", "RT");
+
   protected @NotNull List<LanguageLevel> myVersionsToProcess;
 
   public CompatibilityVisitor(@NotNull List<LanguageLevel> versionsToProcess) {
@@ -282,9 +284,12 @@ public abstract class CompatibilityVisitor extends PyAnnotator {
     else if (level.isOlderThan(LanguageLevel.PYTHON36)) {
       return PYTHON34_PREFIXES;
     }
-    else {
+    else if (level.isOlderThan(LanguageLevel.PYTHON314)) {
       return PYTHON36_PREFIXES;
     }
+    else {
+      return PYTHON314_PREFIXES;
+    }
   }
 
   @Override