diff --git a/tests/test_pandas_typedframe.py b/tests/test_pandas_typedframe.py index c156282..9fcef8c 100644 --- a/tests/test_pandas_typedframe.py +++ b/tests/test_pandas_typedframe.py @@ -224,3 +224,23 @@ def test_multiple_inheritance_2_failure(): def test_multiple_inheritance_2_failure_with_root_overwrite(): with pytest.raises(AssertionError): _ = Down(pd.DataFrame({'root': [True], 'left': [True], 'right': ['string']})) + + +class ListFrame(TypedDataFrame): + schema = { + 'col1': list + } + + +def test_convert_list_attribute_type(): + _ = ListFrame.convert(pd.DataFrame({'col1': [["this"]]})) + + +class DictFrame(TypedDataFrame): + schema = { + 'col1': dict + } + + +def test_convert_dict_attribute_type(): + _ = DictFrame.convert(pd.DataFrame({'col1': [{"this": "that"}]})) diff --git a/tests/test_polars_typedframe.py b/tests/test_polars_typedframe.py index fea772f..4dabf9e 100644 --- a/tests/test_polars_typedframe.py +++ b/tests/test_polars_typedframe.py @@ -47,7 +47,7 @@ def test_base_success_case(): 'datetime_field': [datetime.datetime(2021, 5, 31, 12, 0, 0), datetime.datetime(2021, 6, 1, 12, 0, 0), datetime.datetime(2021, 6, 2, 12, 0, 0)], 'mixin_field': [1, 2, 3], 'new_field': [1, 2, 3]}) - df = df.with_column(pl.col('int_field').cast(pl.Int16)) + df = df.with_columns(pl.col('int_field').cast(pl.Int16)) _ = ChildDataFrame(df) diff --git a/typedframe/pandas_.py b/typedframe/pandas_.py index bacd26a..645d1ec 100644 --- a/typedframe/pandas_.py +++ b/typedframe/pandas_.py @@ -72,6 +72,8 @@ def convert(cls: Type[T], df: pd.DataFrame, add_optional_cols: bool = True) -> T if categories_diff: raise AssertionError(f"For column: {col} there are unknown categories: {categories_diff}") df[col] = pd.Categorical(df[col], categories=expected[col], ordered=True) + elif expected[col] in (list, dict): + df[col] = expected[col](df[col]) elif expected[col] == DATE_TIME_DTYPE: df[col] = pd.to_datetime(df[col]) elif expected[col] == UTC_DATE_TIME_DTYPE: