|
1 | 1 | """Unit tests for pydantic compatibility."""
|
2 | 2 |
|
3 | 3 | # pylint:disable=too-few-public-methods,missing-class-docstring
|
4 |
| -from typing import Annotated, ClassVar, Optional |
| 4 | +from typing import Annotated, Optional |
5 | 5 |
|
6 | 6 | import pandas as pd
|
7 | 7 | import pytest
|
@@ -55,24 +55,28 @@ class SeriesSchemaPydantic(BaseModel):
|
55 | 55 | pa_index: Optional[pa.Index]
|
56 | 56 |
|
57 | 57 |
|
58 |
| -class AnnotatedDfPydantic(BaseModel): |
59 |
| - """Test pydantic model with annotated dataframe model.""" |
| 58 | +if PYDANTIC_V2: |
| 59 | + from pydantic import ConfigDict |
60 | 60 |
|
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) |
66 | 66 |
|
67 |
| - model_config: ClassVar[ConfigDict] = ConfigDict( |
68 |
| - arbitrary_types_allowed=True |
69 |
| - ) |
70 | 67 | df: Annotated[pd.DataFrame, SimpleSchema]
|
71 |
| - else: |
72 | 68 |
|
| 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 |
73 | 75 | class Config:
|
74 | 76 | arbitrary_types_allowed = True
|
75 | 77 |
|
| 78 | + df: Annotated[pd.DataFrame, SimpleSchema] |
| 79 | + |
76 | 80 |
|
77 | 81 | def test_typed_dataframe():
|
78 | 82 | """Test that typed DataFrame is compatible with pydantic."""
|
|
0 commit comments