feat: support tuple[T, ...] syntax

mtshiba · Oct 5, 2024 · d3d3bdb · d3d3bdb
1 parent 02c217e
commit d3d3bdb
Show file tree

Hide file tree

Showing 3 changed files with 37 additions and 10 deletions.
diff --git a/crates/py2erg/convert.rs b/crates/py2erg/convert.rs
@@ -760,6 +760,21 @@ impl ASTConverter {
                             ConstExpr::App(ConstApp::new(obj, None, args))
                         }
                     }
+                    Some("GenericTuple") => {
+                        if args.pos_args.get(1).is_some_and(|arg| matches!(&arg.expr, ConstExpr::Lit(l) if l.is(TokenKind::EllipsisLit))) {
+                            let ty = args.pos_args.remove(0).expr;
+                            let obj = ConstExpr::Accessor(ConstAccessor::Local(
+                                Identifier::private("HomogenousTuple".into()),
+                            ));
+                            let args = ConstArgs::single(ty);
+                            ConstExpr::App(ConstApp::new(obj, None, args))
+                        } else {
+                            let obj = ConstExpr::Accessor(ConstAccessor::Local(
+                                Identifier::private("Tuple".into()),
+                            ));
+                            ConstExpr::App(ConstApp::new(obj, None, args))
+                        }
+                    }
                     Some("Optional") => {
                         let arg = args.pos_args.remove(0).expr;
                         let none = ConstExpr::Accessor(ConstAccessor::Local(Identifier::private(
@@ -975,10 +990,11 @@ impl ASTConverter {
                 ));
                 TypeSpec::poly(acc, ConstArgs::pos_only(vec![key_t, value_t], None))
             }
-            "list" => {
-                let len = ConstExpr::Accessor(ConstAccessor::Local(
-                    self.convert_ident("_".into(), args.location()),
-                ));
+            "List" | "list" => {
+                let len = ConstExpr::Accessor(ConstAccessor::Local(Identifier::private_with_loc(
+                    "_".into(),
+                    pyloc_to_ergloc(args.range()),
+                )));
                 let elem_t = match self.convert_expr_to_const(args) {
                     Some(elem_t) => elem_t,
                     None => {
@@ -998,7 +1014,7 @@ impl ASTConverter {
                     ConstArgs::new(vec![elem_t, len], None, vec![], None, None),
                 )
             }
-            "dict" => {
+            "Dict" | "dict" => {
                 let py_ast::Expr::Tuple(mut tuple) = args else {
                     return Self::gen_dummy_type_spec(args.location());
                 };
@@ -1028,16 +1044,23 @@ impl ASTConverter {
                 ));
                 TypeSpec::poly(acc, ConstArgs::new(vec![dict], None, vec![], None, None))
             }
-            "tuple" => {
-                let py_ast::Expr::Tuple(tuple) = args else {
+            "Tuple" | "tuple" => {
+                let py_ast::Expr::Tuple(mut tuple) = args else {
                     return Self::gen_dummy_type_spec(args.location());
                 };
-                let parens = Self::gen_enclosure_tokens(TokenKind::LParen, tuple.range);
+                // tuple[T, ...] == HomogenousTuple T
+                if tuple.elts.get(1).is_some_and(|ex| matches!(ex, py_ast::Expr::Constant(c) if matches!(c.value, py_ast::Constant::Ellipsis))) {
+                    let acc = ConstAccessor::local(Token::symbol("HomogenousTuple"));
+                    let ty = tuple.elts.remove(0);
+                    let args = ConstArgs::single(self.convert_expr_to_const(ty).unwrap());
+                    return TypeSpec::poly(acc, args);
+                }
                 let tys = tuple
                     .elts
                     .into_iter()
                     .map(|elem| self.convert_type_spec(elem))
                     .collect();
+                let parens = Self::gen_enclosure_tokens(TokenKind::LParen, tuple.range);
                 let tuple = TupleTypeSpec::new(Some(parens), tys);
                 TypeSpec::Tuple(tuple)
             }
@@ -1473,12 +1496,13 @@ impl ASTConverter {
                 Self::mutate_expr(dict)
             }
             py_ast::Expr::Tuple(tuple) => {
+                let (l, r) = Self::gen_enclosure_tokens(TokenKind::LParen, tuple.range);
                 let elements = tuple
                     .elts
                     .into_iter()
                     .map(|ex| PosArg::new(self.convert_expr(ex)))
                     .collect::<Vec<_>>();
-                let elems = Args::pos_only(elements, None);
+                let elems = Args::pos_only(elements, Some((l, r)));
                 Expr::Tuple(Tuple::Normal(NormalTuple::new(elems)))
             }
             py_ast::Expr::Subscript(subs) => {

diff --git a/tests/test.rs b/tests/test.rs
@@ -104,7 +104,7 @@ fn exec_warns() -> Result<(), String> {
 
 #[test]
 fn exec_typespec() -> Result<(), String> {
-    expect("tests/typespec.py", 0, 13)
+    expect("tests/typespec.py", 0, 14)
 }
 
 #[test]

diff --git a/tests/typespec.py b/tests/typespec.py
@@ -9,6 +9,8 @@
 p: Optional[int] = "a" # ERR
 weekdays: Literal[1, 2, 3, 4, 5, 6, 7] = 1 # OK
 weekdays: Literal[1, 2, 3, 4, 5, 6, 7] = 8 # ERR
+_: tuple[int, ...] = (1, 2, 3)
+_: list[tuple[int, ...]] = [(1, 2, 3)]
 _: dict[str, dict[str, Union[int, str]]] = {"a": {"b": 1}}
 _: dict[str, dict[str, list[int]]] = {"a": {"b": [1]}}
 _: dict[str, dict[str, dict[str, int]]] = {"a": {"b": {"c": 1}}}
@@ -22,6 +24,7 @@
 _: dict[str, dict[str, Callable[[int], int]]] = {"a": {"b": print}} # ERR
 _: dict[str, dict[str, Optional[int]]] = {"a": {"b": "c"}} # ERR
 _: dict[str, dict[str, Literal[1, 2]]] = {"a": {"b": 3}} # ERR
+_: list[tuple[int, ...]] = [(1, "a", 3)] # ERR
 
 def f(x: Union[int, str]) -> None:
     pass
-Original file line number
+Diff line change
@@ Expand Up / @@ -104,7 +104,7 @@ fn exec_warns() -> Result<(), String> { @@
     #[test]
     fn exec_typespec() -> Result<(), String> {
-        expect("tests/typespec.py", 0, 13)
+        expect("tests/typespec.py", 0, 14)
     }
     #[test]
@@ Expand Down @@