From 754808412b068e7cca3c716727487bc6869c5bf5 Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveTH@foxmail.com>
Date: Fri, 9 Jun 2023 22:28:36 +0800
Subject: [PATCH] add unpack builder

---
 src/compiler.h    | 18 +++++++++++++-----
 src/expr.h        |  3 ++-
 tests/05_list.py  | 11 ++++++++++-
 tests/06_tuple.py | 11 ++++++++++-
 tests/07_dict.py  | 16 +++++++++++++++-
 tests/08_set.py   | 11 ++++++++++-
 6 files changed, 60 insertions(+), 10 deletions(-)

diff --git a/src/compiler.h b/src/compiler.h
index e2568c85..b55d6c9f 100644
--- a/src/compiler.h
+++ b/src/compiler.h
@@ -343,13 +343,21 @@ class Compiler {
             match_newlines_repl();
             if (curr().type == TK("}")) break;
             EXPR();
-            if(curr().type == TK(":")) parsing_dict = true;
+            int star_level = ctx()->s_expr.top()->star_level();
+            if(star_level==2 || curr().type == TK(":")){
+                parsing_dict = true;
+            }
             if(parsing_dict){
-                consume(TK(":"));
-                EXPR();
                 auto dict_item = make_expr<DictItemExpr>();
-                dict_item->key = ctx()->s_expr.popx();
-                dict_item->value = ctx()->s_expr.popx();
+                if(star_level == 2){
+                    dict_item->key = nullptr;
+                    dict_item->value = ctx()->s_expr.popx();
+                }else{
+                    consume(TK(":"));
+                    EXPR();
+                    dict_item->key = ctx()->s_expr.popx();
+                    dict_item->value = ctx()->s_expr.popx();
+                }
                 items.push_back(std::move(dict_item));
             }else{
                 items.push_back(ctx()->s_expr.popx());
diff --git a/src/expr.h b/src/expr.h
index f2f70669..713a5592 100644
--- a/src/expr.h
+++ b/src/expr.h
@@ -366,7 +366,7 @@ struct SliceExpr: Expr{
 };
 
 struct DictItemExpr: Expr{
-    Expr_ key;
+    Expr_ key;      // maybe nullptr if it is **kwargs
     Expr_ value;
     std::string str() const override { return "DictItem()"; }
 
@@ -374,6 +374,7 @@ struct DictItemExpr: Expr{
 
     void emit(CodeEmitContext* ctx) override {
         if(is_starred()){
+            PK_ASSERT(key == nullptr);
             value->emit(ctx);
         }else{
             value->emit(ctx);
diff --git a/tests/05_list.py b/tests/05_list.py
index 2e5a740e..b7f93e5c 100644
--- a/tests/05_list.py
+++ b/tests/05_list.py
@@ -74,4 +74,13 @@ assert a == [8, 2, 4, 2, 9]
 
 b = [(1, 2), (3, 3), (5, 1)]
 b.sort(key=lambda x:x[1])
-assert b == [(5, 1), (1, 2), (3,3)]
\ No newline at end of file
+assert b == [(5, 1), (1, 2), (3,3)]
+
+# unpack expression
+a = [1, 2, 3]
+b = [*a, 4, 5]
+assert b == [1, 2, 3, 4, 5]
+
+a = []
+b = [*a, 1, 2, 3, *a, *a]
+assert b == [1, 2, 3]
\ No newline at end of file
diff --git a/tests/06_tuple.py b/tests/06_tuple.py
index 7f88ddf1..7392ad32 100644
--- a/tests/06_tuple.py
+++ b/tests/06_tuple.py
@@ -5,4 +5,13 @@ assert b == 2
 a,b = b,a
 assert a == 2
 assert b == 1
-assert len(tup) == 6
\ No newline at end of file
+assert len(tup) == 6
+
+# unpack expression
+a = 1, 2, 3
+b = *a, 4, 5
+assert b == (1, 2, 3, 4, 5)
+
+a = tuple([])
+b = *a, 1, 2, 3, *a, *a
+assert b == (1, 2, 3)
\ No newline at end of file
diff --git a/tests/07_dict.py b/tests/07_dict.py
index def4359c..114e9b8e 100644
--- a/tests/07_dict.py
+++ b/tests/07_dict.py
@@ -52,4 +52,18 @@ assert a == {1: 2, 3: 4}
 assert a.pop(1) == 2
 assert a == {3: 4}
 assert a.pop(3) == 4
-assert a == {}
\ No newline at end of file
+assert a == {}
+
+# unpack expression
+a = {1:2, 3:4}
+b = {**a, 5:6, **a}
+assert b == {1: 2, 3: 4, 5: 6}
+
+a = {}
+b = {**a, 1:2, 3:4}
+assert b == {1: 2, 3: 4}
+
+a = {1:2, 3:4, 7:8}
+b = {**a, 1:5, 3:6}
+c = {**a, **b}
+assert c == {1: 5, 3: 6, 7: 8}
\ No newline at end of file
diff --git a/tests/08_set.py b/tests/08_set.py
index 1f69a1c3..ffa146b4 100644
--- a/tests/08_set.py
+++ b/tests/08_set.py
@@ -76,4 +76,13 @@ assert type({}) is dict
 assert {1,2}.issubset({1,2,3})
 assert {1,2,3}.issuperset({1,2})
 assert {1,2,3}.isdisjoint({4,5,6})
-assert not {1,2,3}.isdisjoint({2,3,4})
\ No newline at end of file
+assert not {1,2,3}.isdisjoint({2,3,4})
+
+# unpack expression
+a = {1, 2, 3}
+b = {*a, 4, 5, *a, *a}
+assert b == {1, 2, 3, 4, 5}
+
+a = set()
+b = {*a, 1, 2, 3, *a, *a}
+assert b == {1, 2, 3}
\ No newline at end of file