Add formatting in number response

Author: ArslanSaleem

pandasai/core/response/base.py

@@ -48,3 +48,6 @@ def to_dict(self) -> dict:
     def to_json(self) -> str:
         """Return a JSON representation."""
         return json.dumps(self.to_dict(), cls=CustomJsonEncoder)
+
+    def __format__(self, fmt):
+        return self.value.__format__(fmt)

tests/unit_tests/response/test_error_response.py

@@ -36,3 +36,17 @@ def test_error_response_with_non_string_value():
     assert response.value == 123
     assert response.last_code_executed == "test_code"
     assert response.error == "error message"
+
+
+def test_error_response_format_alignment():
+    """Test __format__ with string formatting on error message"""
+    response = ErrorResponse("Error!", "test_code", "error message")
+    assert f"{response:>10}" == "    Error!"
+    assert f"{response:<10}" == "Error!    "
+
+
+def test_error_response_format_with_fstring():
+    """Test __format__ in f-string context"""
+    response = ErrorResponse("Failed", "test_code", "error message")
+    result = f"Status: {response:>10}"
+    assert result == "Status:     Failed"

tests/unit_tests/response/test_number_response.py

@@ -26,3 +26,55 @@ def test_number_response_with_string_number():
     response = NumberResponse("123", "test_code")
     assert response.type == "number"
     assert response.value == "123"  # Value remains as string
+
+
+def test_number_response_format_decimal():
+    """Test __format__ with decimal places"""
+    response = NumberResponse(3.14159, "test_code")
+    assert f"{response:.2f}" == "3.14"
+    assert f"{response:.4f}" == "3.1416"
+
+
+def test_number_response_format_with_fstring():
+    """Test __format__ in f-string context"""
+    response = NumberResponse(123.456, "test_code")
+    result = f"Value: {response:.2f}"
+    assert result == "Value: 123.46"
+
+
+def test_number_response_format_function():
+    """Test __format__ with format() function"""
+    response = NumberResponse(42.123, "test_code")
+    assert format(response, ".1f") == "42.1"
+
+
+def test_number_response_format_scientific():
+    """Test __format__ with scientific notation"""
+    response = NumberResponse(1234.5, "test_code")
+    assert f"{response:e}" == "1.234500e+03"
+
+
+def test_number_response_format_percentage():
+    """Test __format__ with percentage"""
+    response = NumberResponse(0.875, "test_code")
+    assert f"{response:.1%}" == "87.5%"
+
+
+def test_number_response_format_padding():
+    """Test __format__ with padding"""
+    response = NumberResponse(42, "test_code")
+    assert f"{response:05d}" == "00042"
+    assert f"{response:>10}" == "        42"
+
+
+def test_number_response_format_integer():
+    """Test __format__ with integer formatting"""
+    response = NumberResponse(42, "test_code")
+    assert f"{response:d}" == "42"
+
+
+def test_number_response_format_with_str_format():
+    """Test __format__ with string .format() method"""
+    response = NumberResponse(99.9, "test_code")
+    result = "Price: ${:.2f}".format(response)
+    assert result == "Price: $99.90"

tests/unit_tests/response/test_string_response.py

@@ -20,3 +20,37 @@ def test_string_response_with_non_string_value():
     assert response.type == "string"
     assert response.value == 123
     assert response.last_code_executed == "test_code"
+
+
+def test_string_response_format_alignment():
+    """Test __format__ with string alignment"""
+    response = StringResponse("hello", "test_code")
+    assert f"{response:>10}" == "     hello"  # Right align
+    assert f"{response:<10}" == "hello     "  # Left align
+    assert f"{response:^10}" == "  hello   "  # Center align
+
+
+def test_string_response_format_with_fstring():
+    """Test __format__ in f-string context"""
+    response = StringResponse("world", "test_code")
+    result = f"Hello {response:>10}!"
+    assert result == "Hello      world!"
+
+
+def test_string_response_format_function():
+    """Test __format__ with format() function"""
+    response = StringResponse("test", "test_code")
+    assert format(response, ">8") == "    test"
+
+
+def test_string_response_format_truncate():
+    """Test __format__ with truncation"""
+    response = StringResponse("hello world", "test_code")
+    assert f"{response:.5}" == "hello"
+
+
+def test_string_response_format_with_str_format():
+    """Test __format__ with string .format() method"""
+    response = StringResponse("Python", "test_code")
+    result = "Language: {:>10}".format(response)
+    assert result == "Language:     Python"