fix(code_env) add numpy in the env by default (#1144)

* fix(code_env) add numpy in the env by default * add: numpy lib to the environment
Sinaptik-AI · May 3, 2024 · 939bbff · 939bbff
1 parent 6c671c4
commit 939bbff
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 0 deletions.
diff --git a/pandasai/helpers/optional.py b/pandasai/helpers/optional.py
@@ -11,6 +11,7 @@
 from typing import TYPE_CHECKING, List
 
 import matplotlib.pyplot as plt
+import numpy as np
 from pandas.util.version import Version
 
 import pandasai.pandas as pd
@@ -56,6 +57,7 @@ def get_environment(additional_deps: List[dict]) -> dict:
     return {
         "pd": pd,
         "plt": plt,
+        "np": np,
         **{
             lib["alias"]: (
                 getattr(import_dependency(lib["module"]), lib["name"])

diff --git a/tests/unit_tests/helpers/test_optional_dependency.py b/tests/unit_tests/helpers/test_optional_dependency.py
@@ -90,3 +90,4 @@ def test_env_for_necessary_deps():
     env = get_environment([])
     assert "pd" in env
     assert "plt" in env
+    assert "np" in env