General support for Starred syntax in the parser + support for few simple cases in blue functions

spy/ast.py

@@ -616,6 +616,12 @@ class List(Expr):
     items: list[Expr]
 
 
+@astnode
+class Starred(Expr):
+    precedence = 17
+    value: Expr
+
+
 @astnode
 class Tuple(Expr):
     precedence = 17

spy/astcompile.py

@@ -736,6 +736,9 @@ def _format_spec_expr(self, item: ast.FormattedExpr) -> ast.Expr:
     def compile_expr_List(self, expr: ast.List) -> ast.Expr:
         return expr.replace(items=[self.compile_expr(item) for item in expr.items])
 
+    def compile_expr_Starred(self, expr: ast.Starred) -> ast.Expr:
+        return expr.replace(value=self.compile_expr(expr.value))
+
     def compile_expr_Tuple(self, expr: ast.Tuple) -> ast.Expr:
         return expr.replace(items=[self.compile_expr(item) for item in expr.items])
 

spy/backend/spy.py

@@ -740,6 +740,10 @@ def fmt_expr_List(self, node: ast.List) -> str:
         items = ", ".join(itemlist)
         return f"[{items}]"
 
+    def fmt_expr_Starred(self, node: ast.Starred) -> str:
+        v = self.fmt_expr(node.value)
+        return f"*{v}"
+
     def fmt_expr_Tuple(self, node: ast.Tuple) -> str:
         itemlist = [self.fmt_expr(it) for it in node.items]
         items = ", ".join(itemlist)

spy/doppler.py

@@ -842,6 +842,14 @@ def shift_expr_List(self, lst: ast.List, wam: W_MetaArg) -> ast.Expr:
             )
         return newlst
 
+    def shift_expr_Starred(self, op: ast.Starred, wam: W_MetaArg) -> ast.Expr:
+        raise SPyError.simple(
+            "W_WIP",
+            "splat expressions are supported only in blue code.",
+            "not supported here",
+            op.loc,
+        )
+
     def shift_expr_Tuple(self, tup: ast.Tuple, wam: W_MetaArg) -> ast.Expr:
         w_opimpl = self.opimpl[tup]
         v_T = make_const(self.vm, tup.loc, wam.w_static_T)
@@ -892,6 +900,7 @@ def shift_expr_Dict(self, dict: ast.Dict, wam: W_MetaArg) -> ast.Expr:
         return newdict
 
     def shift_expr_GetItem(self, op: ast.GetItem, wam: W_MetaArg) -> ast.Expr:
+        self._check_no_starred_args(op.args)
         w_opimpl = self.opimpl[op]
         v = self.shifted_expr[op.value]
         args = [self.shifted_expr[arg] for arg in op.args]
@@ -903,7 +912,21 @@ def shift_expr_GetAttr(self, op: ast.GetAttr, wam: W_MetaArg) -> ast.Expr:
         v_attr = self.shifted_expr[op.attr]
         return self.shift_opimpl(op, w_opimpl, [v, v_attr])
 
+    def _check_no_starred_args(self, args: list[ast.Expr]) -> None:
+        # Splats in call args are expanded eagerly by eval_args_with_starred.
+        # This point is reached only when the call is red.
+        for arg in args:
+            if isinstance(arg, ast.Starred):
+                raise SPyError.simple(
+                    "W_WIP",
+                    "splat arguments (`*expr`) are supported only in calls "
+                    "to blue functions and in blue subscripts",
+                    "not supported here",
+                    arg.loc,
+                )
+
     def shift_expr_Call(self, call: ast.Call, wam: W_MetaArg) -> ast.Expr:
+        self._check_no_starred_args(call.args)
         w_opimpl = self.opimpl[call]
         newfunc = self.shifted_expr[call.func]
         newargs = [self.shifted_expr[arg] for arg in call.args]
@@ -915,6 +938,7 @@ def shift_expr_Call(self, call: ast.Call, wam: W_MetaArg) -> ast.Expr:
             return self.shift_opimpl(call, w_opimpl, [newfunc] + newargs)
 
     def shift_expr_CallMethod(self, op: ast.CallMethod, wam: W_MetaArg) -> ast.Expr:
+        self._check_no_starred_args(op.args)
         w_opimpl = self.opimpl[op]
         v_obj = self.shifted_expr[op.target]
         v_meth = self.shifted_expr[op.method]

