mymaths.py - a maths / stats CLI utility

mymaths.py

#!/usr/bin/env python
"""do basic maths with stuff from stdin"""

import sys
import statistics


def cu(num: str) -> float:
    if not num:
        return None
    num = str(num)
    if num.endswith("m"):
        return float(num[:-1]) * 0.001
    elif num.endswith("ki"):
        return float(num[:-2]) * 1_024
    elif num.endswith("k"):
        return float(num[:-1]) * 1_000
    elif num.endswith("Mi"):
        return float(num[:-2]) * 1024 * 1024
    elif num.endswith("M"):
        return float(num[:-1]) * 1_000_000
    elif num.endswith("Gi"):
        return float(num[:-2]) * 1024 * 1024 * 1024
    elif num.endswith("G"):
        return float(num[:-1]) * 1_000_000_000
    else:
        return float(num)


def ca(nums: list[str]) -> list[float]:
    return [cu(num.strip()) for num in nums if num.strip()]


def f_sum(vals):
    return sum(vals)


def f_sub(vals):
    if len(vals):
        return vals[0] - sum(vals[1:])
    return 0


def f_spread(vals):
    return max(vals) - min(vals)


def f_min(vals):
    return min(vals)


def f_max(vals):
    return max(vals)

def joinvals(fun, **funargs):
    def _w(vals):
        return '; '.join(str(a) for a in fun(vals, **funargs))
    return _w

fmap = {
    "sum": f_sum,
    "add": f_sum,
    "sub": f_sub,
    "diff": f_sub,
    "min": min,
    "max": max,
    "multimode": joinvals(statistics.multimode),
    # stats
    "avg": statistics.mean,
    "mean": statistics.mean,
    "fmean": statistics.fmean,
    "hmean": statistics.harmonic_mean,
    "gmean": statistics.geometric_mean,
    "mode": statistics.mode,
    "variance": statistics.variance,
    "svariance": statistics.variance,
    "pvariance": statistics.pvariance,
    "pstdev": statistics.pstdev,
    "stdev": statistics.stdev,
    "sstdev": statistics.stdev,
    "median": statistics.median,
    "4d": joinvals(statistics.quantiles, n=5),
    "8d": joinvals(statistics.quantiles, n=9),
    "deciles": joinvals(statistics.quantiles, n=11),
    "spread": f_spread,
}

def main(funz, the_in):
    vals = ca(the_in)
    show_headers = len(funz) > 1
    for fun in funz:
        prefix = f"{fun}: " if show_headers else ""
        r = fmap[fun](vals)
        print(f"{prefix}{r}")


if __name__ == "__main__":
    funz = sys.argv[1:]
    if funz[0] == 'all':
        print("All mode")
        funz = list(fmap)
    assert all(fun in fmap for fun in funz)
    main(funz, sys.stdin.readlines())