diff --git a/.travis.yml b/.travis.yml index 220c689..e1e6335 100644 --- a/.travis.yml +++ b/.travis.yml @@ -7,6 +7,7 @@ python: - "3.5" - "3.6" - "3.7" + - "3.8" - "pypy" - "pypy3" diff --git a/README.md b/README.md index ea0bb8c..a0c54f7 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ [![Pypi Entry](https://badge.fury.io/py/goto-statement.svg)](https://pypi.python.org/pypi/goto-statement) A function decorator to use `goto` in Python. -Tested on Python 2.6 through 3.7 and PyPy. +Tested on Python 2.6 through 3.8 and PyPy. [![](https://imgs.xkcd.com/comics/goto.png)](https://xkcd.com/292/) diff --git a/goto.py b/goto.py index f3a86a7..9d40221 100644 --- a/goto.py +++ b/goto.py @@ -34,6 +34,8 @@ def __init__(self): self.have_argument = dis.HAVE_ARGUMENT self.jump_unit = 1 + self.has_loop_blocks = 'SETUP_LOOP' in dis.opmap + @property def argument_bits(self): return self.argument.size * 8 @@ -43,20 +45,23 @@ def argument_bits(self): def _make_code(code, codestring): - args = [ - code.co_argcount, code.co_nlocals, code.co_stacksize, - code.co_flags, codestring, code.co_consts, - code.co_names, code.co_varnames, code.co_filename, - code.co_name, code.co_firstlineno, code.co_lnotab, - code.co_freevars, code.co_cellvars - ] - try: - args.insert(1, code.co_kwonlyargcount) # PY3 - except AttributeError: - pass + return code.replace(co_code=codestring) # new in 3.8+ + except: + args = [ + code.co_argcount, code.co_nlocals, code.co_stacksize, + code.co_flags, codestring, code.co_consts, + code.co_names, code.co_varnames, code.co_filename, + code.co_name, code.co_firstlineno, code.co_lnotab, + code.co_freevars, code.co_cellvars + ] - return types.CodeType(*args) + try: + args.insert(1, code.co_kwonlyargcount) # PY3 + except AttributeError: + pass + + return types.CodeType(*args) def _parse_instructions(code): @@ -86,6 +91,28 @@ def _parse_instructions(code): extended_arg_offset = None yield (dis.opname[opcode], oparg, offset) +def _get_instruction_size(opname, oparg=0): + size = 1 + + extended_arg = oparg >> _BYTECODE.argument_bits + if extended_arg != 0: + size += _get_instruction_size('EXTENDED_ARG', extended_arg) + oparg &= (1 << _BYTECODE.argument_bits) - 1 + + opcode = dis.opmap[opname] + if opcode >= _BYTECODE.have_argument: + size += _BYTECODE.argument.size + + return size + +def _get_instructions_size(ops): + size = 0 + for op in ops: + if isinstance(op, str): + size += _get_instruction_size(op) + else: + size += _get_instruction_size(*op) + return size def _get_instruction_size(opname, oparg=0): size = 1 @@ -128,6 +155,13 @@ def _write_instruction(buf, pos, opname, oparg=0): return pos +def _write_instructions(buf, pos, ops): + for op in ops: + if isinstance(op, str): + pos = _write_instruction(buf, pos, op) + else: + pos = _write_instruction(buf, pos, *op) + return pos def _write_instructions(buf, pos, ops): for op in ops: @@ -144,6 +178,7 @@ def _find_labels_and_gotos(code): block_stack = [] block_counter = 0 + block_exits = [] opname1 = oparg1 = offset1 = None opname2 = oparg2 = offset2 = None @@ -170,9 +205,16 @@ def _find_labels_and_gotos(code): 'SETUP_EXCEPT', 'SETUP_FINALLY', 'SETUP_WITH', 'SETUP_ASYNC_WITH'): block_counter += 1 - block_stack.append(block_counter) + block_stack.append((opname1, block_counter)) + elif not _BYTECODE.has_loop_blocks and opname1 == 'FOR_ITER': + block_counter += 1 + block_stack.append((opname1, block_counter)) + block_exits.append(offset1 + oparg1) elif opname1 == 'POP_BLOCK' and block_stack: block_stack.pop() + elif block_exits and offset1 == block_exits[-1] and block_stack: + block_stack.pop() + block_exits.pop() opname1, oparg1, offset1 = opname2, oparg2, offset2 opname2, oparg2, offset2 = opname3, oparg3, offset3 @@ -206,8 +248,13 @@ def _patch_code(code): raise SyntaxError('Jump into different block') ops = [] - for i in range(len(origin_stack) - target_depth): - ops.append('POP_BLOCK') + + for block, _ in origin_stack[target_depth:]: + if block == 'FOR_ITER': + ops.append('POP_TOP') + else: + ops.append('POP_BLOCK') + ops.append(('JUMP_ABSOLUTE', target // _BYTECODE.jump_unit)) if pos + _get_instructions_size(ops) > end: diff --git a/test_goto.py b/test_goto.py index 7fa82ed..a74f02e 100644 --- a/test_goto.py +++ b/test_goto.py @@ -62,6 +62,18 @@ def func(): assert func() == 0 +def test_jump_out_of_loop_and_survive(): + @with_goto + def func(): + for i in range(10): + for j in range(10): + goto .end + label .end + return (i, j) + + assert func() == (9, 0) + + def test_jump_into_loop(): def func(): for i in range(10): @@ -118,7 +130,6 @@ def func(): assert func() == (0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0) - def test_jump_across_loops(): def func(): for i in range(10): @@ -145,7 +156,6 @@ def func(): assert func() is None - def test_jump_into_try_block(): def func(): try: diff --git a/tox.ini b/tox.ini index 80c937e..3b213b0 100644 --- a/tox.ini +++ b/tox.ini @@ -1,5 +1,5 @@ [tox] -envlist = py26,py27,py32,py33,py34,py35,py36,py37,pypy,pypy3,flake8 +envlist = py26,py27,py32,py33,py34,py35,py36,py37,py38,pypy,pypy3,flake8 skip_missing_interpreters = true [testenv]