PubSub: Type this module

nbxmpp/modules/pubsub.py

@@ -15,13 +15,17 @@
 # You should have received a copy of the GNU General Public License
 # along with this program; If not, see <http://www.gnu.org/licenses/>.
 
+from typing import Optional, Union, Dict, Any
+
 from nbxmpp.namespaces import Namespace
 from nbxmpp.protocol import Node
 from nbxmpp.protocol import Iq
 from nbxmpp.protocol import isResultNode
+from nbxmpp.protocol import JID
 from nbxmpp.structs import StanzaHandler
 from nbxmpp.structs import PubSubEventData
 from nbxmpp.structs import CommonResult
+from nbxmpp.structs import CommonError
 from nbxmpp.structs import PubSubConfigResult
 from nbxmpp.structs import PubSubPublishResult
 from nbxmpp.modules.dataforms import extend_form
@@ -30,9 +34,14 @@ from nbxmpp.util import callback
 from nbxmpp.util import raise_error
 from nbxmpp.modules.base import BaseModule
 
+# TODO: Use better types.
+Client = Any
+Stanza = Any
+Properties = Any
+
 
 class PubSub(BaseModule):
-    def __init__(self, client):
+    def __init__(self, client: Client) -> None:
         BaseModule.__init__(self, client)
 
         self._client = client
@@ -43,7 +52,7 @@ class PubSub(BaseModule):
                           priority=15),
         ]
 
-    def _process_pubsub_base(self, _client, stanza, properties):
+    def _process_pubsub_base(self, _client: Client, stanza: Stanza, properties: Properties) -> None:
         properties.pubsub = True
         event = stanza.getTag('event', namespace=Namespace.PUBSUB_EVENT)
 
@@ -85,7 +94,7 @@ class PubSub(BaseModule):
             properties.pubsub_event = PubSubEventData(node, id_, item)
 
     @call_on_response('_publish_result_received')
-    def publish(self, jid, node, item, id_=None, options=None):
+    def publish(self, jid: JID, node: str, item: Node, id_: Optional[str] = None, options: Optional[Node] = None) -> Iq:
         query = Iq('set', to=jid)
         pubsub = query.addChild('pubsub', namespace=Namespace.PUBSUB)
         publish = pubsub.addChild('publish', {'node': node})
@@ -93,13 +102,13 @@ class PubSub(BaseModule):
         if id_ is not None:
             attrs = {'id': id_}
         publish.addChild('item', attrs, [item])
-        if options:
+        if options is not None:
             publish = pubsub.addChild('publish-options')
             publish.addChild(node=options)
         return query
 
     @callback
-    def _publish_result_received(self, stanza):
+    def _publish_result_received(self, stanza: Stanza) -> Union[PubSubPublishResult, CommonError]:
         if not isResultNode(stanza):
             return raise_error(self._log.warning, stanza)
 
@@ -122,7 +131,7 @@ class PubSub(BaseModule):
         return PubSubPublishResult(jid, node, id_)
 
     @call_on_response('_default_response')
-    def retract(self, jid, node, id_, notify=True):
+    def retract(self, jid: JID, node: str, id_: str, notify: bool = True) -> Iq:
         query = Iq('set', to=jid)
         pubsub = query.addChild('pubsub', namespace=Namespace.PUBSUB)
         attrs = {'node': node}
@@ -133,7 +142,7 @@ class PubSub(BaseModule):
         return query
 
     @call_on_response('_default_response')
-    def set_node_configuration(self, jid, node, form):
+    def set_node_configuration(self, jid: JID, node: str, form: Node) -> Iq:
         self._log.info('Set configuration for %s %s', node, jid)
         query = Iq('set', to=jid)
         pubsub = query.addChild('pubsub', namespace=Namespace.PUBSUB_OWNER)
@@ -143,7 +152,7 @@ class PubSub(BaseModule):
         return query
 
     @call_on_response('_node_configuration_received')
-    def get_node_configuration(self, jid, node):
+    def get_node_configuration(self, jid: JID, node: str) -> Iq:
         self._log.info('Request node configuration')
         query = Iq('get', to=jid)
         pubsub = query.addChild('pubsub', namespace=Namespace.PUBSUB_OWNER)
@@ -151,7 +160,7 @@ class PubSub(BaseModule):
         return query
 
     @callback
-    def _node_configuration_received(self, stanza):
+    def _node_configuration_received(self, stanza: Stanza) -> Union[PubSubConfigResult, CommonError]:
         if not isResultNode(stanza):
             return raise_error(self._log.warning, stanza)
 
@@ -181,13 +190,13 @@ class PubSub(BaseModule):
                            'No valid form type found')
 
     @callback
-    def _default_response(self, stanza):
+    def _default_response(self, stanza: Stanza) -> Union[CommonResult, CommonError]:
         if not isResultNode(stanza):
             return raise_error(self._log.info, stanza)
         return CommonResult(jid=stanza.getFrom())
 
 
-def get_pubsub_request(jid, node, id_=None, max_items=None):
+def get_pubsub_request(jid: JID, node: str, id_: Optional[str] = None, max_items: Optional[str] = None) -> Iq:
     query = Iq('get', to=jid)
     pubsub = query.addChild('pubsub', namespace=Namespace.PUBSUB)
     items = pubsub.addChild('items', {'node': node})
@@ -198,13 +207,13 @@ def get_pubsub_request(jid, node, id_=None, max_items=None):
     return query
 
 
-def get_pubsub_item(stanza):
+def get_pubsub_item(stanza: Node) -> Optional[Node]:
     pubsub_node = stanza.getTag('pubsub')
     items_node = pubsub_node.getTag('items')
     return items_node.getTag('item')
 
 
-def get_pubsub_items(stanza, node=None):
+def get_pubsub_items(stanza: Node, node: Optional[Node] = None) -> Optional[Node]:
     pubsub_node = stanza.getTag('pubsub')
     items_node = pubsub_node.getTag('items')
     if node is not None and items_node.getAttr('node') != node:
@@ -215,7 +224,7 @@ def get_pubsub_items(stanza, node=None):
     return None
 
 
-def get_publish_options(config):
+def get_publish_options(config: Dict[str, str]) -> Node:
     options = Node(Namespace.DATA + ' x', attrs={'type': 'submit'})
     field = options.addChild('field',
                              attrs={'var': 'FORM_TYPE', 'type': 'hidden'})