PR for #155 Improve Python Autofill

source/analysis/top_gdp/topgdp_value.py

@@ -1,19 +1,15 @@
 from pathlib import Path
 import pandas as pd
-from source.lib.JMSLab.autofill import GenerateAutofillMacros
+from source.lib.JMSLab.autofill import AutoFill
 
 def Main():
     instub  = Path("output/derived/wb_clean")
     outstub = Path("output/analysis/top_gdp")
 
     df = pd.read_csv(instub / "gdp_education.csv")
-    TopGDPValue = df['GDP_2010'].sort_values(ascending=False).iloc[0]
+    top_gdp_value = df['GDP_2010'].sort_values(ascending=False).iloc[0]
 
-    GenerateAutofillMacros(
-        ["TopGDPValue"],
-        "{:,.0f}",
-        outstub / "top_gdp.tex"
-    )
+    AutoFill({"TopGDPValue": top_gdp_value}, outstub / "top_gdp.tex", "{:,.0f}")
 
 
 if __name__ == '__main__':

source/lib/JMSLab/autofill.py

@@ -1,56 +1,60 @@
 import inspect
-
-def Autofill(var, format = "{}", namespace = None):
-    newcommand = f"\\newcommand{{{{\\{{}}}}}}{{{{{format}}}}}\n"
-
-    # Look for var in parent frames
-    if namespace is None:
-        parent = inspect.currentframe().f_back
-        while var not in parent.f_locals.keys() and parent.f_back:
-            parent = parent.f_back
-
-        if var not in parent.f_locals.keys() or parent is None:
-            raise Exception(f"Autofill: Variable '{var}' not found")
-
-        namespace = parent.f_locals
-
-    commandname, content = var, namespace[var]
-
-    return newcommand.format(commandname, content)
-
-def GenerateAutofillMacros(autofill_lists, autofill_formats = "{:.2f}", autofill_outfile = "autofill.tex"):
-    '''
-
-    Parameters
-    ----------
-    autofill_lists : TYPE - list
-    autofill_formats : TYPE - str or list
-    autofill_outfile : TYPE - str
-
-    Returns
-    -------
-    .tex file
-
-    '''
-    if type(autofill_lists) != list:
-            raise Exception("Argument 'autofill_lists' must be list")
-
-    nested_list = any(isinstance(i, list) for i in autofill_lists)
-
-    if ((nested_list and type(autofill_formats) is str) 
-        or (not nested_list and type(autofill_formats) is list)):
-        raise Exception("Arguments 'autofill_lists' and 'autofill_formats' are incompatible")
-
-    autofill_file = open(autofill_outfile, 'w')
-
-    if type(autofill_formats) == str:
-        output_macros = ''.join(Autofill(autofill_var, autofill_formats) for autofill_var in autofill_lists)
+from pathlib import Path
+from typing import Any, Literal
+
+
+def AutoFill(
+    macros: dict[str, Any] | list[str],
+    outfile: str | Path,
+    format: str | list[str | None] | None = None,
+    append: bool = False,
+    mode: Literal["math", "text"] = "math",
+) -> None:
+    if isinstance(macros, dict):
+        macro_values = macros
+    elif isinstance(macros, list):
+        caller_frame = inspect.currentframe().f_back # type: ignore
+        macro_values = {}
+        for var in macros:
+            frame = caller_frame
+            found = False
+            while frame is not None:
+                found, value = _LookupVar(var, frame)
+                if found:
+                    break
+                frame = frame.f_back
+            if not found:
+                raise Exception(f"AutoFill: Variable '{var}' not found")
+            macro_values[var] = value
     else:
-        autofill_macros = []
-        for autofill_list, autofill_format in zip(autofill_lists, autofill_formats):
-            autofill_macros.append(''.join(Autofill(autofill_var, autofill_format) for autofill_var in autofill_list))
-        output_macros = ''.join(autofill_macros)
-
-    autofill_file.write(output_macros)
-    autofill_file.close()
+        raise Exception("Argument 'macros' must be a dict or list")
 
+    if isinstance(format, list):
+        if len(format) != len(macro_values):
+            raise Exception("AutoFill: 'format' list length must match number of macros")
+        formats = format
+    else:
+        formats = [format] * len(macro_values)
+
+    output = "".join(
+        _FormatMacro(name, value, fmt, mode)
+        for (name, value), fmt in zip(macro_values.items(), formats)
+    )
+    open_mode = "a" if append else "w"
+    with open(outfile, open_mode) as f:
+        f.write(output)
+
+
+def _LookupVar(name: str, frame) -> tuple[bool, Any]:
+    if name in frame.f_locals:
+        return True, frame.f_locals[name]
+    if name in frame.f_globals:
+        return True, frame.f_globals[name]
+    return False, None
+
+
+def _FormatMacro(name: str, value: Any, format: str | None, mode: str) -> str:
+    formatted = format.format(value) if format is not None else str(value)
+    if mode == "text":
+        return f"\\newcommand{{\\{name}}}{{\\textnormal{{{formatted}}}}}\n"
+    return f"\\newcommand{{\\{name}}}{{{formatted}}}\n"

source/lib/JMSLab/tests/test_autofill.py

@@ -1,69 +1,106 @@
 from unittest import main, TestCase
-from os.path import exists
+import os
+import tempfile
+import shutil
+from pathlib import Path
 
-import tempfile, shutil 
+from ..autofill import AutoFill
 
-from ..autofill import GenerateAutofillMacros
 
-class Test(TestCase):
+class TestOutputContent(TestCase):
 
     def setUp(self):
-        self.tempdir = tempfile.mkdtemp() 
-        self.outfile = self.tempdir + r"output_macros.tex"
-        return None
+        self.tempdir = tempfile.mkdtemp()
+        self.outfile = os.path.join(self.tempdir, "output_macros.tex")
 
     def tearDown(self):
         shutil.rmtree(self.tempdir)
-        return
 
-    def test_file_exists(self):
-        Epsilon = -1.19
-
-        GenerateAutofillMacros(["Epsilon"], autofill_outfile = self.outfile)
-        self.assertTrue(exists(self.outfile))
-        return None
+    def test_dict_output(self):
+        for outfile in [self.outfile, Path(self.outfile)]:
+            AutoFill({"NegativePi": -3.1415, "Pi": 3.1415}, outfile, "{:.2f}")
+            content = Path(outfile).read_text()
+            self.assertEqual(content, "\\newcommand{\\NegativePi}{-3.14}\n\\newcommand{\\Pi}{3.14}\n")
 
-    def test_exception(self):
-        with self.assertRaises(Exception) as context:
-            GenerateAutofillMacros(["Epsilon"], autofill_outfile =  self.outfile)
-
-        self.assertTrue("Autofill: Variable 'Epsilon' not found" in str(context.exception))
-        return None
-
-    def test_output(self):
-        Epsilon = - 1.19
-        MarginalCost = 2.59
-
-        GenerateAutofillMacros([["Epsilon"], ["MarginalCost"]], 
-                               autofill_formats = ["{:.2f}", "\\textnormal{{{:.2f}}}"], 
-                               autofill_outfile = self.outfile)
-        tex_file = open(self.outfile, 'r')
-        content = tex_file.read()
-        self.assertEqual(content, "\\newcommand{\\Epsilon}{-1.19}\n\\newcommand{\\MarginalCost}{\\textnormal{2.59}}\n")
-        tex_file.close()
-        return None
-
-    def test_list_format(self):
-        Epsilon = - 1.19
-        MarginalCost = 2.59
+    def test_list_output(self):
+        NegativePi = -3.1415
+        Pi = 3.1415
+        AutoFill(["NegativePi", "Pi"], self.outfile, "{:.2f}")
+        content = Path(self.outfile).read_text()
+        self.assertEqual(content, "\\newcommand{\\NegativePi}{-3.14}\n\\newcommand{\\Pi}{3.14}\n")
 
