stanza: fix a bunch of type errors

6 files changed, 25 insertions, 17 deletions

diff --git a/slixmpp/stanza/error.py b/slixmpp/stanza/error.py
index 83742a5c..76c6b9cc 100644
--- a/slixmpp/stanza/error.py
+++ b/slixmpp/stanza/error.py
@@ -3,7 +3,7 @@
 # This file is part of Slixmpp.
 # See the file LICENSE for copying permission.
 from __future__ import annotations
-from typing import Optional
+from typing import Optional, Dict, Type, ClassVar
 from slixmpp.xmlstream import ElementBase, ET
 
 
@@ -50,10 +50,10 @@ class Error(ElementBase):
     name = 'error'
     plugin_attrib = 'error'
     interfaces = {'code', 'condition', 'text', 'type',
-                      'gone', 'redirect', 'by'}
+                  'gone', 'redirect', 'by'}
     sub_interfaces = {'text'}
-    plugin_attrib_map = {}
-    plugin_tag_map = {}
+    plugin_attrib_map: ClassVar[Dict[str, Type[ElementBase]]] = {}
+    plugin_tag_map: ClassVar[Dict[str, Type[ElementBase]]] = {}
     conditions = {'bad-request', 'conflict', 'feature-not-implemented',
                   'forbidden', 'gone', 'internal-server-error',
                   'item-not-found', 'jid-malformed', 'not-acceptable',
@@ -66,7 +66,7 @@ class Error(ElementBase):
     condition_ns: str = 'urn:ietf:params:xml:ns:xmpp-stanzas'
     types = {'cancel', 'continue', 'modify', 'auth', 'wait'}
 
-    def setup(self, xml: Optional[ET.Element]=None):
+    def setup(self, xml: Optional[ET.Element] = None):
         """
         Populate the stanza object using an optional XML object.
 
@@ -83,7 +83,9 @@ class Error(ElementBase):
             self['type'] = 'cancel'
             self['condition'] = 'feature-not-implemented'
         if self.parent is not None:
-            self.parent()['type'] = 'error'
+            parent = self.parent()
+            if parent:
+                parent['type'] = 'error'
 
     def get_condition(self) -> str:
         """Return the condition element's name."""
@@ -106,7 +108,7 @@ class Error(ElementBase):
             self.xml.append(ET.Element("{%s}%s" % (self.condition_ns, value)))
         return self
 
-    def del_condition(self) -> None:
+    def del_condition(self) -> Error:
         """Remove the condition element."""
         for child in self.xml:
             if "{%s}" % self.condition_ns in child.tag:
diff --git a/slixmpp/stanza/handshake.py b/slixmpp/stanza/handshake.py
index c58f69aa..70f890be 100644
--- a/slixmpp/stanza/handshake.py
+++ b/slixmpp/stanza/handshake.py
@@ -4,6 +4,7 @@
 # See the file LICENSE for copying permission.
 
 from slixmpp.xmlstream import StanzaBase
+from typing import Optional
 
 
 class Handshake(StanzaBase):
@@ -18,7 +19,7 @@ class Handshake(StanzaBase):
     def set_value(self, value: str):
         self.xml.text = value
 
-    def get_value(self) -> str:
+    def get_value(self) -> Optional[str]:
         return self.xml.text
 
     def del_value(self):
diff --git a/slixmpp/stanza/message.py b/slixmpp/stanza/message.py
index debfb380..2a0f20bd 100644
--- a/slixmpp/stanza/message.py
+++ b/slixmpp/stanza/message.py
@@ -5,6 +5,7 @@
 # See the file LICENSE for copying permission.
 from slixmpp.stanza.rootstanza import RootStanza
 from slixmpp.xmlstream import StanzaBase, ET
+from slixmpp.basexmpp import BaseXMPP
 
 
 ORIGIN_NAME = '{urn:xmpp:sid:0}origin-id'
@@ -61,7 +62,7 @@ class Message(RootStanza):
         """
         StanzaBase.__init__(self, *args, **kwargs)
         if not recv and self['id'] == '':
-            if self.stream is not None and self.stream.use_message_ids:
+            if isinstance(self.stream, BaseXMPP) and self.stream.use_message_ids:
                 self['id'] = self.stream.new_id()
             else:
                 del self['origin_id']
@@ -93,7 +94,7 @@ class Message(RootStanza):
 
         self.xml.attrib['id'] = value
 
-        if self.stream and not self.stream.use_origin_id:
+        if isinstance(self.stream, BaseXMPP) and not self.stream.use_origin_id:
             return None
 
         sub = self.xml.find(ORIGIN_NAME)
diff --git a/slixmpp/stanza/presence.py b/slixmpp/stanza/presence.py
index 1a3b7c78..a0742a95 100644
--- a/slixmpp/stanza/presence.py
+++ b/slixmpp/stanza/presence.py
@@ -5,6 +5,7 @@
 # See the file LICENSE for copying permission.
 from slixmpp.stanza.rootstanza import RootStanza
 from slixmpp.xmlstream import StanzaBase
+from slixmpp.basexmpp import BaseXMPP
 
 
 class Presence(RootStanza):
@@ -69,7 +70,7 @@ class Presence(RootStanza):
         """
         StanzaBase.__init__(self, *args, **kwargs)
         if not recv and self['id'] == '':
-            if self.stream is not None and self.stream.use_presence_ids:
+            if isinstance(self.stream, BaseXMPP) and self.stream.use_presence_ids:
                 self['id'] = self.stream.new_id()
 
     def set_show(self, show: str):
diff --git a/slixmpp/stanza/stream_error.py b/slixmpp/stanza/stream_error.py
index d2036a3c..d0eadd5b 100644
--- a/slixmpp/stanza/stream_error.py
+++ b/slixmpp/stanza/stream_error.py
@@ -4,7 +4,8 @@
 # This file is part of Slixmpp.
 # See the file LICENSE for copying permission.
 from slixmpp.stanza.error import Error
-from slixmpp.xmlstream import StanzaBase
+from slixmpp.xmlstream import StanzaBase, ET
+from typing import Optional, Dict, Union
 
 
 class StreamError(Error, StanzaBase):
@@ -64,17 +65,18 @@ class StreamError(Error, StanzaBase):
         'unsupported-version'}
     condition_ns: str = 'urn:ietf:params:xml:ns:xmpp-streams'
 
