from_api.py

# -*- coding: utf-8 -*-
from uuid import UUID

from arkindex_worker.worker.classification import ClassificationMixin
from arkindex_worker.worker.element import ElementMixin
from arkindex_worker.worker.entity import EntityMixin
from arkindex_worker.worker.metadata import MetaDataMixin
from arkindex_worker.worker.transcription import TranscriptionMixin
from worker_generic_training_dataset import Extractor


class DatasetExtractorFromAPI(
    Extractor,
    ElementMixin,
    ClassificationMixin,
    EntityMixin,
    TranscriptionMixin,
    MetaDataMixin,
):
    def list_classifications(self, element_id: UUID):
        return iter(
            self.api_client.request("RetrieveElement", id=str(element_id))[
                "classifications"
            ]
        )


def main():
    DatasetExtractorFromAPI(
        description="Fill base-worker cache with information about dataset and extract images",
    ).run()


if __name__ == "__main__":
    main()