lint: fix linte

tests/test_memory.py

@@ -1,4 +1,3 @@
-import pytest
 from pandasai.helpers.memory import Memory
 
 
@@ -9,36 +8,32 @@ def test_to_json_empty_memory():
 
 def test_to_json_with_messages():
     memory = Memory()
-    
+
     # Add test messages
     memory.add("Hello", is_user=True)
     memory.add("Hi there!", is_user=False)
     memory.add("How are you?", is_user=True)
-    
+
     expected_json = [
         {"role": "user", "message": "Hello"},
         {"role": "assistant", "message": "Hi there!"},
-        {"role": "user", "message": "How are you?"}
+        {"role": "user", "message": "How are you?"},
     ]
-    
+
     assert memory.to_json() == expected_json
 
 
 def test_to_json_message_order():
     memory = Memory()
-    
+
     # Add messages in specific order
-    messages = [
-        ("Message 1", True),
-        ("Message 2", False),
-        ("Message 3", True)
-    ]
-
+    messages = [("Message 1", True), ("Message 2", False), ("Message 3", True)]
+
     for msg, is_user in messages:
         memory.add(msg, is_user=is_user)
-    
+
     result = memory.to_json()
-    
+
     # Verify order is preserved
     assert len(result) == 3
     assert result[0]["message"] == "Message 1"
@@ -55,13 +50,13 @@ def test_to_openai_messages_with_agent_description():
     memory = Memory(agent_description="I am a helpful assistant")
     memory.add("Hello", is_user=True)
     memory.add("Hi there!", is_user=False)
-    
+
     expected_messages = [
         {"role": "system", "content": "I am a helpful assistant"},
         {"role": "user", "content": "Hello"},
-        {"role": "assistant", "content": "Hi there!"}
+        {"role": "assistant", "content": "Hi there!"},
     ]
-    
+
     assert memory.to_openai_messages() == expected_messages
 
 
@@ -70,11 +65,11 @@ def test_to_openai_messages_without_agent_description():
     memory.add("Hello", is_user=True)
     memory.add("Hi there!", is_user=False)
     memory.add("How are you?", is_user=True)
-    
+
     expected_messages = [
         {"role": "user", "content": "Hello"},
         {"role": "assistant", "content": "Hi there!"},
-        {"role": "user", "content": "How are you?"}
+        {"role": "user", "content": "How are you?"},
     ]
-    
+
     assert memory.to_openai_messages() == expected_messages

tests/unit_tests/dataframe/test_pull.py

@@ -1,13 +1,18 @@
 import os
-import pytest
-from unittest.mock import patch, Mock, mock_open
 from io import BytesIO
+from unittest.mock import Mock, mock_open, patch
 from zipfile import ZipFile
 
 import pandas as pd
+import pytest
+
+from pandasai.data_loader.semantic_layer_schema import (
+    Column,
+    SemanticLayerSchema,
+    Source,
+)
 from pandasai.dataframe.base import DataFrame
-from pandasai.exceptions import PandaAIApiKeyError, DatasetNotFound
-from pandasai.data_loader.semantic_layer_schema import SemanticLayerSchema, Column, Source
+from pandasai.exceptions import DatasetNotFound, PandaAIApiKeyError
 
 
 @pytest.fixture
@@ -41,20 +46,22 @@ def mock_schema():
 
 
 def test_pull_success(mock_env, sample_df, mock_zip_content, mock_schema, tmp_path):
-    with patch("pandasai.dataframe.base.get_pandaai_session") as mock_session, \
-         patch("pandasai.dataframe.base.find_project_root") as mock_root, \
-         patch("pandasai.DatasetLoader.create_loader_from_path") as mock_loader, \
-         patch("builtins.open", mock_open()) as mock_file:
-        
+    with patch("pandasai.dataframe.base.get_pandaai_session") as mock_session, patch(
+        "pandasai.dataframe.base.find_project_root"
+    ) as mock_root, patch(
+        "pandasai.DatasetLoader.create_loader_from_path"
+    ) as mock_loader, patch("builtins.open", mock_open()) as mock_file:
         # Setup mocks
         mock_response = Mock()
         mock_response.status_code = 200
         mock_response.content = mock_zip_content
         mock_session.return_value.get.return_value = mock_response
         mock_root.return_value = str(tmp_path)
-        
+
         mock_loader_instance = Mock()
-        mock_loader_instance.load.return_value = DataFrame(sample_df, schema=mock_schema)
+        mock_loader_instance.load.return_value = DataFrame(
+            sample_df, schema=mock_schema
+        )
         mock_loader.return_value = mock_loader_instance
 
         # Create DataFrame instance and call pull
