ocrd.mets_server

Complexity

Total Complexity

28

Size/Duplication

Total Lines	275
Duplicated Lines	0 %

Importance

Changes

0

20 Methods

Rating	Name	Duplication	Size	Complexity
A	OcrdFileModel.create()	0	3	1
A	ClientSideOcrdMets.unique_identifier()	0	3	1
A	OcrdMetsServer.shutdown()	0	2	1
A	OcrdFileGroupListModel.create()	0	3	1
A	ClientSideOcrdMets.add_agent()	0	2	1
A	OcrdAgentModel.create()	0	3	1
A	OcrdAgentListModel.create()	0	4	1
A	ClientSideOcrdMets.__getattr__()	0	2	1
A	ClientSideOcrdMets.find_files()	0	13	5
A	ClientSideOcrdMets.stop()	0	2	1
A	ClientSideOcrdMets.save()	0	2	1
A	ClientSideOcrdMets.__init__()	0	5	4
A	ClientSideOcrdMets.agents()	0	3	1
A	ClientSideOcrdMets.find_all_files()	0	2	1
A	ClientSideOcrdMets.file_groups()	0	3	1
A	ClientSideOcrdMets.add_file()	0	22	1
A	OcrdFileListModel.create()	0	4	1
A	ClientSideOcrdMets.__str__()	0	2	1
A	OcrdMetsServer.__init__()	0	4	1
B	OcrdMetsServer.startup()	0	91	2

"""
# METS server functionality
"""
import re
from os import environ, _exit
from io import BytesIO
from typing import Any, Dict, Optional, Union, List, Tuple
from urllib.parse import urlparse

from fastapi import FastAPI, Request, File, Form, Response
from fastapi.responses import JSONResponse
from requests import request, Session as requests_session
from requests_unixsocket import Session as requests_unixsocket_session
from pydantic import BaseModel, Field, ValidationError

import uvicorn

from ocrd_models import OcrdMets, OcrdFile, ClientSideOcrdFile, OcrdAgent, ClientSideOcrdAgent
from ocrd_utils import initLogging, getLogger, deprecated_alias

#
# XXX HACKS TODO
#
initLogging()

#
# Models
#

class OcrdFileModel(BaseModel):
    file_grp : str = Field()
    file_id : str = Field()
    mimetype : str = Field()
    page_id : Union[str, None] = Field()
    url : str = Field()
    local_filename : str = Field()

    @staticmethod
    def create(file_grp : str, file_id : str, page_id : Union[str, None], url : str, local_filename : str, mimetype : str):
        return OcrdFileModel(file_grp=file_grp, file_id=file_id, page_id=page_id, mimetype=mimetype, url=url, local_filename=local_filename)

class OcrdAgentModel(BaseModel):
    name : str = Field()
    _type : str = Field()
    role : str = Field()
    otherrole : Optional[str] = Field()
    othertype : str = Field()
    notes : Optional[List[Tuple[Dict[str, str], Optional[str]]]] = Field()

    @staticmethod
    def create(name : str, _type : str, role : str, otherrole : str, othertype : str, notes : List[Tuple[Dict[str, str], Optional[str]]]):
        return OcrdAgentModel(name=name, _type=_type, role=role, otherrole=otherrole, othertype=othertype, notes=notes)


class OcrdFileListModel(BaseModel):
    files : List[OcrdFileModel] = Field()

    @staticmethod
    def create(files : List[OcrdFile]):
        return OcrdFileListModel(
            files=[OcrdFileModel.create(file_grp=f.fileGrp, file_id=f.ID, mimetype=f.mimetype, page_id=f.pageId, url=f.url, local_filename=f.local_filename) for f in files]
        )

class OcrdFileGroupListModel(BaseModel):
    file_groups : List[str] = Field()

    @staticmethod
    def create(file_groups : List[str]):
        return OcrdFileGroupListModel(file_groups=file_groups)

class OcrdAgentListModel(BaseModel):
    agents : List[OcrdAgentModel] = Field()

    @staticmethod
    def create(agents : List[OcrdAgent]):
        return OcrdAgentListModel(
            agents=[OcrdAgentModel(name=a.name, _type=a.type, role=a.role, otherrole=a.otherrole, othertype=a.othertype, notes=a.notes) for a in agents]
        )

#
# Client
#