-        with self.assertRaises(Exception) as context:
-            GenerateAutofillMacros("Epsilon", autofill_outfile = self.outfile)
+    def test_none_format_numeric(self):
+        AutoFill({"NegativePi": -3.1415, "Pi": 3.1415}, self.outfile)
+        content = Path(self.outfile).read_text()
+        self.assertEqual(content, "\\newcommand{\\NegativePi}{-3.1415}\n\\newcommand{\\Pi}{3.1415}\n")
+
+    def test_none_format_text(self):
+        AutoFill({"SampleStart": "January 2010"}, self.outfile)
+        content = Path(self.outfile).read_text()
+        self.assertEqual(content, "\\newcommand{\\SampleStart}{January 2010}\n")
 
-        self.assertTrue("Argument 'autofill_lists' must be list" in str(context.exception))
 
-        with self.assertRaises(Exception) as context:
-            GenerateAutofillMacros(["Epsilon"], autofill_formats = ["{:.2f}", "\\textnormal{{{:.2f}}}"], 
-                                   autofill_outfile = self.outfile)
+class TestModeAndFormat(TestCase):
+
+    def setUp(self):
+        self.tempdir = tempfile.mkdtemp()
+        self.outfile = os.path.join(self.tempdir, "output_macros.tex")
+
+    def tearDown(self):
+        shutil.rmtree(self.tempdir)
 
-        self.assertTrue("Arguments 'autofill_lists' and 'autofill_formats' are incompatible" in str(context.exception))
+    def test_text_mode(self):
+        AutoFill({"Pi": 3.1415}, self.outfile, "{:.2f}", mode="text")
+        content = Path(self.outfile).read_text()
+        self.assertEqual(content, "\\newcommand{\\Pi}{\\textnormal{3.14}}\n")
 
+    def test_format_list(self):
+        AutoFill({"NegativePi": -3.1415, "IntegerPi": 3.1415}, self.outfile, ["{:.2f}", "{:.0f}"])
+        content = Path(self.outfile).read_text()
+        self.assertEqual(content, "\\newcommand{\\NegativePi}{-3.14}\n\\newcommand{\\IntegerPi}{3}\n")
+
+
+class TestFileBehavior(TestCase):
+
+    def setUp(self):
+        self.tempdir = tempfile.mkdtemp()
+        self.outfile = os.path.join(self.tempdir, "output_macros.tex")
+
+    def tearDown(self):
+        shutil.rmtree(self.tempdir)
+
+    def test_append(self):
+        AutoFill({"NegativePi": -3.1415}, self.outfile, "{:.2f}")
+        AutoFill({"Pi": 3.1415}, self.outfile, "{:.2f}", append=True)
+        content = Path(self.outfile).read_text()
+        self.assertEqual(content, "\\newcommand{\\NegativePi}{-3.14}\n\\newcommand{\\Pi}{3.14}\n")
+
+    def test_overwrite_without_append(self):
+        AutoFill({"NegativePi": -3.1415}, self.outfile, "{:.2f}")
+        AutoFill({"Pi": 3.1415}, self.outfile, "{:.2f}")
+        content = Path(self.outfile).read_text()
+        self.assertEqual(content, "\\newcommand{\\Pi}{3.14}\n")
+
+
+class TestErrors(TestCase):
+
+    def setUp(self):
+        self.tempdir = tempfile.mkdtemp()
+        self.outfile = os.path.join(self.tempdir, "output_macros.tex")
+
+    def tearDown(self):
+        shutil.rmtree(self.tempdir)
+
+    def test_list_variable_not_found(self):
         with self.assertRaises(Exception) as context:
-            GenerateAutofillMacros([["Epsilon"], ["MarginalCost"]], autofill_outfile = self.outfile)
+            AutoFill(["NegativePi"], outfile=self.outfile)
+        self.assertIn("AutoFill: Variable 'NegativePi' not found", str(context.exception))
+
+    def test_invalid_macros_type(self):
+        with self.assertRaises(Exception) as context:
+            AutoFill("NegativePi", outfile=self.outfile)
+        self.assertIn("Argument 'macros' must be a dict or list", str(context.exception))
+
 
-        self.assertTrue("Arguments 'autofill_lists' and 'autofill_formats' are incompatible" in str(context.exception))
-        return None
-
 if __name__ == '__main__':
     main()
-