import uuid from src import ModuleManager, utils CAP = utils.irc.Capability(None, "draft/labeled-response-0.2") TAG = utils.irc.MessageTag(None, "draft/label") BATCH = utils.irc.BatchType(None, "draft/labeled-response") CAP_TO_TAG = { "draft/labeled-response-0.2": "draft/label" } class WaitingForLabel(object): def __init__(self, line, events): self.line = line self.events = events self.labels_since = 0 class Module(ModuleManager.BaseModule): @utils.hook("new.server") def new_server(self, event): event["server"]._label_cache = {} @utils.hook("received.cap.ls") @utils.hook("received.cap.new") def on_cap(self, event): if CAP.available(event["capabilities"]): return CAP.copy() @utils.hook("preprocess.send") def raw_send(self, event): available_cap = event["server"].available_capability(CAP) if available_cap: label = TAG.get_value(event["line"].tags) if label == None: tag_key = CAP_TO_TAG[available_cap] label = str(uuid.uuid4()) event["line"].tags[tag_key] = label event["server"]._label_cache[label] = WaitingForLabel(event["line"], event["events"]) @utils.hook("raw.received") def raw_recv(self, event): if not event["line"].command == "BATCH": label = TAG.get_value(event["line"].tags) if not label == None: self._recv(event["server"], label, [event["line"]]) @utils.hook("received.batch.end") def batch_end(self, event): if BATCH.match(event["batch"].type): label = TAG.get_value(event["batch"].tags) self._recv(event["server"], label, event["batch"].get_lines()) def _recv(self, server, label, lines): cached = server._label_cache.pop(label) cached.events.on("labeled-response").call(lines=lines) for label, other_cached in server._label_cache.items(): other_cached.labels_since += 1 if other_cached.labels_since == 10: self.log.warn("%d labels seen while waiting for response to %s", [other_cached.labels_since, label])