Refactor code_gen.py: extract compile_terms and get_latex_equations_str to eliminate duplicate compiler pipeline logic

code/code_gen/code_gen.py

@@ -133,6 +133,36 @@ class VersionInfo(object):
 
 
 
+def compile_terms(terms_raw):
+    """Parses DSL terms spec and runs it through compiler stages 1-4.
+
+    Returns:
+        ir4 (Stage4): The compiled Stage 4 IR object.
+    """
+    parser = Lark(calc_grammar,
+                  parser='lalr',
+                  lexer_callbacks={
+                      'ESCAPED_STRING': tok_to_str,
+                      'INT': tok_to_int,
+                      'BOOL': tok_to_bool
+                  })
+    tree = parser.parse(terms_raw)
+    ir1 = Stage1(tree)
+    ir2 = Stage2(ir1)
+    ir3 = Stage3(ir2)
+    ir4 = Stage4(ir3)
+    return ir4
+
+
+def get_latex_equations_str(ir4):
+    """Generates the LaTeX equations dictionary string from Stage4 IR."""
+    latex_lines = ["{"]
+    for avg in ir4.averaged.values():
+        latex_lines.append('  "{}" : r"${}$",'.format(avg.name, avg.latex))
+    latex_lines.append("}")
+    return "\n".join(latex_lines)
+
+
 def test(terms_raw, report=False, latex=False):
     """Compiles the post-processing term specifications from DSL into targets.
 
@@ -146,39 +176,14 @@ def test(terms_raw, report=False, latex=False):
         latex (bool): If True, prints LaTeX equations of the averaged terms to stdout.
     """
 
-    parser = Lark(calc_grammar,
-            parser='lalr',
-            lexer_callbacks =
-            {
-            'ESCAPED_STRING': tok_to_str,
-            'INT': tok_to_int,
-            'BOOL': tok_to_bool
-            }
-            )
-
-    tree = parser.parse(terms_raw)
-
-    ir1 = Stage1(tree)
-
-    ir2 = Stage2(ir1)
-
-    ir3 = Stage3(ir2)
-
-    ir4 = Stage4(ir3)
-
-    cdict = ir4.print_program()
+    ir4 = compile_terms(terms_raw)
 
     if report:
         ir4.save_ir()
     elif latex:
-        print("{")
-        for avg in ir4.averaged.values():
-            print('"{}" : '.format(avg.name))
-            # print r"\begin{equation}"
-            print('r"${}$" ,'.format(avg.latex))
-            # print r"\end{equation}"
-        print("}")
+        print(get_latex_equations_str(ir4))
     else:
+        cdict = ir4.print_program()
         print(Template(mod_form).render(**cdict))
 
 
@@ -198,24 +203,8 @@ def generate_build_info_module(terms_raw):
     vinfo = VersionInfo(terms_raw)
     build_info_str = vinfo.fparams()
 
-    parser = Lark(calc_grammar,
-                  parser='lalr',
-                  lexer_callbacks={
-                      'ESCAPED_STRING': tok_to_str,
-                      'INT': tok_to_int,
-                      'BOOL': tok_to_bool
-                  })
-    tree = parser.parse(terms_raw)
-    ir1 = Stage1(tree)
-    ir2 = Stage2(ir1)
-    ir3 = Stage3(ir2)
-    ir4 = Stage4(ir3)
-
-    latex_lines = ["{"]
-    for avg in ir4.averaged.values():
-        latex_lines.append(f'  "{avg.name}" : r"${avg.latex}$",')
-    latex_lines.append("}")
-    latex_str = "\n".join(latex_lines)
+    ir4 = compile_terms(terms_raw)
+    latex_str = get_latex_equations_str(ir4)
 
     fortran_build_info = escape_fortran_string(build_info_str)
     fortran_latex = escape_fortran_string(latex_str)