Skip to content

Commit b207b7e

Browse files
committed
Fix annotated pydantic model
Signed-off-by: Charles Swartz <[email protected]>
1 parent 30ba7cc commit b207b7e

File tree

1 file changed

+16
-12
lines changed

1 file changed

+16
-12
lines changed

tests/core/test_pydantic.py

+16-12
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
"""Unit tests for pydantic compatibility."""
22

33
# pylint:disable=too-few-public-methods,missing-class-docstring
4-
from typing import Annotated, ClassVar, Optional
4+
from typing import Annotated, Optional
55

66
import pandas as pd
77
import pytest
@@ -55,24 +55,28 @@ class SeriesSchemaPydantic(BaseModel):
5555
pa_index: Optional[pa.Index]
5656

5757

58-
class AnnotatedDfPydantic(BaseModel):
59-
"""Test pydantic model with annotated dataframe model."""
58+
if PYDANTIC_V2:
59+
from pydantic import ConfigDict
6060

61-
# Required because pandas.DataFrame is not a valid pydantic type. Using
62-
# arbitrary_types_allowed=True essentially adds an isinstance check for
63-
# the annotated type.
64-
if PYDANTIC_V2:
65-
from pydantic import ConfigDict
61+
class AnnotatedDfPydantic(BaseModel): # type: ignore[no-redef]
62+
"""Test pydantic model with annotated dataframe model."""
63+
64+
# arbitrary_types_allowed=True required for pandas.DataFrame
65+
model_config = ConfigDict(arbitrary_types_allowed=True)
6666

67-
model_config: ClassVar[ConfigDict] = ConfigDict(
68-
arbitrary_types_allowed=True
69-
)
7067
df: Annotated[pd.DataFrame, SimpleSchema]
71-
else:
7268

69+
else:
70+
71+
class AnnotatedDfPydantic(BaseModel): # type: ignore[no-redef]
72+
"""Test pydantic model with annotated dataframe model."""
73+
74+
# arbitrary_types_allowed=True required for pandas.DataFrame
7375
class Config:
7476
arbitrary_types_allowed = True
7577

78+
df: Annotated[pd.DataFrame, SimpleSchema]
79+
7680

7781
def test_typed_dataframe():
7882
"""Test that typed DataFrame is compatible with pydantic."""

0 commit comments

Comments
 (0)