XEP-0313: Update the API

- add an iterate() method that makes this plugin more practical
- add a get_fields method to retrieve the available search fields
- add a get_archive_metadata method.

This is a big chunk because git refused to split it further.

4 files changed, 304 insertions, 44 deletions

diff --git a/slixmpp/plugins/xep_0313/__init__.py b/slixmpp/plugins/xep_0313/__init__.py
index 04e65eff..2b7f3273 100644
--- a/slixmpp/plugins/xep_0313/__init__.py
+++ b/slixmpp/plugins/xep_0313/__init__.py
@@ -5,10 +5,10 @@
 # See the file LICENSE for copying permissio
 from slixmpp.plugins.base import register_plugin
 
-from slixmpp.plugins.xep_0313.stanza import Result, MAM
+from slixmpp.plugins.xep_0313.stanza import Result, MAM, Metadata
 from slixmpp.plugins.xep_0313.mam import XEP_0313
 
 
 register_plugin(XEP_0313)
 
-__all__ = ['XEP_0313', 'Result', 'MAM']
+__all__ = ['XEP_0313', 'Result', 'MAM', 'Metadata']
diff --git a/slixmpp/plugins/xep_0313/mam.py b/slixmpp/plugins/xep_0313/mam.py
index 5f2c0bcc..c06360c4 100644
--- a/slixmpp/plugins/xep_0313/mam.py
+++ b/slixmpp/plugins/xep_0313/mam.py
@@ -5,8 +5,17 @@
 # See the file LICENSE for copying permission
 import logging
 
+from asyncio import Future
+from collections.abc import AsyncGenerator
 from datetime import datetime
-from typing import Any, Dict, Callable, Optional, Awaitable
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    Dict,
+    Optional,
+    Tuple,
+)
 
 from slixmpp import JID
 from slixmpp.stanza import Message, Iq
@@ -45,6 +54,9 @@ class XEP_0313(BasePlugin):
         )
         register_stanza_plugin(stanza.MAM, self.xmpp['xep_0059'].stanza.Set)
         register_stanza_plugin(stanza.Fin, self.xmpp['xep_0059'].stanza.Set)
