models.py

"""
This file should contain all relevant data-only information, it should contain minimal to no behavior.
"""

from typing import Dict, List, Optional

from pydantic import BaseModel, Field, model_validator


class Property(BaseModel):
    name: Optional[str] = None
    type: Optional[str] = None
    description: Optional[str] = None
    model_config = {"extra": "ignore"}


def upsert_dict(mapping, kvp):
    k, v = kvp
    mapping[k] = v
    return mapping


class DataModel(BaseModel):
    name: Optional[str] = None
    description: Optional[str] = None
    links: Optional[List[str]] = Field(default={})
    properties: Optional[Dict[str, Property]] = Field(default={})
    model_config = {"extra": "ignore"}

    @model_validator(mode="before")
    def _normalize_properties(cls, node_content):
        props = node_content.get("properties", {})
        if not isinstance(props, list):
            return node_content
        props_to_use = [prop for prop in props if "name" in prop]
        prop_items = [Property(**prop) for prop in props_to_use]
        prop_dict = {
            prop.name: Property(type=prop.type, description=prop.description)
            for prop in prop_items
        }
        return upsert_dict(node_content, ("properties", prop_dict))