xmlstream: add more types

1 files changed, 302 insertions, 168 deletions

diff --git a/slixmpp/xmlstream/xmlstream.py b/slixmpp/xmlstream/xmlstream.py
index 45d814fd..dbebf2c8 100644
--- a/slixmpp/xmlstream/xmlstream.py
+++ b/slixmpp/xmlstream/xmlstream.py
@@ -9,17 +9,24 @@
 # :license: MIT, see LICENSE for more details
 from typing import (
     Any,
+    Dict,
+    Awaitable,
+    Generator,
     Coroutine,
     Callable,
-    Iterable,
     Iterator,
     List,
     Optional,
     Set,
     Union,
     Tuple,
+    TypeVar,
+    NoReturn,
+    Type,
+    cast,
 )
 
+import asyncio
 import functools
 import logging
 import socket as Socket
@@ -27,30 +34,66 @@ import ssl
 import weakref
 import uuid
 
-import asyncio
-from asyncio import iscoroutinefunction, wait, Future
 from contextlib import contextmanager
 import xml.etree.ElementTree as ET
+from asyncio import (
+    AbstractEventLoop,
+    BaseTransport,
+    Future,
+    Task,
+    TimerHandle,
+    Transport,
+    iscoroutinefunction,
+    wait,
+)
 
-from slixmpp.xmlstream import tostring
+from slixmpp.types import FilterString
+from slixmpp.xmlstream.tostring import tostring
 from slixmpp.xmlstream.stanzabase import StanzaBase, ElementBase
 from slixmpp.xmlstream.resolver import resolve, default_resolver
+from slixmpp.xmlstream.handler.base import BaseHandler
+
+T = TypeVar('T')
 
 #: The time in seconds to wait before timing out waiting for response stanzas.
 RESPONSE_TIMEOUT = 30
 
 log = logging.getLogger(__name__)
+
+
 class ContinueQueue(Exception):
     """
     Exception raised in the send queue to "continue" from within an inner loop
     """
 
+
 class NotConnectedError(Exception):
     """
     Raised when we try to send something over the wire but we are not
     connected.
     """
 
+
+_T = TypeVar('_T', str, ElementBase, StanzaBase)
+
+
+SyncFilter = Callable[[StanzaBase], Optional[StanzaBase]]
+AsyncFilter = Callable[[StanzaBase], Awaitable[Optional[StanzaBase]]]
+
+
+Filter = Union[
+    SyncFilter,
+    AsyncFilter,
+]
+
+_FiltersDict = Dict[str, List[Filter]]
+
+Handler = Callable[[Any], Union[
+    Any,
+    Coroutine[Any, Any, Any]
+]]
+
+
 class XMLStream(asyncio.BaseProtocol):
     """
     An XML stream connection manager and event dispatcher.
@@ -78,16 +121,156 @@ class XMLStream(asyncio.BaseProtocol):
     :param int port: The port to use for the connection. Defaults to 0.
     """
 
-    def __init__(self, host='', port=0):
-        # The asyncio.Transport object provided by the connection_made()
-        # callback when we are connected
-        self.transport = None
+    transport: Optional[Transport]
 
-        # The socket that is used internally by the transport object
-        self.socket = None
+    # The socket that is used internally by the transport object
+    socket: Optional[ssl.SSLSocket]
+
+    # The backoff of the connect routine (increases exponentially
+    # after each failure)
+    _connect_loop_wait: float
+
+    parser: Optional[ET.XMLPullParser]
+    xml_depth: int
+    xml_root: Optional[ET.Element]
+
+    force_starttls: Optional[bool]
+    disable_starttls: Optional[bool]
+
+    waiting_queue: asyncio.Queue[Tuple[Union[StanzaBase, str], bool]]
+
+    # A dict of {name: handle}
+    scheduled_events: Dict[str, TimerHandle]
+
+    ssl_context: ssl.SSLContext
+
+    # The event to trigger when the create_connection() succeeds. It can
+    # be "connected" or "tls_success" depending on the step we are at.
+    event_when_connected: str
+
+    #: The list of accepted ciphers, in OpenSSL Format.
+    #: It might be useful to override it for improved security
+    #: over the python defaults.
+    ciphers: Optional[str]
+
+    #: Path to a file containing certificates for verifying the
+    #: server SSL certificate. A non-``None`` value will trigger
+    #: certificate checking.
+    #:
+    #: .. note::
+    #:
+    #:     On Mac OS X, certificates in the system keyring will
+    #:     be consulted, even if they are not in the provided file.
+    ca_certs: Optional[str]
+
+    #: Path to a file containing a client certificate to use for
+    #: authenticating via SASL EXTERNAL. If set, there must also
+    #: be a corresponding `:attr:keyfile` value.
+    certfile: Optional[str]
+
+    #: Path to a file containing the private key for the selected
+    #: client certificate to use for authenticating via SASL EXTERNAL.
+    keyfile: Optional[str]
+
+    # The asyncio event loop
+    _loop: Optional[AbstractEventLoop]
+
+    #: The default port to return when querying DNS records.
+    default_port: int
+
+    #: The domain to try when querying DNS records.
+    default_domain: str
+
+    #: The expected name of the server, for validation.
+    _expected_server_name: str
+    _service_name: str
+
+    #: The desired, or actual, address of the connected server.
+    address: Tuple[str, int]
+
+    #: Enable connecting to the server directly over SSL, in
+    #: particular when the service provides two ports: one for
+    #: non-SSL traffic and another for SSL traffic.
+    use_ssl: bool
+
+    #: If set to ``True``, attempt to use IPv6.
+    use_ipv6: bool
 
