group_manage/ruoyi_common/base/schema_vo.py

# -*- coding: utf-8 -*-
# @Author  : YY

from abc import ABC, abstractmethod
from datetime import datetime
from typing import Type, Dict, Optional, Set, Tuple, TypeVar
from pydantic.alias_generators import to_camel
from pydantic import AliasChoices, AliasGenerator, AliasPath, BaseModel, ConfigDict, Field, create_model
from pydantic.fields import FieldInfo
from pydantic_core import PydanticUndefined
from pydantic import BaseModel

from ruoyi_common.base.model import AllowedExtraModel, BaseEntity, BetOpt, ExtraModel, ExtraOpt, MultiFile, VoAccess, VoValidatorContext
from ruoyi_common.utils.base import DateUtil, get_final_model


T = TypeVar("T")

strict_valid_config = ConfigDict(
    from_attributes = False,
    alias_generator = to_camel,  
    frozen = True,
    extra = "forbid",
    strict = True,
    populate_by_name = False,
    json_encoders = {
        datetime: lambda v: v.strftime(DateUtil.YYYY_MM_DD_HH_MM_SS)
    },
)

query_valid_config = ConfigDict(
    from_attributes = False,
    alias_generator = to_camel,  
    frozen = True,
    extra = "allow",
    strict = True,
    populate_by_name = False,
)


def VoField(
    body=True,
    query=False,
    sort=False, 
    *args, 
    **kwargs
    ):
    vo = VoAccess(
        body=body,
        query=query,
        sort=sort,
    )
    return Field(vo=vo,*args, **kwargs)



class AbcFieldFilter(ABC):
    
    def filter(self, name:str, info:FieldInfo) -> bool:
        """
        过滤字段
        
        Args:
            name(str): 字段名称
            info(FieldInfo): 字段元信息
        
        Returns:
            bool: 是否过滤
        """


class BaseFieldFilter(AbcFieldFilter):
    
    def filter(self, name:str, info:FieldInfo) -> bool:
        """
        过滤字段
        
        Args:
            name(str): 字段名称
            info(FieldInfo): 字段元信息
        
        Returns:
            bool: 是否过滤
        """
        return False
    
    
class VoBodyFieldFilter(BaseFieldFilter):
    
    def __init__(self):
        self.action = "body"
    
    def filter(self, name:str, info:FieldInfo) -> bool:
        """
        过滤字段
        
        Args:
            name(str): 字段名称
            info(FieldInfo): 字段元信息
        
        Returns:
            bool: 是否过滤
        """
        default = True
        if info.json_schema_extra:
            vo_access:VoAccess = info.json_schema_extra.get("vo",default)
            if vo_access:
                perm = getattr(vo_access,self.action,default)
                if perm:
                    is_required = getattr(vo_access,"body_required",False)
                    if is_required:
                        self.change_to_required(info)
                    else:
                        self.change_to_optional(info)
                flag = perm
            else:
                flag = default
        else:
            flag = default
        return flag
    
    @classmethod
    def change_to_required(cls, info:FieldInfo):
        """
        将注解转换为必选项
        
        Args:
            info(FieldInfo): 字段元信息
        """
        info.default = PydanticUndefined
        info.default_factory = None
    
    @classmethod
    def change_to_optional(cls, info:FieldInfo):
        """
        将注解转换为可选项
        
        Args:
            info(FieldInfo): 字段元信息
        """
        if info.is_required:
            info.default = None


class VoQueryFieldFilter(BaseFieldFilter):
    
    extra_opt_cls_list = [BetOpt]
    
    def __init__(self):
        self.action = "query"
        self.sort_fields:Dict[str,FieldInfo] = {}
        self.extra_fields:Dict[str,ExtraOpt] = {}
    
    def filter(self, name:str, info:FieldInfo) -> bool:
        """
        根据权限信息，重置字段
        
        Args:
            name(str): 字段名称
            info(FieldInfo): 字段元信息
        
        Returns:
            Tuple[bool,bool]: 是否支持查询，是否支持排序
        """
        default = False
        if info.json_schema_extra:
            vo_access:VoAccess = info.json_schema_extra.get("vo",False)
            query_perm = getattr(vo_access,self.action,default)
            for extra_opt_cls in self.extra_opt_cls_list:
                if isinstance(query_perm, extra_opt_cls):
                    query_perm.name = name
                    query_perm.info = info
                    self.extra_fields[name] = query_perm
                    query_perm = False
                else:
                    continue
            sort_perm = getattr(vo_access,"sort",default)
            if sort_perm:
                self.sort_fields[name] = info
            return query_perm
        else:
            return default
    

class AbcSchemaFactory(ABC):
    
    @abstractmethod
    def validate_annotation(annotation:Type) -> Optional[Type]:
        """
        检查注解是否有效
        
        Args:    
            annotation(Type): 注解
        
        Returns:
            Optional[Type]: 合法的类型
        """
        pass
    

class BaseSchemaFactory(AbcSchemaFactory):
    
    action = "base"
    model_config = query_valid_config
            
    def __init__(self):
        self.field_filter = None
        self.model_suffix = "Vo"
    
    def validate_annotation(self, annotation:Type) -> Optional[Type]:
        """
        检查注解是否有效
        
        Args:
            annotation(Type): 注解
        
        Returns:
            Optional[Type]: 合法的类型
        """
        bo_model = get_final_model(annotation)
        if issubclass(bo_model, BaseEntity):
            model = self.rebuild_model(model_cls=bo_model)
            return model
        else:
            return None
        
    def rebuild_model(self, model_cls:Type[BaseEntity]) -> Type[BaseEntity]:
        """
        从已有模型类，创建新的模型类
        
        Args:
            model_cls: 已有模型类
        
        Returns:
            Type[BaseModel]: 新的模型类
        """
        
        field_definitions = {}
        for name,info in model_cls.model_fields.items():
            flag = self.rebuild_field(name,info)
            if flag:
                field_definitions[name] = info.annotation,info
        vo_name = model_cls.__name__ + self.action.capitalize() + \
            self.model_suffix

        return create_model(
            vo_name, 
            __base__= model_cls,
            __doc__ = model_cls.__doc__, 
            __module__= model_cls.__module__, 
            **field_definitions
        )
    
    def rebuild_field(self, name:str, info:FieldInfo) -> bool:
        """
        重置字段
        
        Args:
            name(str): 字段名称
            info(FieldInfo): 字段元信息
        
        Returns:
            bool: 是否支持重置
        """
        return True


class BodySchemaFactory(BaseSchemaFactory):
    
    action = "body"
    model_config = query_valid_config
    
    def __init__(self,context:VoValidatorContext):
        super().__init__()
        self.context = context
        self.field_filter = VoBodyFieldFilter()
        
    def validate_annotation(self, annotation: Type) -> Optional[Type[BaseModel]]:
        bo_model = get_final_model(annotation)
        if issubclass(bo_model, BaseModel):
            updated_model = self.rebuild_model(model_cls=bo_model)
            return updated_model
        else:
            if self.context.include_fields or self.context.exclude_fields:
                raise Exception(f"注解{annotation.__name__}不是模型，不支持include和exclude请求条件")
            return None
        
    def rebuild_field(self, name:str, info:FieldInfo) -> bool:
        """
        重置字段
        
        Args:
            name(str): 字段名称
            info(FieldInfo): 字段元信息
        
        Returns:
            bool: 是否支持重置
        """
        if self.context.include_fields and \
            name not in self.context.include_fields:
            return False
        if self.context.exclude_fields and \
            name in self.context.exclude_fields:
            return False
        flag = self.field_filter.filter(name,info)
        return flag


class QuerySchemaFactory(BaseSchemaFactory):
    
    action = "query"
    model_config = query_valid_config
    
    def __init__(
        self,
        context,
        extra_strict_forbid=True,
        extra_allowed_fields:Type[Dict[str,FieldInfo]]=None
    ):
        super().__init__()
        self.context = context
        self.field_filter = VoQueryFieldFilter()
        self.extra_strict_forbid = extra_strict_forbid
        self.extra_allowed_fields = extra_allowed_fields
        self.extra_model = None
    
    def validate_annotation(self, annotation: Type) -> Optional[Type[BaseEntity]]:
        bo_model = get_final_model(annotation)
        if issubclass(bo_model, BaseEntity):
            updated_model = self.rebuild_model(model_cls=bo_model)
            self.model_config = annotation.model_config.copy()
            self._validate_context()
            self.rebuild_extra_model()
            return updated_model
        else:
            if self.context.include_fields or self.context.exclude_fields:
                raise Exception(f"注解{annotation.__name__}不是模型，不支持include和exclude请求条件")
            return None
    
    def _validate_context(self):
        """
        验证上下文信息
        """
        for name,info in self.field_filter.sort_fields.items():
            alias_set = self.get_validate_alias(name,info)
            self.context.include_sort_alias = self.context.include_sort_alias | alias_set
        
    def get_validate_alias(self, name:str,info:FieldInfo) -> Set[str]:
        """
        获取验证别名

        Args:
            name (str): 字段名称
            info (FieldInfo): 字段元信息

        Raises:
            Exception: 模型不支持AliasPath

        Returns:
            Set[str]: 别名集合
        """
        alias_set = set()
        alias = self.get_alias_from_config(name)
        if alias:
            alias_set.add(alias)
        if info.validation_alias:
            if isinstance(info.validation_alias, str):
                alias_set.add(info.validation_alias)
            elif isinstance(info.validation_alias, AliasPath):
                raise Exception(f"模型字段{name}不支持AliasPath")
            elif isinstance(info.validation_alias, AliasChoices):
                alias_set = alias_set | \
                    set(info.validation_alias.choices)
        if "populate_by_name" in self.model_config \
            and self.model_config["populate_by_name"]:
            alias_set.add(name)
        return alias_set
    
    def get_alias_from_config(self,name:str)-> Optional[str]:
        """
        从配置中获取别名
        
        Args:
            name (str): 字段名称
        
        Returns:
            Optional[str]: 别名
        """
        if "generate_alias" in self.model_config:
            generate_alias = self.model_config["generate_alias"]
            if callable(generate_alias):
                alias = generate_alias(name)
                return alias
            elif isinstance(generate_alias, AliasGenerator):
                alias,v_alias,s_alias = generate_alias(name)
                return alias or v_alias
    
    def rebuild_extra_model(self) -> Optional[Type[ExtraModel]]:
        """
        重新构建Extra查询模型
        
        Args:
            fields (Dict[str,ExtraOpt]): 字段元信息
        
        Returns:
            Optional[Type[ExtraModel]]: extra查询模型
        """
        field_defintions = {}
        for name,opt in self.field_filter.extra_fields.items():
            if isinstance(opt, BetOpt):
                min_fieldinfo,max_fieldinfo = self.rebuild_bet_opt(name,opt)
                field_defintions[opt.min] = min_fieldinfo.annotation,min_fieldinfo
                field_defintions[opt.max] = max_fieldinfo.annotation,max_fieldinfo
            elif isinstance(opt, ExtraOpt):
                field_defintions[name] = opt.info.annotation,opt.info
            else:
                continue
        
        if self.extra_allowed_fields:
            for name,info in self.extra_allowed_fields.items():
                if name not in field_defintions:
                    field_defintions[name] = info.annotation,info

        if field_defintions:
            extra_model_cls = ExtraModel if self.extra_strict_forbid else AllowedExtraModel
            self.extra_model = create_model(
                model_name="ExtraModel",
                __base__=extra_model_cls,
                **field_defintions
            )
    
    def rebuild_bet_opt(self, name:str, opt:BetOpt) -> Tuple[FieldInfo,FieldInfo]:
        """
        重新构建BetOpt
        
        Args:
            name (str): 字段名称
            opt (BetOpt): 字段元信息
        
        Returns:
            Tuple[FieldInfo,FieldInfo]: Between查询条件信息
        """
        min_opt = opt.replace(active="min")
        min_fieldinfo = FieldInfo.from_annotation(min_opt.info.annotation)
        min_fieldinfo.json_schema_extra={"vo_opt":min_opt}
        max_opt = opt.replace(active="max")
        max_fieldinfo = FieldInfo.from_annotation(max_opt.info.annotation)
        max_fieldinfo.json_schema_extra={"vo_opt":max_opt}
        return min_fieldinfo,max_fieldinfo

    def rebuild_field(self, name:str, info:FieldInfo) -> bool:
        """
        重置字段
        
        Args:
            name(str): 字段名称
            info(FieldInfo): 字段元信息
        
        Returns:
            bool: 是否支持重置
        """
        if self.context.include_fields and \
            name not in self.context.include_fields:
            return False
        if self.context.exclude_fields and \
            name in self.context.exclude_fields:
            return False
        flag = self.field_filter.filter(name,info)
        return flag


class FormSchemaFactory(AbcSchemaFactory):
    
    action = "form"
    model_config = strict_valid_config
    
    def validate_annotation(self, annotation: Type[BaseModel]) -> Optional[Type[BaseModel]]:
        pass


class PathSchemaFactory(AbcSchemaFactory):
    
    action = "path"
    model_config = query_valid_config
    
    def validate_annotation(self, annotation: Type[BaseModel]) -> Optional[Type[BaseModel]]:
        pass
    

class ArbitrarySchemaFactory(AbcSchemaFactory):
    
    model_config = query_valid_config
    
    def __init__(self):
        super().__init__()
        query_valid_config_copy = query_valid_config.copy()
        query_valid_config_copy.update({
            "arbitrary_types_allowed":True
        })
        self.model_config = query_valid_config_copy
    
    def validate_annotation(self, annotation: Type) -> Optional[Type[BaseEntity]]:
        bo_model = get_final_model(annotation)
        if issubclass(bo_model, BaseEntity):
            return bo_model
        else:
            return None