From dfba2a26df2a7fbee2d2fdf131f1d421e25e8c11 Mon Sep 17 00:00:00 2001
From: Alexander PapaTutuWawa <papatutuwawa@polynom.me>
Date: Sun, 13 Jun 2021 00:10:17 +0200
Subject: [PATCH] fix: Fix inconsistencies in the SubscriptionManager's data
 model

---
 mira/subscription.py | 47 +++++++++++++++++++++++++++++++++-----------
 1 file changed, 35 insertions(+), 12 deletions(-)

diff --git a/mira/subscription.py b/mira/subscription.py
index 80a14ef..bd9c9a3 100644
--- a/mira/subscription.py
+++ b/mira/subscription.py
@@ -21,17 +21,11 @@ import json
 from mira.storage import StorageManager
 
 def append_or_insert(dict_, key, value):
-    if key in dict_ or dict_[key]:
+    if key in dict_:
         dict_[key].append(value)
     else:
         dict_[key] = [value]
 
-def remove_or_delete(dict_, key, value):
-    if value in dict_[key] and len(dict_[key]) == 1:
-        del dict_[key]
-    else:
-        dict_[key].remove(value)
-
 class SubscriptionManager:
     '''
     This class is tasked with providing functions that simplify dealing
@@ -58,6 +52,10 @@ class SubscriptionManager:
         SubscriptionManager.__instance = self
 
     def get_subscriptions_for(self, module, jid):
+        '''
+        Returns a dictionary keyword -> data which represents
+        every subscription a jid has in the context of @module.
+        '''
         if not module in self._subscriptions:
             return []
         if not jid in self._subscriptions[module]:
@@ -81,6 +79,24 @@ class SubscriptionManager:
             tmp.append((jid, data))
         return tmp
 
+    def get_subscriptions_for_keywords(self, module, keywords):
+        '''
+        Returns an array of JIDs that are subscribed to at least one of the keywords
+        of module
+        '''
+        if not module in self._subscriptions:
+            return []
+
+        tmp = []
+        keyword_set = set(keywords)
+        for jid in self._subscriptions[module]:
+            if set(self._subscriptions[module][jid].keys()) & keyword_set:
+                continue
+
+            data = self._subscriptions[module][jid][keyword]['data']
+            tmp.append((jid, data))
+        return tmp
+    
     def get_subscription_keywords(self, module):
         '''Returns a list of subscribed keywords in module'''
         if not module in self._subscriptions:
@@ -118,11 +134,12 @@ class SubscriptionManager:
         '''
         if not module in self._subscriptions:
             self._subscriptions[module] = {}
+        if not jid in self._subscriptions[module]:
+            self._subscriptions[module][jid] = {}
 
-        append_or_insert(self._subscriptions[module], jid, {
-            'keyword': keyword,
+        self._subscriptions[module][jid][keyword] = {
             'data': data
-        })
+        }
         self.__flush()
 
     def append_data_for_subscription(self, module, jid, keyword, item):
@@ -156,7 +173,13 @@ class SubscriptionManager:
         Removes a subscription to @keyword for @jid within the context
         of @module
         '''
-        remove_or_delete(self._subscriptions[module], jid, keyword)
+        del self._subscriptions[module][jid][keyword]
+
+        if not self._subscriptions[module][jid]:
+            del self._subscriptions[module][jid]
+        if not self._subscriptions[module]:
+            del self._subscriptions[module]
+
         self.__flush()
 
     def remove_item_for_subscription(self, module, jid, keyword, item):
@@ -179,4 +202,4 @@ class SubscriptionManager:
         '''
         Write subscription data to disk. Just an interface to StorageManager
         '''
-        self._sm.set_data('_StorageManager', 'subscriptions', self._subscriptions)
+        self._sm.set_data('_SubscriptionManager', 'subscriptions', self._subscriptions)