linting

Author: matteocacciola

pandasai/core/code_generation/code_cleaning.py

@@ -191,7 +191,9 @@ def _remove_make_dirs(self, code: str) -> str:
         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):
+            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)

pandasai/core/response/interactive_chart.py

@@ -22,7 +22,9 @@ def _get_chart(self) -> dict:
 
             return json.loads(self.value)
 
-        raise ValueError("Invalid value type for InteractiveChartResponse. Expected dict or str.")
+        raise ValueError(
+            "Invalid value type for InteractiveChartResponse. Expected dict or str."
+        )
 
     def save(self, path: str):
         img = self._get_chart()
@@ -37,7 +39,9 @@ def get_dict_image(self) -> dict:
 
     def _validate(self):
         if not isinstance(self.value, (dict, str)):
-            raise ValueError("InteractiveChartResponse value must be a dict or a str representing a file path.")
+            raise ValueError(
+                "InteractiveChartResponse value must be a dict or a str representing a file path."
+            )
 
         # if a string, it can be a path to a file or a JSON string
         if isinstance(self.value, str):
@@ -46,4 +50,6 @@ def _validate(self):
             except json.JSONDecodeError:
                 # If it fails, check if it's a valid file path
                 if not os.path.exists(self.value):
-                    raise ValueError("InteractiveChartResponse value must be a valid file path or a JSON string.")
+                    raise ValueError(
+                        "InteractiveChartResponse value must be a valid file path or a JSON string."
+                    )

tests/unit_tests/core/code_generation/test_code_cleaning.py

@@ -205,9 +205,7 @@ def test_replace_output_filenames_with_temp_json_chart_no_json(self):
 
         result = handler._replace_output_filenames_with_temp_json_chart(code)
 
-        self.assertEqual(
-            result, code, f"Expected '{code}', but got '{result}'"
-        )
+        self.assertEqual(result, code, f"Expected '{code}', but got '{result}'")
 
     def test_remove_make_dirs(self):
         handler = self.cleaner
@@ -231,15 +229,11 @@ def test_do_not_remove_make_default_chart_dir(self):
 
         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}'"
-        )
+        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}'"
-        )
+        self.assertEqual(result, code, f"Expected '{code}', but got '{result}'")
 
 
 if __name__ == "__main__":

tests/unit_tests/helpers/test_responses.py

@@ -4,9 +4,8 @@
 from unittest.mock import MagicMock, patch
 
 import pandas as pd
-from PIL import Image
-
 import pytest
+from PIL import Image
 
 from pandasai.core.response import (
     ChartResponse,

tests/unit_tests/response/test_interactive_chart_response.py

@@ -34,7 +34,9 @@ def sample_json():
             "width": 100,
             "height": 100,
             "format": "png",
-            "data": base64.b64encode(io.BytesIO(b"test_image_data").getvalue()).decode("utf-8"),
+            "data": base64.b64encode(io.BytesIO(b"test_image_data").getvalue()).decode(
+                "utf-8"
+            ),
         },
     }