Asynchronous line-based communication with an external program

Question

As part of a larger project, I've written the following code for interacting with git-annex's addurl command in batch mode. addurl takes a series of URLs and output filepaths on stdin, and it outputs a one-line JSON object for each file when that file finishes downloading. git-annex performs the downloads asynchronously, so the order of the input need not match the order of the output, and the caller should input as many URLs as possible before reading any output. It's also necessary for UX, logging, etc. purposes that a "Download completed" message be printed as soon as each download is complete; thus, writing everything at once and then reading everything at once with Popen.communicate() is not an option. Naïvely, you would think this could be done with a basic Popen instance by writing all the lines to the subprocess and then reading the output line by line, but this will encounter blocking issues as soon as you exceed the pipe buffer size (which I believe is 64 kiB on Linux). Hence, I figured this was as good a time as any to finally do something with Python's asynchronous programming.

The below code is written for Python 3.8. Aside from general advice (which is especially needed as this is my first asyncio code), key things I want to know include:

• Is splitting up the reading and writing into separate functions that are run at the same time with asyncio.gather() the best/only way to write to & read from a process at once with asyncio?

• Is asyncio even the right way to go, or is there a completely different (yet still sane, i.e., not involving constant polling) way to communicate with a process asynchronously without risking overflowing the pipe buffer that doesn't involve asyncio or threading?

• How would I make the git-annex process terminate when the Python process receives a Cntrl-C or otherwise hits an error?

• Is it possible to make read_output() into an asynchronous iterator that yields decoded data dicts that are then processed inside download_urls() instead of inside read_output()?

(Note: The main() function here is just something minimal for testing purposes; I am aware that the input it splits apart is almost the same as what feed_input() puts back together again, and I know it does almost nothing with the return value of download_urls(). The calling code for download_urls() in the actual codebase is much different.)

import argparse
import asyncio
import json
import logging
from pathlib import Path
import subprocess
from typing import Dict, Iterable, List, Optional, Tuple

log = logging.getLogger("async_downloader")


async def download_urls(
    repo_path: Path, urls_paths: Iterable[Tuple[str, str]], jobs: int = 10
) -> Dict[str, Optional[str]]:
    """
    :param repo_path: path to a git-annex repository in which to download the
        files
    :param urls_paths: an iterable of (URL, output file path) pairs
    :param jobs: the number of jobs for git-annex to use for downloading
    :returns: a `dict` mapping output file paths to git-annex keys
    """
    process = await asyncio.create_subprocess_exec(
        "git-annex",
        "addurl",
        "--batch",
        "--with-files",
        "--jobs",
        str(jobs),
        "--json",
        "--json-error-messages",
        "--raw",
        stdin=subprocess.PIPE,
        stdout=subprocess.PIPE,
        cwd=repo_path,
    )
    assert process.stdin is not None
    assert process.stdout is not None
    _, (downloaded, failures) = await asyncio.gather(
        feed_input(process.stdin, urls_paths),
        read_output(process.stdout),
    )
    log.info("Downloaded %d files", len(downloaded))
    r = await process.wait()
    if failures:
        raise RuntimeError(f"{failures} files failed to download")
    if r != 0:
        raise RuntimeError(f"git-annex addurl exited with return code {r}")
    return downloaded


async def feed_input(
    fp: asyncio.StreamWriter, urls_paths: Iterable[Tuple[str, str]]
) -> None:
    for url, path in urls_paths:
        fp.write(f"{url} {path}\n".encode("utf-8"))
        await fp.drain()
    fp.close()
    await fp.wait_closed()


async def read_output(fp: asyncio.StreamReader) -> Tuple[Dict[str, Optional[str]], int]:
    downloaded: Dict[str, Optional[str]] = {}
    failures = 0
    async for line in fp:
        data = json.loads(line)
        if not data["success"]:
            log.error(
                "%s: download failed; error messages: %r",
                data["file"],
                data["error-messages"],
            )
            failures += 1
        else:
            key = data.get("key")
            log.info("Finished downloading %s (key = %s)", data["file"], key)
            downloaded[data["file"]] = key
    return (downloaded, failures)


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "repo",
        type=Path,
        help=(
            "Path to a git-annex repository; will be created if it does not"
            " already exist"
        ),
    )
    parser.add_argument(
        "downloads_file",
        type=argparse.FileType("r", encoding="utf-8"),
        help="File containing lines of the form '<URL> <output file path>'",
    )
    args = parser.parse_args()
    logging.basicConfig(
        format="%(asctime)s [%(levelname)-8s] %(name)s %(message)s",
        datefmt="%H:%M:%S%z",
        level=logging.DEBUG,
    )
    urls_paths: List[Tuple[str, str]] = []
    with args.downloads_file:
        for line in args.downloads_file:
            line = line.strip()
            if line and not line.startswith("#"):
                url, path = line.split(maxsplit=1)
                urls_paths.append((url, path))
    if not args.repo.exists():
        subprocess.run(["git", "init", args.repo], check=True)
        subprocess.run(["git-annex", "init"], cwd=args.repo, check=True)
    downloads = asyncio.run(download_urls(args.repo, urls_paths))
    subprocess.run(
        ["git", "commit", "-m", f"Downloaded {len(downloads)} URLs"],
        cwd=args.repo,
        check=True,
    )


if __name__ == "__main__":
    main()