add with and from import

src/compiler.h

@@ -664,7 +664,7 @@ __LISTCOMP:
         return tkmodule;
     }
 
-    // import module1 [as alias1 [, module2 [as alias2 ...]]
+    // import a as b
     void compileRegularImport() {
         do {
             Token tkmodule = compileImportPath();
@@ -678,6 +678,25 @@ __LISTCOMP:
         consumeEndStatement();
     }
 
+    // from a import b as c, d as e
+    void compileFromImport() {
+        Token tkmodule = compileImportPath();
+        consume(TK("import"));
+        do {
+            consume(TK("@id"));
+            Token tkname = parser->previous;
+            int index = getCode()->addName(tkname.str(), NAME_GLOBAL);
+            emitCode(OP_BUILD_ATTR_PTR, index);
+            if (match(TK("as"))) {
+                consume(TK("@id"));
+                tkname = parser->previous;
+            }
+            index = getCode()->addName(tkname.str(), NAME_GLOBAL);
+            emitCode(OP_STORE_NAME_PTR, index);
+        } while (match(TK(",")));
+        consumeEndStatement();
+    }
+
     void parsePrecedence(Precedence precedence) {
         lexToken();
         GrammarFn prefix = rules[parser->previous.type].prefix;
@@ -786,6 +805,21 @@ __LISTCOMP:
             EXPR();
             emitCode(OP_ASSERT);
             consumeEndStatement();
+        } else if(match(TK("with"))){
+            EXPR();
+            consume(TK("as"));
+            consume(TK("@id"));
+            Token tkname = parser->previous;
+            int index = getCode()->addName(
+                tkname.str(),
+                codes.size()>1 ? NAME_LOCAL : NAME_GLOBAL
+            );
+            emitCode(OP_STORE_NAME_PTR, index);
+            emitCode(OP_LOAD_NAME_PTR, index);
+            emitCode(OP_WITH_ENTER);
+            compileBlockBody();
+            emitCode(OP_LOAD_NAME_PTR, index);
+            emitCode(OP_WITH_EXIT);
         } else if(match(TK("label"))){
             if(mode() != EXEC_MODE) syntaxError("'label' is only available in EXEC_MODE");
             consume(TK(".")); consume(TK("@id"));
@@ -926,6 +960,8 @@ __LISTCOMP:
             compileFunction();
         } else if (match(TK("import"))) {
             compileRegularImport();
+        } else if (match(TK("from"))) {
+            compileFromImport();
         } else {
             compileStatement();
         }

src/obj.h

@@ -10,7 +10,7 @@ const _Int _Int_MAX_NEG = -9223372036854775807LL;
 const _Float _FLOAT_INF_POS = INFINITY;
 const _Float _FLOAT_INF_NEG = -INFINITY;
 
-#define PK_VERSION "0.3.5"
+#define PK_VERSION "0.3.6"
 
 class CodeObject;
 class BasePointer;

src/opcodes.h

@@ -63,4 +63,7 @@ OPCODE(GOTO)
 OPCODE(UNARY_REF)           // for &
 OPCODE(UNARY_DEREF)         // for *
 
+OPCODE(WITH_ENTER)
+OPCODE(WITH_EXIT)
+
 #endif

src/parser.h

@@ -12,7 +12,7 @@ constexpr const char* __TOKENS[] = {
     "==", "!=", ">=", "<=",
     "+=", "-=", "*=", "/=", "//=",
     /** KW_BEGIN **/
-    "class", "import", "as", "def", "lambda", "pass", "del",
+    "class", "import", "as", "def", "lambda", "pass", "del", "from", "with",
     "None", "in", "is", "and", "or", "not", "True", "False", "global",
     "goto", "label",      // extended keywords, not available in cpython
     "while", "for", "if", "elif", "else", "break", "continue", "return", "assert", "raise",

src/pocketpy.h

@@ -404,11 +404,18 @@ void __initializeBuiltinFunctions(VM* _vm) {
     _vm->bindMethod("str", "join", [](VM* vm, const pkpy::ArgList& args) {
         vm->__checkArgSize(args, 2, true);
         const _Str& _self = vm->PyStr_AS_C(args[0]);
-        const PyVarList& _list = vm->PyList_AS_C(args[1]);
+        PyVarList* _list;
+        if(args[1]->isType(vm->_tp_list)){
+            _list = &vm->PyList_AS_C(args[1]);
+        }else if(args[1]->isType(vm->_tp_tuple)){
+            _list = &vm->PyTuple_AS_C(args[1]);
+        }else{
+            vm->typeError("can only join a list or tuple");
+        }
         _StrStream ss;
-        for(int i = 0; i < _list.size(); i++){
+        for(int i = 0; i < _list->size(); i++){
             if(i > 0) ss << _self;
-            ss << vm->PyStr_AS_C(vm->asStr(_list[i]));
+            ss << vm->PyStr_AS_C(vm->asStr(_list->operator[](i)));
         }
         return vm->PyStr(ss.str());
     });

src/vm.h

@@ -345,6 +345,18 @@ protected:
                         frame->push(it->second);
                     }
                 } break;
+            case OP_WITH_ENTER:
+            {
+                PyVar obj = frame->popValue(this);
+                PyVar enter_fn = getAttr(obj, "__enter__"_c);
+                call(enter_fn, {});
+            } break;
+            case OP_WITH_EXIT:
+            {
+                PyVar obj = frame->popValue(this);
+                PyVar exit_fn = getAttr(obj, "__exit__"_c);
+                call(exit_fn, {});
+            } break;
             default:
                 systemError(_Str("opcode ") + OP_NAMES[byte.op] + " is not implemented");
                 break;

tests/builtin_ty.py

@@ -48,6 +48,10 @@ seq = ["r","u","n","o","o","b"]
 assert s1.join( seq ) == "r-u-n-o-o-b"
 assert s2.join( seq ) == "runoob"
 
+def test(*seq):
+    return s1.join(seq)
+assert test("r", "u", "n", "o", "o", "b") == "r-u-n-o-o-b"
+
 
 ##num = 6
 ##assert str(num) == '6'   TypeError: 'str' object is not callable

tests/random.py

@@ -1,6 +1,26 @@
-import random
+import random as r, sys as s
 
 for _ in range(100):
-    i = random.randint(1, 10)
+    i = r.randint(1, 10)
     assert i <= 10
     assert i >= 1
+
+
+from sys import version as v
+
+assert type(v) is str
+
+class Context:
+    def __init__(self):
+        self.x = 0
+
+    def __enter__(self):
+        self.x = 1
+
+    def __exit__(self):
+        self.x = 2
+
+with Context() as c:
+    assert c.x == 1
+
+assert c.x == 2