feat: create and enforce blocked file extension list for custom outputs (#377)

* feat: create and enforce blocked file extension list for custom outputs

* test: add coverage for blocked extension checks in plot outputs

* refactor: centralise blocked extension check into utils helper

* style: collapse Records.__init__ docstring to satisfy codacy

---------

Signed-off-by: Hasaan A. <hasaana2005@gmail.com>

Author: ssrhaso

acro/acro.py

@@ -60,12 +60,14 @@ def __init__(self, config: str = "default", suppress: bool = False) -> None:
         Tables.__init__(self, suppress)
         Regression.__init__(self, config)
         self.config: dict[str, Any] = {}
-        self.results: Records = Records()
         self.suppress: bool = suppress
         path: pathlib.Path = pathlib.Path(__file__).with_name(config + ".yaml")
         logger.debug("path: %s", path)
         with open(path, encoding="utf-8") as handle:
             self.config = yaml.load(handle, Loader=yaml.loader.SafeLoader)
+        self.results: Records = Records(
+            blocked_extensions=self.config.get("blocked_extensions", [])
+        )
         logger.info("version: %s", __version__)
         logger.info("config: %s", self.config)
         logger.info("automatic suppression: %s", self.suppress)
@@ -138,7 +140,7 @@ def print_outputs(self) -> str:
         """
         return self.results.print()
 
-    def custom_output(self, filename: str, comment: str = "") -> None:
+    def custom_output(self, filename: str, comment: str = "") -> bool:
         """Add an unsupported output to the results dictionary.
 
         Parameters
@@ -147,8 +149,13 @@ def custom_output(self, filename: str, comment: str = "") -> None:
             The name of the file that will be added to the list of the outputs.
         comment : str
             An optional comment.
+
+        Returns
+        -------
+        bool
+            False if the file extension is blocked, True otherwise.
         """
-        self.results.add_custom(filename, comment)
+        return self.results.add_custom(filename, comment)
 
     def rename_output(self, old: str, new: str) -> None:
         """Rename an output.

acro/acro_stata_parser.py

@@ -376,12 +376,10 @@ def add_custom_output(varlist: list[str]) -> str:
     except IndexError:
         return "syntax error: please pass the name of the output to be added"
 
-    # .gph extension contain data
-    _, file_extension = os.path.splitext(the_output)
-    if file_extension == ".gph":
-        return "Warning: .gph files may not be exported as they contain data."
     comment_str = " ".join(varlist)
-    stata_config.stata_acro.custom_output(the_output, comment_str)
+    if not stata_config.stata_acro.custom_output(the_output, comment_str):
+        _, ext = os.path.splitext(the_output)
+        return f"Warning: {ext} files are not allowed and cannot be exported."
     outcome = f"file {the_output} with comment {comment_str} added to session."
     return outcome
 

acro/acro_tables.py

@@ -578,6 +578,8 @@ def survival_plot(
         summary: str,
     ) -> tuple[Any, str] | None:
         """Create the survival plot according to the status of suppressing."""
+        if utils.is_blocked_extension(filename, self.results.blocked_extensions):
+            return None
         if self.suppress:
             survival_table = _rounded_survival_table(survival_table)
             plot = survival_table.plot(y="rounded_survival_fun", xlim=0, ylim=0)
@@ -703,6 +705,8 @@ def hist(
             The name of the file where the histogram is saved.
         """
         logger.debug("hist()")
+        if utils.is_blocked_extension(filename, self.results.blocked_extensions):
+            return None
         command: str = utils.get_command("hist()", stack())
 
         if isinstance(data, list):  # pragma: no cover
@@ -848,6 +852,8 @@ def pie(
             The path to the saved pie chart file.
         """
         logger.debug("pie()")
+        if utils.is_blocked_extension(filename, self.results.blocked_extensions):
+            return None
         command: str = utils.get_command("pie()", stack())
 
         # COMPUTE PRE-CATEGORY COUNTS

acro/default.yaml

@@ -29,4 +29,13 @@ survival_safe_threshold: 10
 
 # consider zeros to be disclosive
 zeros_are_disclosive: True
+
+################################################################################
+# Blocked file extensions                                                      #
+################################################################################
+# File extensions that are not allowed in custom outputs or plots.
+# Extensions are case-insensitive and must include the leading dot.
+blocked_extensions:
+  - .svg
+  - .gph
 ...

acro/record.py

@@ -15,6 +15,7 @@
 from pandas import DataFrame
 
 from .constants import ARTIFACTS_DIR
+from .utils import is_blocked_extension
 from .version import __version__
 
 logger = logging.getLogger("acro:records")
@@ -210,10 +211,13 @@ def __str__(self) -> str:
 class Records:
     """Stores data related to a collection of output records."""
 
-    def __init__(self) -> None:
+    def __init__(self, blocked_extensions: list[str] | None = None) -> None:
         """Construct a new object for storing multiple records."""
         self.results: dict[str, Record] = {}
         self.output_id: int = 0
+        self.blocked_extensions: list[str] = [
+            ext.lower() for ext in (blocked_extensions or [])
+        ]
 
     def add(
         self,
@@ -323,7 +327,7 @@ def get_index(self, index: int) -> Record:
         key = list(self.results.keys())[index]
         return self.results[key]
 
-    def add_custom(self, filename: str, comment: str | None = None) -> None:
+    def add_custom(self, filename: str, comment: str | None = None) -> bool:
         """Add an unsupported output to the results dictionary.
 
         Parameters
@@ -332,7 +336,14 @@ def add_custom(self, filename: str, comment: str | None = None) -> None:
             The name of the file that will be added to the list of the outputs.
         comment : str | None, default None
             An optional comment.
+
+        Returns
+        -------
+        bool
+            False if the file extension is blocked, True otherwise.
         """
+        if is_blocked_extension(filename, self.blocked_extensions):
+            return False
         if os.path.exists(filename):
             output = Record(
                 uid=f"output_{self.output_id}",
@@ -353,6 +364,7 @@ def add_custom(self, filename: str, comment: str | None = None) -> None:
             logger.info(
                 "WARNING: Unable to add %s because the file does not exist", filename
             )  # pragma: no cover
+        return True
 
     def rename(self, old: str, new: str) -> None:
         """Rename an output.

acro/utils.py

@@ -3,13 +3,27 @@
 from __future__ import annotations
 
 import logging
+import os
 from inspect import FrameInfo, getframeinfo
 
 import pandas as pd
 
 logger = logging.getLogger("acro")
 
 
+def is_blocked_extension(filename: str, blocked_extensions: list[str]) -> bool:
+    """Return True and log a warning if the file's extension is blocked."""
+    _, ext = os.path.splitext(filename)
+    if ext.lower() in blocked_extensions:
+        logger.warning(
+            "Blocked file extension %s. Files with extension %s are not allowed.",
+            filename,
+            ext,
+        )
+        return True
+    return False
+
+
 def get_command(default: str, stack_list: list[FrameInfo]) -> str:
     """Return the calling source line as a string.
 

test/test_initial.py

@@ -506,6 +506,61 @@ def test_custom_output(acro):
     shutil.rmtree(PATH)
 
 
+def test_blocked_extension(acro, tmp_path):
+    """Test that blocked file extensions are rejected in custom output."""
+    # create temporary files with blocked extensions
+    svg_file = tmp_path / "test.svg"
+    svg_file.write_text("<svg></svg>")
+    gph_file = tmp_path / "test.gph"
+    gph_file.write_text("data")
+
+    # blocked extensions should be rejected
+    acro.custom_output(str(svg_file))
+    acro.custom_output(str(gph_file))
+    assert len(acro.results.results) == 0
+
+    # allowed extensions should be accepted
+    txt_file = tmp_path / "test.txt"
+    txt_file.write_text("hello")
+    acro.custom_output(str(txt_file))
+    assert len(acro.results.results) == 1
+
+    # case-insensitive check
+    svg_upper = tmp_path / "test.SVG"
+    svg_upper.write_text("<svg></svg>")
+    acro.custom_output(str(svg_upper))
+    assert len(acro.results.results) == 1
+
+
+def test_blocked_extension_hist(data, acro):
+    """Test that blocked file extensions are rejected for histograms."""
+    result = acro.hist(data, "inc_grants", bins=1, filename="hist.svg")
+    assert result is None
+    assert len(acro.results.results) == 0
+
+
+def test_blocked_extension_pie(data, acro):
+    """Test that blocked file extensions are rejected for pie charts."""
+    result = acro.pie(data, "grant_type", filename="pie.svg")
+    assert result is None
+    assert len(acro.results.results) == 0
+
+
+def test_blocked_extension_survival(acro):
+    """Test that blocked file extensions are rejected for survival plots."""
+    result = acro.survival_plot(
+        survival_table=pd.DataFrame(),
+        survival_func=None,
+        filename="surv.svg",
+        status="pass",
+        sdc={},
+        command="test",
+        summary="test",
+    )
+    assert result is None
+    assert len(acro.results.results) == 0
+
+
 def test_missing(data, acro, monkeypatch):
     """Pivot table and Crosstab with negative values."""
     acro_tables.CHECK_MISSING_VALUES = True

test/test_stata17_interface.py

@@ -472,7 +472,7 @@ def test_stata_custom_output_invalid():
         options="nototals",
         stata_version="17",
     )
-    correct = "Warning: .gph files may not be exported as they contain data."
+    correct = "Warning: .gph files are not allowed and cannot be exported."
     assert ret == correct, f" we got : {ret}\nexpected:{correct}"
     newres = stata_config.stata_acro.results.__dict__
     assert newres == previous, (

test/test_stata_interface.py

@@ -562,7 +562,7 @@ def test_stata_custom_output_invalid():
         options="nototals",
         stata_version="17",
     )
-    correct = "Warning: .gph files may not be exported as they contain data."
+    correct = "Warning: .gph files are not allowed and cannot be exported."
     assert ret == correct, f" we got : {ret}\nexpected:{correct}"
     newres = stata_config.stata_acro.results.__dict__
     assert newres == previous, (