diff --git a/docs/en/_includes/toloka-kit-reference-toc.yaml b/docs/en/_includes/toloka-kit-reference-toc.yaml index 436065d6..73a30f73 100644 --- a/docs/en/_includes/toloka-kit-reference-toc.yaml +++ b/docs/en/_includes/toloka-kit-reference-toc.yaml @@ -1012,6 +1012,9 @@ items: - name: VariantRegistry.register href: ../toloka-kit/reference/toloka.client.primitives.base.VariantRegistry.register.md hidden: true + - name: VariantRegistry.generate_subtype + href: ../toloka-kit/reference/toloka.client.primitives.base.VariantRegistry.generate_subtype.md + hidden: true - name: infinite_overlap items: - name: InfiniteOverlapParametersMixin diff --git a/src/client/primitives/base.py b/src/client/primitives/base.py index 82a11267..7cdb266f 100644 --- a/src/client/primitives/base.py +++ b/src/client/primitives/base.py @@ -8,6 +8,7 @@ ] import inspect +import logging import typing from copy import copy from enum import Enum @@ -17,6 +18,7 @@ import attr import simplejson as json +from ...util._extendable_enum import ExtendableStrEnumMetaclass from .._converter import converter from ..exceptions import SpecClassIdentificationError from ...util._codegen import ( @@ -26,13 +28,20 @@ E = TypeVar('E', bound=Enum) +logger = logging.getLogger(__file__) + class VariantRegistry: - def __init__(self, field: str, enum: Type[E]): + def __init__(self, field: str, enum: Type[E], extendable: Optional[bool] = None): + if extendable is None: + extendable = isinstance(enum, ExtendableStrEnumMetaclass) + if extendable and not isinstance(enum, ExtendableStrEnumMetaclass): + raise ValueError('VariantRegistry could be extendable only if spec_enum is extendable.') self.field: str = field self.enum: Type[E] = enum self.registered_classes: Dict[E, type] = {} + self.extendable = extendable def register(self, type_: type, value: E) -> type: @@ -47,7 +56,19 @@ def register(self, type_: type, value: E) -> type: return type_ - def __getitem__(self, value: E): + def generate_subtype(self, type_: type, value: E) -> type: + if not self.extendable: + raise NotImplementedError('Only extendable VariantRegistry can generate subtype') + + generated_type_name = '_Generated' + value.value.title() + type_.__name__ + BaseTolokaObjectMetaclass(generated_type_name, (type_,), {}, spec_value=value) + logger.info(f'{generated_type_name} class was generated. ' + f'Probably it is a new functionality on the platform.\n' + f'If you want it to be supported by toloka-kit faster you can make feature request here:' + f'https://github.com/Toloka/toloka-kit/issues/new/choose.') + return self.registered_classes[value] + + def __getitem__(self, value: E) -> type: return self.registered_classes[value] @@ -171,13 +192,21 @@ def __new__(cls, *args, **kwargs): return super().__new__(cls) @classmethod - def __init_subclass__(cls, spec_enum: Optional[Union[str, Type[E]]] = None, - spec_field: Optional[str] = None, spec_value=None): + def __init_subclass__( + cls, + spec_enum: Optional[Union[str, Type[E]]] = None, + spec_field: Optional[str] = None, + spec_value=None, + extend_spec: Optional[bool] = None, + ): super().__init_subclass__() # Completing a variant type if spec_value is not None: cls._variant_registry.register(cls, spec_value) + if extend_spec and (spec_enum is None or spec_field is None): + raise ValueError('extend_spec could be True only with spec_enum and spec_field provided') + # Making into a variant type if spec_enum is not None or spec_field is not None: @@ -190,7 +219,7 @@ def __init_subclass__(cls, spec_enum: Optional[Union[str, Type[E]]] = None, # TODO: Possibly make it immutable enum = getattr(cls, spec_enum) if isinstance(spec_enum, str) else spec_enum - cls._variant_registry = VariantRegistry(spec_field, enum) + cls._variant_registry = VariantRegistry(spec_field, enum, extend_spec) # Unexpected fields access @@ -265,7 +294,11 @@ def structure(cls, data: Any): data_field = data.pop(spec_field) try: spec_value = cls._variant_registry.enum(data_field) - spec_class = cls._variant_registry.registered_classes[spec_value] + + if spec_value in cls._variant_registry.registered_classes: + spec_class = cls._variant_registry[spec_value] + else: + spec_class = cls._variant_registry.generate_subtype(cls, spec_value) except Exception: raise SpecClassIdentificationError(spec_field=spec_field, spec_enum=cls._variant_registry.enum.__name__) diff --git a/tests/utils/test_variant_type.py b/tests/utils/test_variant_type.py index 20907d50..dc37e933 100644 --- a/tests/utils/test_variant_type.py +++ b/tests/utils/test_variant_type.py @@ -2,12 +2,14 @@ from typing import Any, Optional import pytest -from toloka.client._converter import converter +from toloka.client.exceptions import SpecClassIdentificationError from toloka.client.primitives.base import BaseTolokaObject +from toloka.client._converter import converter +from toloka.util._extendable_enum import ExtendableStrEnum @unique -class ItemType(Enum): +class ItemType(ExtendableStrEnum): LIST = 'list' SET = 'set' @@ -32,7 +34,7 @@ class SetPopMethodCall(SetMethodCall, spec_value=SetMethod.POP): class ListMethodCall(MethodCall, spec_value=ItemType.LIST, spec_enum='Method', spec_field='method'): @unique - class Method(Enum): + class Method(ExtendableStrEnum): APPEND = 'append' POP = 'pop' @@ -92,6 +94,34 @@ def test_structure_variant(): converter.structure({'type': 'list', 'method': 'pop', 'argument': 'abc'}) +def test_structure_unknown_variant(): + unstructured_with_unknown_high_level_variant = {'type': 'dict', 'method': 'pop'} + unstructured_with_unknown_low_level_variant = {'type': 'list', 'method': 'index'} + unstructured_with_unknown_both_levels_variants = {'type': 'dict', 'method': 'get'} + + method_call_with_unknown_high_level_variant = converter.structure(unstructured_with_unknown_high_level_variant, + MethodCall) + method_call_with_unknown_low_level_variant = converter.structure(unstructured_with_unknown_low_level_variant, + MethodCall) + method_call_with_unknown_both_levels_variants = converter.structure(unstructured_with_unknown_both_levels_variants, + MethodCall) + + assert isinstance(method_call_with_unknown_high_level_variant, MethodCall) + assert isinstance(method_call_with_unknown_low_level_variant, ListMethodCall) + assert isinstance(method_call_with_unknown_both_levels_variants, MethodCall) + + assert converter.unstructure( + method_call_with_unknown_high_level_variant) == unstructured_with_unknown_high_level_variant + assert converter.unstructure( + method_call_with_unknown_low_level_variant) == unstructured_with_unknown_low_level_variant + assert converter.unstructure( + method_call_with_unknown_both_levels_variants) == unstructured_with_unknown_both_levels_variants + + with pytest.raises(SpecClassIdentificationError): + unstructured_unextendable = {'type': 'set', 'method': 'new_method'} + converter.structure(unstructured_unextendable, MethodCall) + + def test_unstructure_variant(): assert {'type': 'set', 'method': 'pop'} == converter.unstructure(SetPopMethodCall()) assert {'type': 'list', 'method': 'pop'} == converter.unstructure(ListPopMethodCall())