spy/parser.py

@@ -928,6 +928,11 @@ def from_py_expr_Attribute(self, py_node: py_ast.Attribute) -> spy.ast.GetAttr:
         attr = spy.ast.StrLiteral(py_node.loc, py_node.attr)
         return spy.ast.GetAttr(py_node.loc, value, attr)
 
+    def from_py_expr_Starred(self, py_node: py_ast.Starred) -> spy.ast.Starred:
+        # `*expr`, e.g. `[*args_m]`, `(*fields.keys(),)`, `f(*args)`.
+        value = self.from_py_expr(py_node.value)
+        return spy.ast.Starred(py_node.loc, value)
+
     def from_py_expr_List(self, py_node: py_ast.List) -> spy.ast.List:
         items = [self.from_py_expr(py_item) for py_item in py_node.elts]
         return spy.ast.List(py_node.loc, items)

spy/tests/compiler/test_blue_generic.py

@@ -188,3 +188,60 @@ def call_via_module_attr() -> i32:
             return mymod.foo[i32](2)
         """)
         assert mod.call_via_module_attr() == 4
+
+    def test_generic_starred_from_blue(self):
+        mod = self.compile("""
+        @blue.generic
+        def add(T):
+            def impl(x: T, y: T) -> T:
+                return x + y
+            return impl
+
+        @blue
+        def make():
+            types = (i32,)
+            return add[*types]
+
+        def foo() -> i32:
+            f = make()
+            return f(1, 2)
+        """)
+        assert mod.foo() == 3
+
+    def test_generic_starred_multiple_args(self):
+        mod = self.compile("""
+        @blue.generic
+        def pick(T, U):
+            def impl(x: T, y: U) -> U:
+                return y
+            return impl
+
+        @blue
+        def make():
+            more = (str,)
+            return pick[i32, *more]
+
+        def foo() -> str:
+            f = make()
+            return f(1, 'hello')
+        """)
+        assert mod.foo() == "hello"
+
+    def test_generic_starred_generic_class(self):
+        mod = self.compile("""
+        @struct
+        class Point[T]:
+            x: T
+            y: T
+
+        @blue
+        def make_point():
+            types = (i32,)
+            return Point[*types]
+
+        def foo() -> i32:
+            P = make_point()
+            p = P(2, 4)
+            return p.x + p.y
+        """)
+        assert mod.foo() == 6

spy/tests/compiler/test_list.py

@@ -1,5 +1,11 @@
 from spy.fqn import FQN
-from spy.tests.support import CompilerTest, expect_errors, no_C, only_interp
+from spy.tests.support import (
+    CompilerTest,
+    expect_errors,
+    no_C,
+    only_doppler,
+    only_interp,
+)
 from spy.vm.b import B
 from spy.vm.object import W_Type
 
@@ -179,3 +185,60 @@ def foo() -> i32:
         """
         mod = self.compile(src)
         assert mod.foo() == 2
+
+    @only_interp
+    def test_splat_red_value_interp(self):
+        src = """
+        def bar(xs: list[i32]) -> list[i32]:
+            return [*xs]
+
+        def foo() -> list[i32]:
+            return bar([1, 2, 3])
+        """
+        mod = self.compile(src)
+        assert mod.foo() == [1, 2, 3]
+
+    @only_doppler
+    def test_splat_red_value_doppler(self):
+        src = """
+        def bar(xs: list[i32]) -> list[i32]:
+            return [*xs]
+
+        def foo() -> list[i32]:
+            return bar([1, 2, 3])
+        """
+        errors = expect_errors(
+            "cannot splat an expression without known value",
+            ("this is not supported", "xs"),
+        )
+        self.compile_raises(src, "foo", errors)
+
+    @only_interp
+    def test_splat_no_fastiter(self):
+        src = """
+        def foo() -> None:
+            x = [*42]
+        """
+        errors = expect_errors(
+            "cannot unpack `i32`: it does not support iteration, so it "
+            "cannot be splatted with `*`",
+            ("this is not supported", "*42"),
+        )
+        self.compile_raises(src, "foo", errors)
+
+    @only_interp
+    def test_splat_in_call_args(self):
+        src = """
+        def bar(x: i32) -> i32:
+            return x
+
+        def foo() -> i32:
+            args = (1,)
+            return bar(*args)
+        """
+        errors = expect_errors(
+            "cannot unpack `tuple[i32]`: it does not support iteration, "
+            "so it cannot be splatted with `*`",
+            ("this is not supported", "*args"),
+        )
+        self.compile_raises(src, "foo", errors)

