From d9fa39df1fb7185df306748071ca9012f13d20e9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sebasti=C3=A1n=20Ram=C3=ADrez?= Date: Tue, 5 May 2020 13:00:55 +0200 Subject: [PATCH] :sparkles: Add support for excluding columns from Pydantic model --- pydantic_sqlalchemy/main.py | 6 ++++-- tests/test_pydantic_sqlalchemy.py | 21 +++++++++++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/pydantic_sqlalchemy/main.py b/pydantic_sqlalchemy/main.py index 6649c1e..a6c4644 100644 --- a/pydantic_sqlalchemy/main.py +++ b/pydantic_sqlalchemy/main.py @@ -1,4 +1,4 @@ -from typing import Type +from typing import Container, Type from pydantic import BaseConfig, BaseModel, create_model from sqlalchemy.inspection import inspect @@ -10,7 +10,7 @@ class OrmConfig(BaseConfig): def sqlalchemy_to_pydantic( - db_model: Type, *, config: Type = OrmConfig + db_model: Type, *, config: Type = OrmConfig, exclude: Container[str] = [] ) -> Type[BaseModel]: mapper = inspect(db_model) fields = {} @@ -20,6 +20,8 @@ def sqlalchemy_to_pydantic( column = attr.columns[0] python_type = column.type.python_type name = attr.key + if name in exclude: + continue default = column.default if default is None and not column.nullable: default = ... diff --git a/tests/test_pydantic_sqlalchemy.py b/tests/test_pydantic_sqlalchemy.py index ace02ed..45c4513 100644 --- a/tests/test_pydantic_sqlalchemy.py +++ b/tests/test_pydantic_sqlalchemy.py @@ -108,3 +108,24 @@ class PydanticUserWithAddresses(PydanticUser): {"emailAddress": "eddy@example.com", "id": 2, "userId": 1}, ], } + + +def test_exclude() -> None: + PydanticUser = sqlalchemy_to_pydantic(User, exclude={"nickname"}) + PydanticAddress = sqlalchemy_to_pydantic(Address, exclude={"user_id"}) + + class PydanticUserWithAddresses(PydanticUser): + addresses: List[PydanticAddress] = [] + + user = db.query(User).first() + pydantic_user_with_addresses = PydanticUserWithAddresses.from_orm(user) + data = pydantic_user_with_addresses.dict(by_alias=True) + assert data == { + "fullname": "Ed Jones", + "id": 1, + "name": "ed", + "addresses": [ + {"email_address": "ed@example.com", "id": 1}, + {"email_address": "eddy@example.com", "id": 2}, + ], + }