diff --git a/python/builtins.py b/python/builtins.py index 57273b95..9bd33ecc 100644 --- a/python/builtins.py +++ b/python/builtins.py @@ -165,6 +165,22 @@ def __f(self, width: int): return '0' * delta + self str.zfill = __f +def __f(self, width: int, fillchar=' '): + delta = width - len(self) + if delta <= 0: + return self + assert len(fillchar) == 1 + return fillchar * delta + self +str.rjust = __f + +def __f(self, width: int, fillchar=' '): + delta = width - len(self) + if delta <= 0: + return self + assert len(fillchar) == 1 + return self + fillchar * delta +str.ljust = __f + ##### list ##### list.__repr__ = lambda self: '[' + ', '.join([repr(i) for i in self]) + ']' list.__json__ = lambda self: '[' + ', '.join([i.__json__() for i in self]) + ']' diff --git a/src/pocketpy.h b/src/pocketpy.h index cb5d6bf4..a307f25e 100644 --- a/src/pocketpy.h +++ b/src/pocketpy.h @@ -103,11 +103,12 @@ inline void init_builtins(VM* _vm) { _vm->bind_builtin_func<3>("pow", [](VM* vm, ArgsView args) { i64 lhs = CAST(i64, args[0]); // assume lhs>=0 i64 rhs = CAST(i64, args[1]); // assume rhs>=0 - i64 mod = CAST(i64, args[2]); // assume mod>0 + i64 mod = CAST(i64, args[2]); // assume mod>0, mod*mod should not overflow i64 res = 1; while(rhs){ - if(rhs & 1) res = res*lhs % mod; - lhs = lhs*lhs % mod; + i64 lhs_mod = lhs % mod; + if(rhs & 1) res = ((res % mod) * lhs_mod) % mod; + lhs = (lhs_mod * lhs_mod) % mod; rhs >>= 1; } return VAR(res); @@ -213,6 +214,22 @@ inline void init_builtins(VM* _vm) { return vm->py_next(args[0]); }); + _vm->bind_builtin_func<1>("bin", [](VM* vm, ArgsView args) { + std::stringstream ss; + i64 x = CAST(i64, args[0]); + if(x < 0){ ss << "-"; x = -x; } + ss << "0b"; + std::string bits; + while(x){ + bits += (x & 1) ? '1' : '0'; + x >>= 1; + } + std::reverse(bits.begin(), bits.end()); + if(bits.empty()) bits = "0"; + ss << bits; + return VAR(ss.str()); + }); + _vm->bind_builtin_func<1>("dir", [](VM* vm, ArgsView args) { std::set names; if(!is_tagged(args[0]) && args[0]->is_attr_valid()){ @@ -295,22 +312,27 @@ inline void init_builtins(VM* _vm) { _vm->bind__pow__(_vm->tp_float, py_number_pow); /************ int ************/ - _vm->bind_constructor<2>("int", [](VM* vm, ArgsView args) { - if (is_type(args[1], vm->tp_float)) return VAR((i64)CAST(f64, args[1])); - if (is_type(args[1], vm->tp_int)) return args[1]; - if (is_type(args[1], vm->tp_bool)) return VAR(_CAST(bool, args[1]) ? 1 : 0); + _vm->bind_constructor<-1>("int", [](VM* vm, ArgsView args) { + if(args.size() == 1+1){ + if (is_type(args[1], vm->tp_float)) return VAR((i64)CAST(f64, args[1])); + if (is_type(args[1], vm->tp_int)) return args[1]; + if (is_type(args[1], vm->tp_bool)) return VAR(_CAST(bool, args[1]) ? 1 : 0); + } + if(args.size() > 1+2) vm->TypeError("int() takes at most 2 arguments"); if (is_type(args[1], vm->tp_str)) { + int base = 10; + if(args.size() == 1+2) base = CAST(i64, args[2]); const Str& s = CAST(Str&, args[1]); try{ size_t parsed = 0; - i64 val = Number::stoi(s.str(), &parsed, 10); + i64 val = Number::stoi(s.str(), &parsed, base); if(parsed != s.length()) throw std::invalid_argument(""); return VAR(val); }catch(std::invalid_argument&){ vm->ValueError("invalid literal for int(): " + s.escape()); } } - vm->TypeError("int() argument must be a int, float, bool or str"); + vm->TypeError("invalid arguments for int()"); return vm->None; }); diff --git a/tests/04_str.py b/tests/04_str.py index f9e2f5e0..9eaf05f1 100644 --- a/tests/04_str.py +++ b/tests/04_str.py @@ -100,4 +100,10 @@ assert b[::-1] == ['!', 'd', 'l', 'r', 'o', 'W', ' ', ',', 'o', 'l', 'l', 'e', ' assert b[::2] == ['H', 'l', 'o', ' ', 'o', 'l', '!'] assert b[2:5:2] == ['l', 'o'] assert b[5:2:-1] == [',', 'o', 'l'] -assert b[5:2:-2] == [',', 'l'] \ No newline at end of file +assert b[5:2:-2] == [',', 'l'] + +a = '123' +assert a.rjust(5) == ' 123' +assert a.rjust(5, '0') == '00123' +assert a.ljust(5) == '123 ' +assert a.ljust(5, '0') == '12300' \ No newline at end of file