Merge pull request #4531 from charliermarsh/charlie/exception-groups

Implement except* syntax

Author: youknowone

ast/Python.asdl

@@ -40,6 +40,7 @@ module Python
 
           | Raise(expr? exc, expr? cause)
           | Try(stmt* body, excepthandler* handlers, stmt* orelse, stmt* finalbody)
+          | TryStar(stmt* body, excepthandler* handlers, stmt* orelse, stmt* finalbody)
           | Assert(expr test, expr? msg)
 
           | Import(alias* names)

ast/src/ast_gen.rs

@@ -139,6 +139,12 @@ pub enum StmtKind<U = ()> {
         orelse: Vec<Stmt<U>>,
         finalbody: Vec<Stmt<U>>,
     },
+    TryStar {
+        body: Vec<Stmt<U>>,
+        handlers: Vec<Excepthandler<U>>,
+        orelse: Vec<Stmt<U>>,
+        finalbody: Vec<Stmt<U>>,
+    },
     Assert {
         test: Box<Expr<U>>,
         msg: Option<Box<Expr<U>>>,
@@ -726,6 +732,17 @@ pub mod fold {
                 orelse: Foldable::fold(orelse, folder)?,
                 finalbody: Foldable::fold(finalbody, folder)?,
             }),
+            StmtKind::TryStar {
+                body,
+                handlers,
+                orelse,
+                finalbody,
+            } => Ok(StmtKind::TryStar {
+                body: Foldable::fold(body, folder)?,
+                handlers: Foldable::fold(handlers, folder)?,
+                orelse: Foldable::fold(orelse, folder)?,
+                finalbody: Foldable::fold(finalbody, folder)?,
+            }),
             StmtKind::Assert { test, msg } => Ok(StmtKind::Assert {
                 test: Foldable::fold(test, folder)?,
                 msg: Foldable::fold(msg, folder)?,

parser/python.lalrpop

@@ -451,6 +451,27 @@ TryStatement: ast::Stmt = {
             },
         }
     },
+    <location:@L> "try" ":" <body:Suite> <handlers:ExceptStarClause+> <else_suite:("else" ":" Suite)?> <finally:("finally" ":" Suite)?> <end_location:@R> => {
+        let orelse = else_suite.map(|s| s.2).unwrap_or_default();
+        let finalbody = finally.map(|s| s.2).unwrap_or_default();
+        let end_location = finalbody
+            .last()
+            .or_else(|| orelse.last())
+            .map(|last| last.end_location)
+            .or_else(|| handlers.last().map(|last| last.end_location))
+            .unwrap();
+        ast::Stmt {
+            custom: (),
+            location,
+            end_location,
+            node: ast::StmtKind::TryStar {
+                body,
+                handlers,
+                orelse,
+                finalbody,
+            },
+        }
+    },
     <location:@L> "try" ":" <body:Suite> <finally:("finally" ":" Suite)> => {
         let handlers = vec![];
         let orelse = vec![];
@@ -470,6 +491,34 @@ TryStatement: ast::Stmt = {
     },
 };
 
+ExceptStarClause: ast::Excepthandler = {
+    <location:@L> "except" "*" <typ:Test<"all">> ":" <body:Suite> => {
+        let end_location = body.last().unwrap().end_location.unwrap();
+        ast::Excepthandler::new(
+            location,
+            end_location,
+            ast::ExcepthandlerKind::ExceptHandler {
+                type_: Some(Box::new(typ)),
+                name: None,
+                body,
+            },
+        )
+    },
+    <location:@L> "except" "*" <x:(Test<"all"> "as" Identifier)> ":" <body:Suite> => {
+        let end_location = body.last().unwrap().end_location.unwrap();
+        ast::Excepthandler::new(
+            location,
+            end_location,
+            ast::ExcepthandlerKind::ExceptHandler {
+                type_: Some(Box::new(x.0)),
+                name: Some(x.2),
+                body,
+            },
+        )
+    },
+};
+
+
 ExceptClause: ast::Excepthandler = {
     <location:@L> "except" <typ:Test<"all">?> ":" <body:Suite> => {
         let end_location = body.last().unwrap().end_location.unwrap();

parser/src/parser.rs

@@ -417,6 +417,37 @@ array[3:5, *idxs_to_select]
         insta::assert_debug_snapshot!(parse_ast);
     }
 
+    #[test]
+    fn test_try() {
+        let parse_ast = parse_program(
+            r#"try:
+    raise ValueError(1)
+except TypeError as e:
+    print(f'caught {type(e)}')
+except OSError as e:
+    print(f'caught {type(e)}')"#,
+            "<test>",
+        )
+        .unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
+    #[test]
+    fn test_try_star() {
+        let parse_ast = parse_program(
+            r#"try:
+    raise ExceptionGroup("eg",
+        [ValueError(1), TypeError(2), OSError(3), OSError(4)])
+except* TypeError as e:
+    print(f'caught {type(e)} with nested {e.exceptions}')
+except* OSError as e:
+    print(f'caught {type(e)} with nested {e.exceptions}')"#,
+            "<test>",
+        )
+        .unwrap();
+        insta::assert_debug_snapshot!(parse_ast);
+    }
+
     #[test]
     fn test_dict_unpacking() {
         let parse_ast = parse_expression(r#"{"a": "b", **c, "d": "e"}"#, "<test>").unwrap();