format

main.py

@@ -11,26 +11,40 @@ from dist import levenshtein_with_wildcard, print_match_summary
 
 description = frontmatter.load("README.md").content
 
+
 def trim(str, n):
     return "\n".join(str.splitlines()[n:])
 
+
 def trim_objdump(str):
     return trim(str, 7)
 
+
 def disassemble_bytes(byte_data, architecture, options):
     with tempfile.NamedTemporaryFile(suffix=".bin", delete=False) as temp_bin_file:
         temp_bin_file.write(byte_data)
         temp_bin_file_name = temp_bin_file.name
 
     disassembly = subprocess.run(
-        ["objdump", "-D", "-b", "binary", "-m", architecture, "-M", options, temp_bin_file_name],
+        [
+            "objdump",
+            "-D",
+            "-b",
+            "binary",
+            "-m",
+            architecture,
+            "-M",
+            options,
+            temp_bin_file_name,
+        ],
         capture_output=True,
-        text=True
+        text=True,
     ).stdout
     disassembly = trim_objdump(disassembly)
 
     return disassembly
 
+
 def compile(compiler, flags, source):
     # Create a temporary file for the C source code
     with tempfile.NamedTemporaryFile(suffix=".c", delete=False) as temp_c_file:
@@ -71,9 +85,7 @@ def compile(compiler, flags, source):
 
     # Disassemble the object file
     disassembly = subprocess.run(
-        ["objdump", "-dr", temp_o_file_name],
-        capture_output=True,
-        text=True
+        ["objdump", "-dr", temp_o_file_name], capture_output=True, text=True
     ).stdout
     disassembly = trim_objdump(disassembly)
 
@@ -86,23 +98,27 @@ def compile(compiler, flags, source):
     # relocs = trim(relocs, 3)
 
     json_relocs = subprocess.run(
-        ["llvm-readobj-19", "--elf-output-style=JSON", "--relocations", temp_o_file_name],
+        [
+            "llvm-readobj-19",
+            "--elf-output-style=JSON",
+            "--relocations",
+            temp_o_file_name,
+        ],
         capture_output=True,
         text=True,
     ).stdout
     json_relocs = json.loads(json_relocs)
     json_relocs = json_relocs[0]["Relocations"]
-    json_relocs = [r["Relocation"] for d in json_relocs for r in d['Relocs']]
+    json_relocs = [r["Relocation"] for d in json_relocs for r in d["Relocs"]]
     # Filter out .text
     json_relocs = [r for r in json_relocs if r["Symbol"]["Name"] != ".text"]
 
-
-
     if result.returncode == 0:
         return json_relocs, compiled_bytes, compile_output, disassembly
     else:
         return None, None, compile_output, disassembly
 
+
 def _reloc_type2size(s):
     match s:
         case "R_X86_64_PC32":
@@ -112,20 +128,31 @@ def _reloc_type2size(s):
         case _:
             assert False, f"Unknown reloc {s}"
 
+
 def _compute_relocs_byte_range(json_relocs):
-        relocs_byte_range = [range(r["Offset"], r["Offset"] + _reloc_type2size(r["Type"]["Name"])) for r in json_relocs]
-        # Flatten relocs_byte_range
-        relocs_byte_range = [i for r in relocs_byte_range for i in r]
-        return relocs_byte_range
+    relocs_byte_range = [
+        range(r["Offset"], r["Offset"] + _reloc_type2size(r["Type"]["Name"]))
+        for r in json_relocs
+    ]
+    # Flatten relocs_byte_range
+    relocs_byte_range = [i for r in relocs_byte_range for i in r]
+    return relocs_byte_range
+
 
 def predict(target_bytes, source, compiler, flags, disasm_arch, disasm_options):
     target_bytes = bytes.fromhex(target_bytes)
-    compiled_relocs, compiled_bytes, compile_output, compiled_disassembly = compile(compiler, flags, source)
+    compiled_relocs, compiled_bytes, compile_output, compiled_disassembly = compile(
+        compiler, flags, source
+    )
     target_disassembly = disassemble_bytes(target_bytes, disasm_arch, disasm_options)
 
     if compiled_bytes is not None:
 
-        reloc_edit_distance, reloc_operations = print_match_summary(target_bytes, compiled_bytes, wildcard_offsets_seq2=_compute_relocs_byte_range(compiled_relocs))
+        reloc_edit_distance, reloc_operations = print_match_summary(
+            target_bytes,
+            compiled_bytes,
+            wildcard_offsets_seq2=_compute_relocs_byte_range(compiled_relocs),
+        )
         print(f"reloc_edit_distance: {reloc_edit_distance}")
         print(f"reloc operations: {reloc_operations}")
 
@@ -138,7 +165,7 @@ def predict(target_bytes, source, compiler, flags, disasm_arch, disasm_options):
             compile_output,
             compiled_disassembly,
             compiled_relocs,
-            target_disassembly
+            target_disassembly,
         )
     else:
         return (
@@ -150,7 +177,7 @@ def predict(target_bytes, source, compiler, flags, disasm_arch, disasm_options):
             compile_output,
             compiled_disassembly,
             compiled_relocs,
-            target_disassembly
+            target_disassembly,
         )
 
 
@@ -172,7 +199,7 @@ def run():
             gr.Textbox(label="Compiler", value="g++"),
             gr.Textbox(label="Compiler Flags", value="-O2"),
             gr.Textbox(label="Architecture (objdump -m)", value="i386"),
-            gr.Textbox(label="Disassembler options (objdump -M)", value="x86-64")
+            gr.Textbox(label="Disassembler options (objdump -M)", value="x86-64"),
         ],
         outputs=[
             gr.Textbox(label="Compiled bytes"),