@@ -64,8 +71,11 @@ def test_pull_success(mock_env, sample_df, mock_zip_content, mock_schema, tmp_pa
         # Verify API call
         mock_session.return_value.get.assert_called_once_with(
             "/datasets/pull",
-            headers={"accept": "application/json", "x-authorization": "Bearer test_api_key"},
-            params={"path": "test/path"}
+            headers={
+                "accept": "application/json",
+                "x-authorization": "Bearer test_api_key",
+            },
+            params={"path": "test/path"},
         )
 
         # Verify file operations
@@ -92,28 +102,35 @@ def test_pull_api_error(mock_env, sample_df, mock_schema):
 
 
 def test_pull_file_exists(mock_env, sample_df, mock_zip_content, mock_schema, tmp_path):
-    with patch("pandasai.dataframe.base.get_pandaai_session") as mock_session, \
-         patch("pandasai.dataframe.base.find_project_root") as mock_root, \
-         patch("pandasai.DatasetLoader.create_loader_from_path") as mock_loader, \
-         patch("builtins.open", mock_open()) as mock_file, \
-         patch("os.path.exists") as mock_exists, \
-         patch("os.makedirs") as mock_makedirs:
-        
+    with patch("pandasai.dataframe.base.get_pandaai_session") as mock_session, patch(
+        "pandasai.dataframe.base.find_project_root"
+    ) as mock_root, patch(
+        "pandasai.DatasetLoader.create_loader_from_path"
+    ) as mock_loader, patch("builtins.open", mock_open()) as mock_file, patch(
+        "os.path.exists"
+    ) as mock_exists, patch("os.makedirs") as mock_makedirs:
         # Setup mocks
         mock_response = Mock()
         mock_response.status_code = 200
         mock_response.content = mock_zip_content
         mock_session.return_value.get.return_value = mock_response
         mock_root.return_value = str(tmp_path)
         mock_exists.return_value = True
-        
+
         mock_loader_instance = Mock()
-        mock_loader_instance.load.return_value = DataFrame(sample_df, schema=mock_schema)
+        mock_loader_instance.load.return_value = DataFrame(
+            sample_df, schema=mock_schema
+        )
         mock_loader.return_value = mock_loader_instance
 
         # Create DataFrame instance and call pull
         df = DataFrame(sample_df, path="test/path", schema=mock_schema)
         df.pull()
 
         # Verify directory creation
-        mock_makedirs.assert_called_with(os.path.dirname(os.path.join(str(tmp_path), "datasets", "test/path", "test.csv")), exist_ok=True)
+        mock_makedirs.assert_called_with(
+            os.path.dirname(
+                os.path.join(str(tmp_path), "datasets", "test/path", "test.csv")
+            ),
+            exist_ok=True,
+        )

tests/unit_tests/helpers/test_logger.py

@@ -1,31 +1,45 @@
 import logging
+
 from pandasai.helpers.logger import Logger
 
 
 def test_verbose_setter():
     # Initialize logger with verbose=False
     logger = Logger(verbose=False)
     assert logger._verbose is False
-    assert not any(isinstance(handler, logging.StreamHandler) for handler in logger._logger.handlers)
+    assert not any(
+        isinstance(handler, logging.StreamHandler)
+        for handler in logger._logger.handlers
+    )
 
     # Set verbose to True
     logger.verbose = True
     assert logger._verbose is True
-    assert any(isinstance(handler, logging.StreamHandler) for handler in logger._logger.handlers)
+    assert any(
+        isinstance(handler, logging.StreamHandler)
+        for handler in logger._logger.handlers
+    )
     assert len(logger._logger.handlers) == 1
 
     # Set verbose to False
     logger.verbose = False
     assert logger._verbose is False
-    assert not any(isinstance(handler, logging.StreamHandler) for handler in logger._logger.handlers)
+    assert not any(
+        isinstance(handler, logging.StreamHandler)
+        for handler in logger._logger.handlers
+    )
     assert len(logger._logger.handlers) == 0
 
     # Set verbose to True again to ensure multiple toggles work
     logger.verbose = True
     assert logger._verbose is True
-    assert any(isinstance(handler, logging.StreamHandler) for handler in logger._logger.handlers)
+    assert any(
+        isinstance(handler, logging.StreamHandler)
+        for handler in logger._logger.handlers
+    )
     assert len(logger._logger.handlers) == 1
 
+
 def test_save_logs_property():
     # Initialize logger with save_logs=False
     logger = Logger(save_logs=False, verbose=False)
@@ -34,22 +48,27 @@ def test_save_logs_property():
     # Enable save_logs
     logger.save_logs = True
     assert logger.save_logs is True
-    assert any(isinstance(handler, logging.FileHandler) for handler in logger._logger.handlers)
+    assert any(
+        isinstance(handler, logging.FileHandler) for handler in logger._logger.handlers
+    )
 
     # Disable save_logs
     logger.save_logs = False
     assert logger.save_logs is False
-    assert not any(isinstance(handler, logging.FileHandler) for handler in logger._logger.handlers)
+    assert not any(
+        isinstance(handler, logging.FileHandler) for handler in logger._logger.handlers
+    )
+
 
 def test_save_logs_property():
     # When logger is initialized with save_logs=True (default), it should have handlers
     logger = Logger(save_logs=True)
     assert logger.save_logs is True
-    
+
     # When logger is initialized with save_logs=False, it should still have handlers if verbose=True
     logger = Logger(save_logs=False, verbose=True)
     assert logger.save_logs is True
-    
+
     # When both save_logs and verbose are False, there should be no handlers
     logger = Logger(save_logs=False, verbose=False)
     logger._logger.handlers = []  # Reset handlers to match the property's expected behavior

tests/unit_tests/helpers/test_session.py

@@ -5,7 +5,7 @@
 import requests
 
 from pandasai.constants import DEFAULT_API_URL
-from pandasai.exceptions import PandaAIApiKeyError, PandaAIApiCallError
+from pandasai.exceptions import PandaAIApiCallError, PandaAIApiKeyError
 from pandasai.helpers.session import Session, get_pandaai_session
 
 
@@ -148,7 +148,7 @@ def test_make_request_error_response(mock_request):
     session = Session(api_key="test-key")
     with pytest.raises(PandaAIApiCallError) as exc_info:
         session.make_request("POST", "/test")
-    
+
     assert str(exc_info.value) == "Bad request"
 
 
@@ -161,7 +161,7 @@ def test_make_request_network_error(mock_request):
     session = Session(api_key="test-key")
     with pytest.raises(PandaAIApiCallError) as exc_info:
         session.make_request("GET", "/test")
-    
+
     assert "Request failed: Network error" in str(exc_info.value)