import re
from enum import Enum
from typing import Callable
from typing import List
from proto.topicData.topicDataRecord_pb2 import TopicDataRecord
[docs]class SubscriptionTokenTYPE(Enum):
TOPIC = 1
REGEX = 2
[docs]class SubscriptionToken:
def __init__(self, id: int, topic: str, callback: Callable[[TopicDataRecord], None], type: SubscriptionTokenTYPE):
self.id: int = id
self.topic: str = topic
self.callback: Callable[[TopicDataRecord], None] = callback
self.type: SubscriptionTokenTYPE = type
[docs]class TopicDataBuffer:
def __init__(self):
self.localTopics: dict[str, TopicDataRecord] = {}
self.dictTopicSubscriptionTokens: dict[str, List[SubscriptionToken]] = {}
self.dictRegexSubscriptionTokens: dict[str, List[SubscriptionToken]] = {}
self.dictTopic2RegexMatches: dict[str, List[str]] = {}
self.currentTokenId = -1
[docs] def publish(self, topicDataRecord: TopicDataRecord):
topic = topicDataRecord.topic
if not topic in self.localTopics:
for regex in self.dictRegexSubscriptionTokens:
if re.match(regex, topic):
if not topic in self.dictTopic2RegexMatches:
self.dictTopic2RegexMatches[topic] = []
self.dictTopic2RegexMatches[topic].append(regex)
self.localTopics[topic] = topicDataRecord
self.notifySubscribers(topicDataRecord)
[docs] def notifySubscribers(self, topicDataRecord):
topic = topicDataRecord.topic
if topic in self.dictTopicSubscriptionTokens:
for token in self.dictTopicSubscriptionTokens[topic]:
token.callback(topicDataRecord)
if topic in self.dictTopic2RegexMatches:
for regex in self.dictTopic2RegexMatches[topic]:
for token in self.dictRegexSubscriptionTokens[regex]:
token.callback(topicDataRecord)
[docs] def subscribeTopic(self, topic: str, callback: Callable[[TopicDataRecord], None]) -> SubscriptionToken:
if not topic in self.dictTopicSubscriptionTokens:
self.dictTopicSubscriptionTokens[topic] = []
token = self.generateToken(topic, callback, SubscriptionTokenTYPE.TOPIC)
self.dictTopicSubscriptionTokens[topic].append(token)
return token
[docs] def subscribeRegex(self, regex: str, callback: Callable[[TopicDataRecord], None]) -> SubscriptionToken:
token = self.generateToken(regex, callback, SubscriptionTokenTYPE.REGEX)
if not regex in self.dictRegexSubscriptionTokens:
self.dictRegexSubscriptionTokens[regex] = []
self.dictRegexSubscriptionTokens[regex].append(token)
for topic in self.localTopics:
if re.match(regex, topic):
if not topic in self.dictTopic2RegexMatches:
self.dictTopic2RegexMatches[topic] = []
self.dictTopic2RegexMatches[topic].append(regex)
return token
[docs] def unsubscribe(self, token: SubscriptionToken) -> bool:
topic = token.topic
existingTopic = False
if topic in self.dictTopicSubscriptionTokens:
self.dictTopicSubscriptionTokens[topic] = list(
filter(lambda a: a.id != token.id, self.dictTopicSubscriptionTokens[topic]))
existingTopic = True
elif topic in self.dictRegexSubscriptionTokens:
self.dictRegexSubscriptionTokens[topic] = list(
filter(lambda a: a.id != token.id, self.dictRegexSubscriptionTokens[topic]))
existingTopic = True
return existingTopic
[docs] def remove(self, topic: str):
if topic in self.localTopics:
del self.localTopics[topic]
[docs] def pull(self, topic: str) -> TopicDataRecord:
return self.localTopics.get(topic)
[docs] def getTopicSubscriptionTokens(self, topic: str) -> List[SubscriptionToken]:
return self.dictTopicSubscriptionTokens.get(topic)
[docs] def getRegexSubscriptionTokens(self, regex: str) -> List[SubscriptionToken]:
return self.dictRegexSubscriptionTokens.get(regex)
[docs] def generateToken(self, topic: str, callback: Callable[[TopicDataRecord], None],
type: SubscriptionTokenTYPE) -> SubscriptionToken:
self.currentTokenId += 1
return SubscriptionToken(self.currentTokenId, topic, callback, type)