diff --git a/src/compiler.h b/src/compiler.h index bbe0305c..e6df0990 100644 --- a/src/compiler.h +++ b/src/compiler.h @@ -730,6 +730,19 @@ __SUBSCR_END: ctx()->emit(OP_YIELD_VALUE, BC_NOARG, kw_line); consume_end_stmt(); break; + case TK("yield from"): + if (contexts.size() <= 1) SyntaxError("'yield from' outside function"); + EXPR_TUPLE(false); + // if yield from present, mark the function as generator + ctx()->co->is_generator = true; + ctx()->emit(OP_GET_ITER, BC_NOARG, kw_line); + ctx()->enter_block(FOR_LOOP); + ctx()->emit(OP_FOR_ITER, BC_NOARG, BC_KEEPLINE); + ctx()->emit(OP_YIELD_VALUE, BC_NOARG, BC_KEEPLINE); + ctx()->emit(OP_LOOP_CONTINUE, BC_NOARG, BC_KEEPLINE); + ctx()->exit_block(); + consume_end_stmt(); + break; case TK("return"): if (contexts.size() <= 1) SyntaxError("'return' outside function"); if(match_end_stmt()){ diff --git a/src/lexer.h b/src/lexer.h index 726a5a36..11d17575 100644 --- a/src/lexer.h +++ b/src/lexer.h @@ -9,7 +9,7 @@ namespace pkpy{ typedef uint8_t TokenIndex; constexpr const char* kTokens[] = { - "is not", "not in", + "is not", "not in", "yield from", "@eof", "@eol", "@sof", "@id", "@num", "@str", "@fstr", "@indent", "@dedent", @@ -231,6 +231,12 @@ struct Lexer { add_token(TK("is not")); return 0; } + }else if(name == "yield"){ + if(strncmp(curr_char, " from", 5) == 0){ + curr_char += 5; + add_token(TK("yield from")); + return 0; + } } add_token(kTokenKwMap.at(name)); } else { diff --git a/tests/45_yield.py b/tests/45_yield.py index d6559d2f..65430e33 100644 --- a/tests/45_yield.py +++ b/tests/45_yield.py @@ -28,3 +28,12 @@ assert a[5] == (1,2) assert a[6] == (2,0) assert a[7] == (2,1) assert a[8] == (2,2) + +def g(): + yield from [1, 2, 3] + +def f(): + yield from g() + +a = [i for i in f()] +assert a == [1, 2, 3] \ No newline at end of file