PY-52930 Support except star

(cherry picked from commit 4d03313cce07e4c3ff4748328cf2e04ba74b8d91)

IJ-MR-24242

GitOrigin-RevId: fa6c1b9077ef78b35c8f3453debcf06bcbc96364

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

@@ -30,4 +30,6 @@ public interface PyExceptPart extends PyElement, StubBasedPsiElement<PyExceptPar
 
   @Nullable
   PyExpression getTarget();
+
+  boolean isStar();
 }

python/python-psi-impl/resources/messages/PyPsiBundle.properties

@@ -156,6 +156,9 @@ ANN.patterns.pattern.does.not.bind.names=Pattern does not bind {0,choice,1#name|
 ANN.patterns.attribute.name.is.repeated=Attribute name ''{0}'' is repeated
 ANN.patterns.name.already.bound=Name ''{0}'' is already bound
 
+ANN.try.except.can.not.have.except.and.star.except=Try statement cannot contain both except and except*
+ANN.exception.group.in.star.except=ExceptionGroup cannot be used in except*
+ANN.continue.break.or.return.in.star.except='break', 'continue' and 'return' cannot appear in an except* block
 
 ### parsing
 PARSE.expected.expression=Expression expected
@@ -807,6 +810,7 @@ INSP.compatibility.basestring.type.not.available.in.py3=basestring type is not a
 INSP.compatibility.new.union.syntax.not.available.in.earlier.version=allow writing union types as X | Y
 INSP.compatibility.feature.support.match.statements=support match statements
 INSP.compatibility.feature.support.parenthesized.context.expressions=support parenthesized context expressions
+INSP.compatibility.feature.support.starred.except.part=support except* part
 
 # PyUnnecessaryBackslashInspection
 INSP.NAME.unnecessary.backslash=Unnecessary backslash

python/python-psi-impl/src/com/jetbrains/python/formatter/PythonFormattingModelBuilder.java

@@ -131,6 +131,7 @@ public class PythonFormattingModelBuilder implements FormattingModelBuilder, Cus
       .aroundInside(ADDITIVE_OPERATIONS, BINARY_EXPRESSION).spaceIf(commonSettings.SPACE_AROUND_ADDITIVE_OPERATORS)
       .aroundInside(STAR_OPERATORS, STAR_PARAMETERS).none()
       .aroundInside(STAR_OPERATORS, STAR_PATTERNS).none()
+      .between(EXCEPT_KEYWORD, MULT).none()
       .around(MULTIPLICATIVE_OPERATIONS).spaceIf(commonSettings.SPACE_AROUND_MULTIPLICATIVE_OPERATORS)
       .around(EXP).spaceIf(pySettings.SPACE_AROUND_POWER_OPERATOR)
       .around(SHIFT_OPERATIONS).spaceIf(commonSettings.SPACE_AROUND_SHIFT_OPERATORS)

python/python-psi-impl/src/com/jetbrains/python/parsing/StatementParsing.java

@@ -784,6 +784,8 @@ public class StatementParsing extends Parsing implements ITokenTypeRemapper {
       while (myBuilder.getTokenType() == PyTokenTypes.EXCEPT_KEYWORD) {
         final SyntaxTreeBuilder.Marker exceptBlock = myBuilder.mark();
         myBuilder.advanceLexer();
+
+        boolean star = matchToken(PyTokenTypes.MULT);
         if (myBuilder.getTokenType() != PyTokenTypes.COLON) {
           if (!getExpressionParser().parseSingleExpression(false)) {
             myBuilder.error(PyPsiBundle.message("PARSE.expected.expression"));
@@ -795,6 +797,9 @@ public class StatementParsing extends Parsing implements ITokenTypeRemapper {
             }
           }
         }
+        else if (star) {
+          myBuilder.error(PyPsiBundle.message("PARSE.expected.expression"));
+        }
         parseColonAndSuite();
         exceptBlock.done(PyElementTypes.EXCEPT_PART);
       }

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

@@ -4,6 +4,7 @@ package com.jetbrains.python.psi.impl;
 import com.intellij.lang.ASTNode;
 import com.intellij.psi.PsiNamedElement;
 import com.jetbrains.python.PyElementTypes;
+import com.jetbrains.python.PyTokenTypes;
 import com.jetbrains.python.PythonDialectsTokenSetProvider;
 import com.jetbrains.python.psi.*;
 import com.jetbrains.python.psi.stubs.PyExceptPartStub;
@@ -17,6 +18,7 @@ import java.util.List;
  * @author dcheryasov
  */
 public class PyExceptPartImpl extends PyBaseElementImpl<PyExceptPartStub> implements PyExceptPart {
+
   public PyExceptPartImpl(ASTNode astNode) {
     super(astNode);
   }
@@ -42,6 +44,11 @@ public class PyExceptPartImpl extends PyBaseElementImpl<PyExceptPartStub> implem
     return childToPsi(PythonDialectsTokenSetProvider.getInstance().getExpressionTokens(), 1);
   }
 
+  @Override
+  public boolean isStar() {
+    return getNode().findChildByType(PyTokenTypes.MULT) != null;
+  }
+
   @Override
   @NotNull
   public PyStatementList getStatementList() {

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

@@ -201,6 +201,10 @@ public class PyTargetExpressionImpl extends PyBaseElementImpl<PyTargetExpression
 
       return PyUnionType.union(collect);
     }
+    if (parent instanceof PyExceptPart && ((PyExceptPart)parent).isStar() &&
+        LanguageLevel.forElement(this).isAtLeast(LanguageLevel.PYTHON311)) {
+      return PyClassTypeImpl.createTypeByQName(this, "ExceptionGroup", false);
+    }
     PyType iterType = getTypeFromIteration(context);
     if (iterType != null) {
       return iterType;

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

@@ -91,6 +91,12 @@ public abstract class CompatibilityVisitor extends PyAnnotator {
                                        new ReplaceExceptPartQuickFix());
       }
     }
+    PsiElement star = PyPsiUtils.getFirstChildOfType(node, PyTokenTypes.MULT);
+    if (star != null) {
+      registerForAllMatchingVersions(level -> level.isOlderThan(LanguageLevel.PYTHON311),
+                                     PyPsiBundle.message("INSP.compatibility.feature.support.starred.except.part"),
+                                     star);
+    }
   }
 
   @Override

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

