use Pydantic to create CSV line

Viewed 2602

We use Pydantic to set the "Domain Models" and use them across the application layer.

import datetime
from decimal import Decimal

from pydantic import BaseModel


class Person(BaseModel):
    name: str
    birth_date: datetime.date
    height: Decimal
    # other 100 fields, most of them Decimal

Now we need to export one of those models to a CSV, first implementation is easy:

import csv
from typing import Iterable


def store(persons: Iterable[Person]):
    fieldnames = list(Person.schema()["properties"].keys())
    
    with open("/tmp/test.csv", "w") as fp:
        writer = csv.DictWriter(fp, fieldnames=fieldnames)
        writer.writeheader()
        for person in persons:
            writer.writerow(person.dict())

But there is more, because we want all the decimals to be limited to two decimal places only, how can we achieve that?

Note: would not put serialization logic in the domain model, but I'm open to create a new model just for the serialization.

1 Answers

You can use a custom json encoder to round all Decimals in your model (Unfortuanatly it doesn't work to use Decimal as the result type in json_encoders as it's not JSON Serializable.

I've also defined a model PersonOut which inherits from Person so you don't have to store the serialization logic in the domain model. In the row writer the model is mapped to the serialization model.

import csv
import json
from decimal import Decimal
from pydantic import BaseModel
from typing import Iterable


class Person(BaseModel):
    """ Person domain model. """
    name: str
    height: Decimal


class PersonOut(Person):
    """ Person model used for serialization. """
    class Config:
        json_encoders = { Decimal: lambda v: float(round(v, 2)) }


def store(persons: Iterable[Person]):
    fieldnames = list(Person.schema()["properties"].keys())
    
    with open("test.csv", "w") as fp:
        writer = csv.DictWriter(fp, fieldnames=fieldnames)
        writer.writeheader()
        for person in persons:
            writer.writerow(json.loads(PersonOut(**person.dict()).json()))

if __name__=="__main__":
    persons = [Person(name='test', height=1.743)]

    store(persons)
Related