diff --git a/slixmpp/plugins/xep_0049/private_storage.py b/slixmpp/plugins/xep_0049/private_storage.py index cb92db51..22286f92 100644 --- a/slixmpp/plugins/xep_0049/private_storage.py +++ b/slixmpp/plugins/xep_0049/private_storage.py @@ -7,12 +7,19 @@ """ import logging +from typing import ( + List, + Optional, + Union, +) +from asyncio import Future -from slixmpp import Iq +from slixmpp import JID +from slixmpp.stanza import Iq from slixmpp.plugins import BasePlugin from slixmpp.xmlstream.handler import Callback from slixmpp.xmlstream.matcher import StanzaPath -from slixmpp.xmlstream import register_stanza_plugin +from slixmpp.xmlstream import register_stanza_plugin, ElementBase from slixmpp.plugins.xep_0049 import stanza, PrivateXML @@ -32,26 +39,25 @@ class XEP_0049(BasePlugin): def register(self, stanza): register_stanza_plugin(PrivateXML, stanza, iterable=True) - def store(self, data, ifrom=None, timeout=None, callback=None, - timeout_callback=None): - iq = self.xmpp.Iq() - iq['type'] = 'set' - iq['from'] = ifrom + def store(self, data: Union[List[ElementBase], ElementBase], ifrom: Optional[JID] = None, **iqkwargs) -> Future: + """Store data in Private XML Storage. + + :param data: An XML element or list of xml element to store. + """ + iq = self.xmpp.make_iq_set(ifrom=ifrom) if not isinstance(data, list): data = [data] - for elem in data: iq['private'].append(elem) - return iq.send(timeout=timeout, callback=callback, - timeout_callback=timeout_callback) + return iq.send(**iqkwargs) - def retrieve(self, name, ifrom=None, timeout=None, callback=None, - timeout_callback=None): - iq = self.xmpp.Iq() - iq['type'] = 'get' - iq['from'] = ifrom + def retrieve(self, name: str, ifrom: Optional[JID] = None, **iqkwargs) -> Future: + """Get previously stored data from Private XML Storage. + + :param name: Name of the payload to retrieve (slixmpp plugin attribute) + """ + iq = self.xmpp.make_iq_get(ifrom=ifrom) iq['private'].enable(name) - return iq.send(timeout=timeout, callback=callback, - timeout_callback=timeout_callback) + return iq.send(**iqkwargs)