diff --git a/royalnet/engineer/command.py b/royalnet/engineer/command.py index 1f843fba..67a60be0 100644 --- a/royalnet/engineer/command.py +++ b/royalnet/engineer/command.py @@ -32,7 +32,7 @@ class Command(c.Conversation): are later completed by a runner. """ - def __init__(self, f: c.ConversationProtocol, *, names: t.List[str] = None, pattern: re.Pattern): + def __init__(self, f: c.ConversationProtocol, *, names: t.List[str] = None, pattern: re.Pattern, lock: bool = True): """ Create a new :class:`.Command` . """ @@ -54,6 +54,11 @@ class Command(c.Conversation): The pattern that should be matched by the command. """ + self.lock: bool = lock + """ + If calling this command should :meth:`~royalnet.engineer.dispenser.Dispenser.lock` the dispenser. + """ + def name(self): """ :return: The main name of the Command. @@ -96,9 +101,17 @@ class Command(c.Conversation): log.debug(f"Match successful, getting capture groups of: {match!r}") message_kwargs: t.Dict[str, str] = match.groupdict() - - log.debug(f"Passing args to function: {message_kwargs!r}") - return await super().run(_sentry=_sentry, _msg=bullet, **base_kwargs, **message_kwargs) + + if self.lock: + log.debug(f"Locking the dispenser...") + + with _sentry.dispenser().lock(self): + log.debug(f"Passing args to function: {message_kwargs!r}") + return await super().run(_sentry=_sentry, _msg=bullet, **base_kwargs, **message_kwargs) + + else: + log.debug(f"Passing args to function: {message_kwargs!r}") + return await super().run(_sentry=_sentry, _msg=bullet, **base_kwargs, **message_kwargs) def help(self) -> t.Optional[str]: """ @@ -116,7 +129,7 @@ class PartialCommand: They can specified later using :meth:`.complete`. """ - def __init__(self, f: c.ConversationProtocol, syntax: str): + def __init__(self, f: c.ConversationProtocol, syntax: str, lock: bool = True): """ Create a new :class:`.PartialCommand` . @@ -133,6 +146,11 @@ class PartialCommand: Part of the pattern from where the arguments should be captured. """ + self.lock: bool = lock + """ + If calling this command should :meth:`~royalnet.engineer.dispenser.Dispenser.lock` the dispenser. + """ + @classmethod def new(cls, *args, **kwargs): """ @@ -167,8 +185,10 @@ class PartialCommand: raise ValueError(f"Name is not alphanumeric: {name!r}") name_regex = f"(?:{'|'.join(names)})" + log.debug(f"Completed pattern: {name_regex!r}") + pattern: re.Pattern = re.compile(pattern.format(name=name_regex, syntax=self.syntax), re.IGNORECASE) - return Command(f=self.f, names=names, pattern=pattern) + return Command(f=self.f, names=names, pattern=pattern, lock=self.lock) def __repr__(self): return f"<{self.__class__.__qualname__} {self.f!r}>" diff --git a/royalnet/engineer/dispenser.py b/royalnet/engineer/dispenser.py index 7e970ae0..ed582df3 100644 --- a/royalnet/engineer/dispenser.py +++ b/royalnet/engineer/dispenser.py @@ -10,6 +10,7 @@ import contextlib from .sentry import SentrySource from .conversation import Conversation +from .exc import LockedDispenserError log = logging.getLogger(__name__) @@ -21,6 +22,13 @@ class Dispenser: A :class:`list` of all the running sentries of this dispenser. """ + self._locked_by: t.List[Conversation] = [] + """ + The conversation that are currently locking this dispenser. + + .. seealso:: :meth:`.lock` + """ + async def put(self, item: t.Any) -> None: """ Insert a new item in the queues of all the running sentries. @@ -54,7 +62,15 @@ class Dispenser: Run the passed conversation. :param conv: The conversation to run. + :raises .LockedDispenserError: If the dispenser is currently :attr:`.locked_by` a :class:`.Conversation`. """ + log.debug(f"Trying to run: {conv!r}") + + if self._locked_by is not None: + log.debug(f"Dispenser is locked by {self._locked_by!r}, refusing to run {conv!r}") + raise LockedDispenserError(self._locked_by, f"The Dispenser is currently locked by {self._locked_by!r} and " + f"cannot start new conversations.") + log.debug(f"Running: {conv}") with self.sentry() as sentry: state = conv(_sentry=sentry, **kwargs) @@ -63,6 +79,26 @@ class Dispenser: while state := await state: log.debug(f"Switched to: {state}") + @contextlib.contextmanager + def lock(self, conv: Conversation): + """ + Lock the dispenser while this :func:`~contextlib.contextmanager` is in scope. + + A locked dispenser will refuse to :meth:`.run` any new conversations, raising :exc:`.exc.LockedDispenserError` + instead. + + :param conv: The conversation that requested the lock. + + .. seealso:: :attr:`._locked_by` + """ + log.debug(f"Adding lock: {conv!r}") + self._locked_by.append(conv) + + yield + + log.debug(f"Clearing lock: {conv!r}") + self._locked_by.remove(conv) + __all__ = ( "Dispenser", diff --git a/royalnet/engineer/exc.py b/royalnet/engineer/exc.py index bfe4c693..fda82e38 100644 --- a/royalnet/engineer/exc.py +++ b/royalnet/engineer/exc.py @@ -61,6 +61,21 @@ class ForbiddenError(FrontendError): """ +class DispenserException(EngineerException): + """ + The base class for errors in :mod:`royalnet.engineer.dispenser`. + """ + + +class LockedDispenserError(DispenserException): + """ + The dispenser couldn't start a new conversation as it is currently locked. + """ + def __init__(self, locked_by, *args): + super().__init__(*args) + self.locked_by = locked_by + + __all__ = ( "EngineerException", "WrenchException", @@ -72,4 +87,6 @@ __all__ = ( "FrontendError", "NotSupportedError", "ForbiddenError", + "DispenserException", + "LockedDispenserError", )