@@ -0,0 +1,97 @@
+package com.jetbrains.python.validation;
+
+import com.intellij.lang.annotation.HighlightSeverity;
+import com.intellij.openapi.util.TextRange;
+import com.intellij.psi.PsiElement;
+import com.intellij.psi.util.PsiTreeUtil;
+import com.jetbrains.python.PyPsiBundle;
+import com.jetbrains.python.PyTokenTypes;
+import com.jetbrains.python.psi.*;
+import com.jetbrains.python.psi.impl.PyPsiUtils;
+import org.jetbrains.annotations.NotNull;
+import org.jetbrains.annotations.Nullable;
+
+public class PyTryExceptAnnotator extends PyAnnotator {
+  @Override
+  public void visitPyTryExceptStatement(@NotNull PyTryExceptStatement node) {
+    boolean haveStar = false;
+    boolean haveNotStar = false;
+    for (PyExceptPart exceptPart: node.getExceptParts()) {
+      var star = PyPsiUtils.getFirstChildOfType(exceptPart, PyTokenTypes.MULT);
+      if (star != null) {
+        haveStar = true;
+      }
+      else {
+        haveNotStar = true;
+      }
+      if (haveNotStar && haveStar) {
+        var exceptKeyword = exceptPart.getFirstChild();
+        int startOffset = exceptKeyword.getTextOffset();
+        int endOffset = exceptKeyword.getTextRange().getEndOffset();
+        if (star != null) {
+          endOffset = star.getTextRange().getEndOffset();
+        }
+        var textRange = new TextRange(startOffset, endOffset);
+        getHolder().newAnnotation(HighlightSeverity.ERROR, PyPsiBundle.message("ANN.try.except.can.not.have.except.and.star.except"))
+            .range(textRange).create();
+        break;
+      }
+    }
+  }
+
+  @Nullable
+  private static PyReferenceExpression tryGetExceptionGroupInExpression(@Nullable PsiElement exceptExpression) {
+    if (exceptExpression instanceof PyReferenceExpression &&
+        "ExceptionGroup".equals(((PyReferenceExpression)exceptExpression).getName())) {
+      return (PyReferenceExpression)exceptExpression;
+    }
+    if (exceptExpression instanceof PyParenthesizedExpression) {
+      return tryGetExceptionGroupInExpression(PyPsiUtils.flattenParens((PyParenthesizedExpression)exceptExpression));
+    }
+    if (exceptExpression instanceof PyTupleExpression) {
+      for (PsiElement child: exceptExpression.getChildren()) {
+        var result = tryGetExceptionGroupInExpression(child);
+        if (result != null) {
+          return result;
+        }
+      }
+    }
+    return null;
+  }
+
+  @Override
+  public void visitPyExceptBlock(@NotNull PyExceptPart node) {
+    if (!node.isStar()) return;
+
+    var exceptClass = node.getExceptClass();
+    var exceptionGroup = tryGetExceptionGroupInExpression(exceptClass);
+    if (exceptionGroup != null) {
+      getHolder().newAnnotation(HighlightSeverity.ERROR, PyPsiBundle.message("ANN.exception.group.in.star.except")).range(exceptionGroup).create();
+    }
+  }
+
+  @Override
+  public void visitPyReturnStatement(@NotNull PyReturnStatement node) {
+    PyExceptPart exceptPart = PsiTreeUtil.getParentOfType(node, PyExceptPart.class, false, PyFunction.class);
+    if (exceptPart != null && exceptPart.isStar()) {
+      getHolder().newAnnotation(HighlightSeverity.ERROR, PyPsiBundle.message("ANN.continue.break.or.return.in.star.except")).create();
+    }
+  }
+
+  private void checkForContinueAndReturn(@NotNull PsiElement node) {
+    PyExceptPart exceptPart = PsiTreeUtil.getParentOfType(node, PyExceptPart.class, false, PyLoopStatement.class);
+    if (exceptPart != null && exceptPart.isStar()) {
+      getHolder().newAnnotation(HighlightSeverity.ERROR, PyPsiBundle.message("ANN.continue.break.or.return.in.star.except")).create();
+    }
+  }
+
+  @Override
+  public void visitPyContinueStatement(@NotNull PyContinueStatement node) {
+    checkForContinueAndReturn(node);
+  }
+
+  @Override
+  public void visitPyBreakStatement(@NotNull PyBreakStatement node) {
+    checkForContinueAndReturn(node);
+  }
+}

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

