rollbot3/lib/rollbot/plugins/injection.py

import inspect

from ..messaging import RollbotMessage, RollbotFailure


class ArgConverter:
    def __init__(self, conv):
        self.conv = conv


Message = ArgConverter(lambda _, __, msg: msg)
Database = ArgConverter(lambda _, db, __: db)
Logger = ArgConverter(lambda cmd, _, __: cmd.logger)
Bot = ArgConverter(lambda cmd, _, __: cmd.bot)
ArgString = ArgConverter(lambda _, __, msg: msg.raw_args)
ArgList = ArgConverter(lambda _, __, msg: msg.arg_list())
Subcommand = ArgConverter(lambda _, __, msg: RollbotMessage.from_subcommand(msg))
Subcommand.ArgString = ArgConverter(lambda _, __, msg: RollbotMessage.from_subcommand(msg).raw_args)
Subcommand.ArgList = ArgConverter(lambda _, __, msg: RollbotMessage.from_subcommand(msg).arg_list())


class Arg(ArgConverter):
    def __init__(self, index, conversion=str, fail_msg=None):
        super().__init__(self._convert)
        self.index = index
        self.conversion = conversion
        self.fail_msg = fail_msg or "Invalid argument: {}"
    
    def _convert(self, cmd, db, msg):
        try:
            arg = msg.arg_list()[self.index]
        except IndexError:
            RollbotFailure.INVALID_ARGUMENTS.with_reason(f"Missing argument {self.index}").raise_exc()
        try:
            return self.conversion(arg)
        except ValueError:
            RollbotFailure.INVALID_ARGUMENTS.with_reason(self.fail_msg.format(arg)).raise_exc()


class OptArg(ArgConverter):
    def __init__(self, index, conversion=str, default=None, fail_msg=None):
        super().__init__(self._convert)
        self.index = index
        self.conversion = conversion
        self.default = default
        self.fail_msg = fail_msg or "Invalid argument: {}"
    
    def _convert(self, cmd, db, msg):
        if msg.raw_args is None:
            return self.default
        try:
            arg = msg.arg_list()[self.index]
        except IndexError:
            return self.default
        try:
            return self.conversion(arg)
        except ValueError:
            RollbotFailure.INVALID_ARGUMENTS.with_reason(self.fail_msg.format(arg)).raise_exc()


class _SubcArg(Arg):
    def _convert(self, cmd, db, msg):
        return super()._convert(self, cmd, db, RollbotMessage.from_subcommand(msg))


class _SubcOptArg(OptArg):
    def _convert(self, cmd, db, msg):
        return super()._convert(self, cmd, db, RollbotMessage.from_subcommand(msg))


Subcommand.Arg = _SubcArg
Subcommand.OptArg = _SubcOptArg


class Config(ArgConverter):
    def __init__(self, key=None):
        if key is None:
            super().__init__(lambda cmd, _, __: cmd.bot.config)
        else:
            super().__init__(lambda cmd, _, __, key=key: cmd.bot.config.get(key))


def _run_converter_function(fn, cmd, db, msg):
    return fn(*[fn.__annotations__[param].conv(cmd, db, msg) for param in inspect.signature(fn).parameters])


class Singleton(ArgConverter):
    def __init__(self, model_cls):
        super().__init__(lambda _, db, msg, model_cls=model_cls: model_cls.get_or_create(db, msg)) 
        self.model_cls = model_cls

    def by(self, fn):
        def do_get(cmd, db, msg):
            key = _run_converter_function(fn, cmd, db, msg)
            return db.query(self.model_cls).get(key) or self.model_cls.create_from_key(key)
        return ArgConverter(do_get)

    def by_all(self, fn):
        def do_get(cmd, db, msg):
            return [db.query(self.model_cls).get(x) or self.model_cls.create_from_key(x) for x in _run_converter_function(fn, cmd, db, msg)]
        return ArgConverter(do_get)


class Query(ArgConverter):
    def __init__(self, model_cls):
        super().__init__(self._convert)
        self.model_cls = model_cls
        self._filters = []
        self._order_by = None
        self._limit = None

    def filter(self, fn):
        self._filters.append(fn)
        return self

    def order_by(self, fn):
        self._order_by = fn
        return self

    def limit(self, n):
        self._limit = n
        return self

    def _convert(self, cmd, db, msg):
        query = db.query(self.model_cls)
        for fn in self._filters:
            query = query.filter(_run_converter_function(fn, cmd, db, msg))
        if self._order_by is not None:
            query = query.order_by(_run_converter_function(self._order_by, cmd, db, msg))
        if self._limit is None:
            return query.all()
        else:
            return query.limit(self._limit)


class Lazy(ArgConverter):
    def __init__(self, child):
        super().__init__(self._convert)
        self.child = child
        self.result = None
        self.run = False

    def _convert(self, cmd, db, msg):
        self.cmd = cmd
        self.db = db
        self.msg = msg
        return self

    def get(self):
        if not self.run:
            self.result = self.child.conv(self.cmd, self.db, self.msg)
        return self.result


class _Executor:
    def __init__(self, cmd, db, msg):
        self.cmd = cmd
        self.db = db
        self.msg = msg

    def run_with_deps(self, fn):
        return _run_converter_function(fn, self.cmd, self.db, self.msg)
        

Injector = ArgConverter(_Executor)