32 lines
870 B
Python
32 lines
870 B
Python
from typing import Dict, ForwardRef, Type, Union
|
|
|
|
from pydantic import BaseModel
|
|
|
|
|
|
class DocsRegistry:
|
|
_registry: Dict[str, Type[BaseModel]] = {}
|
|
|
|
@classmethod
|
|
def register(cls, name: str, doc_type: Type[BaseModel]):
|
|
cls._registry[name] = doc_type
|
|
|
|
@classmethod
|
|
def get(cls, name: str) -> Type[BaseModel]:
|
|
return cls._registry[name]
|
|
|
|
@classmethod
|
|
def evaluate_fr(cls, forward_ref: Union[ForwardRef, Type]):
|
|
"""
|
|
Evaluate forward ref
|
|
|
|
:param forward_ref: ForwardRef - forward ref to evaluate
|
|
:return: Type[BaseModel] - class of the forward ref
|
|
"""
|
|
if (
|
|
isinstance(forward_ref, ForwardRef)
|
|
and forward_ref.__forward_arg__ in cls._registry
|
|
):
|
|
return cls._registry[forward_ref.__forward_arg__]
|
|
else:
|
|
return forward_ref
|