From 79a895f3eeb45df67306d42913d98912f195714a Mon Sep 17 00:00:00 2001 From: Henry Walshaw Date: Thu, 12 Sep 2024 14:04:11 +1000 Subject: [PATCH 1/6] Add failing test to demonstrate issue From the [docs](https://docs.pydantic.dev/latest/concepts/fields/#exclude): > The `exclude` parameter can be used to control which fields should > be excluded from the model when exporting the model. At the moment the values are excluded as expected, but the headers are not correctly filtered. This is worse than we expect because it also breaks the order of the fields rather than leaving the field blank. --- tests/models.py | 5 +++++ tests/test_basemodel_csv_writer.py | 12 +++++++++++- 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/tests/models.py b/tests/models.py index bce8c3c..e0655f0 100644 --- a/tests/models.py +++ b/tests/models.py @@ -51,3 +51,8 @@ def parse_start_date(cls, value): @pydantic.field_validator("end", mode="before") def parse_end_date(cls, value): return datetime.strptime(value, "%d.%m.%Y").date() + + +class ExcludedPassword(pydantic.BaseModel): + username: str = "Wagstaff" + password: str = Field(default="swordfish", exclude=True) diff --git a/tests/test_basemodel_csv_writer.py b/tests/test_basemodel_csv_writer.py index aea54bc..d8155af 100644 --- a/tests/test_basemodel_csv_writer.py +++ b/tests/test_basemodel_csv_writer.py @@ -4,7 +4,7 @@ from pydantic_csv import BasemodelCSVWriter -from .models import NonBaseModelUser, SimpleUser, User +from .models import ExcludedPassword, NonBaseModelUser, SimpleUser, User def test_create_csv_file(users_as_csv_buffer, users_from_csv): @@ -50,3 +50,13 @@ def test_with_wrong_type_in_list(user_list): def test_header_mapping(users_mapped_as_csv_buffer, users_mapped_from_csv): assert users_mapped_as_csv_buffer == users_mapped_from_csv + + +def test_excluded_field(): + output = io.StringIO() + user = ExcludedPassword() + + w = BasemodelCSVWriter(output, [user], ExcludedPassword) + w.write() + + assert output.getvalue() == "username\r\nWagstaff\r\n" From 3afc80f2587d5fcda00f25d767ac4c063ed9309f Mon Sep 17 00:00:00 2001 From: Henry Walshaw Date: Thu, 12 Sep 2024 14:10:42 +1000 Subject: [PATCH 2/6] Add the fix that filters out the excluded fields altogether --- pydantic_csv/basemodel_csv_writer.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/pydantic_csv/basemodel_csv_writer.py b/pydantic_csv/basemodel_csv_writer.py index d7210d5..d58fad8 100644 --- a/pydantic_csv/basemodel_csv_writer.py +++ b/pydantic_csv/basemodel_csv_writer.py @@ -43,10 +43,12 @@ def __init__( self._model = model self._field_mapping: dict[str, str] = {} + fields = {name: field for name, field in self._model.model_fields.items() if not (field.exclude or False)} + if use_alias: - self._fieldnames = [field.alias or name for name, field in self._model.model_fields.items()] + self._fieldnames = [field.alias or name for name, field in fields.items()] else: - self._fieldnames = model.model_fields.keys() + self._fieldnames = fields.keys() self._writer = csv.writer(file_obj, dialect=dialect, **kwargs) From 56d2432299cafd52ed8357743c2f68120bb4180b Mon Sep 17 00:00:00 2001 From: Henry Walshaw Date: Fri, 13 Sep 2024 13:57:58 +1000 Subject: [PATCH 3/6] Add a failing test for computed properties on export --- tests/models.py | 18 +++++++++++++++++- tests/test_basemodel_csv_writer.py | 28 +++++++++++++++++++++++++++- 2 files changed, 44 insertions(+), 2 deletions(-) diff --git a/tests/models.py b/tests/models.py index e0655f0..ce861e8 100644 --- a/tests/models.py +++ b/tests/models.py @@ -2,7 +2,7 @@ from typing import Optional import pydantic -from pydantic import Field +from pydantic import Field, computed_field class User(pydantic.BaseModel): @@ -56,3 +56,19 @@ def parse_end_date(cls, value): class ExcludedPassword(pydantic.BaseModel): username: str = "Wagstaff" password: str = Field(default="swordfish", exclude=True) + + +class ComputedPropertyField(pydantic.BaseModel): + username: str = "Groucho" + + @computed_field + def email(self) -> str: + return f"{self.username.lower()}@marx.bros" + + +class ComputedPropertyWithAlias(pydantic.BaseModel): + username: str = "Harpo" + + @computed_field(alias="e") + def email(self) -> str: + return f"{self.username.lower()}@marx.bros" diff --git a/tests/test_basemodel_csv_writer.py b/tests/test_basemodel_csv_writer.py index d8155af..00a89a0 100644 --- a/tests/test_basemodel_csv_writer.py +++ b/tests/test_basemodel_csv_writer.py @@ -4,7 +4,14 @@ from pydantic_csv import BasemodelCSVWriter -from .models import ExcludedPassword, NonBaseModelUser, SimpleUser, User +from .models import ( + ComputedPropertyField, + ComputedPropertyWithAlias, + ExcludedPassword, + NonBaseModelUser, + SimpleUser, + User, +) def test_create_csv_file(users_as_csv_buffer, users_from_csv): @@ -60,3 +67,22 @@ def test_excluded_field(): w.write() assert output.getvalue() == "username\r\nWagstaff\r\n" + + +@pytest.mark.parametrize( + ("model", "use_alias", "expected_output"), + [ + (ComputedPropertyField, True, "username,email\r\nGroucho,groucho@marx.bros\r\n"), + (ComputedPropertyWithAlias, True, "username,e\r\nHarpo,harpo@marx.bros\r\n"), + (ComputedPropertyField, False, "username,email\r\nGroucho,groucho@marx.bros\r\n"), + (ComputedPropertyWithAlias, False, "username,email\r\nHarpo,harpo@marx.bros\r\n"), + ], +) +def test_computed_property_included(model, use_alias, expected_output): + output = io.StringIO() + user = model() + + w = BasemodelCSVWriter(output, [user], model, use_alias=use_alias) + w.write() + + assert output.getvalue() == expected_output From a1852783d21deb60af63a3bda75c04b49531e37c Mon Sep 17 00:00:00 2001 From: Henry Walshaw Date: Fri, 13 Sep 2024 14:39:29 +1000 Subject: [PATCH 4/6] Deal with computed fields and related alias bug This involves switching to a dict writer, and making user whether to use aliases is passed through to the write method. --- pydantic_csv/basemodel_csv_writer.py | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/pydantic_csv/basemodel_csv_writer.py b/pydantic_csv/basemodel_csv_writer.py index d58fad8..b33f541 100644 --- a/pydantic_csv/basemodel_csv_writer.py +++ b/pydantic_csv/basemodel_csv_writer.py @@ -4,6 +4,7 @@ import csv from collections.abc import Iterable +from itertools import chain from typing import Any import pydantic @@ -43,24 +44,30 @@ def __init__( self._model = model self._field_mapping: dict[str, str] = {} - fields = {name: field for name, field in self._model.model_fields.items() if not (field.exclude or False)} + fields = { + name: field + for name, field in chain(self._model.model_fields.items(), self._model.model_computed_fields.items()) + if not (getattr(field, "exclude", False) or False) + } - if use_alias: + self._use_alias = use_alias + + if self._use_alias: self._fieldnames = [field.alias or name for name, field in fields.items()] else: self._fieldnames = fields.keys() - self._writer = csv.writer(file_obj, dialect=dialect, **kwargs) + self._writer = csv.DictWriter(file_obj, self._fieldnames, dialect=dialect, **kwargs) def _add_to_mapping(self, header: str, fieldname: str) -> None: self._field_mapping[fieldname] = header - def _apply_mapping(self) -> list[str]: - mapped_fields = [] + def _apply_mapping(self) -> dict[str, str]: + mapped_fields = {} for field in self._fieldnames: mapped_item = self._field_mapping.get(field, field) - mapped_fields.append(mapped_item) + mapped_fields[field] = mapped_item return mapped_fields @@ -75,11 +82,12 @@ def write(self, skip_header: bool = False) -> None: Returns: None: well, nothing """ + if not skip_header: if self._field_mapping: - self._fieldnames = self._apply_mapping() - - self._writer.writerow(self._fieldnames) + self._writer.writerow(self._apply_mapping()) + else: + self._writer.writeheader() for item in self._data: if not isinstance(item, self._model): @@ -88,7 +96,7 @@ def write(self, skip_header: bool = False) -> None: f"{self._model.__name__}. All items on the list must be " "instances of the same type" ) - row = item.model_dump().values() + row = item.model_dump(by_alias=self._use_alias) self._writer.writerow(row) def map(self, fieldname: str) -> HeaderMapper: From 21316e7a231fbcdd45971586fdcfc1f159f8405c Mon Sep 17 00:00:00 2001 From: Henry Walshaw Date: Mon, 16 Sep 2024 15:17:31 +1000 Subject: [PATCH 5/6] Bump the version --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 97a3365..575fe4c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "pydantic-csv" -version = "0.1.0" +version = "0.1.1" description = "convert CSV to pydantic.BaseModel and vice versa" authors = ["Nathan Richard "] license = "LICENSE" From 334f40035b3ba5fd097211aad33a7b5d0100ca91 Mon Sep 17 00:00:00 2001 From: Henry Walshaw Date: Tue, 17 Sep 2024 11:51:41 +1000 Subject: [PATCH 6/6] Don't forget the serialization alias! --- pydantic_csv/basemodel_csv_writer.py | 4 +++- pyproject.toml | 2 +- tests/models.py | 1 + tests/test_basemodel_csv_writer.py | 2 +- 4 files changed, 6 insertions(+), 3 deletions(-) diff --git a/pydantic_csv/basemodel_csv_writer.py b/pydantic_csv/basemodel_csv_writer.py index b33f541..eaa6b6c 100644 --- a/pydantic_csv/basemodel_csv_writer.py +++ b/pydantic_csv/basemodel_csv_writer.py @@ -53,7 +53,9 @@ def __init__( self._use_alias = use_alias if self._use_alias: - self._fieldnames = [field.alias or name for name, field in fields.items()] + self._fieldnames = [ + field.alias or getattr(field, "serialization_alias", None) or name for name, field in fields.items() + ] else: self._fieldnames = fields.keys() diff --git a/pyproject.toml b/pyproject.toml index 575fe4c..05f3400 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "pydantic-csv" -version = "0.1.1" +version = "0.1.2" description = "convert CSV to pydantic.BaseModel and vice versa" authors = ["Nathan Richard "] license = "LICENSE" diff --git a/tests/models.py b/tests/models.py index ce861e8..792854a 100644 --- a/tests/models.py +++ b/tests/models.py @@ -56,6 +56,7 @@ def parse_end_date(cls, value): class ExcludedPassword(pydantic.BaseModel): username: str = "Wagstaff" password: str = Field(default="swordfish", exclude=True) + email: str = Field(default="wagstaff@marx.bros", serialization_alias="contact") class ComputedPropertyField(pydantic.BaseModel): diff --git a/tests/test_basemodel_csv_writer.py b/tests/test_basemodel_csv_writer.py index 00a89a0..eac58d5 100644 --- a/tests/test_basemodel_csv_writer.py +++ b/tests/test_basemodel_csv_writer.py @@ -66,7 +66,7 @@ def test_excluded_field(): w = BasemodelCSVWriter(output, [user], ExcludedPassword) w.write() - assert output.getvalue() == "username\r\nWagstaff\r\n" + assert output.getvalue() == "username,contact\r\nWagstaff,wagstaff@marx.bros\r\n" @pytest.mark.parametrize(