1
1
"""Test typing annotations for the model api."""
2
2
# pylint:disable=missing-class-docstring,too-few-public-methods
3
3
import re
4
- from typing import Any , Dict , Optional , Type
4
+ from typing import Any , Dict , NamedTuple , Optional , Type
5
5
6
6
import numpy as np
7
7
import pandas as pd
@@ -487,3 +487,26 @@ def test_init_pandas_dataframe_errors(invalid_data):
487
487
"""Test errors from initializing a pandas.typing.DataFrame with Schema."""
488
488
with pytest .raises (pa .errors .SchemaError ):
489
489
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