Various minor crash/etc fixes

.travis.yml

@@ -7,6 +7,7 @@ python:
   - "3.5"
   - "3.6"
   - "3.7"
+  - "3.8"
   - "pypy"
   - "pypy3"
 

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/)
 

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('<EXCEPT>', previous=exit_block)
+            elif exit_name == 'SETUP_FINALLY':
+                block_stack.push('<FINALLY>', 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] == '<FINALLY>':
+                # in 3.8, only finally blocks are supported, so we must
+                # determine whether it's except/finally ourselves
+                block_stack.replace(top_block,
+                                    ('<EXCEPT>',) + top_block[1:])
+                _, _, setup_block, _ = top_block
+                block_stack.replace(setup_block,
+                                    ('SETUP_EXCEPT',) + setup_block[1:])
+            block_stack.pop_of_type('<EXCEPT>')
+        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] == '<FINALLY>':
+                block_stack.pop_of_type('<FINALLY>')
+        elif opname1 in ('WITH_CLEANUP', 'WITH_CLEANUP_START'):
+            if _BYTECODE.has_setup_with:
+                # temporary block to match END_FINALLY
+                block_stack.push('<FINALLY>')
+            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 == '<EXCEPT>':
+                ops.append('POP_EXCEPT')
+            elif block == '<FINALLY>':
+                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):