+        register_stanza_plugin(Iq, stanza.Metadata)
+        register_stanza_plugin(stanza.Metadata, stanza.Start)
+        register_stanza_plugin(stanza.Metadata, stanza.End)
 
     def retrieve(
             self,
@@ -72,16 +84,10 @@ class XEP_0313(BasePlugin):
         :param bool iterator: Use RSM and iterate over a paginated query
         :param dict rsm: RSM custom options
         """
-        iq = self.xmpp.Iq()
+        iq, stanza_mask = self._pre_mam_retrieve(
+            jid, start, end, with_jid, ifrom
+        )
         query_id = iq['id']
-
-        iq['to'] = jid
-        iq['from'] = ifrom
-        iq['type'] = 'set'
-        iq['mam']['queryid'] = query_id
-        iq['mam']['start'] = start
-        iq['mam']['end'] = end
-        iq['mam']['with'] = with_jid
         amount = 10
         if rsm:
             for key, value in rsm.items():
@@ -90,12 +96,6 @@ class XEP_0313(BasePlugin):
                     amount = value
         cb_data = {}
 
-        stanza_mask = self.xmpp.Message()
-        stanza_mask.xml.remove(stanza_mask.xml.find('{urn:xmpp:sid:0}origin-id'))
-        del stanza_mask['id']
-        del stanza_mask['lang']
-        stanza_mask['from'] = jid
-        stanza_mask['mam_result']['queryid'] = query_id
         xml_mask = str(stanza_mask)
 
         def pre_cb(query: Iq) -> None:
@@ -114,9 +114,11 @@ class XEP_0313(BasePlugin):
                 result['mam']['results'] = results
 
         if iterator:
-            return self.xmpp['xep_0059'].iterate(iq, 'mam', 'results', amount=amount,
-                                                 reverse=reverse, recv_interface='mam_fin',
-                                                 pre_cb=pre_cb, post_cb=post_cb)
+            return self.xmpp['xep_0059'].iterate(
+                iq, 'mam', 'results', amount=amount,
+                reverse=reverse, recv_interface='mam_fin',
+                pre_cb=pre_cb, post_cb=post_cb
+            )
 
         collector = Collector(
             'MAM_Results_%s' % query_id,
@@ -132,26 +134,142 @@ class XEP_0313(BasePlugin):
 
         return iq.send(timeout=timeout, callback=wrapped_cb)
 
-    def get_preferences(self, timeout=None, callback=None):
-        iq = self.xmpp.Iq()
-        iq['type'] = 'get'
+    async def iterate(
+            self,
+            jid: Optional[JID] = None,
+            start: Optional[datetime] = None,
+            end: Optional[datetime] = None,
+            with_jid: Optional[JID] = None,
+            ifrom: Optional[JID] = None,
+            reverse: bool = False,
+            rsm: Optional[Dict[str, Any]] = None,
+            total: Optional[int] = None,
+    ) -> AsyncGenerator:
+        """
+        Iterate over each message of MAM query.
+
+        :param jid: Entity holding the MAM records
+        :param start: MAM query start time
+        :param end: MAM query end time
+        :param with_jid: Filter results on this JID
+        :param ifrom: To change the from address of the query
+        :param reverse: Get the results in reverse order
+        :param rsm: RSM custom options
+        :param total: A number of messages received after which the query
+                      should stop.
+        """
+        iq, stanza_mask = self._pre_mam_retrieve(
+            jid, start, end, with_jid, ifrom
+        )
         query_id = iq['id']
-        iq['mam_prefs']['query_id'] = query_id
-        return iq.send(timeout=timeout, callback=callback)
-
-    def set_preferences(self, jid=None, default=None, always=None, never=None,
-                        ifrom=None, timeout=None, callback=None):
-        iq = self.xmpp.Iq()
-        iq['type'] = 'set'
-        iq['to'] = jid
-        iq['from'] = ifrom
-        iq['mam_prefs']['default'] = default
-        iq['mam_prefs']['always'] = always
-        iq['mam_prefs']['never'] = never
-        return iq.send(timeout=timeout, callback=callback)
-
-    def get_configuration_commands(self, jid, **kwargs):
-        return self.xmpp['xep_0030'].get_items(
-                jid=jid,
-                node='urn:xmpp:mam#configure',
-                **kwargs)
+        amount = 10
+
+        if rsm:
+            for key, value in rsm.items():
+                iq['mam']['rsm'][key] = str(value)
+                if key == 'max':
+                    amount = value
+        cb_data = {}
+
+        def pre_cb(query: Iq) -> None:
+            stanza_mask['mam_result']['queryid'] = query['id']
+            xml_mask = str(stanza_mask)
+            query['mam']['queryid'] = query['id']
+            collector = Collector(
+                'MAM_Results_%s' % query_id,
+                MatchXMLMask(xml_mask))
+            self.xmpp.register_handler(collector)
+            cb_data['collector'] = collector
+
+        def post_cb(result: Iq) -> None:
+            results = cb_data['collector'].stop()
+            if result['type'] == 'result':
+                result['mam']['results'] = results
+
+        iterator = self.xmpp['xep_0059'].iterate(
+            iq, 'mam', 'results', amount=amount,
+            reverse=reverse, recv_interface='mam_fin',
+            pre_cb=pre_cb, post_cb=post_cb
+        )
+        recv_count = 0
+
+        async for page in iterator:
+            messages = [message for message in page['mam']['results']]
+            if reverse:
+                messages.reverse()
+            for message in messages:
+                yield message
+                recv_count += 1
+                if total is not None and recv_count >= total:
+                    break
+            if total is not None and recv_count >= total:
+                break
+
+    def _pre_mam_retrieve(
+            self,
+            jid: Optional[JID] = None,
+            start: Optional[datetime] = None,
+            end: Optional[datetime] = None,
+            with_jid: Optional[JID] = None,
+            ifrom: Optional[JID] = None,
+    ) -> Tuple[Iq, Message]:
+        """Build the IQ and stanza mask for MAM results
+        """
+        iq = self.xmpp.make_iq_set(ito=jid, ifrom=ifrom)
+        query_id = iq['id']
+        iq['mam']['queryid'] = query_id
+        iq['mam']['start'] = start
+        iq['mam']['end'] = end
+        iq['mam']['with'] = with_jid
+
+        stanza_mask = self.xmpp.Message()
+        stanza_mask.xml.remove(
+            stanza_mask.xml.find('{urn:xmpp:sid:0}origin-id')
+        )
+        del stanza_mask['id']
+        del stanza_mask['lang']
+        stanza_mask['from'] = jid
+        stanza_mask['mam_result']['queryid'] = query_id
+
+        return (iq, stanza_mask)
+
+    async def get_fields(self, jid: Optional[JID] = None, **iqkwargs) -> Form:
+        """Get MAM query fields.
+
+        .. versionaddedd:: 1.8.0
+
+        :param jid: JID to retrieve the policy from.
+        :return: The Form of allowed options
+        """
+        ifrom = iqkwargs.pop('ifrom', None)
+        iq = self.xmpp.make_iq_get(ito=jid, ifrom=ifrom)
+        iq.enable('mam')
+        result = await iq.send(**iqkwargs)
+        return result['mam']['form']
+
+    async def get_configuration_commands(self, jid: Optional[JID],
+                                         **discokwargs) -> Future:
+        """Get the list of MAM advanced configuration commands.
+
+        .. versionchanged:: 1.8.0
+
+        :param jid: JID to get the commands from.
+        """
+        if jid is None:
+            jid = self.xmpp.boundjid.bare
+        return await self.xmpp['xep_0030'].get_items(
+            jid=jid,
+            node='urn:xmpp:mam#configure',
+            **discokwargs
+        )
+
+    def get_archive_metadata(self, jid: Optional[JID] = None,
+                             **iqkwargs) -> Future:
+        """Get the archive metadata from a JID.
+
+        :param jid: JID to get the metadata from.
+        """
+        ifrom = iqkwargs.pop('ifrom', None)
+        iq = self.xmpp.make_iq_get(ito=jid, ifrom=ifrom)
+        iq.enable('mam_metadata')
+        return iq.send(**iqkwargs)
diff --git a/slixmpp/plugins/xep_0313/stanza.py b/slixmpp/plugins/xep_0313/stanza.py
index 55c80a35..fbca3c8e 100644
--- a/slixmpp/plugins/xep_0313/stanza.py
+++ b/slixmpp/plugins/xep_0313/stanza.py
@@ -170,12 +170,155 @@ class MAM(ElementBase):
 
 
 class Fin(ElementBase):
+    """A MAM fin element (end of query).
+
+    .. code-block:: xml
+
+        <iq type='result' id='juliet1'>
+          <fin xmlns='urn:xmpp:mam:2'>
+            <set xmlns='http://jabber.org/protocol/rsm'>
+              <first index='0'>28482-98726-73623</first>
+              <last>09af3-cc343-b409f</last>
+            </set>
+          </fin>
+        </iq>
+
+    """
     name = 'fin'
     namespace = 'urn:xmpp:mam:2'
     plugin_attrib = 'mam_fin'
 
+
 class Result(ElementBase):
+    """A MAM result payload.
+
+    .. code-block:: xml
+
+        <message id='aeb213' to='juliet@capulet.lit/chamber'>
+          <result xmlns='urn:xmpp:mam:2' queryid='f27' id='28482-98726-73623'>
+            <forwarded xmlns='urn:xmpp:forward:0'>
+              <delay xmlns='urn:xmpp:delay' stamp='2010-07-10T23:08:25Z'/>
+              <message xmlns='jabber:client' from="witch@shakespeare.lit"
+                       to="macbeth@shakespeare.lit">
+                <body>Hail to thee</body>
+              </message>
+            </forwarded>
+          </result>
+        </message>
+    """
     name = 'result'
     namespace = 'urn:xmpp:mam:2'
     plugin_attrib = 'mam_result'
+    #: Available interfaces:
+    #:
+    #: - ``queryid``: MAM queryid
+    #: - ``id``: ID of the result
     interfaces = {'queryid', 'id'}
+
+
+class Metadata(ElementBase):
+    """Element containing archive metadata
+
+    .. code-block:: xml
+
+        <iq type='result' id='jui8921rr9'>
+          <metadata xmlns='urn:xmpp:mam:2'>
+            <start id='YWxwaGEg' timestamp='2008-08-22T21:09:04Z' />
+            <end id='b21lZ2Eg' timestamp='2020-04-20T14:34:21Z' />
+          </metadata>
+        </iq>
+
+    """
+    name = 'metadata'
+    namespace = 'urn:xmpp:mam:2'
+    plugin_attrib = 'mam_metadata'
+
+
+class Start(ElementBase):
+    """Metadata about the start of an archive.
+
+    .. code-block:: xml
+
+        <iq type='result' id='jui8921rr9'>
+          <metadata xmlns='urn:xmpp:mam:2'>
+            <start id='YWxwaGEg' timestamp='2008-08-22T21:09:04Z' />
+            <end id='b21lZ2Eg' timestamp='2020-04-20T14:34:21Z' />
+          </metadata>
+        </iq>
+
+    """
+    name = 'start'
+    namespace = 'urn:xmpp:mam:2'
+    plugin_attrib = name
+    #: Available interfaces:
+    #:
+    #: - ``id``: ID of the first message of the archive
+    #: - ``timestamp`` (``datetime``): timestamp of the first message of the
+    #:   archive
+    interfaces = {'id', 'timestamp'}
+
+    def get_timestamp(self) -> Optional[datetime]:
+        """Get the timestamp.
+
+        :returns: The timestamp.
+        """
+        stamp = self.xml.attrib.get('timestamp', None)
+        if stamp is not None:
+            return xep_0082.parse(stamp)
+        return stamp
+
+    def set_timestamp(self, value: Union[datetime, str]):
+        """Set the timestamp.
+
+        :param value: Value of the timestamp (either a datetime or a
+                      XEP-0082 timestamp string.
+        """
+        if isinstance(value, str):
+            value = xep_0082.parse(value)
+        value = xep_0082.format_datetime(value)
+        self.xml.attrib['timestamp'] = value
+
+
+class End(ElementBase):
+    """Metadata about the end of an archive.
+
+    .. code-block:: xml
+
+        <iq type='result' id='jui8921rr9'>
+          <metadata xmlns='urn:xmpp:mam:2'>
+            <start id='YWxwaGEg' timestamp='2008-08-22T21:09:04Z' />
+            <end id='b21lZ2Eg' timestamp='2020-04-20T14:34:21Z' />
+          </metadata>
+        </iq>
+
+    """
+    name = 'end'
+    namespace = 'urn:xmpp:mam:2'
+    plugin_attrib = name
+    #: Available interfaces:
+    #:
+    #: - ``id``: ID of the first message of the archive
+    #: - ``timestamp`` (``datetime``): timestamp of the first message of the
+    #:   archive
+    interfaces = {'id', 'timestamp'}
+
+    def get_timestamp(self) -> Optional[datetime]:
+        """Get the timestamp.
+
+        :returns: The timestamp.
+        """
+        stamp = self.xml.attrib.get('timestamp', None)
+        if stamp is not None:
+            return xep_0082.parse(stamp)
+        return stamp
+
+    def set_timestamp(self, value: Union[datetime, str]):
+        """Set the timestamp.
+
+        :param value: Value of the timestamp (either a datetime or a
+                      XEP-0082 timestamp string.
+        """
+        if isinstance(value, str):
+            value = xep_0082.parse(value)
+        value = xep_0082.format_datetime(value)
+        self.xml.attrib['timestamp'] = value
diff --git a/tests/test_stanza_xep_0313.py b/tests/test_stanza_xep_0313.py
index 5c7b42a9..d7bd3080 100644
--- a/tests/test_stanza_xep_0313.py
+++ b/tests/test_stanza_xep_0313.py
@@ -13,7 +13,6 @@ class TestMAM(SlixTest):
     def setUp(self):
         register_stanza_plugin(stanza.MAM, Form)
         register_stanza_plugin(Iq, stanza.MAM)
-        register_stanza_plugin(Iq, stanza.Preferences)
         register_stanza_plugin(Message, stanza.Result)
         register_stanza_plugin(Iq, stanza.Fin)
         register_stanza_plugin(