Quick and dirty stack calculator

stack_calculator.py

#!/usr/bin/env python3

import math
import operator as op
import os
import readline
import statistics
import sys
from functools import reduce


CONSTANTS = {"e": math.e, "pi": math.pi}

OPERATIONS = {
    # Nullary ops
    "q": (0, sys.exit),
    # Unary ops
    "inc":   (1, lambda x: x + 1),
    "dec":   (1, lambda x: x - 1),
    "int":   (1, int),
    "neg":   (1, op.neg),
    "abs":   (1, abs),
    "log":   (1, math.log),
    "floor": (1, math.floor),
    "ceil":  (1, math.ceil),
    "log":   (1, math.log),
    "exp":   (1, math.exp),
    "sqrt":  (1, math.sqrt),
    # Binary ops
    "+":   (2, op.add),
    "-":   (2, op.sub),
    "*":   (2, op.mul),
    "/":   (2, op.truediv),
    "%":   (2, op.mod),
    "//":  (2, op.floordiv),
    "pow": (2, (lambda b, x:
                math.pow(b, x) if isinstance(b, float) and isinstance(x, float)
                else pow(b, x))),
    "logb":(2, math.log),
    # Variadic ops
    "sum": ("whole stack", sum),
    "prod": ("whole stack", lambda xs: reduce(op.mul, xs, 1)),
    "mean": ("whole stack", statistics.mean),
    "median": ("whole stack", statistics.median),
    "stdev": ("whole stack", statistics.stdev),
}

class Calculator:

    def __init__(self):
        self.stack = []
        self.STACK_OPS = {
            # drop, delete
            "d": (0, self.stack.pop),
            # clear
            "c": (0, self.stack.clear),
            # analogous to "x" in vim
            "x": (1, self.stack.pop),
            "dup": (0, lambda: self.stack.append(self.stack[-1])),
            # yank
            "y": (1, lambda ix: self.stack.append(self.stack[ix])),
            # roll
            "r": (1, lambda ix: self.stack.append(self.stack.pop(ix))),
            # swap
            "s": (2, self.swap),
        }

    def mainloop(self):
        clear_screen()
        while True:
            cmd = input("> ").strip()
            clear_screen()
            if cmd:
                self.run(cmd)
            self.print_stack()

    def run(self, cmd):
        result = str2num(cmd)
        if result is not None:
            self.stack.append(result)
            return

        if cmd in OPERATIONS:
            n_args, func = OPERATIONS[cmd]
            if n_args == "whole stack":
                n_args = len(self.stack)
                self.stack[:] = [func(self.stack)]
                return
            if len(self.stack) < n_args:
                sys.stderr.write(f"Need {n_args} arguments in stack for {cmd}\n")
                return
            result = func(*reversed([self.stack.pop() for _ in range(n_args)]))
            self.stack.append(result)
            return

        if cmd in CONSTANTS:
            self.stack.append(CONSTANTS[cmd])
            return

        stack_cmd, *stack_args = cmd.split()
        if stack_cmd in self.STACK_OPS:
            n_args, func = self.STACK_OPS[stack_cmd]
            if len(stack_args) != n_args:
                sys.stderr.write(f"{stack_cmd} takes exactly {n_args} inline arguments\n")
                return
            try:
                indices = [int(arg) for arg in stack_args]
            except ValueError:
                sys.stderr.write("Stack manipulations take integer arguments\n")
                return
            if any(ix >= len(self.stack) or ix < 0 for ix in indices):
                sys.stderr.write("Index out of range\n")
                return
            if not self.stack:
                sys.stderr.write(f"Can't manipulate empty stack\n")
                return
            func(*indices)
            return

        sys.stderr.write(f"{cmd} isn't valid input\n")

    def swap(self, ix1, ix2):
        self.stack[ix1], self.stack[ix2] = self.stack[ix2], self.stack[ix1]

    def print_stack(self):
        for index, element in enumerate(self.stack):
            print(f"{element:>30} [{index:>2}]")


def clear_screen():
    for __ in range(os.get_terminal_size().lines):
        print()


def str2num(string):
    try:
        return int(string)
    except ValueError:
        try:
            return float(string)
        except ValueError:
            return None


if __name__ == "__main__":
    Calculator().mainloop()