customcmd-plugin/customcmd/terms.py

from typing import Any, Awaitable, Callable, Concatenate, ParamSpec, TypeVar, Union, TYPE_CHECKING
from functools import wraps, reduce
import operator

from . import logger
from .errors import CommandValueError

if TYPE_CHECKING:
    from .context import ExecContext

P = ParamSpec("P")
R = TypeVar("R")
T = TypeVar("T")
T1 = TypeVar("T1")
T2 = TypeVar("T2")


TermResult = TypeVar("TermResult")


class ExecutorTerm[TermResult]:
    # May add more here for initial parsing context

    def __init__(self, *args, node_text: str | None = None, **kwargs):
        self.text = node_text

    def __call__(self, ctx: "ExecContext") -> Awaitable[TermResult]:
        return self._execute(ctx)

    async def _execute(self, ctx: "ExecContext") -> Any:
        logger.debug(f"Executing {self} in {ctx}")
        try:
            ctx.termstack.append(self)
            return await self.execute(ctx)
        except Exception:
            logger.exception(f"Unhandled exception running {self} in {ctx}")
            # TODO: Custom execution error class to wrap things
            # TODO: And safe errors to stop execution, e.g. Permission errors or explicit early termination
            raise
        finally:
            ctx.termstack.pop()

    async def execute(self, ctx: "ExecContext") -> Any:
        raise NotImplementedError


class LateTerm(ExecutorTerm[TermResult]):
    """
    A LateTerm is an ExecTerm that runs a given coroutine upon execution 'late'.
    """

    def __init__(
        self,
        coro: Callable[P, Awaitable[TermResult]],
        coroargs,
        corokwargs,
        name: str | None = None,
        **kwargs,
    ):
        self.coro: Callable = coro
        self.coroargs = coroargs
        self.corokwargs = corokwargs
        self.name = name or coro.__name__
        super().__init__(**kwargs)

    async def execute(self, ctx: "ExecContext") -> TermResult:
        return await self.coro(ctx, *self.coroargs, **self.corokwargs)

    def __repr__(self):
        return " ".join(
            (
                "<{self.__class__.__name__}",
                "name={self.name!r}",
                "text={self.text!r}",
                "coroargs={self.coroargs!r}",
                "corokwargs={self.corokwargs!r}",
                "coro={self.coro!r}",
                ">",
            )
        ).format(self=self)


def lazy_term(
    coro: Callable[Concatenate["ExecContext", P], Awaitable[R]],
) -> Callable[Concatenate[str, P], LateTerm[R]]:
    """
    Build a LazyTerm instance builder (args) => LazyTerm wrapping the given coroutine.
    The coroutine should accept a "ExecContext" as the first argument.

    When the builder is called, it will use the arguments it is called with
    to create a LazyTerm which may then later be executed.

    Executing the constructed LazyTerm will execute coro with the given
    execution context, as well as the arguments provided in the builder.
    """

    def term_wrapper(text: str, *args: P.args, **kwargs: P.kwargs):
        return LateTerm(coro, args, kwargs, node_text=text)

    return term_wrapper


@lazy_term
async def ternary_term(
    ctx: "ExecContext",
    condition: ExecutorTerm,
    arg1: ExecutorTerm[T1],
    arg2: ExecutorTerm[T2],
) -> Union[T1, T2]:
    """condition ? arg1 : arg2"""
    if await condition(ctx):
        return await arg1(ctx)
    else:
        return await arg2(ctx)


@lazy_term
async def constant_term(ctx: "ExecContext", constant: T) -> T:
    """Return the provided constant."""
    return constant


@lazy_term
async def cmdarg_term(ctx, argn: int, greedy: bool) -> str:
    args = (ctx.alias, *ctx.args)
    if greedy:
        return " ".join(args[argn:])
    elif argn >= len(args):
        return ''
    else:
        return args[argn]


@lazy_term
async def cascade_term(ctx, *terms: ExecutorTerm):
    for term in terms:
        eval = await term(ctx)
        if eval:
            return eval
    return ""

@lazy_term
async def command_shell(ctx, component_term: ExecutorTerm):
    if isinstance(component_term, ExecutorTerm):
        result = await component_term(ctx)
    else:
        logger.warning(
            f"command_shell LazyTerm stringifying non-term {component_term}"
        )
        result = component_term
    return str(result)

@lazy_term
async def define_equal(ctx, var: str, value: ExecutorTerm):
    ctx.scope[var] = await value.execute(ctx)
    return None

@lazy_term
async def define_walrus(ctx, var: str, value: ExecutorTerm):
    result = await value.execute(ctx)
    ctx.scope[var] = result
    return result

@lazy_term 
async def define_term(ctx, var: str, value: ExecutorTerm):
    ctx.scope[var] = value
    return None

@lazy_term
async def run_function(ctx, var: ExecutorTerm, *args: ExecutorTerm):
    varlue = await var.execute(ctx)
    argues = [await arg.execute(ctx) for arg in args]
    if args and not callable(varlue):
        raise CommandValueError(f"{var} is not callable")
    elif not args:
        return varlue
    else:
        return await varlue(ctx, *argues)

@lazy_term 
async def get_var(ctx, root: str | ExecutorTerm, *attrs: str):
    if isinstance(root, ExecutorTerm):
        value = await root.execute(ctx)
        name = [root.text or '']
    elif root not in ctx.scope:
            return ' '.join((root, *attrs))
    else:
        value = ctx.scope[root]
        name = [root]

    if isinstance(value, ExecutorTerm):
        # ExecTerms only end up in scope from delayed assignment
        value = await value.execute(ctx)

    for attr in attrs:
        try:
            value = getattr(value, '_exec_attr_'+attr)
            name.append(attr)
        except AttributeError:
            raise CommandValueError(f"{'.'.join(name)} has no attribute {attr}")
    return value

@lazy_term 
async def compare(ctx, term1: ExecutorTerm, op: str, term2: ExecutorTerm):
    opmap = {
        '==': operator.eq,
        '!=': operator.ne,
        '>': operator.gt,
        '>=': operator.ge,
        '<': operator.lt,
        '<=': operator.le,
    }
    value1 = await term1.execute(ctx)
    value2 = await term2.execute(ctx)
    oppy = opmap[op]
    return oppy(value1, value2)

@lazy_term
async def binop(ctx, op: str, *terms):
    opmap = {
        '||': operator.or_,
        '&&': operator.and_,
        '+': operator.add,
        '-': operator.sub,
        '*': operator.mul,
        '**': operator.pow,
        '/': operator.truediv,
        '//': operator.floordiv,
        '%': operator.mod,
    }
    # Left associative
    oppy = opmap[op]
    values = [await term.execute(ctx) for term in terms]
    result = reduce(oppy, values)
    return result