aft-may25-2023/util.py

from functools import wraps
from AST import Set, atom
import more_itertools


def powerset(items):
    return map(Set, more_itertools.powerset(items))


def format_set(s: Set[atom]):
    return "{" + ", ".join(sorted(s)) + "}"


def textp(*args):
    whole = "("
    sets = (format_set(arg) for arg in args)
    whole += ", ".join(sets)
    whole += ")"
    return whole


def printp(*args):
    print(textp(*args))


def prints_input(*pos):
    def wrappest(fn):
        @wraps(fn)
        def wrapper(*args, **kwargs):
            output = fn(*args, **kwargs)
            print(
                f"Input of {fn.__name__} pos: {pos}:",
                "  ".join(
                    (format_set(arg) if isinstance(arg, Set) else str(arg))
                    for i, arg in enumerate(args)
                    if i in pos
                ),
            )
            return output

        return wrapper

    return wrappest


def prints_output(fn):
    @wraps(fn)
    def wrapper(*args, **kwargs):
        output = fn(*args, **kwargs)
        print(
            f"Output of {fn.__name__}:",
            format_set(output) if isinstance(output, Set) else output,
        )
        return output

    return wrapper


def disabled_static(return_value):
    "Disables a function by making it return a specified value instead"

    def wrappest(fn):
        @wraps(fn)
        def wrapper(*args, **kwargs):
            return return_value

        return wrapper

    return wrappest