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..93d6cc2 100644 --- a/goto.py +++ b/goto.py @@ -3,7 +3,8 @@ import array import types import functools - +import weakref +import warnings try: _array_to_bytes = array.array.tobytes @@ -34,6 +35,12 @@ def __init__(self): self.have_argument = dis.HAVE_ARGUMENT self.jump_unit = 1 + self.has_loop_blocks = 'SETUP_LOOP' in dis.opmap + self.has_pop_except = 'POP_EXCEPT' in dis.opmap + self.has_setup_with = 'SETUP_WITH' in dis.opmap + self.has_setup_except = 'SETUP_EXCEPT' in dis.opmap + self.has_begin_finally = 'BEGIN_FINALLY' in dis.opmap + @property def argument_bits(self): return self.argument.size * 8 @@ -42,24 +49,35 @@ def argument_bits(self): _BYTECODE = _Bytecode() -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 - ] +# use a weak dictionary in case code objects can be garbage-collected +_patched_code_cache = weakref.WeakKeyDictionary() +try: + _patched_code_cache[_Bytecode.__init__.__code__] = None +except TypeError: + _patched_code_cache = {} # ...unless not supported + +def _make_code(code, codestring): try: - args.insert(1, code.co_kwonlyargcount) # PY3 + return code.replace(co_code=codestring) # new in 3.8+ except AttributeError: - pass + 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 types.CodeType(*args) + return types.CodeType(*args) -def _parse_instructions(code): +def _parse_instructions(code, yield_nones_at_end=0): extended_arg = 0 extended_arg_offset = None pos = 0 @@ -86,6 +104,9 @@ def _parse_instructions(code): extended_arg_offset = None yield (dis.opname[opcode], oparg, offset) + for _ in range(yield_nones_at_end): + yield (None, None, None) + def _get_instruction_size(opname, oparg=0): size = 1 @@ -138,18 +159,87 @@ def _write_instructions(buf, pos, ops): return pos +def _warn_bug(msg): + warnings.warn("Internal error detected" + + " - result of with_goto may be incorrect. (%s)" % msg) + + +class _BlockStack(object): + def __init__(self, labels, gotos): + self.stack = [] + self.block_counter = 0 + self.last_block = None + self.labels = labels + self.gotos = gotos + + def _replace_in_stack(self, stack, old_block, new_block): + for i, block in enumerate(stack): + if block == old_block: + stack[i] = new_block + + def replace(self, old_block, new_block): + self._replace_in_stack(self.stack, old_block, new_block) + + for label in self.labels: + _, _, label_blocks = self.labels[label] + self._replace_in_stack(label_blocks, old_block, new_block) + + for goto in self.gotos: + _, _, _, goto_blocks = goto + self._replace_in_stack(goto_blocks, old_block, new_block) + + def push(self, opname, target_offset=None, previous=None): + self.block_counter += 1 + self.stack.append((opname, target_offset, + previous, self.block_counter)) + + def pop(self): + if self.stack: + self.last_block = self.stack.pop() + return self.last_block + else: + _warn_bug("can't pop block") + + def pop_of_type(self, type): + if self.stack and self.top()[0] != type: + _warn_bug("mismatched block type") + else: + return self.pop() + + def copy_to_list(self): + return list(self.stack) + + def top(self): + return self.stack[-1] if self.stack else None + + def __len__(self): + return len(self.stack) + + def _find_labels_and_gotos(code): labels = {} gotos = [] - block_stack = [] - block_counter = 0 + block_stack = _BlockStack(labels, gotos) opname1 = oparg1 = offset1 = None opname2 = oparg2 = offset2 = None opname3 = oparg3 = offset3 = None - for opname4, oparg4, offset4 in _parse_instructions(code.co_code): + for opname4, oparg4, offset4 in _parse_instructions(code.co_code, 3): + endoffset1 = offset2 + + # check for block exits + while block_stack and offset1 == block_stack.top()[1]: + exit_block = block_stack.pop() + exit_name = exit_block[0] + + if exit_name == 'SETUP_EXCEPT' and _BYTECODE.has_pop_except: + block_stack.push('', previous=exit_block) + elif exit_name == 'SETUP_FINALLY': + block_stack.push('', previous=exit_block) + + # check for special opcodes if opname1 in ('LOAD_GLOBAL', 'LOAD_NAME'): if opname2 == 'LOAD_ATTR' and opname3 == 'POP_TOP': name = code.co_names[oparg1] @@ -160,24 +250,51 @@ def _find_labels_and_gotos(code): )) labels[oparg2] = (offset1, offset4, - tuple(block_stack)) + block_stack.copy_to_list()) elif name == 'goto': gotos.append((offset1, offset4, oparg2, - tuple(block_stack))) - elif opname1 in ('SETUP_LOOP', - 'SETUP_EXCEPT', 'SETUP_FINALLY', - 'SETUP_WITH', 'SETUP_ASYNC_WITH'): - block_counter += 1 - block_stack.append(block_counter) - elif opname1 == 'POP_BLOCK' and block_stack: - block_stack.pop() + block_stack.copy_to_list())) + elif (opname1 in ('SETUP_LOOP', + 'SETUP_EXCEPT', 'SETUP_FINALLY', + 'SETUP_WITH', 'SETUP_ASYNC_WITH')) or \ + (not _BYTECODE.has_loop_blocks and opname1 == 'FOR_ITER'): + block_stack.push(opname1, endoffset1 + oparg1) + elif opname1 == 'POP_EXCEPT': + top_block = block_stack.top() + if not _BYTECODE.has_setup_except and \ + top_block and top_block[0] == '': + # in 3.8, only finally blocks are supported, so we must + # determine whether it's except/finally ourselves + block_stack.replace(top_block, + ('',) + top_block[1:]) + _, _, setup_block, _ = top_block + block_stack.replace(setup_block, + ('SETUP_EXCEPT',) + setup_block[1:]) + block_stack.pop_of_type('') + elif opname1 == 'END_FINALLY': + # Python puts END_FINALLY at the very end of except + # clauses, so we must ignore it in the wrong place. + if block_stack and block_stack.top()[0] == '': + block_stack.pop_of_type('') + elif opname1 in ('WITH_CLEANUP', 'WITH_CLEANUP_START'): + if _BYTECODE.has_setup_with: + # temporary block to match END_FINALLY + block_stack.push('') + else: + # python 2.6 - finally was actually with + last_block = block_stack.last_block + block_stack.replace(last_block, + ('SETUP_WITH',) + last_block[1:]) opname1, oparg1, offset1 = opname2, oparg2, offset2 opname2, oparg2, offset2 = opname3, oparg3, offset3 opname3, oparg3, offset3 = opname4, oparg4, offset4 + if block_stack: + _warn_bug("block stack not empty") + return labels, gotos @@ -187,6 +304,10 @@ def _inject_nop_sled(buf, pos, end): def _patch_code(code): + new_code = _patched_code_cache.get(code) + if new_code is not None: + return new_code + labels, gotos = _find_labels_and_gotos(code) buf = array.array('B', code.co_code) @@ -206,8 +327,26 @@ 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 reversed(origin_stack[target_depth:]): + if block == 'FOR_ITER': + ops.append('POP_TOP') + elif block == '': + ops.append('POP_EXCEPT') + elif block == '': + ops.append('END_FINALLY') + else: + ops.append('POP_BLOCK') + if block in ('SETUP_WITH', 'SETUP_ASYNC_WITH'): + ops.append('POP_TOP') + # END_FINALLY is needed only in pypy, + # but seems logical everywhere + if block in ('SETUP_FINALLY', 'SETUP_WITH', + 'SETUP_ASYNC_WITH'): + ops.append('BEGIN_FINALLY' if + _BYTECODE.has_begin_finally else + ('LOAD_CONST', code.co_consts.index(None))) + ops.append('END_FINALLY') + ops.append(('JUMP_ABSOLUTE', target // _BYTECODE.jump_unit)) if pos + _get_instructions_size(ops) > end: @@ -229,7 +368,10 @@ def _patch_code(code): pos = _write_instructions(buf, pos, ops) _inject_nop_sled(buf, pos, end) - return _make_code(code, _array_to_bytes(buf)) + new_code = _make_code(code, _array_to_bytes(buf)) + + _patched_code_cache[code] = new_code + return new_code def with_goto(func_or_code): diff --git a/test_goto.py b/test_goto.py index 7fa82ed..c07af24 100644 --- a/test_goto.py +++ b/test_goto.py @@ -62,6 +62,31 @@ 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_out_of_loop_and_live(): + @with_goto + def func(): + for i in range(10): + for j in range(10): + for k in range(10): + goto .end + label .end + return (i, j, k) + + assert func() == (9, 0, 0) + + def test_jump_into_loop(): def func(): for i in range(10): @@ -89,6 +114,55 @@ def func(): assert func() == (0, 0) +def test_jump_out_of_nested_4_loops_and_survive(): + @with_goto + def func(): + for i in range(2): + for j in range(2): + for k in range(2): + for m in range(2): + for n in range(2): + goto .end + label .end + return (i, j, k, m, n) + + assert func() == (1, 0, 0, 0, 0) + + +def test_large_jumps_in_diff_orders(): + @with_goto + def func(): + goto .start + + if NonConstFalse: + label .finalle + return (i, j, k, m, n, i1, j1, k1, m1, n1, i2, j2, k2, m2, n2) + + label .start + for i in range(2): + for j in range(2): + for k in range(2): + for m in range(2): + for n in range(2): + goto .end + label .end + for i1 in range(2): + for j1 in range(2): + for k1 in range(2): + for m1 in range(2): + for n1 in range(2): + goto .end2 + label .end2 + for i2 in range(2): + for j2 in range(2): + for k2 in range(2): + for m2 in range(2): + for n2 in range(2): + goto .finalle + + assert func() == (0,) * 15 + + def test_jump_out_of_nested_11_loops(): @with_goto def func(): @@ -130,6 +204,203 @@ def func(): pytest.raises(SyntaxError, with_goto, func) +def test_jump_out_of_while_true_loop(): + @with_goto + def func(): + i = 0 + while True: + i += 1 + goto .out + label .out + return i + + assert func() == 1 + + +def test_jump_out_of_while_true_loop_and_survive(): + @with_goto + def func(): + j = 0 + for i in range(10): + while True: + j += 1 + goto .out + label .out + return i, j + + assert func() == (9, 10) + + +def test_jump_out_of_while_true_loop_and_live(): + @with_goto + def func(): + k = 0 + for i in range(10): + for j in range(10): + while True: + k += 1 + goto .out + label .out + return i, j, k + + assert func() == (9, 0, 10) + + +def test_jump_out_of_while_loop_and_live(): + @with_goto + def func(): + k = 0 + for i in range(10): + for j in range(4): + while k < 5: + k += 1 + goto .out + label .out + return i, j, k + + assert func() == (9, 3, 5) + + +class Context: + def __init__(self): + self.enters = 0 + self.exits = 0 + + def __enter__(self): + self.enters += 1 + return self + + def __exit__(self, a, b, c): + self.exits += 1 + return self + + def data(self): + return (self.enters, self.exits) + + +def test_jump_out_of_with_block(): + @with_goto + def func(): + with Context() as c: + goto .out + label .out + return c.data() + + assert func() == (1, 0) + + +def test_jump_out_of_with_block_and_survive(): + @with_goto + def func(): + c = Context() + for i in range(3): + with c: + goto .out + label .out + return (i, c.data()) + + assert func() == (2, (3, 0)) + + +def test_jump_out_of_with_block_and_live(): + @with_goto + def func(): + c = Context() + for i in range(3): + for j in range(3): + with c: + goto .out + label .out + return (i, j, c.data()) + + assert func() == (2, 0, (3, 0)) + + +def test_jump_into_with_block(): + def func(): + with Context() as c: + label .block + goto .block + return c + + pytest.raises(SyntaxError, with_goto, func) + + +def test_generator(): + @with_goto + def func(): + yield 0 + yield 1 + goto .x + yield 2 + yield 3 + label .x + yield 4 + yield 5 + + assert tuple(func()) == (0, 1, 4, 5) + + +def test_jump_out_of_try_except_block(): + @with_goto + def func(): + try: + rv = None + goto .end + except Exception: + rv = 'except' + label .end + return rv + + assert func() is None + + +def test_jump_out_of_try_except_block_and_live(): + @with_goto + def func(): + for i in range(3): + for j in range(3): + try: + rv = None + goto .end + except Exception: + rv = 'except' + label .end + return (i, j, rv) + + assert func() == (2, 2, None) + + +def test_jump_out_of_try_finally_block(): + @with_goto + def func(): + try: + rv = None + goto .end + finally: + rv = 'finally' + label .end + return rv + + assert func() is None + + +def test_jump_out_of_try_finally_block_and_live(): + @with_goto + def func(): + for i in range(3): + for j in range(3): + try: + rv = None + goto .end + finally: + rv = 'finally' + label .end + return (i, j, rv) + + assert func() == (2, 2, None) + + def test_jump_out_of_try_block(): @with_goto def func(): @@ -146,6 +417,41 @@ def func(): assert func() is None +def test_jump_out_of_try_block_and_survive(): + @with_goto + def func(): + for i in range(10): + try: + rv = None + goto .end + except Exception: + rv = 'except' + finally: + rv = 'finally' + label .end + return (i, rv) + + assert func() == (9, None) + + +def test_jump_out_of_try_block_and_live(): + @with_goto + def func(): + for i in range(3): + for j in range(3): + try: + rv = None + goto .end + except Exception: + rv = 'except' + finally: + rv = 'finally' + label .end + return (i, j, rv) + + assert func() == (2, 2, None) + + def test_jump_into_try_block(): def func(): try: @@ -157,6 +463,195 @@ def func(): pytest.raises(SyntaxError, with_goto, func) +def test_jump_out_of_except_block_wo_finally(): + @with_goto + def func(): + try: + rv = 1 / 0 + except Exception: + rv = 'except' + goto .end + label .end + return rv + + assert func() == 'except' + + +def test_jump_out_of_except_block_wo_finally_and_live(): + @with_goto + def func(): + for i in range(3): + for j in range(3): + try: + rv = 1 / 0 + except Exception: + rv = 'except' + goto .end + label .end + return (i, j, rv) + + assert func() == (2, 0, 'except') + + +def test_jump_out_of_except_block(): + @with_goto + def func(): + try: + rv = 1 / 0 + except Exception: + rv = 'except' + goto .end + finally: + rv = 'finally' + label .end + return rv + + assert func() == 'except' + + +def test_jump_out_of_except_block_and_live(): + @with_goto + def func(): + for i in range(3): + for j in range(3): + try: + rv = 1 / 0 + except Exception: + rv = 'except' + goto .end + finally: + rv = 'finally' + label .end + return (i, j, rv) + + assert func() == (2, 0, 'except') + + +def test_jump_out_of_bare_except_block_and_live(): + @with_goto + def func(): + for i in range(3): + for j in range(3): + try: + rv = 1 / 0 + except: # noqa: E722 + rv = 'except' + goto .end + finally: + rv = 'finally' + label .end + return (i, j, rv) + + assert func() == (2, 0, 'except') + + +"""def test_jump_into_except_block(): + def func(): + try: + pass + except: + label .block + pass + goto .block + + pytest.raises(SyntaxError, with_goto, func)""" + + +def test_jump_out_of_finally_block_w_except(): + @with_goto + def func(): + try: + rv = None + except Exception: + rv = 'except' + finally: + rv = 'finally' + goto .end + rv = 'end' + label .end + return rv + + assert func() == 'finally' + + +def test_jump_out_of_finally_block_w_except_and_live(): + @with_goto + def func(): + for i in range(3): + for j in range(3): + try: + rv = None + except Exception: + rv = 'except' + finally: + rv = 'finally' + goto .end + rv = 'end' + label .end + return i, j, rv + + assert func() == (2, 0, 'finally') + + +def test_jump_out_of_finally_block(): + @with_goto + def func(): + try: + rv = None + finally: + rv = 'finally' + goto .end + rv = 'end' + label .end + return rv + + assert func() == 'finally' + + +def test_jump_out_of_finally_block_and_live(): + @with_goto + def func(): + for i in range(3): + for j in range(3): + try: + rv = None + finally: + rv = 'finally' + goto .end + rv = 'end' + label .end + return i, j, rv + + assert func() == (2, 0, 'finally') + + +def test_jump_out_of_try_in_except_in_finally_and_live(): + @with_goto + def func(): + for i in range(3): + for j in range(3): + try: + rv = None + finally: + rv = 'finally' + try: + rv = 1 / 0 + except Exception: + rv = 'except' + try: + rv = 'try' + goto .end + except Exception: + rv = 'except2' + finally: + rv = 'finally2' + rv = 'end' + label .end + return i, j, rv + + assert func() == (2, 0, 'try') + + def test_jump_to_unknown_label(): def func(): goto .unknown @@ -182,3 +677,15 @@ def func(): assert newfunc is not func assert newfunc.foo == 'bar' + + +def test_code_is_not_copy(): + def outer_func(): + @with_goto + def inner_func(): + goto .test + label .test + return inner_func + + assert outer_func() is not outer_func() + assert outer_func().__code__ is outer_func().__code__ 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]