refactor(Pipelines) : Smart Data Frame Pipeline (Sourcery refactored)

pandasai/smart_datalake/prompt_generation.py

@@ -37,13 +37,9 @@ def execute(self, input: Any, **kwargs) -> Any:
                 "last_code_generated"
             )
 
-        generate_python_code_instruction = (
-            pipeline_context.query_exec_tracker.execute_func(
-                pipeline_context.get_intermediate_value("get_prompt"),
-                "generate_python_code",
-                default_prompt=GeneratePythonCodePrompt,
-                default_values=default_values,
-            )
+        return pipeline_context.query_exec_tracker.execute_func(
+            pipeline_context.get_intermediate_value("get_prompt"),
+            "generate_python_code",
+            default_prompt=GeneratePythonCodePrompt,
+            default_values=default_values,
         )
-
-        return generate_python_code_instruction

tests/pipelines/smart_datalake/test_code_execution.py

@@ -74,8 +74,7 @@ def config(self, llm):
 
     @pytest.fixture
     def context(self, sample_df, config):
-        pipeline_context = PipelineContext([sample_df], config)
-        return pipeline_context
+        return PipelineContext([sample_df], config)
 
     @pytest.fixture
     def logger(self):
@@ -127,15 +126,16 @@ def mock_intermediate_values(key : str):
                 return SkillsManager()
             elif key == "code_manager":
                 return mock_code_manager
+
         context.get_intermediate_value = Mock(side_effect=mock_intermediate_values)
-        
+
         assert isinstance(code_execution, CodeExecution)
 
         result = None
         try:
             result = code_execution.execute(input="Test Code", context=context, logger=logger)
         except Exception as e:
-            assert result == None
+            assert result is None
 
     def test_code_execution_successful_at_retry(self, context, logger):
         # Test Flow : Code Execution Successful with no exceptions

tests/pipelines/smart_datalake/test_code_generator.py

@@ -74,8 +74,7 @@ def config(self, llm):
 
     @pytest.fixture
     def context(self, sample_df, config):
-        pipeline_context = PipelineContext([sample_df], config)
-        return pipeline_context
+        return PipelineContext([sample_df], config)
 
     @pytest.fixture
     def logger(self):

tests/pipelines/smart_datalake/test_result_parsing.py

@@ -72,8 +72,7 @@ def config(self, llm):
 
     @pytest.fixture
     def context(self, sample_df, config):
-        pipeline_context = PipelineContext([sample_df], config)
-        return pipeline_context
+        return PipelineContext([sample_df], config)
 
     @pytest.fixture
     def logger(self):
@@ -111,20 +110,20 @@ def test_result_parsing_unsuccessful_with_exceptions(self, context, logger):
 
         def mock_result_parsing(*args, **kwargs):
             raise Exception("Unit test exception")
-        
+
         context._query_exec_tracker = Mock()
         context.query_exec_tracker.execute_func = Mock(side_effect=mock_result_parsing)
 
         def mock_intermediate_values(key : str):
             if key == "response_parser" :
                 return mock_response_parser
-        
+
         context.get_intermediate_value = Mock(side_effect=mock_intermediate_values)
 
         result = None
         try:
             result = result_parsing.execute(input="Test Result", context=context, logger=logger)
         except Exception as e:
-            assert result == None
+            assert result is None
         assert isinstance(result_parsing, ResultParsing)
 

tests/pipelines/smart_datalake/test_result_validation.py

@@ -72,8 +72,7 @@ def config(self, llm):
 
     @pytest.fixture
     def context(self, sample_df, config):
-        pipeline_context = PipelineContext([sample_df], config)
-        return pipeline_context
+        return PipelineContext([sample_df], config)
 
     @pytest.fixture
     def logger(self):
@@ -96,7 +95,7 @@ def test_result_is_none(self, context, logger):
 
         assert not context.query_exec_tracker.add_step.called
         assert isinstance(result_validation, ResultValidation)
-        assert result == None
+        assert result is None
 
     def test_result_is_not_of_dict_type(self, context, logger):
         # Test Flow : Code Execution Successful with no exceptions