✨ Add support for excluding columns from Pydantic model (tiangolo#3)

tiangolo · web-flow · commit d4ba9b670295 · 2020-05-05T13:03:06.000+02:00
diff --git 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
@@ -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},
+        ],
+    }
