Skip to content

Commit

Permalink
fix: handle multiple annotations when dtype is present in Field
Browse files Browse the repository at this point in the history
  • Loading branch information
brendancooley committed Nov 3, 2023
1 parent ffe343d commit e2bf0d7
Show file tree
Hide file tree
Showing 4 changed files with 14 additions and 1 deletion.
2 changes: 2 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -132,3 +132,5 @@ dmypy.json

# Pyre type checker
.pyre/

.vscode/
3 changes: 2 additions & 1 deletion src/patito/pydantic.py
Original file line number Diff line number Diff line change
Expand Up @@ -218,7 +218,8 @@ def _valid_dtypes( # noqa: C901
f"No valid dtype mapping found for column '{column}'."
)
return [pl.List(dtype) for dtype in item_dtypes]
if "dtype" in props:

if "dtype" in props and 'anyOf' not in props:
return [
props["dtype"],
]
Expand Down
8 changes: 8 additions & 0 deletions tests/test_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -454,3 +454,11 @@ class Model(pt.Model):
assert fields["d"].dtype is not None
assert "unique" in fields["e"]._attributes_set
assert fields["e"].unique is not None


def test_nullable_columns():
class Test(pt.Model):
foo: str | None = pt.Field(dtype=pl.Utf8)

assert Test.nullable_columns == {"foo"}
assert set(Test.valid_dtypes['foo']) == {pl.Utf8, pl.Null}
2 changes: 2 additions & 0 deletions tests/test_polars.py
Original file line number Diff line number Diff line change
Expand Up @@ -253,3 +253,5 @@ class Model(pt.Model):

# Or a list of columns
assert df.drop(["column_1", "column_2"]).columns == []


0 comments on commit e2bf0d7

Please sign in to comment.