@@ -20,10 +20,7 @@ import com.intellij.psi.PsiElement;
 import com.intellij.psi.util.PsiTreeUtil;
 import com.jetbrains.python.PyPsiBundle;
 import com.jetbrains.python.codeInsight.functionTypeComments.psi.PyParameterTypeList;
-import com.jetbrains.python.psi.PyParenthesizedExpression;
-import com.jetbrains.python.psi.PyStarExpression;
-import com.jetbrains.python.psi.PyTupleExpression;
-import com.jetbrains.python.psi.PyYieldExpression;
+import com.jetbrains.python.psi.*;
 import org.jetbrains.annotations.NotNull;
 
 

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

@@ -522,6 +522,7 @@
     <pyAnnotator implementation="com.jetbrains.python.validation.FStringsAnnotator"/>
     <pyAnnotator implementation="com.jetbrains.python.validation.PyHighlightingAnnotator"/>
     <pyAnnotator implementation="com.jetbrains.python.validation.PyPatternAnnotator"/>
+    <pyAnnotator implementation="com.jetbrains.python.validation.PyTryExceptAnnotator"/>
 
     <!--stdlib-->
     <documentationLinkProvider implementation="com.jetbrains.python.codeInsight.stdlib.PyStdlibDocumentationLinkProvider"/>

python/testData/formatter/spaceAfterStarredExcept.py

@@ -0,0 +1,4 @@
+try:
+    raise ExceptionGroup("asdf", [Exception("fdsa")])
+except     *   Exception as ex:
+    pass

python/testData/formatter/spaceAfterStarredExcept_after.py

@@ -0,0 +1,4 @@
+try:
+    raise ExceptionGroup("asdf", [Exception("fdsa")])
+except* Exception as ex:
+    pass

python/testData/highlighting/continueBreakInsideLoopInExceptStarPart.py

@@ -0,0 +1,9 @@
+def f():
+    try:
+        pass
+    except* Exception:
+        while True:
+            break
+
+        for _ in range(42):
+            continue

python/testData/highlighting/continueBreakReturnInExceptStar.py

@@ -0,0 +1,12 @@
+def foo(x):
+    for i in range(0, x):
+        try:
+            ...
+        except* TypeError:
+            if x == 1:
+                <error descr="'break', 'continue' and 'return' cannot appear in an except* block">return 1</error>
+            if x == 2:
+                <error descr="'break', 'continue' and 'return' cannot appear in an except* block">continue</error>
+            if x == 3:
+                <error descr="'break', 'continue' and 'return' cannot appear in an except* block">break</error>
+

python/testData/highlighting/exceptStarAndExceptInTheSameTry.py

@@ -0,0 +1,6 @@
+try:
+    ...
+except ValueError:
+    pass
+<error descr="Try statement cannot contain both except and except*">except*</error> CancelledError:
+    pass

python/testData/highlighting/exceptionGroupInExceptOk.py

@@ -0,0 +1,4 @@
+try:
+    ...
+except ExceptionGroup:
+    pass

python/testData/highlighting/exceptionGroupInExceptStar.py

@@ -0,0 +1,4 @@
+try:
+    ...
+except* <error descr="ExceptionGroup cannot be used in except*">ExceptionGroup</error>:
+    pass

python/testData/highlighting/exceptionGroupInTupleInExceptStar.py

@@ -0,0 +1,4 @@
+try:
+    ...
+except* (TypeError, <error descr="ExceptionGroup cannot be used in except*">ExceptionGroup</error>):
+    pass

python/testData/highlighting/exceptionGroupsStarNoWarning.py

@@ -0,0 +1,6 @@
+try:
+    raise ExceptionGroup("asdf", [Exception("fdsa")])
+except* Exception as ex:
+    print(ex)
+    for e in ex.exceptions:
+        print(e)

python/testData/highlighting/exceptionGroupsStarOlderPythonWarning.py

@@ -0,0 +1,6 @@
+try:
+    raise ExceptionGroup("asdf", [Exception("fdsa")])
+except<error descr="Python version 3.10 does not support except* part">*</error> Exception as ex:
+    print(ex)
+    for e in ex.exceptions:
+        print(e)

python/testData/highlighting/returnInsideFunctionInExceptStarPart.py

@@ -0,0 +1,6 @@
+def f():
+    try:
+        pass
+    except* Exception:
+        def g():
+            return

python/testData/psi/TryExceptStarNoExpression.py

@@ -0,0 +1,4 @@
+try:
+    pass
+except*:
+    pass

python/testData/psi/TryExceptStarNoExpression.txt

@@ -0,0 +1,20 @@
+PyFile:TryExceptStarNoExpression.py
+  PyTryExceptStatement
+    PyTryPart
+      PsiElement(Py:TRY_KEYWORD)('try')
+      PsiElement(Py:COLON)(':')
+      PsiWhiteSpace('\n    ')
+      PyStatementList
+        PyPassStatement
+          PsiElement(Py:PASS_KEYWORD)('pass')
+    PsiWhiteSpace('\n')
+    PyExceptPart
+      PsiElement(Py:EXCEPT_KEYWORD)('except')
+      PsiElement(Py:MULT)('*')
+      PsiErrorElement:Expression expected
+        <empty list>
+      PsiElement(Py:COLON)(':')
+      PsiWhiteSpace('\n    ')
+      PyStatementList
+        PyPassStatement
+          PsiElement(Py:PASS_KEYWORD)('pass')

python/testSrc/com/jetbrains/python/Py3TypeTest.java

@@ -1363,6 +1363,15 @@ public class Py3TypeTest extends PyTestCase {
            "expr: Foo | None");
   }
 
+  // PY-52930
+  public void testExceptionGroupInExceptStar() {
+    doTest("ExceptionGroup",
+           "try:\n" +
+           "    raise ExceptionGroup(\"asdf\", [Exception(\"fdsa\")])\n" +
+           "except* Exception as expr:\n" +
+           "    pass\n");
+  }
+
   /**
    * @see #testRecursiveDictTopDown()
    * @see PyTypeCheckerInspectionTest#testRecursiveDictAttribute()

python/testSrc/com/jetbrains/python/PyFormatterTest.java

@@ -1218,6 +1218,11 @@ public class PyFormatterTest extends PyTestCase {
     doTest();
   }
 
+  // PY-52930
+  public void testSpaceAfterStarredExcept() {
+    doTest();
+  }
+
   // PY-42200
   public void testParenthesizedWithItems() {
     doTest();

python/testSrc/com/jetbrains/python/PythonHighlightingTest.java

@@ -287,6 +287,51 @@ public class PythonHighlightingTest extends PyTestCase {
     doTest(LanguageLevel.PYTHON35, false, false);
   }
 
+  // PY-52930
+  public void testExceptionGroupsStarNoWarning() {
+    doTest(LanguageLevel.getLatest(), false, false);
+  }
+
+  // PY-52930
+  public void testExceptionGroupsStarOlderPythonWarning() {
+    doTest(LanguageLevel.PYTHON310, false, false);
+  }
+
+  // PY-52930
+  public void testExceptionGroupInExceptOk() {
+    doTest(LanguageLevel.getLatest(), false, false);
+  }
+
+  // PY-52930
+  public void testExceptionGroupInExceptStar() {
+    doTest(LanguageLevel.getLatest(), false, false);
+  }
+
+  // PY-52930
+  public void testExceptionGroupInTupleInExceptStar() {
+    doTest(LanguageLevel.getLatest(), false, false);
+  }
+
+  // PY-52930
+  public void testExceptStarAndExceptInTheSameTry() {
+    doTest(LanguageLevel.getLatest(), false, false);
+  }
+
+  // PY-52930
+  public void testContinueBreakReturnInExceptStar() {
+    doTest(LanguageLevel.getLatest(), false, false);
+  }
+
+  // PY-52930
+  public void testContinueBreakInsideLoopInExceptStarPart() {
+    doTest(LanguageLevel.getLatest(), false, false);
+  }
+
+  // PY-52930
+  public void testReturnInsideFunctionInExceptStarPart() {
+    doTest(LanguageLevel.getLatest(), false, false);
+  }
+
   // PY-35961
   public void testUnpackingInNonParenthesizedTuplesInReturnAndYieldBefore38() {
     doTest(LanguageLevel.PYTHON35, false, false);

python/testSrc/com/jetbrains/python/parsing/PythonParsingTest.java

@@ -113,6 +113,11 @@ public class PythonParsingTest extends ParsingTestCase {
     doTest();
   }
 
+  // PY-52930
+  public void testTryExceptStarNoExpression() {
+    doTest();
+  }
+
   public void testPrintAsFunction26() {
     doTest(LanguageLevel.PYTHON26);
   }