corona.py

#!/usr/bin/env python3
import locale
import sys
from operator import itemgetter
from types import MappingProxyType

import requests
import rich
from requests.exceptions import RequestException
from rich.table import Table
from rich.text import Text

API_URL = "https://api.corona-zahlen.org/districts/"

QUERY = "berlin+hamburg+köln"


def format_incidence(value):
    for limit, color in [
        (35, "bright_green"),
        (50, "bright_yellow"),
        (100, "orange3"),
    ]:
        if value < limit:
            break
    else:
        color = "bright_red"

    return Text(f"{value:.01f}", color)


def format_integer(value):
    return locale.format_string("%d", value, grouping=True)


COLUMN_SPECS = MappingProxyType(
    {
        "Landkreis": ("name", str),
        "Inzidenz": ("weekIncidence", format_incidence),
        "Gesamtfälle": ("cases", format_integer),
        "Verstorbene": ("deaths", format_integer)
    }
)


def get_json(url):
    response = requests.get(url)
    response.raise_for_status()
    return response.json()


def get_district_records(query):
    terms = [term.strip().casefold() for term in query.split("+")]
    data = get_json(API_URL)["data"]
    return {
        key: district
        for key, district in data.items()
        for term in terms
        if term in district["name"].casefold()
    }


def get_table(query, specs=COLUMN_SPECS):
    districts = sorted(
        get_district_records(query).values(),
        key=itemgetter("name")
    )
    if not districts:
        raise ValueError(query)
    table = Table(*specs.keys())
    for district in districts:
        table.add_row(*
            (
                formatter(district[key])
                for key, formatter
                in specs.values()
            )
        )
    return table


fail = sys.exit


def main():
    locale.setlocale(locale.LC_ALL, "")
    query = sys.argv[1] if len(sys.argv) > 1 else QUERY
    try:
        rich.print(get_table(query))
    except RequestException as error:
        fail(f"Verbindung fehlgeschlagen: {error}")
    except ValueError as error:
        fail(f"Keine Treffer: {error}")


if __name__ == "__main__":
    main()