pyinfra/pyinfra/server/stream/rest.py

import logging

from flask import jsonify
from funcy import drop

from pyinfra.server.dispatcher.dispatcher import Nothing
from pyinfra.server.stream.queued_stream_function import QueuedStreamFunction

logger = logging.getLogger()


class LazyRestProcessor:
    def __init__(self, queued_stream_function: QueuedStreamFunction, submit_suffix="submit", pickup_suffix="pickup"):
        self.submit_suffix = submit_suffix
        self.pickup_suffix = pickup_suffix
        self.queued_stream_function = queued_stream_function

    def push(self, request):
        self.queued_stream_function.push(request.json)
        return jsonify(replace_suffix(request.base_url, self.submit_suffix, self.pickup_suffix))

    def pop(self):
        result = self.queued_stream_function.pop() or Nothing

        if not valid(result):
            logger.error(f"Received invalid result: {result}")
            result = Nothing

        if result is Nothing:
            resp = jsonify("No more items left")
            resp.status_code = 204

        else:
            resp = jsonify(result)
            resp.status_code = 206

        return resp


def valid(result):
    return isinstance(result, dict) or result is Nothing


def replace_suffix(strn, suf, repl):
    return remove_last_n(strn, len(suf)) + repl


def remove_last_n(strn, n):
    return "".join(reversed(list(drop(n, reversed(strn)))))