From 48bdc16787b8f4554db3540d0bf3989e52ce3e27 Mon Sep 17 00:00:00 2001
From: Andrew Helwer <2n8rn1w1f@mozmail.com>
Date: Wed, 27 Mar 2024 16:00:32 -0400
Subject: [PATCH] Added documentation

Signed-off-by: Andrew Helwer <2n8rn1w1f@mozmail.com>
---
 .github/scripts/unicode_number_set_shim.py | 25 +++++++++++++++++-----
 1 file changed, 20 insertions(+), 5 deletions(-)

diff --git a/.github/scripts/unicode_number_set_shim.py b/.github/scripts/unicode_number_set_shim.py
index 599cace0..846ae61c 100644
--- a/.github/scripts/unicode_number_set_shim.py
+++ b/.github/scripts/unicode_number_set_shim.py
@@ -41,16 +41,26 @@ def build_insertion_point_query(language):
     return language.query('((header_line) name: (identifier) (header_line)) @header (extends) @extends')
 
 def get_required_defs(tree, query):
+    """
+    Gets Nat/Int/Real definitions that are used in the module.
+    """
     return set([name for _, name in query.captures(tree.root_node)])
 
 def get_def_text(required_defs):
-    defs = '\n'
-    for shim in shims:
-        if shim.capture in required_defs:
-            defs += f'LOCAL {shim.unicode} ≜ {shim.ascii}\n'
-    return defs
+    """
+    Builds the definitions to insert into the module.
+    """
+    return '\n' + '\n'.join(
+        f'LOCAL {shim.unicode} ≜ {shim.ascii}'
+        for shim in shims
+        if shim.capture in required_defs
+    )
 
 def get_insertion_point(tree, query):
+    """
+    Find a suitable insertion point in the file: either directly after the
+    header, or directly after the EXTENDS statement if it exists.
+    """
     captures = query.captures(tree.root_node)
     has_extends = any(name for (_, name) in captures if name == 'extends')
     if has_extends:
@@ -61,10 +71,15 @@ def get_insertion_point(tree, query):
         return header.byte_range[1]
 
 def insert_defs(module_path, insertion_point, defs):
+    """
+    Inserts the shim definitions at the given byte offset in the module.
+    """
     def_bytes = bytes(defs, 'utf-8')
     with open(module_path, 'rb+') as module:
         module_bytes = bytearray(module.read())
         module_bytes[insertion_point:insertion_point] = def_bytes
+        module.seek(0)
+        module.truncate()
         module.write(module_bytes)
 
 if __name__ == '__main__':