pypechain: code reuse with format_code and write_code

lib/pypechain/pypechain/run_pypechain.py

@@ -6,6 +6,8 @@
 from dataclasses import asdict
 from pathlib import Path
 
+import black
+from black.mode import Mode  # pylint: disable=no-name-in-module
 from jinja2 import Template
 from pypechain.utilities.abi import (
     get_abi_items,
@@ -21,6 +23,23 @@
 from web3.types import ABIFunction
 
 
+def format_code(code):
+    """Format code with Black on default settings."""
+    while "\n\n" in code:
+        code = code.replace("\n\n", "\n")  # remove all whitespace and let Black sort it out
+    code = code.replace(", )", ")")  # remove trailing comma, it's weird
+    try:
+        return black.format_file_contents(code, fast=False, mode=Mode())
+    except ValueError as exc:
+        raise ValueError(f"cannot format with Black\n code:\n{code}") from exc
+
+
+def write_code(path, code):
+    """save to specified path the provided code."""
+    with open(path, "w", encoding="utf-8") as output_file:
+        output_file.write(code)
+
+
 def main(abi_file_path: str, output_dir: str) -> None:
     """Generates class files for a given abi.
 
@@ -37,6 +56,7 @@ def main(abi_file_path: str, output_dir: str) -> None:
     file_path = Path(abi_file_path)
     filename = file_path.name
     contract_name = os.path.splitext(filename)[0]
+    contract_path = Path(output_dir).joinpath(f"{contract_name}")
 
     # grab the templates
     contract_template, types_template = setup_templates()
@@ -48,13 +68,13 @@ def main(abi_file_path: str, output_dir: str) -> None:
     # TODO: Add more features:
     # TODO:  events
 
-    # Write the renders to a file
-    types_output_file_path = Path(output_dir).joinpath(f"{contract_name}Types.py")
-    contract_output_file_path = Path(output_dir).joinpath(f"{contract_name}Contract.py")
-    with open(contract_output_file_path, "w", encoding="utf-8") as output_file:
-        output_file.write(rendered_contract_code)
-    with open(types_output_file_path, "w", encoding="utf-8") as output_file:
-        output_file.write(rendered_types_code)
+    # Format the generated code using Black
+    formatted_contract_code = format_code(rendered_contract_code)
+    formatted_types_code = format_code(rendered_types_code)
+
+    # Write the code to file
+    write_code(f"{contract_path}Contract.py", formatted_contract_code)
+    write_code(f"{contract_path}Types.py", formatted_types_code)
 
 
 def render_contract_file(contract_name: str, contract_template: Template, abi_file_path: Path) -> str: