From bd97de30ecef4d422e7a0dc3bc8d65062eec8212 Mon Sep 17 00:00:00 2001
From: Hayden <64056131+hay-kot@users.noreply.github.com>
Date: Fri, 4 Jul 2025 12:07:56 -0500
Subject: [PATCH] fix mangles outputs

---
 dev/code-generation/utils/template.py | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 4 deletions(-)

diff --git a/dev/code-generation/utils/template.py b/dev/code-generation/utils/template.py
index a2dd28cb7..32ecf9c47 100644
--- a/dev/code-generation/utils/template.py
+++ b/dev/code-generation/utils/template.py
@@ -1,5 +1,4 @@
 import logging
-import re
 import subprocess
 from dataclasses import dataclass
 from pathlib import Path
@@ -47,12 +46,21 @@ class CodeSlicer:
 
     def push_line(self, string: str) -> None:
         self._next_line = self._next_line or self.start + 1
-        self.text.insert(self._next_line, self.indentation or "" + string + "\n")
+        self.text.insert(self._next_line, (self.indentation or "") + string + "\n")
         self._next_line += 1
 
 
-def get_indentation_of_string(line: str, comment_char: str = "//|#") -> str:
-    return re.sub(rf"{comment_char}.*", "", line).removesuffix("\n")
+def get_indentation_of_string(line: str) -> str:
+    # Extract everything before the comment
+    if "//" in line:
+        indentation = line.split("//")[0]
+    elif "#" in line:
+        indentation = line.split("#")[0]
+    else:
+        indentation = line
+
+    # Keep only the whitespace, remove any non-whitespace characters
+    return "".join(c for c in indentation if c.isspace())
 
 
 def find_start_end(file_text: list[str], gen_id: str) -> tuple[int, int, str | None]: