fix: check if whitelisted lib is actually exists in the additional de…

…ps (#1499)
Sinaptik-AI · Jan 2, 2025 · cfeb071 · cfeb071
1 parent 63df017
commit cfeb071
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 14 deletions.
diff --git 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
@@ -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"