Skip to content

Commit b347ba2

Browse files
author
Sam Goodwin
committed
add tests for complex python types
1 parent 8ab65a5 commit b347ba2

File tree

1 file changed

+24
-1
lines changed

1 file changed

+24
-1
lines changed

tests/core/test_typing.py

+24-1
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
"""Test typing annotations for the model api."""
22
# pylint:disable=missing-class-docstring,too-few-public-methods
33
import re
4-
from typing import Any, Dict, Optional, Type
4+
from typing import Any, Dict, NamedTuple, Optional, Type
55

66
import numpy as np
77
import pandas as pd
@@ -487,3 +487,26 @@ def test_init_pandas_dataframe_errors(invalid_data):
487487
"""Test errors from initializing a pandas.typing.DataFrame with Schema."""
488488
with pytest.raises(pa.errors.SchemaError):
489489
DataFrame[InitSchema](invalid_data)
490+
491+
class ExampleNamedTuple(NamedTuple):
492+
a: int
493+
b: float
494+
495+
class SchemaComplexPythonCollectionTypes(pa.DataFrameModel):
496+
list: list[pa.typing.Int32]
497+
dict: dict[str, pa.typing.Int32]
498+
tuple2: tuple[pa.typing.Int32, pa.typing.Int32]
499+
named_tuple: ExampleNamedTuple
500+
501+
def test_complex_python_collection_types():
502+
"""Test complex python collection types."""
503+
assert isinstance(
504+
DataFrame[SchemaComplexPythonCollectionTypes]({
505+
"list": [[1, 2], [3, 4, 5]],
506+
"dict": [{ "a": 1, "b": 2 }, { "c": 1, "d": 2 }] ,
507+
"tuple2": [[6, 7], [8, 9]],
508+
"named_tuple": [ExampleNamedTuple(1, 2), ExampleNamedTuple(3, 4)]
509+
}),
510+
DataFrame,
511+
)
512+

0 commit comments

Comments
 (0)