diff options
| -rw-r--r-- | slixmpp/plugins/xep_0059/rsm.py | 172 |
1 files changed, 108 insertions, 64 deletions
diff --git a/slixmpp/plugins/xep_0059/rsm.py b/slixmpp/plugins/xep_0059/rsm.py index 0fd6b2f9..00615ee4 100644 --- a/slixmpp/plugins/xep_0059/rsm.py +++ b/slixmpp/plugins/xep_0059/rsm.py @@ -5,9 +5,16 @@ # See the file LICENSE for copying permission. import logging -import slixmpp -from slixmpp import Iq -from slixmpp.plugins import BasePlugin, register_plugin +from collections.abc import AsyncIterator +from typing import ( + Any, + Callable, + Dict, + Optional, +) + +from slixmpp.stanza import Iq +from slixmpp.plugins import BasePlugin from slixmpp.xmlstream import register_stanza_plugin from slixmpp.plugins.xep_0059 import stanza, Set from slixmpp.exceptions import XMPPError @@ -16,41 +23,73 @@ from slixmpp.exceptions import XMPPError log = logging.getLogger(__name__) -class ResultIterator: +class ResultIterator(AsyncIterator): """ An iterator for Result Set Management - """ - def __init__(self, query, interface, results='substanzas', amount=10, - start=None, reverse=False, recv_interface=None, - pre_cb=None, post_cb=None): + Example: + + .. code-block:: python + + q = Iq() + q['to'] = 'pubsub.example.com' + q['disco_items']['node'] = 'blog' + async for i in ResultIterator(q, 'disco_items', '10'): + print(i['disco_items']['items']) + + """ + #: Template for the RSM query + query: Iq + #: Substanza of the query to send, e.g. "disco_items" + interface: str + #: Stanza interface on the query results providing the retrieved + #: elements (used to count them) + results: str + #: From which item id to start + start: Optional[str] + #: Amount of elements to retrieve for each page + amount: int + #: If True, page backwards through the results + reverse: bool + #: Callback to run before sending the stanza + pre_cb: Optional[Callable[[Iq], None]] + #: Callback to run after receiving the reply + post_cb: Optional[Callable[[Iq], None]] + #: Optional dict of Iq options (timeout, etc…) for Iq.send() + iq_options: Dict[str, Any] + + def __init__(self, query: Iq, interface: str, results: str = 'substanzas', + amount: int = 10, + start: Optional[str] = None, reverse: bool = False, + recv_interface: Optional[str] = None, + pre_cb: Optional[Callable[[Iq], None]] = None, + post_cb: Optional[Callable[[Iq], None]] = None, + iq_options: Optional[Dict[str, Any]] = None): """ - Arguments: - query -- The template query - interface -- The substanza of the query to send, for example disco_items - recv_interface -- The substanza of the query to receive, for example disco_items - results -- The query stanza's interface which provides a + :param query: The template query + :param interface: The substanza of the query to send, for example + disco_items + :param recv_interface: The substanza of the query to receive, for + example disco_items + :param results: The query stanza's interface which provides a countable list of query results. - amount -- The max amounts of items to request per iteration - start -- From which item id to start - reverse -- If True, page backwards through the results - pre_cb -- Callback to run before sending the stanza - post_cb -- Callback to run after receiving the reply - - Example: - q = Iq() - q['to'] = 'pubsub.example.com' - q['disco_items']['node'] = 'blog' - for i in ResultIterator(q, 'disco_items', '10'): - print i['disco_items']['items'] - + :param amount: The max amounts of items to request per iteration + :param start: From which item id to start + :param reverse: If True, page backwards through the results + :param pre_cb: Callback to run before sending the stanza + :param post_cb: Callback to run after receiving the reply + :param iq_options: Optional dict of parameters for Iq.send """ self.query = query self.amount = amount self.start = start + if iq_options is None: + self.iq_options = {} + else: + self.iq_options = iq_options self.interface = interface - if recv_interface: + if recv_interface is not None: self.recv_interface = recv_interface else: self.recv_interface = interface @@ -63,10 +102,10 @@ class ResultIterator: def __aiter__(self): return self - async def __anext__(self): + async def __anext__(self) -> Iq: return await self.next() - async def next(self): + async def next(self) -> Iq: """ Return the next page of results from a query. @@ -76,20 +115,19 @@ class ResultIterator: """ if self._stop: raise StopAsyncIteration - if self.query[self.interface]['rsm']['before'] is None: - self.query[self.interface]['rsm']['before'] = self.reverse self.query['id'] = self.query.stream.new_id() self.query[self.interface]['rsm']['max'] = str(self.amount) - if self.start and self.reverse: - self.query[self.interface]['rsm']['before'] = self.start - elif self.start: - self.query[self.interface]['rsm']['after'] = self.start + if self.start: + if self.reverse: + self.query[self.interface]['rsm']['before'] = self.start + else: + self.query[self.interface]['rsm']['after'] = self.start try: if self.pre_cb: self.pre_cb(self.query) - r = await self.query.send() + r = await self.query.send(**self.iq_options) if not r[self.recv_interface]['rsm']['first'] and \ not r[self.recv_interface]['rsm']['last']: @@ -118,7 +156,7 @@ class ResultIterator: class XEP_0059(BasePlugin): """ - XEP-0050: Result Set Management + XEP-0059: Result Set Management """ name = 'xep_0059' @@ -139,34 +177,40 @@ class XEP_0059(BasePlugin): def session_bind(self, jid): self.xmpp['xep_0030'].add_feature(Set.namespace) - def iterate(self, stanza, interface, results='substanzas', amount=10, reverse=False, - recv_interface=None, pre_cb=None, post_cb=None): + def iterate(self, stanza: Iq, interface: str, results: str = 'substanzas', + amount: int = 10, reverse: bool = False, + recv_interface: Optional[str] = None, + pre_cb: Optional[Callable[[Iq], None]] = None, + post_cb: Optional[Callable[[Iq], None]] = None, + iq_options: Optional[Dict[str, Any]] = None + ) -> ResultIterator: """ Create a new result set iterator for a given stanza query. - Arguments: - stanza -- A stanza object to serve as a template for - queries made each iteration. For example, a - basic disco#items query. - interface -- The name of the substanza to which the - result set management stanza should be - appended in the query stanza. For example, - for disco#items queries the interface - 'disco_items' should be used. - recv_interface -- The name of the substanza from which the - result set management stanza should be - read in the result stanza. If unspecified, - it will be set to the same value as the - ``interface`` parameter. - pre_cb -- Callback to run before sending each stanza e.g. - setting the MAM queryid and starting a stanza - collector. - post_cb -- Callback to run after receiving each stanza e.g. - stopping a MAM stanza collector in order to - gather results. - results -- The name of the interface containing the - query results (typically just 'substanzas'). + :param stanza: A stanza object to serve as a template for + queries made each iteration. For example, a + basic disco#items query. + :param interface: The name of the substanza to which the + result set management stanza should be + appended in the query stanza. For example, + for disco#items queries the interface + 'disco_items' should be used. + :param recv_interface: The name of the substanza from which the + result set management stanza should be + read in the result stanza. If unspecified, + it will be set to the same value as the + ``interface`` parameter. + :param pre_cb: Callback to run before sending each stanza e.g. + setting the MAM queryid and starting a stanza + collector. + :param post_cb: Callback to run after receiving each stanza e.g. + stopping a MAM stanza collector in order to + gather results. + :param results: The name of the interface containing the + query results (typically just 'substanzas'). + :param iq_options: Optional dict of parameters for Iq.send """ - return ResultIterator(stanza, interface, results, amount, reverse=reverse, - recv_interface=recv_interface, pre_cb=pre_cb, - post_cb=post_cb) + return ResultIterator(stanza, interface, results, amount, + reverse=reverse, recv_interface=recv_interface, + pre_cb=pre_cb, post_cb=post_cb, + iq_options=iq_options) |