version with pydantic inheritance passing all the tests
This commit is contained in:
collerek
196
ormar/models/metaclass.py
Normal file
196
ormar/models/metaclass.py
Normal file
@ -0,0 +1,196 @@
|
||||
from typing import Any, Dict, List, Optional, TYPE_CHECKING, Tuple, Type, Union
|
||||
|
||||
import databases
|
||||
import pydantic
|
||||
import sqlalchemy
|
||||
from pydantic import BaseConfig, create_model, Extra
|
||||
from pydantic.fields import ModelField, FieldInfo
|
||||
|
||||
from ormar import ForeignKey, ModelDefinitionError # noqa I100
|
||||
from ormar.fields import BaseField
|
||||
from ormar.fields.foreign_key import ForeignKeyField
|
||||
from ormar.queryset import QuerySet
|
||||
from ormar.relations import RelationshipManager
|
||||
|
||||
if TYPE_CHECKING: # pragma no cover
|
||||
from ormar import Model
|
||||
|
||||
relationship_manager = RelationshipManager()
|
||||
|
||||
|
||||
class ModelMeta:
|
||||
tablename: str
|
||||
table: sqlalchemy.Table
|
||||
metadata: sqlalchemy.MetaData
|
||||
database: databases.Database
|
||||
columns: List[sqlalchemy.Column]
|
||||
pkname: str
|
||||
model_fields: Dict[str, Union[BaseField, ForeignKey]]
|
||||
_orm_relationship_manager: RelationshipManager
|
||||
|
||||
|
||||
def parse_pydantic_field_from_model_fields(object_dict: dict) -> Dict[str, Tuple]:
|
||||
pydantic_fields = {
|
||||
field_name: (
|
||||
base_field.__type__,
|
||||
... if base_field.is_required else base_field.default_value,
|
||||
)
|
||||
for field_name, base_field in object_dict.items()
|
||||
if isinstance(base_field, BaseField)
|
||||
}
|
||||
return pydantic_fields
|
||||
|
||||
|
||||
def register_relation_on_build(table_name: str, field: ForeignKey, name: str) -> None:
|
||||
child_relation_name = (
|
||||
field.to.get_name(title=True)
|
||||
+ "_"
|
||||
+ (field.related_name or (name.lower() + "s"))
|
||||
)
|
||||
reverse_name = child_relation_name
|
||||
relation_name = name.lower().title() + "_" + field.to.get_name()
|
||||
relationship_manager.add_relation_type(
|
||||
relation_name, reverse_name, field, table_name
|
||||
)
|
||||
|
||||
|
||||
def expand_reverse_relationships(model: Type["Model"]) -> None:
|
||||
for model_field in model.Meta.model_fields.values():
|
||||
if issubclass(model_field, ForeignKeyField):
|
||||
child_model_name = model_field.related_name or model.get_name() + "s"
|
||||
parent_model = model_field.to
|
||||
child = model
|
||||
if (
|
||||
child_model_name not in parent_model.__fields__
|
||||
and child.get_name() not in parent_model.__fields__
|
||||
):
|
||||
register_reverse_model_fields(parent_model, child, child_model_name)
|
||||
|
||||
|
||||
def register_reverse_model_fields(
|
||||
model: Type["Model"], child: Type["Model"], child_model_name: str
|
||||
) -> None:
|
||||
# model.__fields__[child_model_name] = ModelField(
|
||||
# name=child_model_name,
|
||||
# type_=Optional[Union[List[child], child]],
|
||||
# model_config=child.__config__,
|
||||
# class_validators=child.__validators__,
|
||||
# )
|
||||
model.Meta.model_fields[child_model_name] = ForeignKey(
|
||||
child, name=child_model_name, virtual=True
|
||||
)
|
||||
|
||||
|
||||
def sqlalchemy_columns_from_model_fields(
|
||||
name: str, object_dict: Dict, table_name: str
|
||||
) -> Tuple[Optional[str], List[sqlalchemy.Column], Dict[str, BaseField]]:
|
||||
columns = []
|
||||
pkname = None
|
||||
model_fields = {
|
||||
field_name: field
|
||||
for field_name, field in object_dict['__annotations__'].items()
|
||||
if issubclass(field, BaseField)
|
||||
}
|
||||
for field_name, field in model_fields.items():
|
||||
if field.primary_key:
|
||||
if pkname is not None:
|
||||
raise ModelDefinitionError("Only one primary key column is allowed.")
|
||||
if field.pydantic_only:
|
||||
raise ModelDefinitionError('Primary key column cannot be pydantic only')
|
||||
pkname = field_name
|
||||
if not field.pydantic_only:
|
||||
columns.append(field.get_column(field_name))
|
||||
if issubclass(field, ForeignKeyField):
|
||||
register_relation_on_build(table_name, field, name)
|
||||
|
||||
return pkname, columns, model_fields
|
||||
|
||||
|
||||
def populate_pydantic_default_values(attrs: Dict) -> Dict:
|
||||
for field, type_ in attrs["__annotations__"].items():
|
||||
if issubclass(type_, BaseField):
|
||||
if type_.name is None:
|
||||
type_.name = field
|
||||
def_value = type_.default_value()
|
||||
curr_def_value = attrs.get(field, 'NONE')
|
||||
print(field, curr_def_value, 'def val', type_.nullable)
|
||||
if curr_def_value == 'NONE' and isinstance(def_value, FieldInfo):
|
||||
attrs[field] = def_value
|
||||
elif curr_def_value == 'NONE' and type_.nullable:
|
||||
print(field, 'defsults tp none')
|
||||
attrs[field] = FieldInfo(default=None)
|
||||
return attrs
|
||||
|
||||
|
||||
def get_pydantic_base_orm_config() -> Type[BaseConfig]:
|
||||
class Config(BaseConfig):
|
||||
orm_mode = True
|
||||
arbitrary_types_allowed = True
|
||||
# extra = Extra.allow
|
||||
|
||||
return Config
|
||||
|
||||
|
||||
class ModelMetaclass(pydantic.main.ModelMetaclass):
|
||||
def __new__(mcs: type, name: str, bases: Any, attrs: dict) -> type:
|
||||
|
||||
attrs['Config'] = get_pydantic_base_orm_config()
|
||||
new_model = super().__new__( # type: ignore
|
||||
mcs, name, bases, attrs
|
||||
)
|
||||
|
||||
if hasattr(new_model, 'Meta'):
|
||||
|
||||
if attrs.get("__abstract__"):
|
||||
return new_model
|
||||
|
||||
annotations = attrs.get("__annotations__") or new_model.__annotations__
|
||||
attrs["__annotations__"]= annotations
|
||||
attrs = populate_pydantic_default_values(attrs)
|
||||
|
||||
print(attrs)
|
||||
|
||||
tablename = name.lower() + "s"
|
||||
new_model.Meta.tablename = new_model.Meta.tablename or tablename
|
||||
|
||||
# sqlalchemy table creation
|
||||
|
||||
pkname, columns, model_fields = sqlalchemy_columns_from_model_fields(
|
||||
name, attrs, new_model.Meta.tablename
|
||||
)
|
||||
|
||||
if hasattr(new_model.Meta, "model_fields") and not pkname:
|
||||
model_fields = new_model.Meta.model_fields
|
||||
for fieldname, field in new_model.Meta.model_fields.items():
|
||||
if field.primary_key:
|
||||
pkname=fieldname
|
||||
columns = new_model.Meta.table.columns
|
||||
|
||||
if not hasattr(new_model.Meta, "table"):
|
||||
new_model.Meta.table = sqlalchemy.Table(new_model.Meta.tablename, new_model.Meta.metadata, *columns)
|
||||
|
||||
new_model.Meta.columns = columns
|
||||
new_model.Meta.pkname = pkname
|
||||
|
||||
if not pkname:
|
||||
raise ModelDefinitionError("Table has to have a primary key.")
|
||||
|
||||
# pydantic model creation
|
||||
new_model.Meta.pydantic_fields = parse_pydantic_field_from_model_fields(attrs)
|
||||
new_model.Meta.pydantic_model = create_model(
|
||||
name, __config__=get_pydantic_base_orm_config(), **new_model.Meta.pydantic_fields
|
||||
)
|
||||
|
||||
new_model.Meta.model_fields = model_fields
|
||||
print(attrs, 'before super')
|
||||
print(new_model.Meta.__dict__)
|
||||
new_model = super().__new__( # type: ignore
|
||||
mcs, name, bases, attrs
|
||||
)
|
||||
expand_reverse_relationships(new_model)
|
||||
|
||||
new_model.Meta._orm_relationship_manager = relationship_manager
|
||||
new_model.objects = QuerySet(new_model)
|
||||
|
||||
# breakpoint()
|
||||
return new_model
|
||||
Reference in New Issue
Block a user