cleaning the eventual creation of dirs

Author: matteocacciola

pandasai/core/code_generation/code_cleaning.py

@@ -138,6 +138,8 @@ def clean_code(self, code: str) -> str:
         code = self._replace_output_filenames_with_temp_chart(code)
         code = self._replace_output_filenames_with_temp_json_chart(code)
 
+        code = self._remove_make_dirs(code)
+
         # If plt.show or fig.show is in the code, remove that line
         code = re.sub(r"[a-z].show\(\)", "", code)
 
@@ -180,3 +182,16 @@ def _replace_output_filenames_with_temp_json_chart(self, code: str) -> str:
             lambda m: f"{m.group(1)}{chart_path}{m.group(1)}",
             code,
         )
+
+    def _remove_make_dirs(self, code: str) -> str:
+        """
+        Remove any directory creation commands from the code.
+        """
+        # Remove lines that create directories, except for the default chart directory DEFAULT_CHART_DIRECTORY
+        code_lines = code.splitlines()
+        cleaned_lines = []
+        for line in code_lines:
+            if DEFAULT_CHART_DIRECTORY not in line and ("os.makedirs(" in line or "os.mkdir(" in line):
+                continue
+            cleaned_lines.append(line)
+        return "\n".join(cleaned_lines)

tests/unit_tests/core/code_generation/test_code_cleaning.py

@@ -5,6 +5,7 @@
 from unittest.mock import MagicMock
 
 from pandasai.agent.state import AgentState
+from pandasai.constants import DEFAULT_CHART_DIRECTORY
 from pandasai.core.code_generation.code_cleaning import CodeCleaner
 from pandasai.dataframe.base import DataFrame
 from pandasai.exceptions import MaliciousQueryError
@@ -201,14 +202,45 @@ def test_replace_output_filenames_with_temp_json_chart_no_json(self):
         handler = self.cleaner
 
         code = "some text without json"
-        expected_code = "some text without json"  # No change should occur
 
         result = handler._replace_output_filenames_with_temp_json_chart(code)
 
+        self.assertEqual(
+            result, code, f"Expected '{code}', but got '{result}'"
+        )
+
+    def test_remove_make_dirs(self):
+        handler = self.cleaner
+
+        code = "os.makedirs('/some/path')\nplt.show()\nfig.show()"
+        expected_code = "plt.show()\nfig.show()"  # Should remove the os.makedirs line
+        result = handler._remove_make_dirs(code)
+        self.assertEqual(
+            result, expected_code, f"Expected '{expected_code}', but got '{result}'"
+        )
+
+        code = "os.mkdir('/some/path')\nplt.show()\nfig.show()"
+        expected_code = "plt.show()\nfig.show()"  # Should remove the os.mkdir line
+        result = handler._remove_make_dirs(code)
         self.assertEqual(
             result, expected_code, f"Expected '{expected_code}', but got '{result}'"
         )
 
+    def test_do_not_remove_make_default_chart_dir(self):
+        handler = self.cleaner
+
+        code = f"os.makedirs({DEFAULT_CHART_DIRECTORY}')\nplt.show()\nfig.show()"
+        result = handler._remove_make_dirs(code)
+        self.assertEqual(
+            result, code, f"Expected '{code}', but got '{result}'"
+        )
+
+        code = f"os.mkdir({DEFAULT_CHART_DIRECTORY}')\nplt.show()\nfig.show()"
+        result = handler._remove_make_dirs(code)
+        self.assertEqual(
+            result, code, f"Expected '{code}', but got '{result}'"
+        )
+
 
 if __name__ == "__main__":
     unittest.main()