subprocess-util/receive_inform.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import socket
from pathlib import Path
from typing import IO, AnyStr, Callable

from unix_sock_input import accept_loop_until_command_received


def receive_inform(in_pipe: IO[AnyStr],
                   socket_file: Path,
                   chunk_file_tmpl: Path,
                   get_chunk_file: Callable[[Path, int], Path],
                   ) -> None:
    """
    :param get_chunk_file:
    :param in_pipe:
    :param chunk_file_tmpl:
    :param socket_file: Create a UNIX socket and wait for messages.
    :return:
    """
    print(f'Listening on socket {socket_file}')
    sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
    sock.bind(str(socket_file))
    sock.listen(1)

    ct = 1
    commands = [b'OK\n', b'EOF\n']
    while True:
        command = accept_loop_until_command_received(sock, commands)
        if command not in commands:
            raise ValueError("Invalid state")

        chunk_file = get_chunk_file(chunk_file_tmpl, ct)
        chunk = chunk_file.read_bytes()
        in_pipe.write(chunk)
        # in_pipe.flush()  # TODO: is this required?
        chunk_file.unlink(missing_ok=False)

        if command == b'OK\n':
            ct += 1
        elif command == b'EOF\n':
            break
        else:
            raise ValueError("Invalid state")

    print(f'Closing socket {socket_file}')
    sock.close()
    socket_file.unlink(missing_ok=False)

    in_pipe.flush()

    # TODO: Has this any effect? On stdin probably yes!
    in_pipe.close()