diff options
| -rw-r--r-- | slixmpp/plugins/xep_0059/rsm.py | 81 |
1 files changed, 55 insertions, 26 deletions
diff --git a/slixmpp/plugins/xep_0059/rsm.py b/slixmpp/plugins/xep_0059/rsm.py index a02d94e0..8f7f48f0 100644 --- a/slixmpp/plugins/xep_0059/rsm.py +++ b/slixmpp/plugins/xep_0059/rsm.py @@ -19,23 +19,27 @@ from slixmpp.exceptions import XMPPError log = logging.getLogger(__name__) -class ResultIterator(): +class ResultIterator: """ An iterator for Result Set Managment """ def __init__(self, query, interface, results='substanzas', amount=10, - start=None, reverse=False): + start=None, reverse=False, recv_interface=None, + pre_cb=None, post_cb=None): """ Arguments: query -- The template query - interface -- The substanza of the query, for example disco_items + 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 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() @@ -49,17 +53,23 @@ class ResultIterator(): self.amount = amount self.start = start self.interface = interface + if recv_interface: + self.recv_interface = recv_interface + else: + self.recv_interface = interface + self.pre_cb = pre_cb + self.post_cb = post_cb self.results = results self.reverse = reverse self._stop = False - def __iter__(self): + def __aiter__(self): return self - def __next__(self): - return self.next() + async def __anext__(self): + return await self.next() - def next(self): + async def next(self): """ Return the next page of results from a query. @@ -68,7 +78,7 @@ class ResultIterator(): of items. """ if self._stop: - raise StopIteration + raise StopAsyncIteration 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) @@ -79,28 +89,32 @@ class ResultIterator(): self.query[self.interface]['rsm']['after'] = self.start try: - r = self.query.send(block=True) - - if not r[self.interface]['rsm']['first'] and \ - not r[self.interface]['rsm']['last']: - raise StopIteration - - if r[self.interface]['rsm']['count'] and \ - r[self.interface]['rsm']['first_index']: - count = int(r[self.interface]['rsm']['count']) - first = int(r[self.interface]['rsm']['first_index']) - num_items = len(r[self.interface][self.results]) + if self.pre_cb: + self.pre_cb(self.query) + r = await self.query.send() + + if not r[self.recv_interface]['rsm']['first'] and \ + not r[self.recv_interface]['rsm']['last']: + raise StopAsyncIteration + + if r[self.recv_interface]['rsm']['count'] and \ + r[self.recv_interface]['rsm']['first_index']: + count = int(r[self.recv_interface]['rsm']['count']) + first = int(r[self.recv_interface]['rsm']['first_index']) + num_items = len(r[self.recv_interface][self.results]) if first + num_items == count: self._stop = True if self.reverse: - self.start = r[self.interface]['rsm']['first'] + self.start = r[self.recv_interface]['rsm']['first'] else: - self.start = r[self.interface]['rsm']['last'] + self.start = r[self.recv_interface]['rsm']['last'] + if self.post_cb: + self.post_cb(r) return r except XMPPError: - raise StopIteration + raise StopAsyncIteration class XEP_0059(BasePlugin): @@ -127,7 +141,8 @@ class XEP_0059(BasePlugin): def session_bind(self, jid): self.xmpp['xep_0030'].add_feature(Set.namespace) - def iterate(self, stanza, interface, results='substanzas'): + def iterate(self, stanza, interface, results='substanzas', + recv_interface=None, pre_cb=None, post_cb=None): """ Create a new result set iterator for a given stanza query. @@ -137,9 +152,23 @@ class XEP_0059(BasePlugin): basic disco#items query. interface -- The name of the substanza to which the result set management stanza should be - appended. For example, for disco#items queries - the interface 'disco_items' should be used. + 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'). """ - return ResultIterator(stanza, interface, results) + return ResultIterator(stanza, interface, results, + recv_interface=recv_interface, pre_cb=pre_cb, + post_cb=post_cb) |