argos-httping/httping.py

#! /usr/bin/env python3
# -*- coding: utf-8 -*-

""" Simulates httping for `argos <https://github.com/p-e-w/argos>`"""
# <bitbar.title>httping</bitbar.title>
# <bitbar.version>v1.0</bitbar.version>
# <bitbar.author>Tobias Schmidl</bitbar.author>
# <bitbar.author.github>schtobia</bitbar.author.github>
# <bitbar.desc>Short description of what your plugin does.</bitbar.desc>
# <bitbar.image></bitbar.image>
# <bitbar.dependencies>python,requests</bitbar.dependencies>
# <bitbar.abouturl>https://github.com/schtobia/argos-httping/blob/master/httping.py</bitbar.abouturl>

from __future__ import annotations

import datetime  # for the whole ms → s and vice versa stuff
import statistics  # for median and stdev
import sys  # for redirecting print to stderr
from typing import Generic, List, Sequence, TypeVar, Optional, Tuple, Dict, Union

import requests  # for the central HTTP HEAD request

__author__ = "Tobias Schmidl"
__copyright__ = "Copyright 2018, Tobias Schmidl"
__license__ = "MIT"
__version__ = "0.0.1"
__maintainer__ = "Tobias Schmidl"

MAX_PING: datetime.timedelta = datetime.timedelta(milliseconds=1000)  # in seconds
TARGETS: List[str] = ["https://www.google.de", "https://www.bing.com", "https://www.wikipedia.org"]

T = TypeVar('T')
Bucket = Tuple[int, int]
BucketList = List[Bucket]

MAX_PING_IN_MS: int = int(MAX_PING.total_seconds() * 1000)


def in_bucket(element: int, bucket: Bucket) -> bool:
    """ Returns True if element is in bucket """
    return bucket[0] <= element <= bucket[1]


def generate_buckets(min_val: int, max_val: int, step_width: int = 2) -> BucketList:
    """ Generates a list of partitioned lists between min_val and max_val."""
    return [(x, x + step_width - 1) for x in range(min_val, max_val, step_width)]


class Themes:
    """Maintains a list of themes and generates buckets between 0 and MAX_PING_IN_MS based on the selected theme."""
    # Themes are from http://colorbrewer2.org/
    purple_green: List[str] = ["#762a83", "#9970ab", "#c2a5cf", "#a6dba0", "#5aae61", "#1b7837"]
    red_green: List[str] = ["#d73027", "#fc8d59", "#fee08b", "#d9ef8b", "#91cf60", "#1a9850"]
    original: List[str] = ["#acacac", "#ff0101", "#cc673b", "#ce8458", "#6bbb15", "#0ed812"]
    selected_colors: List[str] = []
    buckets: BucketList = []

    def __init__(self, selected_colors):
        self.selected_colors = selected_colors
        self.buckets = generate_buckets(0, int(MAX_PING_IN_MS), int(MAX_PING_IN_MS / (len(selected_colors) - 1)))

    def colorize(self, input_text):
        """Colorized the current input according to the selected theme."""
        for index, current_bucket in enumerate(self.buckets):
            if in_bucket(input_text, current_bucket):
                return self.selected_colors[len(self.selected_colors) - 1 - index]
        return self.selected_colors[0]


CURRENT_THEME: Themes = Themes(Themes.red_green)


def httping(targets: List[str]) -> Tuple[Dict[str, Optional[float]], Optional[float], Optional[float]]:
    """Pings the supplied targets and returns a mean and a std deviation over all target responses."""
    responses: Dict[str, Optional[float]] = {}
    for target in targets:
        try:
            responses[target] = requests.head(target, timeout=MAX_PING.total_seconds()).elapsed.total_seconds() * 1000
        except requests.exceptions.RequestException as ex:
            print(str(ex), file=sys.stderr)
            responses[target] = None
    response_times = [value for key, value in responses.items() if isinstance(value, float)]

    if len(response_times) > 1:
        return responses, statistics.median(response_times), statistics.stdev(response_times)
    elif response_times:
        return responses, response_times[0], 0
    else:
        return responses, None, None


if __name__ == "__main__":
    responses, median, stddev = httping(TARGETS)
    if isinstance(median, float) and median != MAX_PING_IN_MS:
        print("⦿ " + str(round(median, 2)) + "±" + str(round(stddev, 2)) + " | color=" +
              CURRENT_THEME.colorize(median))  # type: ignore
    else:
        print("☠ | color=" + CURRENT_THEME.selected_colors[0])
    print("---")
    for current_host, current_response in responses.items():
        if isinstance(current_response, float):
            print(current_host + ": " + str(round(current_response, 2)) + " | color=" +
                  CURRENT_THEME.colorize(current_response))
        else:
            print(current_host + ": ☠ | color=" + CURRENT_THEME.selected_colors[0])