Использование в классах Pydantic полей типа numpy.array с поддержкой сериализации и десериализации в json

import json
from pydantic import BaseModel, field_serializer, field_validator
from typing import List
import numpy as np
 
class TestDataClass(BaseModel):
    id: int
    matrix: np.ndarray
    @field_serializer("matrix")   # Кастомный сериализатор для поля нестандартного типа (numpy array)
    def serialize_matrix(self, dt: np.ndarray, _info):
        ls = dt.tolist()
        dmp = json.dumps(ls)
        return dmp
 
    @field_validator("matrix", mode="before")  # Валидатор для поля нестандартного типа (numpy array), позволяющий 
                                                                    # присваивать ему значение типа List (вызывается в parse_raw)
    def check_matrix(cls, value):
        if isinstance(value, np.ndarray):
            return value
        elif isinstance(value, List):
            return np.array(value)
        elif isinstance(value, str):
            vl = json.loads(value)
            rv = np.array(vl)
            return rv
        else:
            raise ValueError("matrix may be numpy array or list")
 
    class Config:   # Разрешение использования нестандартного типа (numpy array)
        arbitrary_types_allowed = True

    to1 = TestDataClass(id=3802, matrix=np.array([[0, 0, 1], [0, 1, 0], [1, 0, 1]]))
 
    to1js = to1.json()
 
    print(json.dumps(to1js))
 
    to2 = TestDataClass.parse_raw(to1js)
 
    print(to2.stretch_id)

class TestDataClass(ExtendedBaseModel):
    id: int
    matrix: np.ndarray

import json
from pydantic import BaseModel, model_validator, model_serializer, Field
from pydantic._internal._model_construction import ModelMetaclass
import typing
from typing import List
import types
import inspect
import numpy as np
 
"""
Определяем расширение для типа BaseModel из библиотеки pydantic в виде класса-наследника ExtendedBaseModel(BaseModel),
которое позволяет определять в модели атрибуты нестандартных типов, в частности, типа numpy.array
и выполнять сериализацию и десериализацию объектов с этими атрибутами.
 
Для атрибутов типа numpy.array можно также указать дополнительные ограничения:
 
dimension_limits - определяется разрешенная размерность массива (количество измерений), а также верхний и нижний предел
                   для каждого измерения. 
                   Задаётся в виде списка пар [<мин.значение>, <макс.значения>] для каждого измерения:
                
                   Например:
                   
                   dimension_limits = [[3,3],[3,3]]     -  поле может быть только матрицей строгой размерности 3 на 3
                   
                   dimension_limits = [[2,None]]     -  поле может быть одномерным массивом минимум из двух элементов
 
elem_type - задаётся ограничение на тип элементов массива, например:
                   elem_type = float  
                   
Если эти ограничения заданы для определённых полей - то при инициализации значения этих полей (типа numpy.array) 
проверяются на соответcтвие этип ограничениям, и в случае нарушения генерируется ошибка
 
 
Если необходимо задать ограничения dimension_limits и/или elem_type, поле в классе должно описываться с помощью
конструкции Field, например:
 
class SomeClass(ExtendedBaseModel):
    ...
    array_field: np.ndarray = Field(dimension_limits = [[3,3],[3,3]], elem_type = float)
"""
 
 
""" 
Класс для хранения дополнительной информации о составе полей модели BaseModel в виде статических полей
ExtendedBaseModel будет содержать эти статические поля, т.к. от будет также и наследником ExtendedBaseModelInfo
через множественное наследование 
"""
class ExtendedBaseModelInfo(object):
    need_custom_json_serializing = False # Нужна ли кастомная сериализация (присутствуют ли поля нестандартных типов)
    list_of_basemodel_fields = []  # Список полей модели, которые в свою очередь являются списками объектов типа BaseModel
 
dimension_limits_error = "'dimension_limits' must be a list of pairs [<minimum size>, <maximum size>] defining limitations for array dimensions, where <minimum size> and <maximum size> are integer value or None"
 
"""
Чтобы поля, определённые в ExtendedBaseModelInfo, инициализировались сразу при определении классов-наследников ExtendedBaseModel
- используем магию метакласса (https://habr.com/ru/articles/145835/)
 
Т.К. для BaseBodel в pydantic уже определён метакласс - наш иетакласс должен быть его наследником    
 
"""
 