class ClientSideOcrdMets():
    """
    Partial substitute for :py:class:`ocrd_models.ocrd_mets.OcrdMets` which provides for
    :py:meth:`ocrd_models.ocrd_mets.OcrdMets.find_files`,
    :py:meth:`ocrd_models.ocrd_mets.OcrdMets.find_all_files`, and
    :py:meth:`ocrd_models.ocrd_mets.OcrdMets.add_agent`,
    :py:meth:`ocrd_models.ocrd_mets.OcrdMets.agents`,
    :py:meth:`ocrd_models.ocrd_mets.OcrdMets.add_file` to query via HTTP a
    :py:class:`ocrd.mets_server.OcrdMetsServer`.
    """

    def __init__(self, url):
        protocol = 'tcp' if url.startswith('http://') else 'uds'
        self.log = getLogger(f'ocrd.mets_client.{protocol}')
        self.url = url if protocol == 'tcp' else f'http+unix://{url.replace("/", "%2F")}'
        self.session = requests_session() if protocol == 'tcp' else requests_unixsocket_session()

    def __getattr__(self, name):
        raise NotImplementedError(f"ClientSideOcrdMets has no access to '{name}' - try without METS server")

    def __str__(self):
        return f'<ClientSideOcrdMets[url={self.url}]>'

    @deprecated_alias(ID="file_id")
    @deprecated_alias(pageId="page_id")
    @deprecated_alias(fileGrp="file_grp")
    def find_files(self, **kwargs):
        if 'pageId' in kwargs:
            kwargs['page_id'] = kwargs.pop('pageId')
        if 'ID' in kwargs:
            kwargs['file_id'] = kwargs.pop('ID')
        if 'fileGrp' in kwargs:
            kwargs['file_grp'] = kwargs.pop('fileGrp')
        r = self.session.request('GET', f'{self.url}/file', params={**kwargs})
        for f in r.json()['files']:
            yield ClientSideOcrdFile(None, ID=f['file_id'], pageId=f['page_id'], fileGrp=f['file_grp'], url=f['url'], local_filename=f['local_filename'], mimetype=f['mimetype'])

    def find_all_files(self, *args, **kwargs):
        return list(self.find_files(*args, **kwargs))

    def add_agent(self, *args, **kwargs):
        return self.session.request('POST', f'{self.url}/agent', json=OcrdAgentModel.create(**kwargs).dict())

    @property
    def agents(self):
        return [ClientSideOcrdAgent(None, **agent_dict) for agent_dict in self.session.request('GET', f'{self.url}/agent').json()['agents']]

    @property
    def unique_identifier(self):
        return self.session.request('GET', f'{self.url}/unique_identifier').text

    @property
    def file_groups(self):
        return self.session.request('GET', f'{self.url}/file_groups').json()['file_groups']

    @deprecated_alias(pageId="page_id")
    @deprecated_alias(ID="file_id")
    def add_file(self, file_grp, content=None, file_id=None, url=None, local_filename=None, mimetype=None, page_id=None, **kwargs):
        self.session.request(
            'POST',
            f'{self.url}/file',
            data=OcrdFileModel.create(
                file_id=file_id,
                file_grp=file_grp,
                page_id=page_id,
                mimetype=mimetype,
                url=url,
                local_filename=local_filename).dict(),
        )
        return ClientSideOcrdFile(
                None,
                ID=file_id,
                fileGrp=file_grp,
                url=url,
                pageId=page_id,
                mimetype=mimetype,
                local_filename=local_filename)


    def save(self):
        self.session.request('PUT', self.url)

    def stop(self):
        self.session.request('DELETE', self.url)

#
# Server
#

class OcrdMetsServer():

    def __init__(self, workspace, url):
        self.workspace = workspace
        self.url = url
        self.log = getLogger('ocrd.workspace_client')

    def shutdown(self):
        _exit(0)

    def startup(self):

        workspace = self.workspace

        app = FastAPI(
            title="OCR-D METS Server",
            description="Providing simultaneous write-access to mets.xml for OCR-D",
        )

        @app.exception_handler(ValidationError)
        async def exception_handler_validation_error(request: Request, exc: ValidationError):
            return JSONResponse(status_code=400, content=exc.errors())

        @app.exception_handler(FileExistsError)
        async def exception_handler_file_exists(request: Request, exc: FileExistsError):
            return JSONResponse(status_code=400, content=str(exc))

        @app.exception_handler(re.error)
        async def exception_handler_invalid_regex(request: Request, exc: re.error):
            return JSONResponse(status_code=400, content=f'invalid regex: {exc}')

        @app.get("/file", response_model=OcrdFileListModel)
        async def find_files(
            file_grp : Union[str, None] = None,
            file_id : Union[str, None] = None,
            page_id : Union[str, None] = None,
            mimetype : Union[str, None] = None,
        ):
            """
            Find files in the mets
            """
            found = workspace.mets.find_all_files(fileGrp=file_grp, ID=file_id, pageId=page_id, mimetype=mimetype)
            return OcrdFileListModel.create(found)

        @app.put('/')
        def save():
            return workspace.save_mets()

        @app.post('/file', response_model=OcrdFileModel)
        async def add_file(
            file_grp : str = Form(),
            file_id : str = Form(),
            page_id : Union[str, None] = Form(),
            mimetype : str = Form(),
            url : str = Form(),
            local_filename : str = Form(),
        ):
            """
            Add a file
            """
            # Validate
            file_resource = OcrdFileModel.create(file_grp=file_grp, file_id=file_id, page_id=page_id, mimetype=mimetype, url=url, local_filename=local_filename)
            # Add to workspace
            kwargs = file_resource.dict()
            workspace.add_file(**kwargs)
            return file_resource

        @app.get('/file_groups', response_model=OcrdFileGroupListModel)
        async def file_groups():
            return {'file_groups': workspace.mets.file_groups}

        @app.post('/agent', response_model=OcrdAgentModel)
        async def add_agent(agent : OcrdAgentModel):
            kwargs = agent.dict()
            workspace.mets.add_agent(**kwargs)
            return agent

        @app.get('/agent', response_model=OcrdAgentListModel)
        async def agents():
            return OcrdAgentListModel.create(workspace.mets.agents)

        @app.get('/unique_identifier', response_model=str)
        async def unique_identifier():
            return Response(content=workspace.mets.unique_identifier, media_type='text/plain')

        @app.delete('/')
        async def stop():
            """
            Stop the server
            """
            getLogger('ocrd_models.ocrd_mets').info('Shutting down')
            workspace.save_mets()
            # os._exit because uvicorn catches SystemExit raised by sys.exit
            _exit(0)

        if self.url.startswith('http'):
            parsed = urlparse(self.url)
            uvicorn_kwargs = {'host': parsed.hostname, 'port': parsed.port}
        else:
            uvicorn_kwargs = {'uds': self.url}
        uvicorn.run(app, **uvicorn_kwargs)