From a744ac4a58c7b355003f6ffcbeabf6b5d92b0d2a Mon Sep 17 00:00:00 2001
From: Vaughn Kottler <vaughnkottler@gmail.com>
Date: Thu, 15 Aug 2024 22:01:18 -0700
Subject: [PATCH] Initial common method

---
 ifgen/common/python.py              | 21 ++++++++++++++++----
 ifgen/enum/python/__init__.py       |  7 ++-----
 ifgen/generation/python/__init__.py | 30 ++++++++++++++++++++++++++++-
 3 files changed, 48 insertions(+), 10 deletions(-)

diff --git a/ifgen/common/python.py b/ifgen/common/python.py
index f7f7da5..e8bddf7 100644
--- a/ifgen/common/python.py
+++ b/ifgen/common/python.py
@@ -5,13 +5,26 @@
 
 # internal
 from ifgen.generation.interface import GenerateTask
+from ifgen.generation.python import python_function, python_imports
 
 
 def create_python_common(task: GenerateTask) -> None:
     """Create a Python module that aggregates runtime registration tasks."""
 
     with task.boilerplate() as writer:
-        # Add registration function for all generated enumerations. Some base
-        # class for structs that registers the generated enumerations?
-        writer.empty()
-        writer.write("# todo")
+        # Write imports.
+        python_imports(
+            writer, third_party={"runtimepy.enum.registry": ["EnumRegistry"]}
+        )
+
+        with python_function(
+            writer,
+            "register_enums",
+            "Register generated enumerations.",
+            params="registry: EnumRegistry",
+        ):
+            writer.write("# iterate over all custom enums, register each")
+            writer.empty()
+            writer.write("del registry")
+
+        writer.write("# base class for structs that auto registers enums?")
diff --git a/ifgen/enum/python/__init__.py b/ifgen/enum/python/__init__.py
index 6d86807..231bca0 100644
--- a/ifgen/enum/python/__init__.py
+++ b/ifgen/enum/python/__init__.py
@@ -18,14 +18,11 @@ def python_enum_header(task: GenerateTask, writer: IndentedFileWriter) -> None:
         writer, third_party={"runtimepy.enum.registry": ["RuntimeIntEnum"]}
     )
 
-    writer.empty()
-    writer.empty()
-
     with python_class(
         writer,
         task.name,
         task.resolve_description() or "No description.",
         parents=["RuntimeIntEnum"],
+        final_empty=0,
     ):
-        # Add members.
-        writer.write("# todo")
+        writer.write("# register members")
diff --git a/ifgen/generation/python/__init__.py b/ifgen/generation/python/__init__.py
index 9709d26..a1a5785 100644
--- a/ifgen/generation/python/__init__.py
+++ b/ifgen/generation/python/__init__.py
@@ -26,13 +26,17 @@ def write_imports(
 
 
 def python_imports(
-    writer: IndentedFileWriter, third_party: PythonImports = None
+    writer: IndentedFileWriter,
+    third_party: PythonImports = None,
+    final_empty: int = 2,
 ) -> None:
     """Write Python-module imports."""
 
     if third_party:
         write_imports(writer, "third-party", third_party)
 
+    writer.empty(count=final_empty)
+
 
 @contextmanager
 def python_class(
@@ -40,6 +44,7 @@ def python_class(
     name: str,
     docstring: str,
     parents: list[str] = None,
+    final_empty: int = 2,
 ) -> Iterator[None]:
     """Write class definition contents."""
 
@@ -53,3 +58,26 @@ def python_class(
         writer.write(f'"""{docstring}"""')
         writer.empty()
         yield
+
+    writer.empty(count=final_empty)
+
+
+@contextmanager
+def python_function(
+    writer: IndentedFileWriter,
+    name: str,
+    docstring: str,
+    params: str = "",
+    return_type: str = "None",
+    final_empty: int = 2,
+) -> Iterator[None]:
+    """Write a Python function."""
+
+    writer.write(f"def {name}({params}) -> {return_type}:")
+
+    with writer.indented():
+        writer.write(f'"""{docstring}"""')
+        writer.empty()
+        yield
+
+    writer.empty(count=final_empty)