fix: make seaborn as an optional dependency (#1500)

* fix: make seaborn as an optional dependency * fix: linting errors
Sinaptik-AI · Jan 2, 2025 · 63df017 · 63df017
1 parent d63f918
commit 63df017
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/pandasai/helpers/dataframe_serializer.py b/pandasai/helpers/dataframe_serializer.py
@@ -1,8 +1,6 @@
 import json
 from enum import Enum
 
-import yaml
-
 import pandasai.pandas as pd
 
 
@@ -160,6 +158,9 @@ def convert_df_to_json_str(self, df: pd.DataFrame, extras: dict) -> str:
 
     def convert_df_to_yml(self, df: pd.DataFrame, extras: dict) -> str:
         json_df = self.convert_df_to_json(df, extras)
+
+        import yaml
+
         yml_str = yaml.dump(json_df, sort_keys=False, allow_unicode=True)
         if "is_direct_sql" in extras and extras["is_direct_sql"]:
             return f"<table>\n{yml_str}\n</table>\n"

diff --git a/pandasai/helpers/optional.py b/pandasai/helpers/optional.py
@@ -19,7 +19,6 @@
 from pandasai.safe_libs.restricted_matplotlib import RestrictedMatplotlib
 from pandasai.safe_libs.restricted_numpy import RestrictedNumpy
 from pandasai.safe_libs.restricted_pandas import RestrictedPandas
-from pandasai.safe_libs.restricted_seaborn import RestrictedSeaborn
 
 if TYPE_CHECKING:
     import types
@@ -80,6 +79,8 @@ def get_environment(additional_deps: List[dict]) -> dict:
 
     for lib in additional_deps:
         if lib["name"] == "seaborn":
+            from pandasai.safe_libs.restricted_seaborn import RestrictedSeaborn
+
             env["sns"] = RestrictedSeaborn()
 
         if lib["name"] == "datetime":