chainable.py - Implement pipe operator + currying

chainable.py

#!/usr/bin/env python3


from typing import Any
from collections.abc import Callable

SHOULD_LOG = False


class Chainable:
    """
    This class is used to wrap a function and its arguments in a chainable manner.

    It is also used to overload the right shift operator (>>) to chain the functions together.
    """

    def __init__(
        self,
        func,
        *args,
        **kwargs,
    ):
        self.func = func
        self.func_name = func.__name__
        self.args = args
        self.kwargs = kwargs
        self.retval = func(*args, **kwargs)
        self.__is_magic_shell_func__ = getattr(func, "__is_magic_shell_func__", False)

    def __call__(self, *args, **kwargs):
        new_args = args + self.args
        self.retval = self.func(*new_args, **kwargs)
        cur = Chainable(self.func, *new_args, **kwargs)
        cur.retval = self.retval
        if SHOULD_LOG:
            print(cur)
        return cur

    def __rshift__(self, other):
        left = self
        right = other
        if not isinstance(right, Callable):
            raise ValueError("Right-hand side must be a callable")

        if getattr(right, "__is_magic_shell_func__", False):
            return right(stdin=left.retval)

        return right(left.retval)

    def __repr__(self):
        args_str = ", ".join(map(str, self.args)).replace("\n", r"\n")
        kwargs_str = ", ".join([f"{k}={v}" for k, v in self.kwargs.items()])
        sep = ", " if args_str and kwargs_str else ""
        string = f"{self.func_name}"
        if args_str:
            string += f"({args_str})"
        if kwargs_str:
            string = string[:-1]
            string += f"{sep}{kwargs_str})"
        return string

    @property
    def value(self):
        return self.retval


# Make a decorator to wrap the function with Chainable class
def chainable(func):
    if func.__name__ == "sh":
        func.__is_magic_shell_func__ = True
    else:
        func.__is_magic_shell_func__ = False

    def wrapper(*args, **kwargs):
        return Chainable(func, *args, **kwargs)

    return wrapper


@chainable
def add(*args) -> int | float:
    return sum(args)


@chainable
def sub(*args) -> int | float:
    if len(args) == 0:
        raise ValueError("At least one argument is required")
    return args[0] - sum(args[1:])


@chainable
def mul(*args) -> int | float:
    retval = 1
    for arg in args:
        retval *= arg
    return retval


@chainable
def div(*args) -> int | float:
    if len(args) == 0:
        raise ValueError("At least one argument is required")
    retval = args[0]
    for arg in args[1:]:
        retval /= arg
    return retval


@chainable
def pow(*args) -> int | float:
    if len(args) == 0:
        raise ValueError("At least one argument is required")
    retval = args[0]
    for arg in args[1:]:
        retval **= arg
    return retval


def toValue(x: Any) -> Any:
    return x


from subprocess import Popen, PIPE


class ShellError(Exception):
    pass


@chainable
def sh(cmd: str, *, stdin=None, **kwargs):
    proc = Popen(
        cmd,
        shell=True,
        stdout=PIPE,
        stderr=PIPE,
        stdin=PIPE,
        encoding="utf-8",
    )
    if stdin is None:
        stdout, stderr = proc.communicate()
    else:
        stdout, stderr = proc.communicate(input=stdin)

    if proc.returncode != 0 or stderr:
        raise ShellError(stderr)

    return stdout.strip()


if __name__ == "__main__":
    assert add(1, 2, 4).value == 7
    assert add(1)(2)(4).value == 7
    assert add(1) >> add(2) >> add(4) >> toValue == 7
    assert add(1)(2) >> add(2)(2) >> toValue == 7
    assert add(1)(2) >> add(2)(2) >> sub(2) >> toValue == 5
    assert add(4) >> add(4) >> sub(5) >> mul(2) >> toValue == 6
    assert add(4) >> add(4) >> sub(5) >> mul(2) >> div(3) >> toValue == 2

    assert sh("echo 'a b c'") >> toValue == "a b c"
    assert sh("echo 'a b c'") >> sh("wc -w") >> toValue == "3"
    assert sh("echo 'foo'") >> sh("awk '{print toupper($0)}'") >> toValue == "FOO"

    print("All tests passed!")

system calls with Unix style pipes directly in python using the >> operator behave as expected :)

Playing with syntax is fun :)