Python automation utils

Here are some utility functions that have been useful when writing automation scripts.

requests

I like using requests or httpx, but if you add a dependency to your script, now you’ll need to create an environment, install the depdencies, etc. just to run it. I like keeping automation scripts in a single file I can copy around. Instead of using the requests, module, you can use this. This function has been adapted from the urlrequest function in fastai/fastcore.

import urllib.request
from urllib.request import Request
from typing import Optional, Callable
from urllib.parse import urlencode
import json


# Adapted from fastai/fastcore
def request(
    url,
    verb: str = "GET",
    headers: Optional[dict] = None,
    route: Optional[dict] = None,
    query=None,
    data=None,
    json_data=True,
):
    "`Request` for `url` with optional route params replaced by `route`, plus `query` string, and post `data`"

    dumper_func: Callable[..., str] = json.dumps if json_data else urlencode
    if route:
        url = url.format(**route)
    if query:
        url += "?" + urlencode(query)
    if isinstance(data, dict):
        data = dumper_func(data).encode("ascii")
    req = Request(url, headers=headers or {}, data=data or None, method=verb.upper())

    return urllib.request.urlopen(req).read()

Example usage:

import urllib.error

headers = {"Foo": "Bar"}

try:
    data = request(
        "http://example.com/{foo}/1",
        "POST",
        headers=headers,
        route={"foo": "3"},
        query={"q": "4"},
        data={"d": "5"},
    )
except urllib.error.HTTPError as e:
    print(f"Error code: {e.code}")
    raise

run subprocesses in parallel

When you’re automating something, it’s very common to run tasks using subprocess. To run multiple tasks in parallel, wrap each subprocess call inside a function and pass a list of functions this par_run function. Each function will be run in a separate thread.

from threading import Thread
from typing import Callable, List

def par_run(callables: List[Callable]):
    ths: List[Thread] = []
    for func in callables:
        t = Thread(target=func, daemon=True)
        t.start()
        ths.append(t)

    for t in ths:
        t.join()

running a subprocess

Instead of using subprocess directly, this function lets you pass both a string or a list of strings. It also logs the command to the standard output. This is useful because it let’s you copy/paste the command if something goes wrong.

from typing import Union, List
import sys
import shlex
import subprocess


def e(*args, **kwargs):
    """
    Print to stderr
    """
    print(*args, **kwargs, file=sys.stderr, flush=True)


def c(
    cmd: Union[str, List[str]], hide=False, capture=True, **kwargs
) -> Union[str, int]:
    """
    Run command. The command can be a string or a list of strings.

    Args:
        hide: Hide full command and do not print it in a copy/paste-able format
        capture: Capture command output as text and return it
        **kargs: Keyword arguments to forward to the `subprocess.run` call. If `capture=True`,
                 the `capture_output` and `text` arguments will be overridden.
    """

    if isinstance(cmd, str):
        for char in "$()`":
            if char in cmd:
                raise ValueError(f"Command shouldn't include the character '{char}'")

        _cmd = shlex.split(cmd)
    else:
        _cmd = cmd

    _cmd: List[str] = [str(x) for x in _cmd]
    if not hide:
        e(f"Running: {shlex.join(_cmd)}")

    if capture:
        kwargs["capture_output"] = True
        kwargs["text"] = True
    p = subprocess.run(_cmd, **kwargs)

    try:
        p.check_returncode()
        return p.stdout if capture else p.returncode
    except subprocess.CalledProcessError:
        e(p.stdout)
        e(p.stderr)
        raise

Example usage:

import json

c("echo foo")
c(["echo", "foo"])
c("aws s3 ls")

json.loads(c("aws sts get-caller-identity"))

Asking the user to select a value from a list

This is useful when the automation script requires input from the user to select a value from a list of options.

from typing import List
from typing import TypeVar
import sys
import pprint

T = TypeVar("T")


def interactive_select_value_from_list(l: List[T]) -> T:
    """
    Select a value from a list interactively (user input).
    """

    ret = None

    while not ret:
        print("Select an index number (integer) from the list below:")
        pprint.pprint({idx: v for idx, v in enumerate(l)})
        sel = input("Index: ")
        try:
            idx = int(sel)
        except ValueError:
            print(f"You must select an index number, you selected: '{sel}'")
            continue
        if idx < 0:
            print(f"The index can't be a negative number")
            continue
        try:
            ret = l[idx]
        except IndexError:
            print(f"The index you selected ({idx}) is not in the list")
            continue

    return ret

rand[om]

Python automation utils

requests

run subprocesses in parallel

running a subprocess

Asking the user to select a value from a list