69 lines
2.4 KiB
Python
69 lines
2.4 KiB
Python
|
from typing import Container, Optional, Type
|
||
|
|
|
||
|
|
from pydantic import BaseConfig, BaseModel, create_model
|
||
|
|
from sqlalchemy.inspection import inspect
|
||
|
|
from sqlalchemy.orm.properties import ColumnProperty
|
||
|
|
from hashlib import md5
|
||
|
|
|
||
|
|
|
||
|
|
# 文件md
|
||
|
|
def file_md5(body):
|
||
|
|
md = md5()
|
||
|
|
md.update(body)
|
||
|
|
return md.hexdigest()
|
||
|
|
|
||
|
|
|
||
|
|
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
|
||
|
|
|
||
|
|
|
||
|
|
def get_sqlalchemy_model_fields(db_model):
|
||
|
|
mapper = inspect(db_model)
|
||
|
|
fields = {}
|
||
|
|
for attr in mapper.attrs:
|
||
|
|
if isinstance(attr, ColumnProperty):
|
||
|
|
if attr.columns:
|
||
|
|
name = attr.key
|
||
|
|
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)
|
||
|
|
return fields
|