Efficient treatment of constants in TemplateRegistry.
authorBenjamin Braatz <bb@bbraatz.eu>
Fri, 3 Sep 2021 14:12:31 +0000 (16:12 +0200)
committerBenjamin Braatz <bb@bbraatz.eu>
Fri, 3 Sep 2021 14:12:31 +0000 (16:12 +0200)
controlpi/messagebus.py

index 3e216db458d11aa54f77ea016a8b75555d5f8626..d88707a947b61dde77662c89e0e76e07aaba1fb5 100644 (file)
@@ -643,11 +643,11 @@ class MessageTemplate(Dict[str, JSONSchema]):
         return True
 
 
-class MessageTemplateRegistry:
+class TemplateRegistry:
     """Manage a collection of message templates with registered clients.
 
-    A new MessageTemplateRegistry is created by:
-    >>> r = MessageTemplateRegistry()
+    A new TemplateRegistry is created by:
+    >>> r = TemplateRegistry()
 
     Client names (strings) can be registered for message templates, which
     are mappings from keys to JSON schemas:
@@ -751,84 +751,48 @@ class MessageTemplateRegistry:
     def __init__(self) -> None:
         """Initialise an empty registry.
 
-        >>> r = MessageTemplateRegistry()
+        >>> r = TemplateRegistry()
         """
         self._clients: List[str] = []
-        self._children: Dict[str, Dict[str, MessageTemplateRegistry]] = {}
+        self._constants: Dict[str, Dict[str, TemplateRegistry]] = {}
+        # First key is the message key, second key is the constant string
+        self._schemas: Dict[str, Dict[str, TemplateRegistry]] = {}
+        # First key is the message key, second key is the JSON schema string
         self._templates: Dict[str, List[MessageTemplate]] = {}
 
     def insert(self, template: MessageTemplate, client: str) -> None:
         """Register a client for a template.
 
-        >>> r = MessageTemplateRegistry()
-        >>> r.insert({'k1': {'const': 'v1'}, 'k2': {'const': 'v1'}}, 'C 1')
-        >>> r.insert({'k1': {'const': 'v1'}, 'k2': {'const': 'v2'}}, 'C 2')
-        >>> r.insert({'k1': {'const': 'v2'}, 'k2': {'const': 'v1'}}, 'C 3')
-        >>> r.insert({'k1': {'const': 'v2'}, 'k2': {'const': 'v2'}}, 'C 4')
+        >>> r = TemplateRegistry()
+        >>> r.insert({'k1': {'const': 'v1'}, 'k2': {'type': 'integer'}}, 'C 1')
+        >>> r.insert({'k1': {'const': 'v1'}, 'k2': {'type': 'string'}}, 'C 2')
+        >>> r.insert({'k1': {'type': 'integer'}, 'k2': {'const': 'v1'}}, 'C 3')
+        >>> r.insert({'k1': {'type': 'integer'}, 'k2': {'const': 'v2'}}, 'C 4')
         >>> r.insert({}, 'C 5')