class ExtendedBaseModelMetaClass(ModelMetaclass):
    def __new__(mcs, *args, **kwargs) -> type:
 
        # Класс - наследник ExtendedBaseModel создаётся с помощью метода родительского метакласса ModelMetaclass
        ebm_class = super(ExtendedBaseModelMetaClass, mcs).__new__(mcs, *args, **kwargs)
        cls_name = args[0]
        if (cls_name != "ExtendedBaseModel") and (len(ebm_class.model_fields) > 0):
            # Перехватываем его после создания, когда в нём уже определены описания полей в списке model_fields
 
            # Проходим по списку полей
            for fln in ebm_class.model_fields:
 
                if ebm_class.model_fields[fln].exclude != True:
 
                    # Определяем, не является ли поле списком объектов, являющихся наследниками BaseModel
                    list_of_basemodel = False
                    if isinstance(ebm_class.model_fields[fln].annotation, types.GenericAlias) or isinstance(
                            ebm_class.model_fields[fln].annotation, typing._GenericAlias):
 
                        """ Как работает магтческий метод __reduce__ - не вполне понятно, но в случае, если тип поля (в annotation)
                            - это types.GenericAlias или typing._GenericAlias, вызов __reduce__ от этого типа
                            позволяет получить информацию об основном типе (напр. list или List), и о типе элемента 
                        """
                        rdc = ebm_class.model_fields[fln].annotation.__reduce__()
                        if len(rdc) == 2:
                            if len(rdc[1]) == 2:
                                if (rdc[1][0] is list) or (rdc[1][0] is List):
                                    if inspect.isclass(rdc[1][1]) == True:
                                        if issubclass(rdc[1][1], BaseModel):
                                            list_of_basemodel = True
                                            if issubclass(rdc[1][1], ExtendedBaseModel):
                                                ebm_class.need_custom_json_serializing = True
                                    elif isinstance(rdc[1][1], tuple):
                                        for te in rdc[1][1]:
                                            if inspect.isclass(te) == True:
                                                if issubclass(te, BaseModel):
                                                    list_of_basemodel = True
                                                    if issubclass(te, ExtendedBaseModel):
                                                        ebm_class.need_custom_json_serializing = True
                                                    break
                    if list_of_basemodel == True:
                        ebm_class.list_of_basemodel_fields.append(fln)
                    elif (ebm_class.model_fields[fln].annotation is np.ndarray) or (issubclass(ebm_class.model_fields[fln].annotation, ExtendedBaseModel)):
                        ebm_class.need_custom_json_serializing = True
 
                # Контролируем правильность значений в "dimension_limits" и/или "elem_type", если они заданы при определении поля
                if ebm_class.model_fields[fln].json_schema_extra != None:
                    if "dimension_limits" in ebm_class.model_fields[fln].json_schema_extra:
 
                        dls = ebm_class.model_fields[fln].json_schema_extra["dimension_limits"]
 
                        if isinstance(dls, list) == False:
                            raise ValueError("Class '"+cls_name+"', field '"+fln+"': "+dimension_limits_error)
 
                        if len(dls) < 1:
                            raise ValueError("Class '"+cls_name+"', field '"+fln+"': "+"'dimension_limits' list must contain at least one dimension limitation descriptor [<minimum size>, <maximum size>]")
 
                        for dl in dls:
                            if isinstance(dl, list) == False:
                                raise ValueError("Class '"+cls_name+"', field '"+fln+"': "+dimension_limits_error)
                            if len(dl) != 2:
                                raise ValueError("Class '"+cls_name+"', field '"+fln+"': "+dimension_limits_error)
 
                            if ((dl[0] != None) and (isinstance(dl[0], int) == False)) or (
                                    (dl[1] != None) and (isinstance(dl[1], int) == False)):
                                raise ValueError("Class '"+cls_name+"', field '"+fln+"': "+dimension_limits_error)
 
                            if (dl[0] != None) and (dl[1] != None):
                                if dl[0] > dl[1]:
                                    raise ValueError("Class '"+cls_name+"', field '"+fln+"': "+"'dimension_limits' list must contain limitation descriptors [<minimum size>, <maximum size>] where minimum size <= maximum size")
                    if "elem_type" in ebm_class.model_fields[fln].json_schema_extra:
                        if isinstance(ebm_class.model_fields[fln].json_schema_extra["elem_type"], type) == False:
                            raise ValueError("Class '"+cls_name+"', field '"+fln+"': 'elem_type' must  be a type")
 
        return ebm_class
 
 
