impl `dictcomp` and `setcomp`

python/dict.py

@@ -97,4 +97,19 @@ class dict:
             if type(k) is not str:
                 raise TypeError('json keys must be strings, got ' + repr(k) )
             a.append(k.__json__()+': '+v.__json__())
-        return '{'+ ', '.join(a) + '}'
+        return '{'+ ', '.join(a) + '}'
+    
+    def __eq__(self, __o: object) -> bool:
+        if type(__o) is not dict:
+            return False
+        if len(self) != len(__o):
+            return False
+        for k in self.keys():
+            if k not in __o:
+                return False
+            if self[k] != __o[k]:
+                return False
+        return True
+
+    def __ne__(self, __o: object) -> bool:
+        return not self.__eq__(__o)

src/ceval.h

@@ -80,11 +80,6 @@ PyVar VM::run_frame(Frame* frame){
             frame->push(VAR(ss.str()));
         } continue;
         case OP_LOAD_EVAL_FN: frame->push(builtins->attr(m_eval)); continue;
-        case OP_LIST_APPEND: {
-            PyVar obj = frame->pop_value(this);
-            List& list = CAST(List&, frame->top_1());
-            list.push_back(std::move(obj));
-        } continue;
         case OP_BEGIN_CLASS: {
             auto& name = frame->co->names[byte.arg];
             PyVar clsBase = frame->pop_value(this);
@@ -206,6 +201,20 @@ PyVar VM::run_frame(Frame* frame){
             PyVar obj = call(builtins->attr("set"), one_arg(list));
             frame->push(obj);
         } continue;
+        case OP_LIST_APPEND: {
+            PyVar obj = frame->pop_value(this);
+            List& list = CAST(List&, frame->top_1());
+            list.push_back(std::move(obj));
+        } continue;
+        case OP_MAP_ADD: {
+            PyVar value = frame->pop_value(this);
+            PyVar key = frame->pop_value(this);
+            call(frame->top_1(), __setitem__, two_args(key, value));
+        } continue;
+        case OP_SET_ADD: {
+            PyVar obj = frame->pop_value(this);
+            call(frame->top_1(), "add", one_arg(obj));
+        } continue;
         case OP_DUP_TOP_VALUE: frame->push(frame->top_value(this)); continue;
         case OP_UNARY_STAR: {
             if(byte.arg > 0){   // rvalue

src/compiler.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include "codeobject.h"
 #include "parser.h"
 #include "error.h"
 #include "ceval.h"
@@ -528,28 +529,12 @@ private:
         consume(TK(")"));
     }
 
-    void exprList() {
-        int _patch = emit(OP_NO_OP);
-        int _body_start = co()->codes.size();
-        int ARGC = 0;
-        do {
-            match_newlines(mode()==REPL_MODE);
-            if (peek() == TK("]")) break;
-            EXPR(); ARGC++;
-            match_newlines(mode()==REPL_MODE);
-            if(ARGC == 1 && match(TK("for"))) goto __LISTCOMP;
-        } while (match(TK(",")));
-        match_newlines(mode()==REPL_MODE);
-        consume(TK("]"));
-        emit(OP_BUILD_LIST, ARGC);
-        return;
-
-__LISTCOMP:
+    void _consume_comp(Opcode op0, Opcode op1, int _patch, int _body_start){
         int _body_end_return = emit(OP_JUMP_ABSOLUTE, -1);
         int _body_end = co()->codes.size();
         co()->codes[_patch].op = OP_JUMP_ABSOLUTE;
         co()->codes[_patch].arg = _body_end;
-        emit(OP_BUILD_LIST, 0);
+        emit(op0, 0);
         EXPR_FOR_VARS();consume(TK("in"));EXPR_TUPLE();
         match_newlines(mode()==REPL_MODE);
         
@@ -572,23 +557,44 @@ __LISTCOMP:
             int ifpatch = emit(OP_POP_JUMP_IF_FALSE);
             emit(OP_JUMP_ABSOLUTE, _body_start);
             patch_jump(_body_end_return);
-            emit(OP_LIST_APPEND);
+            emit(op1);
             patch_jump(ifpatch);
         }else{
             emit(OP_JUMP_ABSOLUTE, _body_start);
             patch_jump(_body_end_return);
-            emit(OP_LIST_APPEND);
+            emit(op1);
         }
 
         emit(OP_LOOP_CONTINUE, -1, true);
         co()->_exit_block();
         match_newlines(mode()==REPL_MODE);
+    }
+
+    void exprList() {
+        int _patch = emit(OP_NO_OP);
+        int _body_start = co()->codes.size();
+        int ARGC = 0;
+        do {
+            match_newlines(mode()==REPL_MODE);
+            if (peek() == TK("]")) break;
+            EXPR(); ARGC++;
+            match_newlines(mode()==REPL_MODE);
+            if(ARGC == 1 && match(TK("for"))){
+                _consume_comp(OP_BUILD_LIST, OP_LIST_APPEND, _patch, _body_start);
+                consume(TK("]"));
+                return;
+            }
+        } while (match(TK(",")));
+        match_newlines(mode()==REPL_MODE);
         consume(TK("]"));
+        emit(OP_BUILD_LIST, ARGC);
     }
 
     void exprMap() {
+        int _patch = emit(OP_NO_OP);
+        int _body_start = co()->codes.size();
         bool parsing_dict = false;
-        int size = 0;
+        int ARGC = 0;
         do {
             match_newlines(mode()==REPL_MODE);
             if (peek() == TK("}")) break;
@@ -598,13 +604,19 @@ __LISTCOMP:
                 consume(TK(":"));
                 EXPR();
             }
-            size++;
+            ARGC++;
             match_newlines(mode()==REPL_MODE);
+            if(ARGC == 1 && match(TK("for"))){
+                if(parsing_dict) _consume_comp(OP_BUILD_MAP, OP_MAP_ADD, _patch, _body_start);
+                else _consume_comp(OP_BUILD_SET, OP_SET_ADD, _patch, _body_start);
+                consume(TK("}"));
+                return;
+            }
         } while (match(TK(",")));
         consume(TK("}"));
 
-        if(size == 0 || parsing_dict) emit(OP_BUILD_MAP, size);
-        else emit(OP_BUILD_SET, size);
+        if(ARGC == 0 || parsing_dict) emit(OP_BUILD_MAP, ARGC);
+        else emit(OP_BUILD_SET, ARGC);
     }
 
     void exprCall() {

src/obj.h

@@ -136,7 +136,7 @@ inline bool is_type(const PyVar& obj, Type type) noexcept {
 }
 
 inline bool is_both_int_or_float(const PyVar& a, const PyVar& b) noexcept {
-    return ((a.bits | b.bits) & 0b11) != 0b00;
+    return a.is_tagged() && b.is_tagged();
 }
 
 inline bool is_both_int(const PyVar& a, const PyVar& b) noexcept {

src/opcodes.h

@@ -29,6 +29,8 @@ OPCODE(BUILD_TUPLE_REF)
 OPCODE(BUILD_STRING)
 
 OPCODE(LIST_APPEND)
+OPCODE(MAP_ADD)
+OPCODE(SET_ADD)
 OPCODE(IMPORT_NAME)
 OPCODE(PRINT_EXPR)
 

tests/07_dict.py

@@ -20,6 +20,8 @@ tinydict.update(tinydict2)
 updated_dict = {'Name': 'circle', 'Age': 7, 'Sex': 'female'}
 for k,v in tinydict.items():
     assert updated_dict[k] == v
+assert len(tinydict) == 3
+assert tinydict == updated_dict
 
 dishes = {'eggs': 2, 'sausage': 1, 'bacon': 1, 'spam': 500}
 keys = dishes.keys()
@@ -29,8 +31,15 @@ assert sorted(values) == sorted([2, 1, 1, 500])
 
 d={1:"a",2:"b",3:"c"}
 result=[]
-for kv in d.items():
-    k = kv[0]; v=kv[1]
+for k,v in d.items():
     result.append(k)
     result.append(v)
-assert result == [1, 'a', 2, 'b', 3, 'c']
+assert len(result) == 6
+assert set(result) == set([1, 'a', 2, 'b', 3, 'c'])
+
+# test __eq__
+d1 = {1:2, 3:4}
+d2 = {3:4, 1:2}
+d3 = {1:2, 3:4, 5:6}
+assert d1 == d2
+assert d1 != d3

tests/23_dictcomp.py

@@ -0,0 +1,8 @@
+a = {i: i**2 for i in range(10)}
+assert a == {0: 0, 1: 1, 2: 4, 3: 9, 4: 16, 5: 25, 6: 36, 7: 49, 8: 64, 9: 81}
+
+a = {i: i**2 for i in range(10) if i % 2 == 0}
+assert a == {0: 0, 2: 4, 4: 16, 6: 36, 8: 64}
+
+b = {k:v for k,v in a.items()}
+assert b == a

tests/23_setcomp.py

@@ -0,0 +1,11 @@
+a = {i for i in range(10)}
+assert a == set(range(10))
+
+a = {i for i in range(10) if i % 2 == 0}
+assert a == {0, 2, 4, 6, 8}
+
+a = {i**3 for i in range(10) if i % 2 == 0}
+assert a == {0, 8, 64, 216, 512}
+
+a = {(i,i+1) for i in range(5)}
+assert a == {(0, 1), (1, 2), (2, 3), (3, 4), (4, 5)}