MAM: Use tasks

nbxmpp/modules/mam.py

@@ -18,19 +18,18 @@
 
 from nbxmpp.protocol import JID
 from nbxmpp.protocol import Iq
-from nbxmpp.protocol import isResultNode
 from nbxmpp.protocol import Node
 from nbxmpp.namespaces import Namespace
 from nbxmpp.structs import MAMQueryData
 from nbxmpp.structs import MAMPreferencesData
-from nbxmpp.structs import CommonResult
-from nbxmpp.util import call_on_response
-from nbxmpp.util import callback
-from nbxmpp.util import raise_error
+from nbxmpp.task import iq_request_task
+from nbxmpp.errors import StanzaError
+from nbxmpp.errors import MalformedStanzaError
+from nbxmpp.modules.base import BaseModule
 from nbxmpp.modules.rsm import parse_rsm
 from nbxmpp.modules.dataforms import SimpleDataForm
 from nbxmpp.modules.dataforms import create_field
-from nbxmpp.modules.base import BaseModule
+from nbxmpp.modules.util import process_response
 
 
 class MAM(BaseModule):
@@ -40,7 +39,7 @@ class MAM(BaseModule):
         self._client = client
         self.handlers = []
 
-    @call_on_response('_query_result')
+    @iq_request_task
     def make_query(self,
                    jid,
                    queryid=None,
@@ -50,166 +49,155 @@ class MAM(BaseModule):
                    after=None,
                    max_=70):
 
-        iq = Iq(typ='set', to=jid, queryNS=Namespace.MAM_2)
-        if queryid is not None:
-            iq.getQuery().setAttr('queryid', queryid)
-
-        payload = [
-            self._make_query_form(start, end, with_),
-            self._make_rsm_query(max_, after)
-        ]
-
-        iq.setQueryPayload(payload)
-        return iq
-
-    @staticmethod
-    def _make_query_form(start, end, with_):
-        fields = [
-            create_field(typ='hidden', var='FORM_TYPE', value=Namespace.MAM_2)
-        ]
-
-        if start:
-            fields.append(create_field(
-                typ='text-single',
-                var='start',
-                value=start.strftime('%Y-%m-%dT%H:%M:%SZ')))
-
-        if end:
-            fields.append(create_field(
-                typ='text-single',
-                var='end',
-                value=end.strftime('%Y-%m-%dT%H:%M:%SZ')))
-
-        if with_:
-            fields.append(create_field(
-                typ='jid-single',
-                var='with',
-                value=with_))
-
-        return SimpleDataForm(type_='submit', fields=fields)
-
-    @staticmethod
-    def _make_rsm_query(max_, after):
-        rsm_set = Node('set', attrs={'xmlns': Namespace.RSM})
-        if max_ is not None:
-            rsm_set.setTagData('max', max_)
-        if after is not None:
-            rsm_set.setTagData('after', after)
-        return rsm_set
-
-    @callback
-    def _query_result(self, stanza):
-        if not isResultNode(stanza):
-            return raise_error(self._log.info, stanza)
-
-        jid = stanza.getFrom()
-        fin = stanza.getTag('fin', namespace=Namespace.MAM_2)
+        _task = yield
+
+        response = yield _make_request(jid, queryid,
+                                       start, end, with_, after, max_)
+        if response.isError():
+            raise StanzaError(response)
+
+        jid = response.getFrom()
+        fin = response.getTag('fin', namespace=Namespace.MAM_2)
         if fin is None:
-            return raise_error(self._log.warning,
-                               stanza,
-                               'stanza-malformed',
-                               'No fin node found')
+            raise MalformedStanzaError('fin node missing', response)
 
         rsm = parse_rsm(fin)
         if rsm is None:
-            return raise_error(self._log.warning,
-                               stanza,
-                               'stanza-malformed',
-                               'rsm set missing')
+            raise MalformedStanzaError('rsm set missing', response)
 
         complete = fin.getAttr('complete') == 'true'
         if not complete:
             if rsm.first is None or rsm.last is None:
-                return raise_error(self._log.warning,
-                                   stanza,
-                                   'stanza-malformed',
-                                   'missing first or last element')
+                raise MalformedStanzaError('first or last element missing',
+                                           response)
 
-        return MAMQueryData(jid=jid,
-                            complete=complete,
-                            rsm=rsm)
+        yield MAMQueryData(jid=jid,
+                           complete=complete,
+                           rsm=rsm)
 
-    @call_on_response('_preferences_result')
+    @iq_request_task
     def request_preferences(self):
-        iq = Iq('get', queryNS=Namespace.MAM_2)
-        iq.setQuery('prefs')
-        return iq
+        _task = yield
 
-    @callback
-    def _preferences_result(self, stanza):
-        if not isResultNode(stanza):
-            return raise_error(self._log.info, stanza)
+        response = yield _make_pref_request()
+        if response.isError():
+            raise StanzaError(response)
 
-        prefs = stanza.getTag('prefs', namespace=Namespace.MAM_2)
+        prefs = response.getTag('prefs', namespace=Namespace.MAM_2)
         if prefs is None:
-            return raise_error(self._log.warning,
-                               stanza,
-                               'stanza-malformed',
-                               'No prefs node found')
+            raise MalformedStanzaError('prefs node missing', response)
 
         default = prefs.getAttr('default')
         if default is None:
-            return raise_error(self._log.warning,
-                               stanza,
-                               'stanza-malformed',
-                               'No default attr found')
+            raise MalformedStanzaError('default attr missing', response)
 
         always_node = prefs.getTag('always')
         if always_node is None:
-            return raise_error(self._log.warning,
-                               stanza,
-                               'stanza-malformed',
-                               'No always node found')
+            raise MalformedStanzaError('always node missing', response)
 
-        always = self._get_preference_jids(always_node)
+        always = _get_preference_jids(always_node)
 
         never_node = prefs.getTag('never')
         if never_node is None:
-            return raise_error(self._log.warning,
-                               stanza,
-                               'stanza-malformed',
-                               'No never node found')
-
-        never = self._get_preference_jids(never_node)
-        return MAMPreferencesData(default=default,
-                                  always=always,
-                                  never=never)
-
-    def _get_preference_jids(self, node):
-        jids = []
-        for item in node.getTags('jid'):
-            jid = item.getData()
-            if not jid:
-                continue
-
-            try:
-                jid = JID.from_string(jid)
-            except Exception:
-                self._log.warning('Invalid jid found in preferences: %s',
-                                  jid)
-            jids.append(jid)
-        return jids
-
-    @call_on_response('_default_response')
+            raise MalformedStanzaError('never node missing', response)
+
+        never = _get_preference_jids(never_node)
+        yield MAMPreferencesData(default=default,
+                                 always=always,
+                                 never=never)
+
+    @iq_request_task
     def set_preferences(self, default, always, never):
+        _task = yield
+
         if default not in ('always', 'never', 'roster'):
             raise ValueError('Wrong default preferences type')
 
-        iq = Iq(typ='set')
-        prefs = iq.addChild(name='prefs',
-                            namespace=Namespace.MAM_2,
-                            attrs={'default': default})
-        always_node = prefs.addChild(name='always')
-        never_node = prefs.addChild(name='never')
-        for jid in always:
-            always_node.addChild(name='jid').setData(jid)
-
-        for jid in never:
-            never_node.addChild(name='jid').setData(jid)
-        return iq
-
-    @callback
-    def _default_response(self, stanza):
-        if not isResultNode(stanza):
-            return raise_error(self._log.info, stanza)
-        return CommonResult(jid=stanza.getFrom())
+        response = yield _make_set_pref_request(default, always, never)
+        yield process_response(response)
+
+
+def _make_query_form(start, end, with_):
+    fields = [
+        create_field(typ='hidden', var='FORM_TYPE', value=Namespace.MAM_2)
+    ]
+
+    if start:
+        fields.append(create_field(
+            typ='text-single',
+            var='start',
+            value=start.strftime('%Y-%m-%dT%H:%M:%SZ')))
+
+    if end:
+        fields.append(create_field(
+            typ='text-single',
+            var='end',
+            value=end.strftime('%Y-%m-%dT%H:%M:%SZ')))
+
+    if with_:
+        fields.append(create_field(
+            typ='jid-single',
+            var='with',
+            value=with_))
+
+    return SimpleDataForm(type_='submit', fields=fields)
+
+
+def _make_rsm_query(max_, after):
+    rsm_set = Node('set', attrs={'xmlns': Namespace.RSM})
+    if max_ is not None:
+        rsm_set.setTagData('max', max_)
+    if after is not None:
+        rsm_set.setTagData('after', after)
+    return rsm_set
+
+
+def _make_request(jid, queryid, start, end, with_, after, max_):
+    iq = Iq(typ='set', to=jid, queryNS=Namespace.MAM_2)
+    if queryid is not None:
+        iq.getQuery().setAttr('queryid', queryid)
+
+    payload = [
+        _make_query_form(start, end, with_),
+        _make_rsm_query(max_, after)
+    ]
+
+    iq.setQueryPayload(payload)
+    return iq
+
+
+def _make_pref_request():
+    iq = Iq('get', queryNS=Namespace.MAM_2)
+    iq.setQuery('prefs')
+    return iq
+
+
+def _get_preference_jids(node):
+    jids = []
+    for item in node.getTags('jid'):
+        jid = item.getData()
+        if not jid:
+            continue
+
+        try:
+            jid = JID.from_string(jid)
+        except Exception:
+            continue
+
+        jids.append(jid)
+    return jids
+
+
+def _make_set_pref_request(default, always, never):
+    iq = Iq(typ='set')
+    prefs = iq.addChild(name='prefs',
+                        namespace=Namespace.MAM_2,
+                        attrs={'default': default})
+    always_node = prefs.addChild(name='always')
+    never_node = prefs.addChild(name='never')
+    for jid in always:
+        always_node.addChild(name='jid').setData(jid)
+
+    for jid in never:
+        never_node.addChild(name='jid').setData(jid)
+    return iq