-
-        Implementation details:
-        -----------------------
-        The tree nodes on the way to a registered object are used/created
-        in the order given in the message template, which can be used to
-        design more efficient lookups (e.g., putting rarer key-value pairs
-        earlier in the template).
-        >>> r._clients
-        ['C 5']
-        >>> r._children.keys()
-        dict_keys(['k1'])
-        >>> r._children['k1'].keys()
-        dict_keys(['{"const": "v1"}', '{"const": "v2"}'])
-        >>> r._children['k1']['{"const": "v1"}']._clients
-        []
-        >>> r._children['k1']['{"const": "v1"}']._children.keys()
-        dict_keys(['k2'])
-        >>> r._children['k1']['{"const": "v1"}']._children['k2'].keys()
-        dict_keys(['{"const": "v1"}', '{"const": "v2"}'])
-        >>> (r._children['k1']['{"const": "v1"}']
-        ...   ._children['k2']['{"const": "v1"}'])._clients
-        ['C 1']
-        >>> (r._children['k1']['{"const": "v1"}']
-        ...   ._children['k2']['{"const": "v1"}'])._children.keys()
-        dict_keys([])
-        >>> (r._children['k1']['{"const": "v1"}']
-        ...   ._children['k2']['{"const": "v2"}'])._clients
-        ['C 2']
-        >>> (r._children['k1']['{"const": "v1"}']
-        ...   ._children['k2']['{"const": "v2"}'])._children.keys()
-        dict_keys([])
-        >>> r._children['k1']['{"const": "v2"}']._clients
-        []
-        >>> r._children['k1']['{"const": "v2"}']._children.keys()
-        dict_keys(['k2'])
-        >>> r._children['k1']['{"const": "v2"}']._children['k2'].keys()
-        dict_keys(['{"const": "v1"}', '{"const": "v2"}'])
-        >>> (r._children['k1']['{"const": "v2"}']
-        ...   ._children['k2']['{"const": "v1"}'])._clients
-        ['C 3']
-        >>> (r._children['k1']['{"const": "v2"}']
-        ...   ._children['k2']['{"const": "v1"}'])._children.keys()
-        dict_keys([])
-        >>> (r._children['k1']['{"const": "v2"}']
-        ...   ._children['k2']['{"const": "v2"}'])._clients
-        ['C 4']
-        >>> (r._children['k1']['{"const": "v2"}']
-        ...   ._children['k2']['{"const": "v2"}'])._children.keys()
-        dict_keys([])
         """
         if not template:
             self._clients.append(client)
         else:
             key, schema = next(iter(template.items()))
-            schema_string = json.dumps(schema)
             reduced_template = MessageTemplate({k: template[k]
                                                 for k in template
                                                 if k != key})
-            if key not in self._children:
-                self._children[key] = {}
-            if schema_string not in self._children[key]:
-                self._children[key][schema_string] = MessageTemplateRegistry()
-            self._children[key][schema_string].insert(reduced_template, client)
+            if (isinstance(schema, dict) and len(schema) == 1 and
+                    'const' in schema and isinstance(schema['const'], str)):
+                value = schema['const']
+                if key not in self._constants:
+                    self._constants[key] = {}
+                if value not in self._constants[key]:
+                    self._constants[key][value] = TemplateRegistry()
+                self._constants[key][value].insert(reduced_template, client)
+            else:
+                schema_string = json.dumps(schema)
+                if key not in self._schemas:
+                    self._schemas[key] = {}
+                if schema_string not in self._schemas[key]:
+                    self._schemas[key][schema_string] = TemplateRegistry()
+                self._schemas[key][schema_string].insert(reduced_template,
+                                                         client)
         if client not in self._templates:
             self._templates[client] = []
         self._templates[client].append(template)
@@ -836,91 +800,78 @@ class MessageTemplateRegistry:
     def delete(self, client: str) -> bool:
         """Unregister a client from all templates.
 
-        >>> r = MessageTemplateRegistry()
-        >>> r.insert({'k1': {'const': 'v1'}, 'k2': {'const': 'v1'}}, 'C 1')
-        >>> r.insert({'k1': {'const': 'v1'}, 'k2': {'const': 'v2'}}, 'C 2')
-        >>> r.insert({'k1': {'const': 'v2'}, 'k2': {'const': 'v1'}}, 'C 3')
-        >>> r.insert({'k1': {'const': 'v2'}, 'k2': {'const': 'v2'}}, 'C 4')
+        >>> r = TemplateRegistry()
+        >>> r.insert({'k1': {'const': 'v1'}, 'k2': {'type': 'integer'}}, 'C 1')
+        >>> r.insert({'k1': {'const': 'v1'}, 'k2': {'type': 'string'}}, 'C 2')
+        >>> r.insert({'k1': {'type': 'integer'}, 'k2': {'const': 'v1'}}, 'C 3')
+        >>> r.insert({'k1': {'type': 'integer'}, 'k2': {'const': 'v2'}}, 'C 4')
         >>> r.insert({}, 'C 5')
         >>> r.delete('C 3')
         True
         >>> r.delete('C 4')
         True
