Add support and test for pd.Series

Author: EkberHasanov

pydantic_extra_types/pandas_types.py

@@ -0,0 +1,43 @@
+from typing import Any, List, Tuple, TypeVar, Union
+
+import pandas as pd
+from pydantic import GetCoreSchemaHandler
+from pydantic_core import core_schema
+
+T = TypeVar('T', str, bytes, bool, int, float, complex, pd.Timestamp, pd.Timedelta, pd.Period)
+
+
+class Series:
+    def __init__(self, value: Any) -> None:
+        self.value = pd.Series(value)
+
+    @classmethod
+    def __get_pydantic_core_schema__(
+        cls, source: type[Any], handler: GetCoreSchemaHandler
+    ) -> core_schema.AfterValidatorFunctionSchema:
+        return core_schema.general_after_validator_function(
+            cls._validate,
+            core_schema.any_schema(),
+        )
+
+    @classmethod
+    def _validate(cls, __input_value: Any, _: core_schema.ValidationInfo) -> 'Series':
+        if isinstance(__input_value, pd.Series):
+            return cls(__input_value)
+        return cls(pd.Series(__input_value))
+
+    def __repr__(self) -> str:
+        return repr(self.value)
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self.value, name)
+
+    def __eq__(self, __value: object) -> bool:
+        return isinstance(__value, pd.Series) or isinstance(__value, Series)
+
+    def __add__(self, other: Union['Series', List[Any], Tuple[Any], T]) -> 'Series':
+        if isinstance(other, Series):
+            result_val = self.value + other.value
+        else:
+            result_val = self.value + other
+        return Series(result_val)

requirements/linting.txt

@@ -50,4 +50,5 @@ virtualenv==20.23.0
     # via pre-commit
 
 # The following packages are considered to be unsafe in a requirements file:
+pandas-stubs==2.0.2.230605
 # setuptools

tests/test_pandas_types.py

@@ -0,0 +1,89 @@
+import pandas as pd
+import pytest
+
+from pydantic_extra_types.pandas_types import Series
+
+
+@pytest.mark.parametrize(
+    'data, expected',
+    [
+        ([1, 2, 3], [1, 2, 3]),
+        ([], []),
+        ([10, 20, 30, 40], [10, 20, 30, 40]),
+    ],
+)
+def test_series_creation(data, expected):
+    s = Series(data)
+    assert isinstance(s, Series)
+    assert isinstance(s.value, pd.Series)
+    assert s.value.tolist() == expected
+
+
+def test_series_repr():
+    data = [1, 2, 3]
+    s = Series(data)
+    assert repr(s) == repr(pd.Series(data))
+
+
+def test_series_attribute_access():
+    data = [1, 2, 3]
+    s = Series(data)
+    assert s.sum() == pd.Series(data).sum()
+
+
+def test_series_equality():
+    data = [1, 2, 3]
+    s1 = Series(data)
+    s2 = Series(data)
+    assert s1 == s2
+    assert s2 == pd.Series(data)
+
+
+def test_series_addition():
+    data1 = [1, 2, 3]
+    data2 = [4, 5, 6]
+    s1 = Series(data1)
+    s2 = Series(data2)
+    s3 = s1 + s2
+    assert isinstance(s3, Series)
+    assert s3.value.tolist() == [5, 7, 9]
+
+
+@pytest.mark.parametrize(
+    'data, other, expected',
+    [
+        ([1, 2, 3], [4, 5, 6], [5, 7, 9]),
+        ([10, 20, 30], (1, 2, 3), [11, 22, 33]),
+        ([5, 10, 15], pd.Series([1, 2, 3]), [6, 12, 18]),
+    ],
+)
+def test_series_addition_with_types(data, other, expected):
+    s = Series(data)
+    result = s + other
+    assert isinstance(result, Series)
+    assert result.value.tolist() == expected
+
+
+@pytest.mark.parametrize(
+    'data, other',
+    [
+        ([1, 2, 3], 'invalid'),  # Invalid type for addition
+        ([1, 2, 3], {'a': 1, 'b': 2}),  # Invalid type for addition
+    ],
+)
+def test_series_addition_invalid_type_error(data, other) -> None:
+    s = Series(data)
+    with pytest.raises(TypeError):
+        s + other
+
+
+@pytest.mark.parametrize(
+    'data, other',
+    [
+        ([1, 2, 3], []),
+    ],
+)
+def test_series_addition_invalid_value_error(data, other) -> None:
+    s = Series(data)
+    with pytest.raises(ValueError):
+        s + other