From 62cc5d2b3018a9d8719f3c8ebafe1eaabda9be79 Mon Sep 17 00:00:00 2001 From: Stijn de Gooijer Date: Mon, 24 Jun 2024 16:18:39 +0200 Subject: [PATCH] Add namedtuple test --- .../polars/_utils/construction/dataframe.py | 2 +- py-polars/tests/unit/datatypes/test_parse.py | 16 +++++++++++++++- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/py-polars/polars/_utils/construction/dataframe.py b/py-polars/polars/_utils/construction/dataframe.py index 676445aa80711..7f15c38cb90de 100644 --- a/py-polars/polars/_utils/construction/dataframe.py +++ b/py-polars/polars/_utils/construction/dataframe.py @@ -650,7 +650,7 @@ def _sequence_of_tuple_to_pydf( orient: Orientation | None, infer_schema_length: int | None, ) -> PyDataFrame: - # infer additional meta information if named tuple + # infer additional meta information if namedtuple if is_namedtuple(first_element.__class__): if schema is None: schema = first_element._fields # type: ignore[attr-defined] diff --git a/py-polars/tests/unit/datatypes/test_parse.py b/py-polars/tests/unit/datatypes/test_parse.py index f7861b6e4fa76..1e376d27690e1 100644 --- a/py-polars/tests/unit/datatypes/test_parse.py +++ b/py-polars/tests/unit/datatypes/test_parse.py @@ -1,7 +1,7 @@ from __future__ import annotations from datetime import date, datetime -from typing import TYPE_CHECKING, Any, ForwardRef, Optional, Union +from typing import TYPE_CHECKING, Any, ForwardRef, NamedTuple, Optional, Union import pytest @@ -108,3 +108,17 @@ def test_parse_union_type_into_dtype(input: Any, expected: PolarsDataType) -> No def test_parse_union_type_into_dtype_invalid(input: Any) -> None: with pytest.raises(TypeError): _parse_union_type_into_dtype(input) + + +def test_parse_dtype_namedtuple_fields() -> None: + # Utilizes ForwardRef parsing + + class MyTuple(NamedTuple): + a: str + b: int + c: str | None = None + + schema = {c: parse_into_dtype(a) for c, a in MyTuple.__annotations__.items()} + + expected = pl.Schema({"a": pl.String(), "b": pl.Int64(), "c": pl.String()}) + assert schema == expected