Source code for luna_handlers.classes.schemas.attribute_upgrade

"""
Module contains schemas for attribute upgrade handler
"""

from pydantic import Field, constr, root_validator
from vlutils.regexps import UUID_REGEXP_STR

from classes.schemas import types
from classes.schemas.base_schema import BaseSchema, HandlerSettings
from classes.schemas.types import OptionalNotNullable
from crutches_on_wheels.cow.errors.errors import Error
from crutches_on_wheels.cow.errors.exception import VLException

SamplesIDs = types.userConstrainedList(
    constr(regex=UUID_REGEXP_STR), minItems=1, maxItemsGetter=lambda: HandlerSettings.receivedImagesLimit
)


class FaceUpgradePolicy(BaseSchema):
    """Face attribute upgrade policy"""

    _descriptorVersion: int

    # whether to extract basic attributes
    extractBasicAttributes: types.Int01 = 0
    # descriptor quality score threshold
    fdScoreThreshold: types.StrictFloat01 = 0.0
    # whether to extract face descriptor
    extractDescriptor: types.Int01 = 1
    # descriptor version
    descriptorVersion: types.FaceDescriptorVersion = Field(
        default_factory=lambda: HandlerSettings.faceDescriptorVersion
    )

    @root_validator(pre=False, skip_on_failure=True)
    def validatePolicy(cls, values):
        """
        Validate attributes are select for extract
        Raises:
            VLException(Error.NotSelectedAttributesForExtract, 400, False) if no attributes select for extract
        """
        isExtractBasicAttributes = values["extractBasicAttributes"]
        isExtractDescriptor = values["extractDescriptor"]
        if not isExtractDescriptor and not isExtractBasicAttributes:
            raise VLException(Error.NotSelectedAttributesForExtract, 400, False)
        return values


class BodyUpgradePolicy(BaseSchema):
    """Body attribute upgrade policy"""

    # descriptor version
    descriptorVersion: types.BodyDescriptorVersion = Field(
        default_factory=lambda: HandlerSettings.bodyDescriptorVersion
    )


class FaceUpgradeData(BaseSchema):
    """Face upgrade data schema"""

    # sample ids
    samples: SamplesIDs
    # face attribute upgrade policy
    policy: FaceUpgradePolicy = FaceUpgradePolicy()


class BodyUpgradeData(BaseSchema):
    """Body upgrade data schema"""

    # sample ids
    samples: SamplesIDs
    # body attribute upgrade policy
    policy: BodyUpgradePolicy = BodyUpgradePolicy()


class SamplesForUpgrade(BaseSchema):
    """Samples for upgrade schema"""

    # face id
    faceId: constr(regex=UUID_REGEXP_STR) = OptionalNotNullable()
    # face upgrade data
    upgradeData: FaceUpgradeData = OptionalNotNullable()

    # event id
    eventId: constr(regex=UUID_REGEXP_STR) = OptionalNotNullable()
    # face upgrade data
    faceUpgradeData: FaceUpgradeData = OptionalNotNullable()
    # body upgrade data
    bodyUpgradeData: BodyUpgradeData = OptionalNotNullable()

    @root_validator(pre=False, skip_on_failure=True)
    def validateOneOf(cls, values):
        """
        Validate one of request value:

            'face_id' and 'upgrade_data' or
            'event_id' and 'face_upgrade_data' and/or 'body_upgrade_data'
        """
        existingKeys = {key for key, value in values.items() if value}
        isValidForFace = existingKeys == {"faceId", "upgradeData"}
        isValidForEvent = (
            "eventId" in existingKeys
            and ("faceUpgradeData" in existingKeys or "bodyUpgradeData" in existingKeys)
            and not (existingKeys.intersection({"faceId", "samples"}))
        )

        if isValidForFace == isValidForEvent:
            raise ValueError(
                "Required one of: `face_id` with `upgrade_data` or "
                "`event_id` with `face_upgrade_data` and/or `body_upgrade_data` in request"
            )

        return values