-
-        Implementation details:
-        -----------------------
-        If parts of the tree become superfluous by the deletion of the
-        client, they are also completely removed to reduce the lookup
-        effort and keep the tree clean.
-        >>> r._clients
-        ['C 5']
-        >>> r._children.keys()
-        dict_keys(['k1'])
-        >>> r._children['k1'].keys()
-        dict_keys(['{"const": "v1"}'])
-        >>> r._children['k1']['{"const": "v1"}']._clients
-        []
-        >>> r._children['k1']['{"const": "v1"}']._children.keys()
-        dict_keys(['k2'])
-        >>> r._children['k1']['{"const": "v1"}']._children['k2'].keys()
-        dict_keys(['{"const": "v1"}', '{"const": "v2"}'])
-        >>> (r._children['k1']['{"const": "v1"}']
-        ...   ._children['k2']['{"const": "v1"}'])._clients
-        ['C 1']
-        >>> (r._children['k1']['{"const": "v1"}']
-        ...   ._children['k2']['{"const": "v1"}'])._children.keys()
-        dict_keys([])
-        >>> (r._children['k1']['{"const": "v1"}']
-        ...   ._children['k2']['{"const": "v2"}'])._clients
-        ['C 2']
-        >>> (r._children['k1']['{"const": "v1"}']
-        ...   ._children['k2']['{"const": "v2"}'])._children.keys()
-        dict_keys([])
         """
         self._clients = [c for c in self._clients if c != client]
-        new_children: Dict[str, Dict[str, MessageTemplateRegistry]] = {}
-        for key in self._children:
-            new_children[key] = {}
-            for schema in self._children[key]:
-                if self._children[key][schema].delete(client):
-                    new_children[key][schema] = self._children[key][schema]
-            if not new_children[key]:
-                del new_children[key]
-        self._children = new_children
+        new_constants: Dict[str, Dict[str, TemplateRegistry]] = {}
+        for key in self._constants:
+            new_constants[key] = {}
+            for value in self._constants[key]:
+                if self._constants[key][value].delete(client):
+                    new_constants[key][value] = self._constants[key][value]
+            if not new_constants[key]:
+                del new_constants[key]
+        self._constants = new_constants
+        new_schemas: Dict[str, Dict[str, TemplateRegistry]] = {}
+        for key in self._schemas:
+            new_schemas[key] = {}
+            for schema in self._schemas[key]:
+                if self._schemas[key][schema].delete(client):
+                    new_schemas[key][schema] = self._schemas[key][schema]
+            if not new_schemas[key]:
+                del new_schemas[key]
+        self._schemas = new_schemas
         if client in self._templates:
             del self._templates[client]
-        if self._clients or self._children:
+        if self._clients or self._constants or self._schemas:
             return True
         return False
 
     def check(self, client: str, message: Message) -> bool:
         """Get if a client has a registered template matching a message.
 
-        >>> r = MessageTemplateRegistry()
+        >>> r = TemplateRegistry()
         >>> r.insert({'k1': {'const': 'v1'}}, 'Client 1')
-        >>> for m in [{'k1': 'v1', 'k2': 'v1'}, {'k1': 'v1', 'k2': 'v2'},
-        ...           {'k1': 'v2', 'k2': 'v1'}, {'k1': 'v2', 'k2': 'v2'}]:
+        >>> for m in [{'k1': 'v1', 'k2': 'v1'}, {'k1': 'v1', 'k2': 2},
+        ...           {'k1': 'v2', 'k2': 'v1'}, {'k1': 'v2', 'k2': 2}]:
         ...     print(f"{m}: {r.check('Client 1', m)}")
         {'k1': 'v1', 'k2': 'v1'}: True
-        {'k1': 'v1', 'k2': 'v2'}: True
+        {'k1': 'v1', 'k2': 2}: True
         {'k1': 'v2', 'k2': 'v1'}: False
-        {'k1': 'v2', 'k2': 'v2'}: False
-        >>> r.insert({'k2': {'const': 'v2'}}, 'Client 2')
-        >>> for m in [{'k1': 'v1', 'k2': 'v1'}, {'k1': 'v1', 'k2': 'v2'},
-        ...           {'k1': 'v2', 'k2': 'v1'}, {'k1': 'v2', 'k2': 'v2'}]:
+        {'k1': 'v2', 'k2': 2}: False
+        >>> r.insert({'k2': {'type': 'integer'}}, 'Client 2')
+        >>> for m in [{'k1': 'v1', 'k2': 'v1'}, {'k1': 'v1', 'k2': 2},
+        ...           {'k1': 'v2', 'k2': 'v1'}, {'k1': 'v2', 'k2': 2}]:
         ...     print(f"{m}: {r.check('Client 2', m)}")
         {'k1': 'v1', 'k2': 'v1'}: False
