revng-revng/python/revng/model/_common/base.py

#
# This file is distributed under the MIT License. See LICENSE.md for details.
#
from typing import Dict

import yaml
from pydantic import BaseModel
from pydantic.main import ModelMetaclass


def get_monkey_patching_metaclass(
    substitutions: Dict[str, type],
):
    """Returns a metaclass which transparently substitutes some types with others. This way pydantic uses the substitute
    types when creating the validators for its types, without the need to modify the autogenerated classes.
    """

    class MonkeyPatchingMetaClass(ModelMetaclass):
        @staticmethod
        def __new__(mcs, clsname, bases, namespace):
            substitution = substitutions.get(clsname)
            if substitution is not None:
                return substitution

            created_class = super(MonkeyPatchingMetaClass, mcs).__new__(
                mcs, clsname, bases, namespace
            )

            if "__root__" in namespace:

                def yaml_representer(dumper: yaml.dumper.Dumper, instance):
                    classname = instance.__root__.__class__.__name__
                    tag = f"!{classname}"
                    return dumper.represent_mapping(
                        tag,
                        {k: v for k, v in instance.__root__._iter(exclude_none=True)},
                    )

            else:

                def yaml_representer(dumper: yaml.dumper.Dumper, instance):
                    return dumper.represent_dict(
                        {k: v for k, v in instance._iter(exclude_none=True)},
                    )

            yaml.add_representer(
                created_class,
                yaml_representer,
            )

            def yaml_constructor(loader, node):
                mapping = loader.construct_mapping(node, deep=True)
                return created_class(**mapping)

            tag = f"!{clsname}"
            yaml.add_constructor(tag, yaml_constructor)

            return created_class

    return MonkeyPatchingMetaClass


def get_monkey_patching_base_class(
    substitutions: Dict[str, type],
):
    """Returns a base class which transparently substitutes some model types with others"""
    MonkeyPatchingMetaClass = get_monkey_patching_metaclass(
        substitutions,
    )

    class MonkeyPatchingBaseClass(BaseModel, metaclass=MonkeyPatchingMetaClass):
        class Config:
            # Allows enums to be converted to strings when calling dict() on a model instance
            use_enum_values = True

    return MonkeyPatchingBaseClass