From d4ae2727a14b427a7da9c93c1c46b722291736d4 Mon Sep 17 00:00:00 2001
From: blueloveTH <blueloveTH@foxmail.com>
Date: Thu, 16 Mar 2023 23:32:23 +0800
Subject: [PATCH] add `collections`

---
 python/builtins.py      |  2 +-
 python/collections.py   | 87 +++++++++++++++++++++++++++++++++++++++++
 src/compiler.h          |  4 +-
 src/pocketpy.h          | 13 +++---
 tests/70_collections.py | 16 ++++++++
 5 files changed, 113 insertions(+), 9 deletions(-)
 create mode 100644 python/collections.py
 create mode 100644 tests/70_collections.py

diff --git a/python/builtins.py b/python/builtins.py
index 0e4e5641..2da2fc9a 100644
--- a/python/builtins.py
+++ b/python/builtins.py
@@ -113,7 +113,7 @@ def str::strip(self, chars=None):
 
 ##### list #####
 
-list.__new__ = lambda obj: [i for i in obj]
+list.__new__ = lambda iterable: [x for x in iterable]
 list.__repr__ = lambda self: '[' + ', '.join([repr(i) for i in self]) + ']'
 tuple.__repr__ = lambda self: '(' + ', '.join([repr(i) for i in self]) + ')'
 list.__json__ = lambda self: '[' + ', '.join([i.__json__() for i in self]) + ']'
diff --git a/python/collections.py b/python/collections.py
new file mode 100644
index 00000000..b913686b
--- /dev/null
+++ b/python/collections.py
@@ -0,0 +1,87 @@
+class _LinkedListNode:
+    def __init__(self, prev, next, value) -> None:
+        self.prev = prev
+        self.next = next
+        self.value = value
+
+class deque:
+    def __init__(self, iterable=None) -> None:
+        self.head = _LinkedListNode(None, None, None)
+        self.tail = _LinkedListNode(None, None, None)
+        self.head.next = self.tail
+        self.tail.prev = self.head
+        self.size = 0
+        if iterable is not None:
+            for value in iterable:
+                self.append(value)
+
+    def append(self, value):
+        node = _LinkedListNode(self.tail.prev, self.tail, value)
+        self.tail.prev.next = node
+        self.tail.prev = node
+        self.size += 1
+    
+    def appendleft(self, value):
+        node = _LinkedListNode(self.head, self.head.next, value)
+        self.head.next.prev = node
+        self.head.next = node
+        self.size += 1
+
+    def pop(self):
+        assert self.size > 0
+        node = self.tail.prev
+        node.prev.next = self.tail
+        self.tail.prev = node.prev
+        self.size -= 1
+        return node.value
+    
+    def popleft(self):
+        assert self.size > 0
+        node = self.head.next
+        node.next.prev = self.head
+        self.head.next = node.next
+        self.size -= 1
+        return node.value
+    
+    def copy(self):
+        new_list = deque()
+        for value in self:
+            new_list.append(value)
+        return new_list
+    
+    def __len__(self):
+        return self.size
+    
+    def __iter__(self):
+        node = self.head.next
+        while node is not self.tail:
+            yield node.value
+            node = node.next
+
+    def __repr__(self) -> str:
+        return f"deque({list(self)})"
+    
+    def __eq__(self, __o: object) -> bool:
+        if not isinstance(__o, deque):
+            return False
+        if len(self) != len(__o):
+            return False
+        t1, t2 = self.head.next, __o.head.next
+        while t1 is not self.tail:
+            if t1.value != t2.value:
+                return False
+            t1, t2 = t1.next, t2.next
+        return True
+    
+    def __ne__(self, __o: object) -> bool:
+        return not self == __o
+
+
+def Counter(iterable):
+    a = {}
+    for x in iterable:
+        if x in a:
+            a[x] += 1
+        else:
+            a[x] = 1
+    return a
\ No newline at end of file
diff --git a/src/compiler.h b/src/compiler.h
index 3957def1..8781d97d 100644
--- a/src/compiler.h
+++ b/src/compiler.h
@@ -1067,7 +1067,9 @@ __LISTCOMP:
             _compile_f_args(func, true);
             consume(TK(")"));
         }
-        if(match(TK("->"))) consume(TK("@id")); // eat type hints
+        if(match(TK("->"))){
+            if(!match(TK("None"))) consume(TK("@id"));
+        }
         func.code = make_sp<CodeObject>(parser->src, func.name.str());
         this->codes.push(func.code);
         compile_block_body();
diff --git a/src/pocketpy.h b/src/pocketpy.h
index 7f3605a2..3f5cab94 100644
--- a/src/pocketpy.h
+++ b/src/pocketpy.h
@@ -133,6 +133,10 @@ void init_builtins(VM* _vm) {
         return VAR("0x" + ss.str());
     });
 
+    _vm->bind_builtin_func<1>("iter", [](VM* vm, Args& args) {
+        return vm->asIter(args[0]);
+    });
+
     _vm->bind_builtin_func<1>("dir", [](VM* vm, Args& args) {
         std::set<StrName> names;
         if(args[0]->is_attr_valid()){
@@ -713,12 +717,6 @@ void add_module_random(VM* vm){
     vm->_exec(code, mod);
 }
 
-void add_module_functools(VM* vm){
-    PyVar mod = vm->new_module("functools");
-    CodeObject_ code = vm->compile(kPythonLibs["functools"], "functools.py", EXEC_MODE);
-    vm->_exec(code, mod);
-}
-
 void VM::post_init(){
     init_builtins(this);
     add_module_sys(this);
@@ -730,8 +728,9 @@ void VM::post_init(){
     add_module_random(this);
     add_module_io(this);
     add_module_os(this);
-    add_module_functools(this);
     add_module_c(this);
+    _lazy_modules["functools"] = kPythonLibs["functools"];
+    _lazy_modules["collections"] = kPythonLibs["collections"];
 
     CodeObject_ code = compile(kPythonLibs["builtins"], "<builtins>", EXEC_MODE);
     this->_exec(code, this->builtins);
diff --git a/tests/70_collections.py b/tests/70_collections.py
new file mode 100644
index 00000000..477b16d3
--- /dev/null
+++ b/tests/70_collections.py
@@ -0,0 +1,16 @@
+from collections import Counter, deque
+
+q = deque()
+q.append(1)
+q.append(2)
+q.appendleft(3)
+q.append(4)
+
+assert len(q) == 4
+
+assert q == deque([3, 1, 2, 4])
+assert q.popleft() == 3
+assert q.pop() == 4
+
+assert len(q) == 2
+assert q == deque([1, 2])
\ No newline at end of file