class ExtendedBaseModel(BaseModel, ExtendedBaseModelInfo, metaclass=ExtendedBaseModelMetaClass):
    class Config:  # Разрешение использования нестандартных типов (numpy.array и пр)
        arbitrary_types_allowed = True
 
 
    # Кастомный валидатор свойств
    @model_validator(mode='before')
    @classmethod
    def validate_cust(cls, dt: any):
 
        """ При создании объекта конструктором класса (например:  p = Person(name="max", age=34) )
            в dt передаётся словарь со значениями атрибутов, например:
            {
                "name": "max",
                "age": 34
            }
        """
        data = dt
 
 
        """ А при создании объекта парсингом из json-строки  (например: p = <Имя класса>.parse_raw(<json-строка>) )
            в dt передаётся содержимое этой json-строки (тип str)
            и её надо загрузить в словарь с помощью стандартного метода json.loads
        """
        if isinstance(data, str):
            data = json.loads(data)
 
        if isinstance(data, dict):
            # Проходим по описаниям полей класса
            for fldn in cls.model_fields:
                # и если тип поля - нестандартный (например, numpy.array) - выполняем его кастомную инициализацию и валидацию
                if cls.model_fields[fldn].annotation is np.ndarray:
                    if fldn in data:
                        dim = None
                        elt = None
                        if cls.model_fields[fldn].json_schema_extra != None:
                            if "dimension_limits" in cls.model_fields[fldn].json_schema_extra:
                                dim = cls.model_fields[fldn].json_schema_extra["dimension_limits"]
                            if "elem_type" in cls.model_fields[fldn].json_schema_extra:
                                elt = cls.model_fields[fldn].json_schema_extra["elem_type"]
                        vl = data[fldn]
                        if isinstance(vl, list):
                            if elt != None:
                                data[fldn] = np.array(vl, dtype=elt)
                            else:
                                data[fldn] = np.array(vl)
                        elif isinstance(vl, str):
                            if elt != None:
                                data[fldn] = np.array(json.loads(vl), dtype=elt)
                            else:
                                data[fldn] = np.array(json.loads(vl))
                        # Проверка на соответствие типа элемента массива, если он задан
                        elif (isinstance(vl, np.ndarray) == True) and (elt != None):
                            if vl.dtype != elt:
                                raise ValueError(fldn + " field must be a numpy array of type '"+str(elt)+"'")
                        elif (isinstance(vl, np.ndarray) == False) and (vl != None):
                            raise ValueError(fldn +" field may be numpy array or list or json string for list")
                        # Проверка на соответствие размерности, если она задана
                        if dim != None:
                            rdim = data[fldn].shape
                            if len(rdim) != len(dim):
                                raise ValueError("the number of dimensions in "+ fldn + " bust be "+str(len(dim))+" as defined in 'dimension_limits'" )
                            dn = 0
                            for td in rdim:
                                if ((dim[dn][0] != None) and (td < dim[dn][0])) or ((dim[dn][1] != None) and (td > dim[dn][1])):
                                    raise ValueError("shape of " + fldn + " numpy array "+str(rdim)+" violates limitations defined in 'dimension_limits' "+str(dim))
                                dn = dn + 1
 
        return data
 
    # Кастомный сериализатор
    @model_serializer(mode='wrap')
    def serialize_cust(self, std, _info):
        if _info.mode == "json":
            """
            Если в классе есть нестандартные поля, требующие кастомной сериализации - генерируем json-строку "вручную", проходя по полям объекта
            """
            if self.need_custom_json_serializing == True:
                rv = "{"
                first_fld = True
                for fln in self.model_fields:
 
                    if self.model_fields[fln].exclude == True:
                        continue
 
                    if first_fld == False:
                        rv = rv + ","
                    first_fld = False
                    rv = rv + "\"" + fln + "\":"
 
                    flv = self.__getattribute__(fln)
                    if flv is None:
                        rv = rv + "null"
                    else:
                        # Если поле - это список объектов-наследников BaseModel
                        # - проходим по нему и вызываем метод .json() для каждого члена списка
                        if fln in self.list_of_basemodel_fields:
                            rv = rv + "["
                            first_list_elem = True
                            for le in flv:
                                if first_list_elem == False:
                                    rv = rv + ","
                                rv = rv + le.json()
                                first_list_elem = False
                            rv = rv + "]"
                        # для поля типа numpy.array преобразуем его значение обратно в список и сериализуем с помощью json
                        elif self.model_fields[fln].annotation is np.ndarray:
                            rv = rv + json.dumps(flv.tolist())
                        # Для поля, являющегося наследником BaseModel - вызываем метод .json()
                        elif isinstance(flv, BaseModel):
                            rv = rv + flv.json()
                        # Прочие поля стандартных типов сериализуем с помощью стандартной библиотеки json
                        else:
                            rv = rv + json.dumps(flv)
 
                rv = rv + "}"
                return rv
            else:
                # Если в классе нет полей нестандартных типов - используем стандартную серивлизацию от BaseModel
                return std(self)
        else:
            return std(self)

class Person(ExtendedBaseModel):
    name: str
    age: int
 
 
class TestDataClass(ExtendedBaseModel):
    id: int
    matrix: np.ndarray = Field(dimension_limits = [[3,3],[3,3]], elem_type = float)
    authors: List[Person]
 
if __name__ == '__main__':
 
    pr = Person(name="max", age=34)
 
    to1 = TestDataClass(id=1, matrix= [[0, 0, 1], [0, 1, 0], [1, 0, 1]], authors=[pr])
 
    to1js = to1.json()
 
    print(to1js)
 
    to2 = TestDataClass.parse_raw(to1js)
 
    print(to2.id)

import numpy as np
from pydantic import BaseModel
 
 
class Trixie(BaseModel):
    class Config:
        @staticmethod
        def ndarray_encode(value: np.ndarray):
            return {
                '__shape': value.shape,
                '__value': value.tolist(),
            }
 
        arbitrary_types_allowed = True
        json_encoders = {
            np.ndarray: ndarray_encode,
        }
 
 
class TestDataClass(Trixie):
    id: int
    matrix: np.ndarray
 
 
if __name__ == '__main__':
    to1 = TestDataClass(id=3802, matrix=np.array([[0, 0, 1], [0, 1, 0], [1, 0, 1]]))
    to1js = to1.model_dump_json()
    print(to1js)

                '__value': value.flatten().tolist(),