fix https://github.com/blueloveTH/pocketpy/issues/80

docs/modules/timeit.md

@@ -0,0 +1,8 @@
+---
+icon: package
+label: timeit
+---
+
+### `timeit.timeit(f, number)`
+
+Returns the time taken to execute the given function `f` `number` times.

python/_long.py

@@ -0,0 +1,300 @@
+from c import sizeof
+
+if sizeof('void_p') == 4:
+    PyLong_SHIFT = 28//2
+elif sizeof('void_p') == 8:
+    PyLong_SHIFT = 60//2
+else:
+    raise NotImplementedError
+
+PyLong_BASE = 2 ** PyLong_SHIFT
+PyLong_MASK = PyLong_BASE - 1
+PyLong_DECIMAL_SHIFT = 4
+PyLong_DECIMAL_BASE = 10 ** PyLong_DECIMAL_SHIFT
+
+def ulong_fromint(x: int):
+    # return a list of digits and sign
+    if x == 0: return [0], 1
+    sign = 1 if x > 0 else -1
+    if sign < 0: x = -x
+    res = []
+    while x:
+        res.append(x & PyLong_MASK)
+        x >>= PyLong_SHIFT
+    return res, sign
+
+def ulong_cmp(a: list, b: list) -> int:
+    # return 1 if a>b, -1 if a<b, 0 if a==b
+    if len(a) > len(b): return 1
+    if len(a) < len(b): return -1
+    for i in range(len(a)-1, -1, -1):
+        if a[i] > b[i]: return 1
+        if a[i] < b[i]: return -1
+    return 0
+
+def ulong_pad_(a: list, size: int):
+    # pad leading zeros to have `size` digits
+    delta = size - len(a)
+    if delta > 0:
+        a.extend([0] * delta)
+
+def ulong_unpad_(a: list):
+    # remove leading zeros
+    while len(a)>1 and a[-1]==0:
+        a.pop()
+
+def ulong_add(a: list, b: list) -> list:
+    res = [0] * max(len(a), len(b))
+    ulong_pad_(a, len(res))
+    ulong_pad_(b, len(res))
+    carry = 0
+    for i in range(len(res)):
+        carry += a[i] + b[i]
+        res[i] = carry & PyLong_MASK
+        carry >>= PyLong_SHIFT
+    if carry > 0:
+        res.append(carry)
+    return res
+
+def ulong_sub(a: list, b: list) -> list:
+    # a >= b
+    res = []
+    borrow = 0
+    for i in range(len(b)):
+        tmp = a[i] - b[i] - borrow
+        if tmp < 0:
+            tmp += PyLong_BASE
+            borrow = 1
+        else:
+            borrow = 0
+        res.append(tmp)
+    for i in range(len(b), len(a)):
+        tmp = a[i] - borrow
+        if tmp < 0:
+            tmp += PyLong_BASE
+            borrow = 1
+        else:
+            borrow = 0
+        res.append(tmp)
+    ulong_unpad_(res)
+    return res
+
+def ulong_divmodi(a: list, b: int):
+    # b > 0
+    res = []
+    carry = 0
+    for i in range(len(a)-1, -1, -1):
+        carry <<= PyLong_SHIFT
+        carry += a[i]
+        res.append(carry // b)
+        carry %= b
+    res.reverse()
+    ulong_unpad_(res)
+    return res, carry
+
+def ulong_floordivi(a: list, b: int):
+    # b > 0
+    return ulong_divmodi(a, b)[0]
+
+def ulong_muli(a: list, b: int):
+    # b >= 0
+    res = [0] * len(a)
+    carry = 0
+    for i in range(len(a)):
+        carry += a[i] * b
+        res[i] = carry & PyLong_MASK
+        carry >>= PyLong_SHIFT
+    if carry > 0:
+        res.append(carry)
+    return res
+
+def ulong_mul(a: list, b: list):
+    res = [0] * (len(a) + len(b))
+    for i in range(len(a)):
+        carry = 0
+        for j in range(len(b)):
+            carry += res[i+j] + a[i] * b[j]
+            res[i+j] = carry & PyLong_MASK
+            carry >>= PyLong_SHIFT
+        res[i+len(b)] = carry
+    ulong_unpad_(res)
+    return res
+
+def ulong_powi(a: list, b: int):
+    # b >= 0
+    if b == 0: return [1]
+    res = [1]
+    while b:
+        if b & 1:
+            res = ulong_mul(res, a)
+        a = ulong_mul(a, a)
+        b >>= 1
+    return res
+
+def ulong_repr(x: list) -> str:
+    res = []
+    while len(x)>1 or x[0]>0:   # non-zero
+        x, r = ulong_divmodi(x, PyLong_DECIMAL_BASE)
+        res.append(str(r).zfill(PyLong_DECIMAL_SHIFT))
+    res.reverse()
+    s = ''.join(res)
+    if len(s) == 0: return '0'
+    if len(s) > 1: s = s.lstrip('0')
+    return s
+
+def ulong_fromstr(s: str):
+    res = [0]
+    base = [1]
+    if s[0] == '-':
+        sign = -1
+        s = s[1:]
+    else:
+        sign = 1
+    s = s[::-1]
+    for c in s:
+        c = ord(c) - 48
+        assert 0 <= c <= 9
+        res = ulong_add(res, ulong_muli(base, c))
+        base = ulong_muli(base, 10)
+    return res, sign
+
+class long:
+    def __init__(self, x):
+        if type(x) is tuple:
+            self.digits, self.sign = x
+        elif type(x) is int:
+            self.digits, self.sign = ulong_fromint(x)
+        elif type(x) is str:
+            self.digits, self.sign = ulong_fromstr(x)
+        else:
+            raise TypeError('expected int or str')
+
+    def __add__(self, other):
+        if type(other) is int:
+            other = long(other)
+        else:
+            assert type(other) is long
+        if self.sign == other.sign:
+            return long(ulong_add(self.digits, other.digits), self.sign)
+        else:
+            cmp = ulong_cmp(self.digits, other.digits)
+            if cmp == 0:
+                return long(0)
+            if cmp > 0:
+                return long(ulong_sub(self.digits, other.digits), self.sign)
+            else:
+                return long(ulong_sub(other.digits, self.digits), other.sign)
+            
+    def __radd__(self, other):
+        return self.__add__(other)
+    
+    def __sub__(self, other):
+        if type(other) is int:
+            other = long(other)
+        else:
+            assert type(other) is long
+        if self.sign != other.sign:
+            return long(ulong_add(self.digits, other.digits), self.sign)
+        else:
+            cmp = ulong_cmp(self.digits, other.digits)
+            if cmp == 0:
+                return long(0)
+            if cmp > 0:
+                return long(ulong_sub(self.digits, other.digits), self.sign)
+            else:
+                return long(ulong_sub(other.digits, self.digits), -other.sign)
+            
+    def __rsub__(self, other):
+        return self.__sub__(other)
+    
+    def __mul__(self, other):
+        if type(other) is int:
+            return long((
+                ulong_muli(self.digits, abs(other)),
+                self.sign * (1 if other >= 0 else -1)
+            ))
+        elif type(other) is long:
+            return long((
+                ulong_mul(self.digits, other.digits),
+                self.sign * other.sign
+            ))
+        raise TypeError('unsupported operand type(s) for *')
+    
+    def __rmul__(self, other):
+        return self.__mul__(other)
+    
+    #######################################################
+    def __divmod__(self, other: int):
+        assert type(other) is int and other > 0
+        assert self.sign == 1
+        q, r = ulong_divmodi(self.digits, other)
+        return long((q, 1)), r
+
+    def __floordiv__(self, other: int):
+        return self.__divmod__(other)[0]
+
+    def __mod__(self, other: int):
+        return self.__divmod__(other)[1]
+
+    def __pow__(self, other: int):
+        assert type(other) is int and other >= 0
+        if self.sign == -1 and other & 1:
+            sign = -1
+        else:
+            sign = 1
+        return long((ulong_powi(self.digits, other), sign))
+    
+    def __lshift__(self, other: int):
+        # TODO: optimize
+        assert type(other) is int and other >= 0
+        x = self.digits.copy()
+        for _ in range(other):
+            x = ulong_muli(x, 2)
+        return long((x, self.sign))
+    
+    def __rshift__(self, other: int):
+        # TODO: optimize
+        assert type(other) is int and other >= 0
+        x = self.digits.copy()
+        for _ in range(other):
+            x = ulong_floordivi(x, 2)
+        return long((x, self.sign))
+    
+    def __and__(self, other):
+        raise NotImplementedError
+    
+    def __or__(self, other):
+        raise NotImplementedError
+    
+    def __xor__(self, other):
+        raise NotImplementedError
+    
+    def __neg__(self):
+        return long(self.digits, -self.sign)
+    
+    def __cmp__(self, other):
+        if type(other) is int:
+            other = long(other)
+        else:
+            assert type(other) is long
+        if self.sign > other.sign:
+            return 1
+        elif self.sign < other.sign:
+            return -1
+        else:
+            return ulong_cmp(self.digits, other.digits)
+        
+    def __eq__(self, other):
+        return self.__cmp__(other) == 0
+    def __lt__(self, other):
+        return self.__cmp__(other) < 0
+    def __le__(self, other):
+        return self.__cmp__(other) <= 0
+    def __gt__(self, other):
+        return self.__cmp__(other) > 0
+    def __ge__(self, other):
+        return self.__cmp__(other) >= 0
+            
+    def __repr__(self):
+        prefix = '-' if self.sign < 0 else ''
+        return prefix + ulong_repr(self.digits) + 'L'

python/builtins.py

@@ -158,6 +158,13 @@ def __f(self, chars=None):
     return self[i:j+1]
 str.strip = __f
 
+def __f(self, width: int):
+    delta = width - len(self)
+    if delta <= 0:
+        return self
+    return '0' * delta + self
+str.zfill = __f
+
 ##### list #####
 list.__repr__ = lambda self: '[' + ', '.join([repr(i) for i in self]) + ']'
 list.__json__ = lambda self: '[' + ', '.join([i.__json__() for i in self]) + ']'
@@ -235,3 +242,5 @@ def help(obj):
 
 
 del __f
+
+from _long import long

src/compiler.h

@@ -221,9 +221,11 @@ class Compiler {
         std::vector<Expr_> items;
         items.push_back(ctx()->s_expr.popx());
         do {
+            if(curr().brackets_level) match_newlines_repl();
             if(!is_expression()) break;
             EXPR();
             items.push_back(ctx()->s_expr.popx());
+            if(curr().brackets_level) match_newlines_repl();
         } while(match(TK(",")));
         ctx()->s_expr.push(make_expr<TupleExpr>(
             std::move(items)

src/lexer.h

@@ -56,6 +56,7 @@ struct Token{
   const char* start;
   int length;
   int line;
+  int brackets_level;
   TokenValue value;
 
   Str str() const { return Str(start, length);}
@@ -151,11 +152,11 @@ struct Lexer {
         // https://docs.python.org/3/reference/lexical_analysis.html#indentation
         if(spaces > indents.top()){
             indents.push(spaces);
-            nexts.push_back(Token{TK("@indent"), token_start, 0, current_line});
+            nexts.push_back(Token{TK("@indent"), token_start, 0, current_line, brackets_level});
         } else if(spaces < indents.top()){
             while(spaces < indents.top()){
                 indents.pop();
-                nexts.push_back(Token{TK("@dedent"), token_start, 0, current_line});
+                nexts.push_back(Token{TK("@dedent"), token_start, 0, current_line, brackets_level});
             }
             if(spaces != indents.top()){
                 return false;
@@ -262,6 +263,7 @@ struct Lexer {
             token_start,
             (int)(curr_char - token_start),
             current_line - ((type == TK("@eol")) ? 1 : 0),
+            brackets_level,
             value
         };
         // handle "not in", "is not", "yield from"
@@ -526,7 +528,7 @@ struct Lexer {
         this->src = src;
         this->token_start = src->source.c_str();
         this->curr_char = src->source.c_str();
-        this->nexts.push_back(Token{TK("@sof"), token_start, 0, current_line});
+        this->nexts.push_back(Token{TK("@sof"), token_start, 0, current_line, brackets_level});
         this->indents.push(0);
     }
 

src/pocketpy.h

@@ -115,10 +115,15 @@ inline void init_builtins(VM* _vm) {
     });
 
     _vm->bind_builtin_func<2>("divmod", [](VM* vm, ArgsView args) {
-        i64 lhs = CAST(i64, args[0]);
+        if(is_int(args[0])){
+            i64 lhs = _CAST(i64, args[0]);
             i64 rhs = CAST(i64, args[1]);
             auto res = std::div(lhs, rhs);
             return VAR(Tuple({VAR(res.quot), VAR(res.rem)}));
+        }else{
+            DEF_SNAME(__divmod__);
+            return vm->call_method(args[0], __divmod__, args[1]);
+        }
     });
 
     _vm->bind_builtin_func<1>("eval", [](VM* vm, ArgsView args) {
@@ -1048,6 +1053,19 @@ inline void init_builtins(VM* _vm) {
     Generator::register_class(_vm, _vm->builtins);
 }
 
+inline void add_module_timeit(VM* vm){
+    PyObject* mod = vm->new_module("timeit");
+    vm->bind_func<2>(mod, "timeit", [](VM* vm, ArgsView args) {
+        PyObject* f = args[0];
+        i64 iters = CAST(i64, args[1]);
+        auto now = std::chrono::system_clock::now();
+        for(i64 i=0; i<iters; i++) vm->call(f);
+        auto end = std::chrono::system_clock::now();
+        f64 elapsed = std::chrono::duration_cast<std::chrono::milliseconds>(end - now).count() / 1000.0;
+        return VAR(elapsed);
+    });
+}
+
 inline void add_module_time(VM* vm){
     PyObject* mod = vm->new_module("time");
     vm->bind_func<0>(mod, "time", [](VM* vm, ArgsView args) {
@@ -1279,29 +1297,7 @@ inline void add_module_gc(VM* vm){
 
 inline void VM::post_init(){
     init_builtins(this);
-#if !DEBUG_NO_BUILTIN_MODULES
-    add_module_sys(this);
-    add_module_traceback(this);
-    add_module_time(this);
-    add_module_json(this);
-    add_module_math(this);
-    add_module_re(this);
-    add_module_dis(this);
-    add_module_c(this);
-    add_module_gc(this);
-    add_module_random(this);
-    add_module_base64(this);
 
-    for(const char* name: {"this", "functools", "collections", "heapq", "bisect", "pickle"}){
-        _lazy_modules[name] = kPythonLibs[name];
-    }
-
-    CodeObject_ code = compile(kPythonLibs["builtins"], "<builtins>", EXEC_MODE);
-    this->_exec(code, this->builtins);
-    code = compile(kPythonLibs["_set"], "<set>", EXEC_MODE);
-    this->_exec(code, this->builtins);
-
-    // property is defined in builtins.py so we need to add it after builtins is loaded
     _t(tp_object)->attr().set("__class__", property(CPP_LAMBDA(vm->_t(args[0]))));
     _t(tp_type)->attr().set("__base__", property([](VM* vm, ArgsView args){
         const PyTypeInfo& info = vm->_all_types[OBJ_GET(Type, args[0])];
@@ -1338,6 +1334,35 @@ inline void VM::post_init(){
         return VAR(MappingProxy(args[0]));
     }));
 
+#if !DEBUG_NO_BUILTIN_MODULES
+    add_module_sys(this);
+    add_module_traceback(this);
+    add_module_time(this);
+    add_module_json(this);
+    add_module_math(this);
+    add_module_re(this);
+    add_module_dis(this);
+    add_module_c(this);
+    add_module_gc(this);
+    add_module_random(this);
+    add_module_base64(this);
+    add_module_timeit(this);
+
+    for(const char* name: {"this", "functools", "collections", "heapq", "bisect", "pickle", "_long"}){
+        _lazy_modules[name] = kPythonLibs[name];
+    }
+
+    try{
+        CodeObject_ code = compile(kPythonLibs["builtins"], "<builtins>", EXEC_MODE);
+        this->_exec(code, this->builtins);
+        code = compile(kPythonLibs["_set"], "<set>", EXEC_MODE);
+        this->_exec(code, this->builtins);
+    }catch(Exception& e){
+        std::cerr << e.summary() << std::endl;
+        std::cerr << "failed to load builtins module!!" << std::endl;
+        exit(1);
+    }
+
     if(enable_os){
         add_module_io(this);
         add_module_os(this);

tests/04_str.py

@@ -60,6 +60,9 @@ seq = ["r","u","n","o","o","b"]
 assert s1.join( seq ) == "r-u-n-o-o-b"
 assert s2.join( seq ) == "runoob"
 
+assert 'x'.zfill(5) == '0000x'
+assert '568'.zfill(1) == '568'
+
 def test(*seq):
     return s1.join(seq)
 assert test("r", "u", "n", "o", "o", "b") == "r-u-n-o-o-b"

tests/99_bugs2.py

@@ -0,0 +1,15 @@
+def g(x):
+    return x
+def f(x):
+    return x
+
+assert (g(1), 2) == (1, 2)
+assert (
+    g(1),
+    2
+) == (1, 2)
+
+assert f((
+    g(1),
+    2
+)) == (1, 2)