history_plot.py

"""
Take history.jsonl from history_get.py and plot it

Usage:

# Prep
pip install -r history-requirements.txt
python3 history_get.py  # creates/updates history.jsonl

# All available data
python3 history_plot.py && open history.png

# Subset
python3 history_plot.py -v 2.6 2.7 3.3 3.4 3.5 3.6 3.7 && open history.png
"""

import argparse
import hashlib
from pprint import pprint  # noqa: F401

from natsort import natsorted  # pip install natsort
from termcolor import colored  # pip install termcolor

from history_get import load_jsonlines
from template import get_eols


def dopplr(name):
    """
    Take the MD5 digest of a name,
    convert it to hex and take the
    first 6 characters as an RGB value.
    """
    # Tweak "2.8" because it's too close in colour to "3.5"
    if name == "2.8":
        name = "python 2.8"

    return "#" + hashlib.sha224(name.encode()).hexdigest()[:6]


def make_chart(dates, totals):
    # x: list of dates
    # y: totals for each version
    import datetime as dt

    import matplotlib.dates as mdates
    import matplotlib.pyplot as plt  # pip install matplotlib
    import matplotlib.ticker as plticker

    # "2020-01-26 22:35:22+02:00" -> "2020-01-26"
    dates = [date.split()[0] for date in dates]
    # Real dates so matplotlib can place ticks sensibly
    x = [dt.date.fromisoformat(date) for date in dates]

    fig, ax = plt.subplots()

    print("Plot...")
    eol = get_eols()
    for version, v in natsorted(totals.items()):
        print(version)

        if version in eol and eol[version] in dates:
            # Add a vertical line to zero at EOL
            eol_pos = dates.index(eol[version])
            # dates.insert(eol_pos, eol[version])
            totals[version][eol_pos] = 0
            # breakpoint()

        ax.plot(x, v, label=version, color=dopplr(version))

    ax.set_ylim(ymin=0, ymax=360)

    # Shrink current axis by 20% so legend is outside chart
    box = ax.get_position()
    ax.set_position((box.x0, box.y0, box.width * 0.9, box.height))

    # Put a legend to the right of the current axis
    ax.legend(
        loc="center left",
        bbox_to_anchor=(1, 0.5),
    )

    # One labelled tick per year, minor ticks per quarter
    ax.xaxis.set_major_locator(mdates.YearLocator())
    ax.xaxis.set_major_formatter(mdates.DateFormatter("%Y"))
    ax.xaxis.set_minor_locator(mdates.MonthLocator(bymonth=(4, 7, 10)))
    loc = plticker.MultipleLocator(base=60)
    ax.yaxis.set_major_locator(loc)

    plt.suptitle("Dropped Python versions")
    plt.title("By the top 360 packages downloaded from PyPI", fontsize=10)
    plt.ylabel("Packages")

    outfile = "history.png"
    print(colored(outfile, "green"))
    plt.savefig(outfile, dpi=96 * 2.5)


def main():
    parser = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("-v", "--versions", nargs="+", help="Show only these versions")
    args = parser.parse_args()

    # Each line looks like:
    # {
    #  "date": "2020-01-26 23:35:22+02:00",
    #  "drop_totals": {"2.0": 348, "2.1": 348, "2.2": 348, "2.3": 348, "2.4": 342, ...}
    # }
    lines = load_jsonlines("history.jsonl")

    #  First 3-4 days of data is junk, ditch it
    lines = [line for line in lines if line["date"] > "2017-10-11 20:44:23+03:00"]

    # Ditch first few days for some due to bug fixes causing a sudden dip
    for line in lines:
        if line["date"] <= "2017-10-24 23:44:27+03:00":
            line["drop_totals"].pop("2.6", None)
            line["drop_totals"].pop("3.2", None)
        if line["date"] <= "2018-09-06 13:35:20+03:00":
            line["drop_totals"].pop("3.4", None)

    # Sort by date
    lines = sorted(lines, key=lambda k: k["date"])
    print("Number of lines:", len(lines))

    # First get a list of all versions, as not every data point has every version
    if args.versions:
        all_versions = args.versions
    else:
        all_versions = set()
        for line in lines:
            all_versions.update(line["drop_totals"].keys())
        all_versions = sorted(all_versions)
    print("All versions: ", all_versions)

    print("Prep data...")
    dates = []
    totals = {}
    for line in lines:
        dates.append(line["date"])

        for version in all_versions:
            try:
                version_total = line["drop_totals"][version]
            except KeyError:
                version_total = None

            try:
                totals[version].append(version_total)
            except KeyError:
                totals[version] = [version_total]

    make_chart(dates, totals)


if __name__ == "__main__":
    main()