bytecode.py

import dis
import io
import re
import sys
import textwrap
from op_codes import op_codes_37


def fix_spacing(unformatted_decompiled_python):
    decompiled_python = re.sub(r"[ \t]+", " ", unformatted_decompiled_python)
    decompiled_python = re.sub(r"\n ", "\n", decompiled_python)
    return decompiled_python


def get_byte_code(python_code):
    code_object = compile(python_code, "f.py", "exec")

    original_stdout = sys.stdout
    string_output = io.StringIO()
    sys.stdout = string_output
    dis.dis(code_object)
    sys.stdout = original_stdout
    disassembled_pyc = string_output.getvalue()
    return fix_spacing(textwrap.dedent(str(disassembled_pyc)).strip())


def extract_text_between_parentheses(input_string):
    start_index = input_string.find("(")
    end_index = input_string.rfind(")")

    if start_index != -1 and end_index != -1 and start_index < end_index:
        return input_string[start_index + 1 : end_index]
    else:
        return ""


def reformat_bytecode_line(input_string):
    if input_string.startswith("Disassembly of <"):
        return None

    for code in op_codes_37:
        if code in input_string:
            reformatted_string = extract_text_between_parentheses(input_string)
            return f"{code} {reformatted_string}".strip()

    return input_string


def format_bytecode(bytecode: str):
    formatted_strings = []
    for line in bytecode.split("\n"):
        reformatted_string = reformat_bytecode_line(line)
        if reformatted_string is not None:
            formatted_strings.append(reformatted_string)
    return "\n".join(formatted_strings).strip()


if __name__ == "__main__":
    file_path = sys.argv[1]
    print(f"Converting python source code to byte code {file_path}")
    with open(file_path, "r") as file:
        file_contents = file.read()
        print(file_contents)
        byte_code = get_byte_code(file_contents)

        print(format_bytecode(bytecode=byte_code))