-        {'k1': 'v1', 'k2': 'v2'}: True
+        {'k1': 'v1', 'k2': 2}: True
         {'k1': 'v2', 'k2': 'v1'}: False
-        {'k1': 'v2', 'k2': 'v2'}: True
+        {'k1': 'v2', 'k2': 2}: True
         """
         if client in self._clients:
             return True
-        for key in self._children:
+        for key in self._constants:
+            if (key in message and isinstance(message[key], str) and
+                    message[key] in self._constants[key]):
+                value = message[key]
+                assert isinstance(value, str)
+                child = self._constants[key][value]
+                if child.check(client, message):
+                    return True
+        for key in self._schemas:
             if key in message:
-                for schema_string in self._children[key]:
+                for schema_string in self._schemas[key]:
                     if validate(schema_string, message[key]):
-                        child = self._children[key][schema_string]
+                        child = self._schemas[key][schema_string]
                         if child.check(client, message):
                             return True
         return False
@@ -928,26 +879,35 @@ class MessageTemplateRegistry:
     def get(self, message: Message) -> List[str]:
         """Get all clients registered for templates matching a message.
 
-        >>> r = MessageTemplateRegistry()
+        >>> r = TemplateRegistry()
         >>> r.insert({'k1': {'const': 'v1'}}, 'Client 1')
-        >>> r.insert({'k2': {'const': 'v2'}}, 'Client 2')
-        >>> for m in [{'k1': 'v1', 'k2': 'v1'}, {'k1': 'v1', 'k2': 'v2'},
-        ...           {'k1': 'v2', 'k2': 'v1'}, {'k1': 'v2', 'k2': 'v2'}]:
+        >>> r.insert({'k2': {'type': 'integer'}}, 'Client 2')
+        >>> for m in [{'k1': 'v1', 'k2': 'v1'}, {'k1': 'v1', 'k2': 2},
+        ...           {'k1': 'v2', 'k2': 'v1'}, {'k1': 'v2', 'k2': 2}]:
         ...     print(f"{m}: {r.get(m)}")
         {'k1': 'v1', 'k2': 'v1'}: ['Client 1']
-        {'k1': 'v1', 'k2': 'v2'}: ['Client 1', 'Client 2']
+        {'k1': 'v1', 'k2': 2}: ['Client 1', 'Client 2']
         {'k1': 'v2', 'k2': 'v1'}: []
-        {'k1': 'v2', 'k2': 'v2'}: ['Client 2']
+        {'k1': 'v2', 'k2': 2}: ['Client 2']
         """
         result = []
         for client in self._clients:
             if client not in result:
                 result.append(client)
-        for key in self._children:
+        for key in self._constants:
+            if (key in message and isinstance(message[key], str) and
+                    message[key] in self._constants[key]):
+                value = message[key]
+                assert isinstance(value, str)
+                child = self._constants[key][value]
+                for client in child.get(message):
+                    if client not in result:
+                        result.append(client)
+        for key in self._schemas:
             if key in message:
-                for schema_string in self._children[key]:
+                for schema_string in self._schemas[key]:
                     if validate(schema_string, message[key]):
-                        child = self._children[key][schema_string]
+                        child = self._schemas[key][schema_string]
                         for client in child.get(message):
                             if client not in result:
                                 result.append(client)
@@ -956,7 +916,7 @@ class MessageTemplateRegistry:
     def get_templates(self, client: str) -> List[MessageTemplate]:
         """Get all templates for a client.
 
-        >>> r = MessageTemplateRegistry()
+        >>> r = TemplateRegistry()
         >>> r.insert({'k1': {'const': 'v1'}}, 'Client 1')
         >>> r.get_templates('Client 1')
         [{'k1': {'const': 'v1'}}]
@@ -1099,8 +1059,8 @@ class MessageBus:
         """
         self._queue: asyncio.Queue = asyncio.Queue()
         self._plugins: Dict[str, str] = {}
-        self._send_reg: MessageTemplateRegistry = MessageTemplateRegistry()
-        self._recv_reg: MessageTemplateRegistry = MessageTemplateRegistry()
+        self._send_reg: TemplateRegistry = TemplateRegistry()
+        self._recv_reg: TemplateRegistry = TemplateRegistry()
         self._callbacks: Dict[str, MessageCallback] = {}
 
     def register(self, client: str, plugin: str,