proxy.py

import functools
from enum import Enum
from typing import List, Optional, Union


class OpType(Enum):
    ATTR = 1
    ITEM = 2
    CALL = 3
    FOREACH = 4


class Operation:
    op_type: OpType


class AttrOperation(Operation):
    op_type = OpType.ATTR

    def __init__(self, attribute: str):
        self.attribute = attribute

    def apply(self, obj):
        return getattr(obj, self.attribute)


class ItemOperation(Operation):
    op_type = OpType.ITEM

    def __init__(self, key: Union[int, str]):
        self.key = key

    def apply(self, obj):
        return obj[self.key]


class CallOperation(Operation):
    op_type = OpType.CALL

    def __init__(self, args, kwargs):
        self.args = args
        self.kwargs = kwargs

    def apply(self, obj):
        return obj(*self.args, **self.kwargs)


class ForeachOperation(Operation):
    op_type = OpType.FOREACH

    def apply(self, obj, next_proxy):
        return [next_proxy.apply(item) for item in obj]


class Applier:
    def __init__(self, ops: List[Operation], child: Optional["Applier"] = None):
        self.ops = ops
        self.child = child

    @classmethod
    def build(cls, ops: List[Operation]):
        partial_ops: List[Operation] = []
        for idx, op in enumerate(ops):
            if op.op_type == OpType.FOREACH:
                return cls(partial_ops, cls.build(ops[idx + 1 :]))
            partial_ops.append(op)
        return cls(partial_ops)

    def apply(self, obj):
        current_obj = obj
        for op in self.ops:
            # XXX: You might wonder why this isn't done with a
            # method/abstracted. This is because we are trying to avoid the
            # overhead of recursive calls
            if isinstance(op, AttrOperation):
                current_obj = getattr(current_obj, op.attribute)
            elif isinstance(op, ItemOperation):
                current_obj = current_obj[op.key]
            elif isinstance(op, CallOperation):
                current_obj = current_obj(*op.args, **op.kwargs)

        if self.child:
            return [self.child.apply(item) for item in current_obj]

        return current_obj


class Proxy:
    def __init__(self, operations: Optional[List[Operation]] = None):
        self._operations = operations or []

    def __add_op(self, operation) -> "Proxy":
        return type(self)(operations=self._operations + [operation])

    def attr(self, attribute: str) -> "Proxy":
        return self.__add_op(AttrOperation(attribute))

    __getattr__ = attr

    def __getitem__(self, key) -> "Proxy":
        return self.__add_op(ItemOperation(key))

    def __call__(self, *args, **kwargs) -> "Proxy":
        return self.__add_op(CallOperation(args, kwargs))

    def foreach(self):
        return self.__add_op(ForeachOperation())

    def apply(self, obj):
        return self.applier.apply(obj)

    @functools.cached_property
    def applier(self):
        return self.build_applier()

    def build_applier(self) -> Applier:
        return Applier.build(self._operations)

    def __repr__(self):
        repr_str = "Proxy.start"
        for op in self._operations:
            if isinstance(op, AttrOperation):
                repr_str += f".{op.attribute}"
            elif isinstance(op, ItemOperation):
                repr_str += f"[{repr(op.key)}]"
            elif isinstance(op, CallOperation):
                args_str = ", ".join(repr(arg) for arg in op.args)
                kwargs_str = ", ".join(
                    f"{key}={repr(value)}" for key, value in op.kwargs.items()
                )
                all_args_str = ", ".join(filter(None, [args_str, kwargs_str]))
                repr_str += f"({all_args_str})"
            elif isinstance(op, ForeachOperation):
                repr_str += ".foreach()"
        return repr_str


proxy = Proxy()
each = proxy.foreach()