-
-
Notifications
You must be signed in to change notification settings - Fork 69
/
main.py
38 lines (33 loc) · 1.39 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
from typing import Container, Optional, Type
from pydantic import BaseConfig, BaseModel, create_model
from sqlalchemy.inspection import inspect
from sqlalchemy.orm.properties import ColumnProperty
class OrmConfig(BaseConfig):
orm_mode = True
def sqlalchemy_to_pydantic(
db_model: Type, *, config: Type = OrmConfig, exclude: Container[str] = []
) -> Type[BaseModel]:
mapper = inspect(db_model)
fields = {}
for attr in mapper.attrs:
if isinstance(attr, ColumnProperty):
if attr.columns:
name = attr.key
if name in exclude:
continue
column = attr.columns[0]
python_type: Optional[type] = None
if hasattr(column.type, "impl"):
if hasattr(column.type.impl, "python_type"):
python_type = column.type.impl.python_type
elif hasattr(column.type, "python_type"):
python_type = column.type.python_type
assert python_type, f"Could not infer python_type for {column}"
default = None
if column.default is None and not column.nullable:
default = ...
fields[name] = (python_type, default)
pydantic_model = create_model(
db_model.__name__, __config__=config, **fields # type: ignore
)
return pydantic_model