From cfeb071308fc857ea3d54c191e8e73a6f5638b80 Mon Sep 17 00:00:00 2001
From: Arslan Saleem <khan.arslan38@gmail.com>
Date: Thu, 2 Jan 2025 17:32:29 +0100
Subject: [PATCH] fix: check if whitelisted lib is actually exists in the
 additional deps (#1499)

---
 pandasai/pipelines/chat/code_cleaning.py      | 27 +++++++++----------
 .../smart_datalake/test_code_cleaning.py      |  4 +++
 2 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/pandasai/pipelines/chat/code_cleaning.py b/pandasai/pipelines/chat/code_cleaning.py
index a185aeec4..96dbd91f9 100644
--- a/pandasai/pipelines/chat/code_cleaning.py
+++ b/pandasai/pipelines/chat/code_cleaning.py
@@ -627,23 +627,22 @@ def _check_imports(self, node: Union[ast.Import, ast.ImportFrom]):
         if library == "pandas":
             return
 
-        if (
-            library
-            in WHITELISTED_LIBRARIES + self._config.custom_whitelisted_dependencies
-        ):
-            for alias in node.names:
-                self._additional_dependencies.append(
-                    {
-                        "module": module,
-                        "name": alias.name,
-                        "alias": alias.asname or alias.name,
-                    }
-                )
-            return
+        whitelisted_libs = (
+            WHITELISTED_LIBRARIES + self._config.custom_whitelisted_dependencies
+        )
 
-        if library not in WHITELISTED_LIBRARIES:
+        if library not in whitelisted_libs:
             raise BadImportError(
                 f"The library '{library}' is not in the list of whitelisted libraries. "
                 "To learn how to whitelist custom dependencies, visit: "
                 "https://docs.pandas-ai.com/custom-whitelisted-dependencies#custom-whitelisted-dependencies"
             )
+
+        for alias in node.names:
+            self._additional_dependencies.append(
+                {
+                    "module": module,
+                    "name": alias.name,
+                    "alias": alias.asname or alias.name,
+                }
+            )
diff --git a/tests/unit_tests/pipelines/smart_datalake/test_code_cleaning.py b/tests/unit_tests/pipelines/smart_datalake/test_code_cleaning.py
index d632f6cc3..9506a06bb 100644
--- a/tests/unit_tests/pipelines/smart_datalake/test_code_cleaning.py
+++ b/tests/unit_tests/pipelines/smart_datalake/test_code_cleaning.py
@@ -309,7 +309,11 @@ def test_custom_whitelisted_dependencies(
         code_cleaning._config.custom_whitelisted_dependencies = ["my_custom_library"]
         output = code_cleaning.execute(code, context=context, logger=logger)
 
+        print(code_cleaning._additional_dependencies)
         assert output.output == "my_custom_library.do_something()"
+        assert (
+            code_cleaning._additional_dependencies[0]["module"] == "my_custom_library"
+        )
         assert isinstance(output, LogicUnitOutput)
         assert output.success
         assert output.message == "Code Cleaned Successfully"