XEP-0191: Add type hints and switch to default args

1 files changed, 34 insertions, 23 deletions

diff --git a/slixmpp/plugins/xep_0191/blocking.py b/slixmpp/plugins/xep_0191/blocking.py
index 5425544d..53c279df 100644
--- a/slixmpp/plugins/xep_0191/blocking.py
+++ b/slixmpp/plugins/xep_0191/blocking.py
@@ -8,7 +8,15 @@
 
 import logging
 
-from slixmpp import Iq
+from asyncio import Future
+from typing import (
+    List,
+    Optional,
+    Set,
+    Union,
+)
+
+from slixmpp.stanza import Iq
 from slixmpp.plugins import BasePlugin
 from slixmpp.xmlstream.handler import Callback
 from slixmpp.xmlstream.matcher import StanzaPath
@@ -18,6 +26,12 @@ from slixmpp.plugins.xep_0191 import stanza, Block, Unblock, BlockList
 
 log = logging.getLogger(__name__)
 
+BlockedJIDs = Union[
+    JID,
+    Set[JID],
+    List[JID]
+]
+
 
 class XEP_0191(BasePlugin):
 
@@ -45,42 +59,39 @@ class XEP_0191(BasePlugin):
         self.xmpp.remove_handler('Blocked Contact')
         self.xmpp.remove_handler('Unblocked Contact')
 
-    def get_blocked(self, ifrom=None, timeout=None, callback=None,
-                          timeout_callback=None):
-        iq = self.xmpp.Iq()
-        iq['type'] = 'get'
-        iq['from'] = ifrom
+    def get_blocked(self, ifrom: Optional[JID] = None, **iqkwargs) -> Future:
+        """Get the list of blocked JIDs."""
+        iq = self.xmpp.make_iq_get(ifrom=ifrom)
         iq.enable('blocklist')
-        return iq.send(timeout=timeout, callback=callback,
-                       timeout_callback=timeout_callback)
+        return iq.send(**iqkwargs)
 
-    def block(self, jids, ifrom=None, timeout=None, callback=None,
-                          timeout_callback=None):
-        iq = self.xmpp.Iq()
-        iq['type'] = 'set'
-        iq['from'] = ifrom
+    def block(self, jids: BlockedJIDs,
+              ifrom: Optional[JID] = None, **iqkwargs) -> Future:
+        """Block a JID or a list of JIDs.
 
+        :param jids: JID(s) to block.
+        """
+        iq = self.xmpp.make_iq_set(ifrom=ifrom)
         if not isinstance(jids, (set, list)):
             jids = [jids]
 
         iq['block']['items'] = jids
-        return iq.send(timeout=timeout, callback=callback,
-                       timeout_callback=timeout_callback)
+        return iq.send(**iqkwargs)
 
-    def unblock(self, jids=None, ifrom=None, timeout=None, callback=None,
-                      timeout_callback=None):
-        iq = self.xmpp.Iq()
-        iq['type'] = 'set'
-        iq['from'] = ifrom
+    def unblock(self, jids: BlockedJIDs, ifrom: Optional[JID] = None, **iqkwargs) -> Future:
+        """Unblock a JID or a list of JIDs.
 
+        :param jids: JID(s) to unblock.
+        """
         if jids is None:
-            jids = []
+            raise ValueError("jids cannot be empty.")
+        iq = self.xmpp.make_iq_set(ifrom=ifrom)
+
         if not isinstance(jids, (set, list)):
             jids = [jids]
 
         iq['unblock']['items'] = jids
-        return iq.send(timeout=timeout, callback=callback,
-                       timeout_callback=timeout_callback)
+        return iq.send(**iqkwargs)
 
     def _handle_blocked(self, iq):
         self.xmpp.event('blocked', iq)