From 9903692811655a2d47636a1cd9917eda7f778cf4 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Tue, 26 Jun 2018 03:56:45 +0100
Subject: [PATCH] master: update scan_code_imports to cope with wordcode

Constant-sized opcodes were introduced as an optimization in Python 3.6.
See https://bugs.python.org/issue26647
---
 mitogen/master.py    | 41 ++++++++++++++++++++++++++++-------------
 tests/master_test.py | 19 +++++++++++++------
 2 files changed, 41 insertions(+), 19 deletions(-)

diff --git a/mitogen/master.py b/mitogen/master.py
index 46b1e84b..159bf427 100644
--- a/mitogen/master.py
+++ b/mitogen/master.py
@@ -105,6 +105,33 @@ LOAD_CONST = dis.opname.index('LOAD_CONST')
 IMPORT_NAME = dis.opname.index('IMPORT_NAME')
 
 
+if sys.version_info < (3, 0):
+    def iter_opcodes(co):
+        # Yield `(op, oparg)` tuples from the code object `co`.
+        ordit = imap(ord, co.co_code)
+        nextb = ordit.next
+        return ((c, (None
+                     if c < dis.HAVE_ARGUMENT else
+                     (nextb() | (nextb() << 8))))
+                    for c in ordit)
+elif sys.version_info < (3, 6):
+    def iter_opcodes(co):
+        # Yield `(op, oparg)` tuples from the code object `co`.
+        ordit = iter(co.co_code)
+        nextb = ordit.__next__
+        return ((c, (None
+                     if c < dis.HAVE_ARGUMENT else
+                     (nextb() | (nextb() << 8))))
+                    for c in ordit)
+else:
+    def iter_opcodes(co):
+        # Yield `(op, oparg)` tuples from the code object `co`.
+        ordit = iter(co.co_code)
+        nextb = ordit.__next__
+        # https://github.com/abarnert/cpython/blob/c095a32f/Python/wordcode.md
+        return ((c, nextb()) for c in ordit)
+
+
 def scan_code_imports(co):
     """Given a code object `co`, scan its bytecode yielding any
     ``IMPORT_NAME`` and associated prior ``LOAD_CONST`` instructions
@@ -120,19 +147,7 @@ def scan_code_imports(co):
         * `namelist`: for `ImportFrom`, the list of names to be imported from
           `modname`.
     """
-    # Yield `(op, oparg)` tuples from the code object `co`.
-    if mitogen.core.PY3:
-        ordit = iter(co.co_code)
-        nextb = ordit.__next__
-    else:
-        ordit = imap(ord, co.co_code)
-        nextb = ordit.next
-
-    opit = ((c, (None
-                 if c < dis.HAVE_ARGUMENT else
-                 (nextb() | (nextb() << 8))))
-            for c in ordit)
-
+    opit = iter_opcodes(co)
     opit, opit2, opit3 = itertools.tee(opit, 3)
     try:
         next(opit2)
diff --git a/tests/master_test.py b/tests/master_test.py
index cf16d6c5..19a9b414 100644
--- a/tests/master_test.py
+++ b/tests/master_test.py
@@ -9,15 +9,22 @@ import mitogen.master
 class ScanCodeImportsTest(unittest2.TestCase):
     func = staticmethod(mitogen.master.scan_code_imports)
 
+    if mitogen.core.PY3:
+        level = 0
+    else:
+        level = -1
+
+    SIMPLE_EXPECT = [
+        (level, 'inspect', ()),
+        (level, 'unittest2', ()),
+        (level, 'testlib', ()),
+        (level, 'mitogen.master', ()),
+    ]
+
     def test_simple(self):
         source_path = inspect.getsourcefile(ScanCodeImportsTest)
         co = compile(open(source_path).read(), source_path, 'exec')
-        self.assertEquals(list(self.func(co)), [
-            (-1, 'inspect', ()),
-            (-1, 'unittest2', ()),
-            (-1, 'testlib', ()),
-            (-1, 'mitogen.master', ()),
-        ])
+        self.assertEquals(list(self.func(co)), self.SIMPLE_EXPECT)
 
 
 if __name__ == '__main__':