diff --git a/ircrobots/interface.py b/ircrobots/interface.py index 2039ec0..586f798 100644 --- a/ircrobots/interface.py +++ b/ircrobots/interface.py @@ -33,10 +33,15 @@ class SendPriority(IntEnum): DEFAULT = MEDIUM class SentLine(object): - def __init__(self, priority: int, line: Line): + def __init__(self, + id: int, + priority: int, + line: Line): + self.id = id self.priority = priority self.line = line self.future: Future = Future() + def __lt__(self, other: "SentLine") -> bool: return self.priority < other.priority @@ -61,9 +66,11 @@ class IServer(Server): params: ConnectionParams desired_caps: Set[ICapability] - def send_raw(self, line: str, priority=SendPriority.DEFAULT) -> Future: + def send_raw(self, line: str, priority=SendPriority.DEFAULT + ) -> Awaitable[SentLine]: pass - def send(self, line: Line, priority=SendPriority.DEFAULT) -> Future: + def send(self, line: Line, priority=SendPriority.DEFAULT + ) -> Awaitable[SentLine]: pass def wait_for(self, response: IMatchResponse) -> Awaitable[Line]: diff --git a/ircrobots/ircv3.py b/ircrobots/ircv3.py index 27cb72b..bdb27bf 100644 --- a/ircrobots/ircv3.py +++ b/ircrobots/ircv3.py @@ -34,7 +34,15 @@ class Capability(ICapability): alias=self.alias, depends_on=self.depends_on[:]) -CAP_SASL = Capability("sasl") +CAP_SASL = Capability("sasl") +CAP_ECHO = Capability("echo-message") +CAP_LABEL = Capability("labeled-response", "draft/labeled-response-0.2") + +LABEL_TAG = { + "draft/labeled-response-0.2": "draft/label", + "labeled-response": "label" +} + CAPS: List[ICapability] = [ Capability("multi-prefix"), Capability("chghost"), diff --git a/ircrobots/server.py b/ircrobots/server.py index d8f1aef..566f6e8 100644 --- a/ircrobots/server.py +++ b/ircrobots/server.py @@ -6,7 +6,7 @@ from asyncio_throttle import Throttler from ircstates import Emit, Channel, NUMERIC_NAMES from irctokens import build, Line, tokenise -from .ircv3 import CAPContext, CAP_SASL +from .ircv3 import CAPContext, CAP_ECHO, CAP_SASL, CAP_LABEL, LABEL_TAG from .sasl import SASLContext, SASLResult from .matching import ResponseOr, Numerics, Numeric, ParamAny, ParamFolded from .asyncs import MaybeAwait @@ -33,6 +33,7 @@ class Server(IServer): self.sasl_state = SASLResult.NONE + self._sent_count: int = 0 self._wait_for: List[Tuple["Future[Line]", IMatchResponse]] = [] self._write_queue: PriorityQueue[SentLine] = PriorityQueue() self.desired_caps: Set[ICapability] = set([]) @@ -48,12 +49,27 @@ class Server(IServer): return hostmask def send_raw(self, line: str, priority=SendPriority.DEFAULT - ) -> Future: + ) -> Awaitable[SentLine]: return self.send(tokenise(line), priority) - def send(self, line: Line, priority=SendPriority.DEFAULT) -> Future: - prio_line = SentLine(priority, line) - self._write_queue.put_nowait(prio_line) - return prio_line.future + def send(self, line: Line, priority=SendPriority.DEFAULT + ) -> Awaitable[SentLine]: + sent_line = SentLine(self._sent_count, priority, line) + self._sent_count += 1 + + label = self.cap_available(CAP_LABEL) + if not label is None: + tag = LABEL_TAG[label] + if line.tags is None or not tag in line.tags: + if line.tags is None: + line.tags = {} + line.tags[tag] = str(sent_line.id) + + self._write_queue.put_nowait(sent_line) + + async def _assure() -> SentLine: + await sent_line.future + return sent_line + return MaybeAwait(_assure) def set_throttle(self, rate: int, time: float): self.throttle.rate_limit = rate @@ -164,7 +180,7 @@ class Server(IServer): await self._writer.drain() for line in lines: - line.future.set_result(None) + line.future.set_result(line) await self.line_send(line.line) return [l.line for l in lines] diff --git a/requirements.txt b/requirements.txt index cddd6db..0a89b21 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,5 @@ anyio ==1.3.0 asyncio-throttle ==1.0.1 dataclasses ==0.6 -ircstates ==0.8.8 +ircstates ==0.8.9 async_stagger ==0.3.0