spy/tests/compiler/test_metafunc.py

@@ -1,6 +1,12 @@
 import pytest
 
-from spy.tests.support import CompilerTest, expect_errors, no_C
+from spy.errors import SPyError
+from spy.fqn import FQN
+from spy.tests.support import (
+    CompilerTest,
+    expect_errors,
+    no_C,
+)
 from spy.vm.b import B
 
 
@@ -302,3 +308,137 @@ def foo() -> tuple[str, str]:
         result = mod.foo()
         assert result[0] == "x + 2"
         assert result[1].endswith("test.spy")
+
+    def test_variadic_starred(self):
+        src = """
+        from operator import OpSpec
+
+        @blue.metafunc
+        def foo(*args_m):
+            m_x = args_m[0]
+            m_y = args_m[1]
+
+            T = m_x.static_type
+            assert m_y.static_type == T
+
+            def impl_str(x: T, y: T) -> T:
+                return x + y
+
+            return OpSpec(impl_str, [*args_m])
+
+        def test_i32(x: i32, y: i32) -> i32:
+            return foo(x, y)
+
+        def test_str(x: str, y: str) -> str:
+            return foo(x, y)
+        """
+
+        mod = self.compile(src)
+        assert mod.test_i32(1, 2) == 3
+        assert mod.test_str("S", "Py") == "SPy"
+
+    def test_splat_of_a_red_value_inside_a_blue_function(self):
+        src = """
+        @blue
+        def make():
+            # cannot yet use a tuple, missing __fastiter__
+            t = [1, 2, 3]
+            l = [*t]  # red
+            l.append(4)
+            t2 = (*l,)
+
+            def func(i: i32) -> i32:
+                return t2[0] + i
+            return func
+
+        def foo(i: i32) -> i32:
+            return make()(i)
+        """
+        mod = self.compile(src)
+        assert mod.foo(2) == 3
+
+    def test_splat_call_args_blue_func(self):
+        src = """
+        @blue
+        def bfunc(a, b, c):
+            return a + b + c
+
+        @blue
+        def make():
+            t = [2, 3]
+            return bfunc(*t, 4), bfunc(1, *t), bfunc(*[1, 1], 1)
+
+        def foo() -> tuple[i32, i32, i32]:
+            return make()
+        """
+        mod = self.compile(src)
+        assert mod.foo() == (9, 6, 3)
+
+    def test_splat_call_args_empty(self):
+        src = """
+        @blue
+        def bfunc(a):
+            return a
+
+        @blue
+        def make():
+            empty = [*[]]
+            return bfunc(1, *empty)
+
+        def foo() -> i32:
+            return make()
+        """
+        mod = self.compile(src)
+        assert mod.foo() == 1
+
+    def test_splat_call_args_dict_keys(self):
+        src = """
+        @blue
+        def bfunc(a, b):
+            return a + b
+
+        @blue
+        def make():
+            d = {1: "x", 2: "y"}
+            return bfunc(*d.keys())
+
+        def foo() -> i32:
+            return make()
+        """
+        mod = self.compile(src)
+        assert mod.foo() == 3
+
+    def test_splat_call_args_method(self):
+        src = """
+        @blue
+        def make():
+            l = [1]
+            more = [2, 3]
+            l.extend([*more])
+            return len(l)
+
+        def foo() -> i32:
+            return make()
+        """
+        mod = self.compile(src)
+        assert mod.foo() == 3
+
+    def test_splat_call_args_not_iterable(self):
+        src = """
+        @blue
+        def bfunc(a):
+            return a
+
+        @blue
+        def make():
+            return bfunc(*42)
+
+        def foo() -> i32:
+            return make()
+        """
+        errors = expect_errors(
+            "cannot unpack `i32`: it does not support iteration, "
+            "so it cannot be splatted with `*`",
+            ("this is not supported", "*42"),
+        )
+        self.compile_raises(src, "foo", errors)