forked from HoloTech/customcmd-plugin
255 lines
6.8 KiB
Python
255 lines
6.8 KiB
Python
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
|
|
|
|
|
|
__all__ = (
|
|
"constant_term",
|
|
"define_equal",
|
|
"define_walrus",
|
|
"define_term",
|
|
"binop",
|
|
"compare",
|
|
"cmdarg_term",
|
|
"command_shell",
|
|
"get_var",
|
|
"ternary_term",
|
|
"run_function",
|
|
"cascade_term",
|
|
'ExecutorTerm',
|
|
'LateTerm',
|
|
'lazy_term',
|
|
)
|
|
|
|
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.this_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.this_scope[var] = result
|
|
return result
|
|
|
|
@lazy_term
|
|
async def define_term(ctx, var: str, value: ExecutorTerm):
|
|
ctx.this_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
|