chore: type hint issue (#537)

Sinaptik-AI · Sep 8, 2023 · 1b5e942 · 1b5e942
1 parent 9e6a71a
commit 1b5e942
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/pandasai/helpers/code_manager.py b/pandasai/helpers/code_manager.py
@@ -46,12 +46,12 @@ def __init__(
         if self._config.middlewares is not None:
             self.add_middlewares(*self._config.middlewares)
 
-    def add_middlewares(self, *middlewares: List[Middleware]):
+    def add_middlewares(self, *middlewares: Optional[Middleware]):
         """
         Add middlewares to PandasAI instance.
 
         Args:
-            *middlewares: A list of middlewares
+            *middlewares: Middlewares to be added
 
         """
         self._middlewares.extend(middlewares)

diff --git a/pandasai/smart_dataframe/__init__.py b/pandasai/smart_dataframe/__init__.py
@@ -177,12 +177,12 @@ def __setitem__(self, key, value):
     def __repr__(self):
         return self._df.__repr__()
 
-    def add_middlewares(self, *middlewares: List[Middleware]):
+    def add_middlewares(self, *middlewares: Optional[Middleware]):
         """
         Add middlewares to PandasAI instance.
 
         Args:
-            *middlewares: A list of middlewares
+            *middlewares: Middlewares to be added
 
         """
         self._dl.add_middlewares(*middlewares)

diff --git a/pandasai/smart_datalake/__init__.py b/pandasai/smart_datalake/__init__.py
@@ -169,12 +169,12 @@ def _load_llm(self, llm: LLM):
 
         self._llm = llm
 
-    def add_middlewares(self, *middlewares: List[Middleware]):
+    def add_middlewares(self, *middlewares: Optional[Middleware]):
         """
         Add middlewares to PandasAI instance.
 
         Args:
-            *middlewares: A list of middlewares
+            *middlewares: Middlewares to be added
         """
         self._code_manager.add_middlewares(*middlewares)