Add a coroutine parameter to iq.send() to return a coroutine

(instead of exposing a different send_coroutine method)
This commit is contained in:
mathieui 2015-02-23 17:20:47 +01:00
parent 06358d0665
commit 1450d36377
No known key found for this signature in database
GPG key ID: C59F84CEEFD616E3

View file

@ -160,11 +160,11 @@ class Iq(RootStanza):
return new_iq return new_iq
@asyncio.coroutine @asyncio.coroutine
def send_coroutine(self, timeout=None): def _send_coroutine(self, matcher=None, timeout=None):
"""Send an <iq> stanza over the XML stream. """Send an <iq> stanza over the XML stream.
Blocks (with asyncio) until a the reply is received. Blocks (with asyncio) until a the reply is received.
Use with yield from iq.send_coroutine(). Use with yield from iq.send(coroutine=True).
Overrides StanzaBase.send Overrides StanzaBase.send
@ -174,15 +174,6 @@ class Iq(RootStanza):
response before an IqTimeout is raised response before an IqTimeout is raised
""" """
if self.stream.session_bind_event.is_set():
matcher = MatchIDSender({
'id': self['id'],
'self': self.stream.boundjid,
'peer': self['to']
})
else:
matcher = MatcherId(self['id'])
future = asyncio.Future() future = asyncio.Future()
def callback(result): def callback(result):
@ -217,7 +208,7 @@ class Iq(RootStanza):
raise IqError(result) raise IqError(result)
return result return result
def send(self, callback=None, timeout=None, timeout_callback=None): def send(self, callback=None, timeout=None, timeout_callback=None, coroutine=False):
"""Send an <iq> stanza over the XML stream. """Send an <iq> stanza over the XML stream.
A callback handler can be provided that will be executed when the Iq A callback handler can be provided that will be executed when the Iq
@ -237,6 +228,8 @@ class Iq(RootStanza):
function. Will be executed when the timeout expires function. Will be executed when the timeout expires
before a response has been received with the before a response has been received with the
originally-sent IQ stanza. originally-sent IQ stanza.
coroutine -- This function will return a coroutine if this argument
is True.
""" """
if self.stream.session_bind_event.is_set(): if self.stream.session_bind_event.is_set():
matcher = MatchIDSender({ matcher = MatchIDSender({
@ -247,33 +240,36 @@ class Iq(RootStanza):
else: else:
matcher = MatcherId(self['id']) matcher = MatcherId(self['id'])
if callback is not None and self['type'] in ('get', 'set'): if not coroutine:
handler_name = 'IqCallback_%s' % self['id'] if callback is not None and self['type'] in ('get', 'set'):
if asyncio.iscoroutinefunction(callback): handler_name = 'IqCallback_%s' % self['id']
constr = CoroutineCallback if asyncio.iscoroutinefunction(callback):
constr = CoroutineCallback
else:
constr = Callback
if timeout_callback:
self.callback = callback
self.timeout_callback = timeout_callback
self.stream.schedule('IqTimeout_%s' % self['id'],
timeout,
self._fire_timeout,
repeat=False)
handler = constr(handler_name,
matcher,
self._handle_result,
once=True)
else:
handler = constr(handler_name,
matcher,
callback,
once=True)
self.stream.register_handler(handler)
StanzaBase.send(self)
return handler_name
else: else:
constr = Callback return StanzaBase.send(self)
if timeout_callback:
self.callback = callback
self.timeout_callback = timeout_callback
self.stream.schedule('IqTimeout_%s' % self['id'],
timeout,
self._fire_timeout,
repeat=False)
handler = constr(handler_name,
matcher,
self._handle_result,
once=True)
else:
handler = constr(handler_name,
matcher,
callback,
once=True)
self.stream.register_handler(handler)
StanzaBase.send(self)
return handler_name
else: else:
return StanzaBase.send(self) return self._send_coroutine(timeout=timeout, matcher=matcher)
def _handle_result(self, iq): def _handle_result(self, iq):
# we got the IQ, so don't fire the timeout # we got the IQ, so don't fire the timeout