Move print to verbose flag

PolicyEngine · Dec 17, 2024 · 597574a · 597574a
1 parent a43ac8c
commit 597574a
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 8 deletions.
diff --git a/policyengine/outputs/macro/comparison/budget/breakdown.py b/policyengine/outputs/macro/comparison/budget/breakdown.py
@@ -8,6 +8,7 @@ def breakdown(
     provisions: List[dict] | None = None,
     provision_names: List[str] | None = None,
     count_years: int = 5,
+    verbose: bool = False,
 ) -> pd.DataFrame:
     """Calculate the budgetary impact of the reform by provision.
 
@@ -16,6 +17,7 @@ def breakdown(
         provisions (List[dict], optional): The provisions to include in the breakdown. Defaults to None.
         provision_names (List[str], optional): The names of the provisions to include in the breakdown. Defaults to None.
         count_years (int, optional): The number of years to include in the breakdown. Defaults to 5.
+        verbose (bool, optional): Whether to print the budgetary impacts. Defaults to False.
 
     Returns:
         pd.DataFrame: A DataFrame (long) containing the budgetary impact of the reform by provision.
@@ -54,9 +56,10 @@ def breakdown(
             year_items.append(year)
             budget_items.append(difference)
 
-            print(
-                f"Year: {year}, Provision: {key_focus}, Budgetary impact: {round(difference/1e6)}"
-            )
+            if verbose:
+                print(
+                    f"Year: {year}, Provision: {key_focus}, Budgetary impact: {round(difference/1e9)}"
+                )
 
     return pd.DataFrame(
         {

diff --git a/policyengine/utils/budget.py b/policyengine/utils/budget.py
@@ -8,8 +8,21 @@ def budget_breakdown(
     years: list,
     baseline: dict,
     options: dict,
+    verbose: bool = False,
 ):
-    """Create a table of per-provision budgetary impacts for a given reform."""
+    """Create a table of per-provision budgetary impacts for a given reform.
+
+    Args:
+        subreforms (list): A list of reform dictionaries.
+        subreform_names (list): A list of names for the provisions in the reform.
+        years (list): A list of years to include in the breakdown.
+        baseline (dict): The baseline reform.
+        options (dict): The simulation options.
+        verbose (bool, optional): Whether to print the budgetary impacts. Defaults to False.
+    
+    Returns:
+        pd.DataFrame: A DataFrame containing the budgetary impact of the reform by provision.
+    """
     subreform_items = []
     year_items = []
     budget_items = []
@@ -35,10 +48,10 @@ def budget_breakdown(
             subreform_items.append(key_focus)
             year_items.append(year)
             budget_items.append(difference)
-
-            print(
-                f"Year: {year}, Provision: {key_focus}, Budgetary impact: {round(difference/1e6)}"
-            )
+            if verbose:
+                print(
+                    f"Year: {year}, Provision: {key_focus}, Budgetary impact: {round(difference/1e9)}"
+                )
 
     df = pd.DataFrame(
         {