-    def get_see_other_host(self) -> str:
+    def get_see_other_host(self) -> Union[str, Dict[str, str]]:
         ns = self.condition_ns
         return self._get_sub_text('{%s}see-other-host' % ns, '')
 
-    def set_see_other_host(self, value: str) -> None:
+    def set_see_other_host(self, value: str) -> Optional[ET.Element]:
         if value:
             del self['condition']
             ns = self.condition_ns
             return self._set_sub_text('{%s}see-other-host' % ns, value)
         elif self['condition'] == 'see-other-host':
             del self['condition']
+        return None
 
     def del_see_other_host(self) -> None:
         self._del_sub('{%s}see-other-host' % self.condition_ns)
diff --git a/slixmpp/stanza/stream_features.py b/slixmpp/stanza/stream_features.py
index 7362f17b..a3e05a1a 100644
--- a/slixmpp/stanza/stream_features.py
+++ b/slixmpp/stanza/stream_features.py
@@ -3,7 +3,8 @@
 # Copyright (C) 2010  Nathanael C. Fritz
 # This file is part of Slixmpp.
 # See the file LICENSE for copying permission.
-from slixmpp.xmlstream import StanzaBase
+from slixmpp.xmlstream import StanzaBase, ElementBase
+from typing import ClassVar, Dict, Type
 
 
 class StreamFeatures(StanzaBase):
@@ -15,8 +16,8 @@ class StreamFeatures(StanzaBase):
     namespace = 'http://etherx.jabber.org/streams'
     interfaces = {'features', 'required', 'optional'}
     sub_interfaces = interfaces
-    plugin_tag_map = {}
-    plugin_attrib_map = {}
+    plugin_attrib_map: ClassVar[Dict[str, Type[ElementBase]]] = {}
+    plugin_tag_map: ClassVar[Dict[str, Type[ElementBase]]] = {}
 
     def setup(self, xml):
         StanzaBase.setup(self, xml)