advent22/api/advent22_api/core/webdav.py

import functools
import logging
import re
from io import BytesIO

import requests
from asyncify import asyncify
from cachetools import TTLCache, cachedmethod
from cachetools.keys import hashkey
from webdav3.client import Client as WebDAVclient

from .settings import SETTINGS

_logger = logging.getLogger(__name__)


def davkey(name, _, *args, **kwargs):
    """Return a cache key for use with cached methods."""

    return hashkey(name, *args, **kwargs)


class WebDAV:
    class __WebDAVclient(WebDAVclient):
        def execute_request(
            self,
            action,
            path,
            data=None,
            headers_ext=None,
        ) -> requests.Response:
            res = super().execute_request(action, path, data, headers_ext)

            # the "Content-Length" header can randomly be missing on txt files,
            # this should fix that (probably serverside bug)
            if action == "download" and "Content-Length" not in res.headers:
                res.headers["Content-Length"] = str(len(res.text))

            return res

    _webdav_client = __WebDAVclient(
        {
            "webdav_hostname": SETTINGS.webdav.url,
            "webdav_login": SETTINGS.webdav.username,
            "webdav_password": SETTINGS.webdav.password,
        }
    )

    _cache = TTLCache(
        ttl=SETTINGS.webdav.cache_ttl,
        maxsize=SETTINGS.webdav.cache_size,
    )

    @classmethod
    @asyncify
    @cachedmethod(
        cache=lambda cls: cls._cache,
        key=functools.partial(davkey, "list_files"),
    )
    def list_files(
        cls,
        directory: str = "",
        *,
        regex: re.Pattern[str] = re.compile(""),
    ) -> list[str]:
        """
        List files in directory `directory` matching RegEx `regex`
        """

        _logger.debug(f"list_files {directory!r}")
        ls = cls._webdav_client.list(directory)

        return [path for path in ls if regex.search(path)]

    @classmethod
    @asyncify
    @cachedmethod(
        cache=lambda cls: cls._cache,
        key=functools.partial(davkey, "exists"),
    )
    def exists(cls, path: str) -> bool:
        """
        `True` iff there is a WebDAV resource at `path`
        """

        _logger.debug(f"file_exists {path!r}")
        return cls._webdav_client.check(path)

    @classmethod
    @asyncify
    @cachedmethod(
        cache=lambda cls: cls._cache,
        key=functools.partial(davkey, "read_bytes"),
    )
    def read_bytes(cls, path: str) -> bytes:
        """
        Load WebDAV file from `path` as bytes
        """

        _logger.debug(f"read_bytes {path!r}")
        buffer = BytesIO()
        cls._webdav_client.download_from(buffer, path)
        buffer.seek(0)

        return buffer.read()

    @classmethod
    async def read_str(cls, path: str, encoding="utf-8") -> str:
        """
        Load WebDAV file from `path` as string
        """

        _logger.debug(f"read_str {path!r}")
        return (await cls.read_bytes(path)).decode(encoding=encoding).strip()

    @classmethod
    @asyncify
    def write_bytes(cls, path: str, buffer: bytes) -> None:
        """
        Write bytes from `buffer` into WebDAV file at `path`
        """

        _logger.debug(f"write_bytes {path!r}")
        cls._webdav_client.upload_to(buffer, path)

        # invalidate cache entry
        cls._cache.pop(hashkey("read_bytes", path))

    @classmethod
    async def write_str(cls, path: str, content: str, encoding="utf-8") -> None:
        """
        Write string from `content` into WebDAV file at `path`
        """

        _logger.debug(f"write_str {path!r}")
        await cls.write_bytes(path, content.encode(encoding=encoding))