-        # The backoff of the connect routine (increases exponentially
-        # after each failure)
+    #: If set to ``True``, allow using the ``dnspython`` DNS library
+    #: if available. If set to ``False``, the builtin DNS resolver
+    #: will be used, even if ``dnspython`` is installed.
+    use_aiodns: bool
+
+    #: Use CDATA for escaping instead of XML entities. Defaults
+    #: to ``False``.
+    use_cdata: bool
+
+    #: The default namespace of the stream content, not of the
+    #: stream wrapper it
+    default_ns: str
+
+    default_lang: Optional[str]
+    peer_default_lang: Optional[str]
+
+    #: The namespace of the enveloping stream element.
+    stream_ns: str
+
+    #: The default opening tag for the stream element.
+    stream_header: str
+
+    #: The default closing tag for the stream element.
+    stream_footer: str
+
+    #: If ``True``, periodically send a whitespace character over the
+    #: wire to keep the connection alive. Mainly useful for connections
+    #: traversing NAT.
+    whitespace_keepalive: bool
+
+    #: The default interval between keepalive signals when
+    #: :attr:`whitespace_keepalive` is enabled.
+    whitespace_keepalive_interval: int
+
+    #: Flag for controlling if the session can be considered ended
+    #: if the connection is terminated.
+    end_session_on_disconnect: bool
+
+    #: A mapping of XML namespaces to well-known prefixes.
+    namespace_map: dict
+
+    __root_stanza: List[Type[StanzaBase]]
+    __handlers: List[BaseHandler]
+    __event_handlers: Dict[str, List[Tuple[Handler, bool]]]
+    __filters: _FiltersDict
+
+    # Current connection attempt (Future)
+    _current_connection_attempt: Optional[Future[None]]
+
+    #: A list of DNS results that have not yet been tried.
+    _dns_answers: Optional[Iterator[Tuple[str, str, int]]]
+
+    #: The service name to check with DNS SRV records. For
+    #: example, setting this to ``'xmpp-client'`` would query the
+    #: ``_xmpp-client._tcp`` service.
+    dns_service: Optional[str]
+
+    #: The reason why we are disconnecting from the server
+    disconnect_reason: Optional[str]
+
+    #: An asyncio Future being done when the stream is disconnected.
+    disconnected: Future[bool]
+
+    # If the session has been started or not
+    _session_started: bool
+    # If we want to bypass the send() check (e.g. unit tests)
+    _always_send_everything: bool
+
+    _run_out_filters: Optional[Future]
+    __slow_tasks: List[Task]
+    __queued_stanzas: List[Tuple[Union[StanzaBase, str], bool]]
+
+    def __init__(self, host: str = '', port: int = 0):
+        self.transport = None
+        self.socket = None
         self._connect_loop_wait = 0
 
         self.parser = None
@@ -106,126 +289,60 @@ class XMLStream(asyncio.BaseProtocol):
         self.ssl_context.check_hostname = False
         self.ssl_context.verify_mode = ssl.CERT_NONE
 
-        # The event to trigger when the create_connection() succeeds. It can
-        # be "connected" or "tls_success" depending on the step we are at.
         self.event_when_connected = "connected"
 
-        #: The list of accepted ciphers, in OpenSSL Format.
-        #: It might be useful to override it for improved security
-        #: over the python defaults.
         self.ciphers = None
 
-        #: Path to a file containing certificates for verifying the
-        #: server SSL certificate. A non-``None`` value will trigger
-        #: certificate checking.
-        #:
-        #: .. note::
-        #:
-        #:     On Mac OS X, certificates in the system keyring will
-        #:     be consulted, even if they are not in the provided file.
         self.ca_certs = None
 
-        #: Path to a file containing a client certificate to use for
-        #: authenticating via SASL EXTERNAL. If set, there must also
-        #: be a corresponding `:attr:keyfile` value.
-        self.certfile = None
-
-        #: Path to a file containing the private key for the selected
-        #: client certificate to use for authenticating via SASL EXTERNAL.
         self.keyfile = None
 
-        self._der_cert = None
-
-        # The asyncio event loop
         self._loop = None
 
-        #: The default port to return when querying DNS records.
         self.default_port = int(port)
-
-        #: The domain to try when querying DNS records.
         self.default_domain = ''
 
-        #: The expected name of the server, for validation.
         self._expected_server_name = ''
         self._service_name = ''
 
-        #: The desired, or actual, address of the connected server.
         self.address = (host, int(port))
 
-        #: Enable connecting to the server directly over SSL, in
-        #: particular when the service provides two ports: one for
-        #: non-SSL traffic and another for SSL traffic.
         self.use_ssl = False
-
-        #: If set to ``True``, attempt to use IPv6.
         self.use_ipv6 = True
 
-        #: If set to ``True``, allow using the ``dnspython`` DNS library
-        #: if available. If set to ``False``, the builtin DNS resolver
-        #: will be used, even if ``dnspython`` is installed.
         self.use_aiodns = True
-
-        #: Use CDATA for escaping instead of XML entities. Defaults
-        #: to ``False``.
         self.use_cdata = False
 
-        #: The default namespace of the stream content, not of the
-        #: stream wrapper itself.
         self.default_ns = ''
 
         self.default_lang = None
         self.peer_default_lang = None
 
-        #: The namespace of the enveloping stream element.
         self.stream_ns = ''
-
-        #: The default opening tag for the stream element.
         self.stream_header = "<stream>"
-
-        #: The default closing tag for the stream element.
         self.stream_footer = "</stream>"
 
-        #: If ``True``, periodically send a whitespace character over the
-        #: wire to keep the connection alive. Mainly useful for connections
-        #: traversing NAT.
         self.whitespace_keepalive = True
-
-        #: The default interval between keepalive signals when
-        #: :attr:`whitespace_keepalive` is enabled.
         self.whitespace_keepalive_interval = 300
 
-        #: Flag for controlling if the session can be considered ended
-        #: if the connection is terminated.
         self.end_session_on_disconnect = True
-
-        #: A mapping of XML namespaces to well-known prefixes.
         self.namespace_map = {StanzaBase.xml_ns: 'xml'}
 
         self.__root_stanza = []
         self.__handlers = []
         self.__event_handlers = {}
-        self.__filters = {'in': [], 'out': [], 'out_sync': []}
+        self.__filters = {
+            'in': [], 'out': [], 'out_sync': []
+        }
 
-        # Current connection attempt (Future)
         self._current_connection_attempt = None
 
-        #: A list of DNS results that have not yet been tried.
-        self._dns_answers: Optional[Iterator[Tuple[str, str, int]]] = None
-
-        #: The service name to check with DNS SRV records. For
-        #: example, setting this to ``'xmpp-client'`` would query the
-        #: ``_xmpp-client._tcp`` service.
+        self._dns_answers = None
         self.dns_service = None
 
-        #: The reason why we are disconnecting from the server
         self.disconnect_reason = None
-
-        #: An asyncio Future being done when the stream is disconnected.
-        self.disconnected: Future = Future()
-
-        # If the session has been started or not
+        self.disconnected = Future()
         self._session_started = False
-        # If we want to bypass the send() check (e.g. unit tests)
         self._always_send_everything = False
 
         self.add_event_handler('disconnected', self._remove_schedules)
@@ -234,21 +351,21 @@ class XMLStream(asyncio.BaseProtocol):
         self.add_event_handler('session_start', self._set_session_start)
         self.add_event_handler('session_resumed', self._set_session_start)
 
-        self._run_out_filters: Optional[Future] = None
-        self.__slow_tasks: List[Future] = []
-        self.__queued_stanzas: List[Tuple[StanzaBase, bool]] = []
+        self._run_out_filters = None
+        self.__slow_tasks = []
+        self.__queued_stanzas = []
 
     @property
-    def loop(self):
+    def loop(self) -> AbstractEventLoop:
         if self._loop is None:
             self._loop = asyncio.get_event_loop()
         return self._loop
 
     @loop.setter
-    def loop(self, value):
+    def loop(self, value: AbstractEventLoop) -> None:
         self._loop = value
 
-    def new_id(self):
+    def new_id(self) -> str:
         """Generate and return a new stream ID in hexadecimal form.
 
         Many stanzas, handlers, or matchers may require unique
@@ -257,7 +374,7 @@ class XMLStream(asyncio.BaseProtocol):
         """
         return uuid.uuid4().hex
 
-    def _set_session_start(self, event):
+    def _set_session_start(self, event: Any) -> None:
         """
         On session start, queue all pending stanzas to be sent.
         """
@@ -266,17 +383,17 @@ class XMLStream(asyncio.BaseProtocol):
             self.waiting_queue.put_nowait(stanza)
         self.__queued_stanzas = []
 
-    def _set_disconnected(self, event):
+    def _set_disconnected(self, event: Any) -> None:
         self._session_started = False
 
-    def _set_disconnected_future(self):
+    def _set_disconnected_future(self) -> None:
         """Set the self.disconnected future on disconnect"""
         if not self.disconnected.done():
             self.disconnected.set_result(True)
         self.disconnected = asyncio.Future()
 
-    def connect(self, host='', port=0, use_ssl=False,
-                force_starttls=True, disable_starttls=False):
+    def connect(self, host: str = '', port: int = 0, use_ssl: Optional[bool] = False,
+                force_starttls: Optional[bool] = True, disable_starttls: Optional[bool] = False) -> None:
         """Create a new socket and connect to the server.
 
         :param host: The name of the desired server for the connection.
@@ -327,7 +444,7 @@ class XMLStream(asyncio.BaseProtocol):
             loop=self.loop,
         )
 
-    async def _connect_routine(self):
+    async def _connect_routine(self) -> None:
         self.event_when_connected = "connected"
 
         if self._connect_loop_wait > 0:
@@ -345,6 +462,7 @@ class XMLStream(asyncio.BaseProtocol):
             # and try (host, port) as a last resort
             self._dns_answers = None
 
+        ssl_context: Optional[ssl.SSLContext]
         if self.use_ssl:
             ssl_context = self.get_ssl_context()
         else:
@@ -373,7 +491,7 @@ class XMLStream(asyncio.BaseProtocol):
                 loop=self.loop,
             )
 
-    def process(self, *, forever=True, timeout=None):
+    def process(self, *, forever: bool = True, timeout: Optional[int] = None) -> None:
         """Process all the available XMPP events (receiving or sending data on the
         socket(s), calling various registered callbacks, calling expired
         timers, handling signal events, etc).  If timeout is None, this
@@ -386,12 +504,12 @@ class XMLStream(asyncio.BaseProtocol):
             else:
                 self.loop.run_until_complete(self.disconnected)
         else:
-            tasks = [asyncio.sleep(timeout, loop=self.loop)]
+            tasks: List[Future[bool]] = [asyncio.sleep(timeout, loop=self.loop)]
             if not forever:
                 tasks.append(self.disconnected)
             self.loop.run_until_complete(asyncio.wait(tasks, loop=self.loop))
 
-    def init_parser(self):
+    def init_parser(self) -> None:
         """init the XML parser. The parser must always be reset for each new
         connexion
         """
@@ -399,11 +517,13 @@ class XMLStream(asyncio.BaseProtocol):
         self.xml_root = None
         self.parser = ET.XMLPullParser(("start", "end"))
 
-    def connection_made(self, transport):
+    def connection_made(self, transport: BaseTransport) -> None:
         """Called when the TCP connection has been established with the server
         """
         self.event(self.event_when_connected)
-        self.transport = transport
+        self.transport = cast(Transport, transport)
+        if self.transport is None:
+            raise ValueError("Transport cannot be none")
         self.socket = self.transport.get_extra_info(
             "ssl_object",
             default=self.transport.get_extra_info("socket")
@@ -413,7 +533,7 @@ class XMLStream(asyncio.BaseProtocol):
         self.send_raw(self.stream_header)
         self._dns_answers = None
 
-    def data_received(self, data):
+    def data_received(self, data: bytes) -> None:
         """Called when incoming data is received on the socket.
 
         We feed that data to the parser and the see if this produced any XML
@@ -467,18 +587,18 @@ class XMLStream(asyncio.BaseProtocol):
             self.send(error)
             self.disconnect()
 
-    def is_connecting(self):
+    def is_connecting(self) -> bool:
         return self._current_connection_attempt is not None
 
-    def is_connected(self):
+    def is_connected(self) -> bool:
         return self.transport is not None
 
-    def eof_received(self):
+    def eof_received(self) -> None:
         """When the TCP connection is properly closed by the remote end
         """
         self.event("eof_received")
 
-    def connection_lost(self, exception):
+    def connection_lost(self, exception: Optional[BaseException]) -> None:
         """On any kind of disconnection, initiated by us or not.  This signals the
         closure of the TCP connection
         """
@@ -493,9 +613,9 @@ class XMLStream(asyncio.BaseProtocol):
             self._reset_sendq()
             self.event('session_end')
         self._set_disconnected_future()
-        self.event("disconnected", self.disconnect_reason or exception and exception.strerror)
+        self.event("disconnected", self.disconnect_reason or exception)
 
-    def cancel_connection_attempt(self):
+    def cancel_connection_attempt(self) -> None:
         """
         Immediately cancel the current create_connection() Future.
         This is useful when a client using slixmpp tries to connect
@@ -506,7 +626,7 @@ class XMLStream(asyncio.BaseProtocol):
             self._current_connection_attempt.cancel()
             self._current_connection_attempt = None
 
-    def disconnect(self, wait: Union[float, int] = 2.0, reason: Optional[str] = None, ignore_send_queue: bool = False) -> Future:
+    def disconnect(self, wait: Union[float, int] = 2.0, reason: Optional[str] = None, ignore_send_queue: bool = False) -> Future[None]:
         """Close the XML stream and wait for an acknowldgement from the server for
         at most `wait` seconds.  After the given number of seconds has
         passed without a response from the server, or when the server
@@ -526,7 +646,7 @@ class XMLStream(asyncio.BaseProtocol):
         # `disconnect(wait=True)` for ages. This doesn't mean anything to the
         # schedule call below. It would fortunately be converted to `1` later
         # down the call chain. Praise the implicit casts lord.
-        if wait == True:
+        if wait is True:
             wait = 2.0
 
         if self.transport:
@@ -545,11 +665,11 @@ class XMLStream(asyncio.BaseProtocol):
         else:
             self._set_disconnected_future()
             self.event("disconnected", reason)
-            future = Future()
+            future: Future[None] = Future()
             future.set_result(None)
             return future
 
-    async def _consume_send_queue_before_disconnecting(self, reason: Optional[str], wait: float):
+    async def _consume_send_queue_before_disconnecting(self, reason: Optional[str], wait: float) -> None:
         """Wait until the send queue is empty before disconnecting"""
         try:
             await asyncio.wait_for(
@@ -561,7 +681,7 @@ class XMLStream(asyncio.BaseProtocol):
         self.disconnect_reason = reason
         await self._end_stream_wait(wait)
 
-    async def _end_stream_wait(self, wait: Union[int, float] = 2, reason: Optional[str] = None):
+    async def _end_stream_wait(self, wait: Union[int, float] = 2, reason: Optional[str] = None) -> None:
         """
         Run abort() if we do not received the disconnected event
         after a waiting time.
@@ -578,7 +698,7 @@ class XMLStream(asyncio.BaseProtocol):
             # that means the disconnect has already been handled
             pass
 
-    def abort(self):
+    def abort(self) -> None:
         """
         Forcibly close the connection
         """
@@ -588,26 +708,26 @@ class XMLStream(asyncio.BaseProtocol):
             self.transport.abort()
             self.event("killed")
 
-    def reconnect(self, wait=2.0, reason="Reconnecting"):
+    def reconnect(self, wait: Union[int, float] = 2.0, reason: str = "Reconnecting") -> None:
         """Calls disconnect(), and once we are disconnected (after the timeout, or
         when the server acknowledgement is received), call connect()
         """
         log.debug("reconnecting...")
-        async def handler(event):
+        async def handler(event: Any) -> None:
             # We yield here to allow synchronous handlers to work first
             await asyncio.sleep(0, loop=self.loop)
             self.connect()
         self.add_event_handler('disconnected', handler, disposable=True)
         self.disconnect(wait, reason)
 
-    def configure_socket(self):
+    def configure_socket(self) -> None:
         """Set timeout and other options for self.socket.
 
         Meant to be overridden.
         """
         pass
 
-    def configure_dns(self, resolver, domain=None, port=None):
+    def configure_dns(self, resolver: Any, domain: Optional[str] = None, port: Optional[int] = None) -> None:
         """
         Configure and set options for a :class:`~dns.resolver.Resolver`
         instance, and other DNS related tasks. For example, you
@@ -624,7 +744,7 @@ class XMLStream(asyncio.BaseProtocol):
         """
         pass
 
-    def get_ssl_context(self):
+    def get_ssl_context(self) -> ssl.SSLContext:
         """
         Get SSL context.
         """
@@ -644,12 +764,14 @@ class XMLStream(asyncio.BaseProtocol):
 
         return self.ssl_context
 
-    async def start_tls(self):
+    async def start_tls(self) -> bool:
         """Perform handshakes for TLS.
 
         If the handshake is successful, the XML stream will need
         to be restarted.
         """
+        if self.transport is None:
+            raise ValueError("Transport should not be None")
         self.event_when_connected = "tls_success"
         ssl_context = self.get_ssl_context()
         try:
@@ -685,7 +807,7 @@ class XMLStream(asyncio.BaseProtocol):
             self.connection_made(transp)
         return True
 
-    def _start_keepalive(self, event):
+    def _start_keepalive(self, event: Any) -> None:
         """Begin sending whitespace periodically to keep the connection alive.
 
         May be disabled by setting::
@@ -702,11 +824,11 @@ class XMLStream(asyncio.BaseProtocol):
                       args=(' ',),
                       repeat=True)
 
-    def _remove_schedules(self, event):
+    def _remove_schedules(self, event: Any) -> None:
         """Remove some schedules that become pointless when disconnected"""
         self.cancel_schedule('Whitespace Keepalive')
 
-    def start_stream_handler(self, xml):
+    def start_stream_handler(self, xml: ET.Element) -> None:
         """Perform any initialization actions, such as handshakes,
         once the stream header has been sent.
 
@@ -714,7 +836,7 @@ class XMLStream(asyncio.BaseProtocol):
         """
         pass
 
-    def register_stanza(self, stanza_class):
+    def register_stanza(self, stanza_class: Type[StanzaBase]) -> None:
         """Add a stanza object class as a known root stanza.
 
         A root stanza is one that appears as a direct child of the stream's
@@ -732,7 +854,7 @@ class XMLStream(asyncio.BaseProtocol):
         """
         self.__root_stanza.append(stanza_class)
 
-    def remove_stanza(self, stanza_class):
+    def remove_stanza(self, stanza_class: Type[StanzaBase]) -> None:
         """Remove a stanza from being a known root stanza.
 
         A root stanza is one that appears as a direct child of the stream's
@@ -744,7 +866,7 @@ class XMLStream(asyncio.BaseProtocol):
         """
         self.__root_stanza.remove(stanza_class)
 
-    def add_filter(self, mode, handler, order=None):
+    def add_filter(self, mode: FilterString, handler: Callable[[StanzaBase], Optional[StanzaBase]], order: Optional[int] = None) -> None:
         """Add a filter for incoming or outgoing stanzas.
 
         These filters are applied before incoming stanzas are
@@ -766,11 +888,11 @@ class XMLStream(asyncio.BaseProtocol):
         else:
             self.__filters[mode].append(handler)
 
-    def del_filter(self, mode, handler):
+    def del_filter(self, mode: str, handler: Callable[[StanzaBase], Optional[StanzaBase]]) -> None:
         """Remove an incoming or outgoing filter."""
         self.__filters[mode].remove(handler)
 
-    def register_handler(self, handler, before=None, after=None):
+    def register_handler(self, handler: BaseHandler, before: Optional[BaseHandler] = None, after: Optional[BaseHandler] = None) -> None:
         """Add a stream event handler that will be executed when a matching
         stanza is received.
 
@@ -782,7 +904,7 @@ class XMLStream(asyncio.BaseProtocol):
             self.__handlers.append(handler)
             handler.stream = weakref.ref(self)
 
-    def remove_handler(self, name):
+    def remove_handler(self, name: str) -> bool:
         """Remove any stream event handlers with the given name.
 
         :param name: The name of the handler.
@@ -831,9 +953,9 @@ class XMLStream(asyncio.BaseProtocol):
         try:
             return next(self._dns_answers)
         except StopIteration:
-            return
+            return None
 
-    def add_event_handler(self, name, pointer, disposable=False):
+    def add_event_handler(self, name: str, pointer: Callable[..., Any], disposable: bool = False) -> None:
         """Add a custom event handler that will be executed whenever
         its event is manually triggered.
 
@@ -847,7 +969,7 @@ class XMLStream(asyncio.BaseProtocol):
             self.__event_handlers[name] = []
         self.__event_handlers[name].append((pointer, disposable))
 
-    def del_event_handler(self, name, pointer):
+    def del_event_handler(self, name: str, pointer: Callable[..., Any]) -> None:
         """Remove a function as a handler for an event.
 
         :param name: The name of the event.
@@ -858,21 +980,21 @@ class XMLStream(asyncio.BaseProtocol):
 
         # Need to keep handlers that do not use
         # the given function pointer
-        def filter_pointers(handler):
+        def filter_pointers(handler: Tuple[Callable[..., Any], bool]) -> bool:
             return handler[0] != pointer
 
         self.__event_handlers[name] = list(filter(
             filter_pointers,
             self.__event_handlers[name]))
 
-    def event_handled(self, name):
+    def event_handled(self, name: str) -> int:
         """Returns the number of registered handlers for an event.
 
         :param name: The name of the event to check.
         """
         return len(self.__event_handlers.get(name, []))
 
-    async def event_async(self, name: str, data: Any = {}):
+    async def event_async(self, name: str, data: Any = {}) -> None:
         """Manually trigger a custom event, but await coroutines immediately.
 
         This event generator should only be called in situations when
@@ -908,7 +1030,7 @@ class XMLStream(asyncio.BaseProtocol):
                 except Exception as e:
                     self.exception(e)
 
-    def event(self, name: str, data: Any = {}):
+    def event(self, name: str, data: Any = {}) -> None:
         """Manually trigger a custom event.
         Coroutine handlers are wrapped into a future and sent into the
         event loop for their execution, and not awaited.
@@ -928,7 +1050,7 @@ class XMLStream(asyncio.BaseProtocol):
             # If the callback is a coroutine, schedule it instead of
             # running it directly
             if iscoroutinefunction(handler_callback):
-                async def handler_callback_routine(cb):
+                async def handler_callback_routine(cb: Callable[[ElementBase], Any]) -> None:
                     try:
                         await cb(data)
                     except Exception as e:
@@ -957,8 +1079,9 @@ class XMLStream(asyncio.BaseProtocol):
                 except ValueError:
                     pass
 
-    def schedule(self, name, seconds, callback, args=tuple(),
-                 kwargs={}, repeat=False):
+    def schedule(self, name: str, seconds: int, callback: Callable[..., None],
+            args: Tuple[Any, ...] = tuple(),
+            kwargs: Dict[Any, Any] = {}, repeat: bool = False) -> None:
         """Schedule a callback function to execute after a given delay.
 
         :param name: A unique name for the scheduled callback.
@@ -986,21 +1109,21 @@ class XMLStream(asyncio.BaseProtocol):
         # canceling scheduled_events[name]
         self.scheduled_events[name] = handle
 
-    def cancel_schedule(self, name):
+    def cancel_schedule(self, name: str) -> None:
         try:
             handle = self.scheduled_events.pop(name)
             handle.cancel()
         except KeyError:
             log.debug("Tried to cancel unscheduled event: %s" % (name,))
 
-    def _safe_cb_run(self, name, cb):
+    def _safe_cb_run(self, name: str, cb: Callable[[], None]) -> None:
         log.debug('Scheduled event: %s', name)
         try:
             cb()
         except Exception as e:
             self.exception(e)
 
-    def _execute_and_reschedule(self, name, cb, seconds):
+    def _execute_and_reschedule(self, name: str, cb: Callable[[], None], seconds: int) -> None:
         """Simple method that calls the given callback, and then schedule itself to
         be called after the given number of seconds.
         """
@@ -1009,7 +1132,7 @@ class XMLStream(asyncio.BaseProtocol):
                                       name, cb, seconds)
         self.scheduled_events[name] = handle
 
-    def _execute_and_unschedule(self, name, cb):
+    def _execute_and_unschedule(self, name: str, cb: Callable[[], None]) -> None:
         """
         Execute the callback and remove the handler for it.
         """
@@ -1018,7 +1141,7 @@ class XMLStream(asyncio.BaseProtocol):
         if name in self.scheduled_events:
             del self.scheduled_events[name]
 
-    def incoming_filter(self, xml):
+    def incoming_filter(self, xml: ET.Element) -> ET.Element:
         """Filter incoming XML objects before they are processed.
 
         Possible uses include remapping namespaces, or correcting elements
@@ -1028,7 +1151,7 @@ class XMLStream(asyncio.BaseProtocol):
         """
         return xml
 
-    def _reset_sendq(self):
+    def _reset_sendq(self) -> None:
         """Clear sending tasks on session end"""
         # Cancel all pending slow send tasks
         log.debug('Cancelling %d slow send tasks', len(self.__slow_tasks))
@@ -1042,8 +1165,8 @@ class XMLStream(asyncio.BaseProtocol):
 
     async def _continue_slow_send(
             self,
-            task: asyncio.Task,
-            already_used: Set[Callable[[ElementBase], Optional[StanzaBase]]]
+            task: asyncio.Task[Optional[StanzaBase]],
+            already_used: Set[Filter]
     ) -> None:
         """
         Used when an item in the send queue has taken too long to process.
@@ -1062,12 +1185,14 @@ class XMLStream(asyncio.BaseProtocol):
             if iscoroutinefunction(filter):
                 data = await filter(data)
             else:
+                filter = cast(SyncFilter, filter)
                 data = filter(data)
             if data is None:
                 return
 
-        if isinstance(data, ElementBase):
+        if isinstance(data, StanzaBase):
             for filter in self.__filters['out_sync']:
+                filter = cast(SyncFilter, filter)
                 data = filter(data)
                 if data is None:
                     return
@@ -1077,19 +1202,21 @@ class XMLStream(asyncio.BaseProtocol):
         else:
             self.send_raw(data)
 
-    async def run_filters(self):
+    async def run_filters(self) -> NoReturn:
         """
         Background loop that processes stanzas to send.
         """
         while True:
+            data: Optional[Union[StanzaBase, str]]
             (data, use_filters) = await self.waiting_queue.get()
             try:
-                if isinstance(data, ElementBase):
+                if isinstance(data, StanzaBase):
                     if use_filters:
                         already_run_filters = set()
                         for filter in self.__filters['out']:
                             already_run_filters.add(filter)
                             if iscoroutinefunction(filter):
+                                filter = cast(AsyncFilter, filter)
                                 task = asyncio.create_task(filter(data))
                                 completed, pending = await wait(
                                     {task},
@@ -1108,19 +1235,24 @@ class XMLStream(asyncio.BaseProtocol):
                                         "Slow coroutine, rescheduling filters"
                                     )
                                 data = task.result()
-                            else:
+                            elif isinstance(data, StanzaBase):
+                                filter = cast(SyncFilter, filter)
                                 data = filter(data)
                             if data is None:
                                 raise ContinueQueue('Empty stanza')
 
-                if isinstance(data, ElementBase):
+                if isinstance(data, StanzaBase):
                     if use_filters:
                         for filter in self.__filters['out_sync']:
+                            filter = cast(SyncFilter, filter)
                             data = filter(data)
                             if data is None:
                                 raise ContinueQueue('Empty stanza')
-                    str_data = tostring(data.xml, xmlns=self.default_ns,
-                                        stream=self, top_level=True)
+                    if isinstance(data, StanzaBase):
+                        str_data = tostring(data.xml, xmlns=self.default_ns,
+                                            stream=self, top_level=True)
+                    else:
+                        str_data = data
                     self.send_raw(str_data)
                 else:
                     self.send_raw(data)
@@ -1130,10 +1262,10 @@ class XMLStream(asyncio.BaseProtocol):
                 log.error('Exception raised in send queue:', exc_info=True)
             self.waiting_queue.task_done()
 
-    def send(self, data, use_filters=True):
+    def send(self, data: Union[StanzaBase, str], use_filters: bool = True) -> None:
         """A wrapper for :meth:`send_raw()` for sending stanza objects.
 
-        :param data: The :class:`~slixmpp.xmlstream.stanzabase.ElementBase`
+        :param data: The :class:`~slixmpp.xmlstream.stanzabase.StanzaBase`
                      stanza to send on the stream.
         :param bool use_filters: Indicates if outgoing filters should be
                                  applied to the given stanza data. Disabling
@@ -1156,15 +1288,15 @@ class XMLStream(asyncio.BaseProtocol):
                 return
         self.waiting_queue.put_nowait((data, use_filters))
 
-    def send_xml(self, data):
+    def send_xml(self, data: ET.Element) -> None:
         """Send an XML object on the stream
 
         :param data: The :class:`~xml.etree.ElementTree.Element` XML object
                      to send on the stream.
         """
-        return self.send(tostring(data))
+        self.send(tostring(data))
 
-    def send_raw(self, data):
+    def send_raw(self, data: Union[str, bytes]) -> None:
         """Send raw data across the stream.
 
         :param string data: Any bytes or utf-8 string value.
@@ -1176,7 +1308,8 @@ class XMLStream(asyncio.BaseProtocol):
             data = data.encode('utf-8')
         self.transport.write(data)
 
-    def _build_stanza(self, xml, default_ns=None):
+    def _build_stanza(self, xml: ET.Element,
+                      default_ns: Optional[str] = None) -> StanzaBase:
         """Create a stanza object from a given XML object.
 
         If a specialized stanza type is not found for the XML, then
@@ -1201,7 +1334,7 @@ class XMLStream(asyncio.BaseProtocol):
             stanza['lang'] = self.peer_default_lang
         return stanza
 
-    def _spawn_event(self, xml):
+    def _spawn_event(self, xml: ET.Element) -> None:
         """
         Analyze incoming XML stanzas and convert them into stanza
         objects if applicable and queue stream events to be processed
@@ -1215,9 +1348,10 @@ class XMLStream(asyncio.BaseProtocol):
 
         # Convert the raw XML object into a stanza object. If no registered
         # stanza type applies, a generic StanzaBase stanza will be used.
-        stanza = self._build_stanza(xml)
+        stanza: Optional[StanzaBase] = self._build_stanza(xml)
         for filter in self.__filters['in']:
             if stanza is not None:
+                filter = cast(SyncFilter, filter)
                 stanza = filter(stanza)
         if stanza is None:
             return
@@ -1244,7 +1378,7 @@ class XMLStream(asyncio.BaseProtocol):
         if not handled:
             stanza.unhandled()
 
-    def exception(self, exception):
+    def exception(self, exception: Exception) -> None:
         """Process an unknown exception.
 
         Meant to be overridden.
@@ -1253,7 +1387,7 @@ class XMLStream(asyncio.BaseProtocol):
         """
         pass
 
-    async def wait_until(self, event: str, timeout=30) -> Any:
+    async def wait_until(self, event: str, timeout: Union[int, float] = 30) -> Any:
         """Utility method to wake on the next firing of an event.
         (Registers a disposable handler on it)
 
@@ -1261,9 +1395,9 @@ class XMLStream(asyncio.BaseProtocol):
         :param int timeout: Timeout
         :raises: :class:`asyncio.TimeoutError` when the timeout is reached
         """
-        fut = asyncio.Future()
+        fut: Future[Any] = asyncio.Future()
 
-        def result_handler(event_data):
+        def result_handler(event_data: Any) -> None:
             if not fut.done():
                 fut.set_result(event_data)
             else:
@@ -1280,19 +1414,19 @@ class XMLStream(asyncio.BaseProtocol):
         return await asyncio.wait_for(fut, timeout)
 
     @contextmanager
-    def event_handler(self, event: str, handler: Callable):
+    def event_handler(self, event: str, handler: Callable[..., Any]) -> Generator[None, None, None]:
         """
         Context manager that adds then removes an event handler.
         """
         self.add_event_handler(event, handler)
         try:
             yield
-        except Exception as exc:
+        except Exception:
             raise
         finally:
             self.del_event_handler(event, handler)
 
-    def wrap(self, coroutine: Coroutine[Any, Any, Any]) -> Future:
+    def wrap(self, coroutine: Coroutine[None, None, T]) -> Future[T]:
         """Make a Future out of a coroutine with the current loop.
 
         :param coroutine: The coroutine to wrap.