From abcd10628a91723ccb2cb08ed7b275583112a496 Mon Sep 17 00:00:00 2001
From: rocky <rb@dustyfeet.com>
Date: Sun, 20 Nov 2016 21:11:38 -0500
Subject: [PATCH] Add --linemaps: shows line number correspondences

---
 uncompyle6/linenumbers.py | 59 +++++++++++++++++++++++++++++++++++----
 uncompyle6/main.py        | 10 ++++---
 2 files changed, 59 insertions(+), 10 deletions(-)

diff --git a/uncompyle6/linenumbers.py b/uncompyle6/linenumbers.py
index 933ca5af..dbf30342 100644
--- a/uncompyle6/linenumbers.py
+++ b/uncompyle6/linenumbers.py
@@ -1,14 +1,61 @@
+from collections import deque, namedtuple
+
+from xdis.code import iscode
 from xdis.load import load_file, load_module
-from xdis.bytecode import findlinestarts, offset2line
+from xdis.main import get_opcode
+from xdis.bytecode import Bytecode, findlinestarts, offset2line
 
 def line_number_mapping(pyc_filename, src_filename):
-    (version, timestamp, magic_int, code_obj1, is_pypy,
+    (version, timestamp, magic_int, code1, is_pypy,
      source_size) = load_module(pyc_filename)
     try:
-        code_obj2 = load_file(src_filename)
+        code2 = load_file(src_filename)
     except SyntaxError as e:
         return str(e)
 
-    linestarts_orig = findlinestarts(code_obj1)
-    linestarts_uncompiled = list(findlinestarts(code_obj2))
-    return [[line, offset2line(offset, linestarts_uncompiled)] for offset, line in linestarts_orig]
+    queue = deque([code1, code2])
+
+    mappings = []
+
+    opc = get_opcode(version, is_pypy)
+    number_loop(queue, mappings, opc)
+    return sorted(mappings, key=lambda x: x[1])
+
+
+def number_loop(queue, mappings, opc):
+    while len(queue) > 0:
+        code1 = queue.popleft()
+        code2 = queue.popleft()
+        assert code1.co_name == code2.co_name
+        linestarts_orig = findlinestarts(code1)
+        linestarts_uncompiled = list(findlinestarts(code2))
+        mappings += [[line, offset2line(offset, linestarts_uncompiled)] for offset, line in linestarts_orig]
+        bytecode1 = Bytecode(code1, opc)
+        bytecode2 = Bytecode(code2, opc)
+        instr2s = bytecode2.get_instructions(code2)
+        seen = set([code1.co_name])
+        for instr in bytecode1.get_instructions(code1):
+            next_code1 = None
+            if iscode(instr.argval):
+                next_code1 = instr.argval
+            if next_code1:
+                next_code2 = None
+                while not next_code2:
+                    try:
+                        instr2 = next(instr2s)
+                        if iscode(instr2.argval):
+                            next_code2 = instr2.argval
+                            pass
+                    except StopIteration:
+                        break
+                    pass
+                if next_code2:
+                    assert next_code1.co_name == next_code2.co_name
+                    if next_code1.co_name not in seen:
+                        seen.add(next_code1.co_name)
+                        queue.append(next_code1)
+                        queue.append(next_code2)
+                        pass
+                    pass
+            pass
+        pass
diff --git a/uncompyle6/main.py b/uncompyle6/main.py
index 11102507..9c3de62a 100644
--- a/uncompyle6/main.py
+++ b/uncompyle6/main.py
@@ -126,7 +126,7 @@ def main(in_base, out_base, files, codes, outfile=None,
                 prefix = os.path.basename(filename)
                 if prefix.endswith('.py'):
                     prefix = prefix[:-len('.py')]
-                junk, outfile = tempfile.mkstemp(suffix=".pyc",
+                junk, outfile = tempfile.mkstemp(suffix=".py",
                                              prefix=prefix)
                 # Unbuffer output
                 sys.stdout = os.fdopen(sys.stdout.fileno(), 'w', 0)
@@ -168,9 +168,11 @@ def main(in_base, out_base, files, codes, outfile=None,
             if outfile:
                 if do_linemaps:
                     mapping = line_number_mapping(infile, outfile)
-                    print("\n\n## Line number correspondences", infile)
-                    for m in mapping:
-                        print("## %s" % m, infile)
+                    outstream.write("\n\n## Line number correspondences\n")
+                    import pprint
+                    s = pprint.pformat(mapping, indent=2, width=80)
+                    s2 = '##' + '\n##'.join(s.split("\n")) + "\n"
+                    outstream.write(s2)
                 outstream.close()
 
                 if do_verify: