chore: Update `__exit__` signatures

Co-authored-by: Marcelo Trylesinski <marcelotryle@gmail.com>

9 files changed, 106 insertions, 13 deletions

diff --git a/kombu/asynchronous/http/base.py b/kombu/asynchronous/http/base.py
index 2bd28076..5feb8f3a 100644
--- a/kombu/asynchronous/http/base.py
+++ b/kombu/asynchronous/http/base.py
@@ -2,6 +2,7 @@
 
 import sys
 from http.client import responses
+from typing import TYPE_CHECKING, Optional, Type
 
 from vine import Thenable, maybe_promise, promise
 
@@ -10,6 +11,9 @@ from kombu.utils.compat import coro
 from kombu.utils.encoding import bytes_to_str
 from kombu.utils.functional import maybe_list, memoize
 
+if TYPE_CHECKING:
+    from types import TracebackType
+
 __all__ = ('Headers', 'Response', 'Request')
 
 PYPY = hasattr(sys, 'pypy_version_info')
@@ -253,5 +257,10 @@ class BaseClient:
     def __enter__(self):
         return self
 
-    def __exit__(self, *exc_info):
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional['TracebackType']
+    ) -> None:
         self.close()
diff --git a/kombu/asynchronous/timer.py b/kombu/asynchronous/timer.py
index 21ad37c1..2ee26c61 100644
--- a/kombu/asynchronous/timer.py
+++ b/kombu/asynchronous/timer.py
@@ -7,6 +7,7 @@ from datetime import datetime
 from functools import total_ordering
 from time import monotonic
 from time import time as _time
+from typing import TYPE_CHECKING, Optional, Type
 from weakref import proxy as weakrefproxy
 
 from vine.utils import wraps
@@ -18,6 +19,9 @@ try:
 except ImportError:  # pragma: no cover
     utc = None
 
+if TYPE_CHECKING:
+    from types import TracebackType
+
 __all__ = ('Entry', 'Timer', 'to_timestamp')
 
 logger = get_logger(__name__)
@@ -101,7 +105,12 @@ class Timer:
     def __enter__(self):
         return self
 
-    def __exit__(self, *exc_info):
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional['TracebackType']
+    ) -> None:
         self.stop()
 
     def call_at(self, eta, fun, args=(), kwargs=None, priority=0):
diff --git a/kombu/compat.py b/kombu/compat.py
index 1fa3f631..df37c3b1 100644
--- a/kombu/compat.py
+++ b/kombu/compat.py
@@ -4,10 +4,14 @@ See https://pypi.org/project/carrot/ for documentation.
 """
 
 from itertools import count
+from typing import TYPE_CHECKING, Optional, Type
 
 from . import messaging
 from .entity import Exchange, Queue
 
+if TYPE_CHECKING:
+    from types import TracebackType
+
 __all__ = ('Publisher', 'Consumer')
 
 # XXX compat attribute
@@ -65,7 +69,12 @@ class Publisher(messaging.Producer):
     def __enter__(self):
         return self
 
-    def __exit__(self, *exc_info):
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional['TracebackType']
+    ) -> None:
         self.close()
 
     @property
@@ -127,7 +136,12 @@ class Consumer(messaging.Consumer):
     def __enter__(self):
         return self
 
-    def __exit__(self, *exc_info):
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional['TracebackType']
+    ) -> None:
         self.close()
 
     def __iter__(self):
diff --git a/kombu/connection.py b/kombu/connection.py
index e2dd69b6..7cdef505 100644
--- a/kombu/connection.py
+++ b/kombu/connection.py
@@ -5,6 +5,7 @@ import socket
 from contextlib import contextmanager
 from itertools import count, cycle
 from operator import itemgetter
+from typing import TYPE_CHECKING, Optional, Type
 
 try:
     from ssl import CERT_NONE
@@ -13,6 +14,7 @@ except ImportError:  # pragma: no cover
     CERT_NONE = None
     ssl_available = False
 
+
 # jython breaks on relative import for .exceptions for some reason
 # (Issue #112)
 from kombu import exceptions
@@ -25,6 +27,9 @@ from .utils.functional import dictfilter, lazy, retry_over_time, shufflecycle
 from .utils.objects import cached_property
 from .utils.url import as_url, maybe_sanitize_url, parse_url, quote, urlparse
 
+if TYPE_CHECKING:
+    from types import TracebackType
+
 __all__ = ('Connection', 'ConnectionPool', 'ChannelPool')
 
 logger = get_logger(__name__)
@@ -828,7 +833,12 @@ class Connection:
     def __enter__(self):
         return self
 
-    def __exit__(self, *args):
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional['TracebackType']
+    ) -> None:
         self.release()
 
     @property
diff --git a/kombu/messaging.py b/kombu/messaging.py
index 0bed52c5..ebba43ee 100644
--- a/kombu/messaging.py
+++ b/kombu/messaging.py
@@ -1,6 +1,7 @@
 """Sending and receiving messages."""
 
 from itertools import count
+from typing import TYPE_CHECKING, Optional, Type
 
 from .common import maybe_declare
 from .compression import compress
@@ -10,6 +11,9 @@ from .exceptions import ContentDisallowed
 from .serialization import dumps, prepare_accept_content
 from .utils.functional import ChannelPromise, maybe_list
 
+if TYPE_CHECKING:
+    from types import TracebackType
+
 __all__ = ('Exchange', 'Queue', 'Producer', 'Consumer')
 
 
@@ -236,7 +240,12 @@ class Producer:
     def __enter__(self):
         return self
 
-    def __exit__(self, *exc_info):
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional['TracebackType']
+    ) -> None:
         self.release()
 
     def release(self):
@@ -435,7 +444,12 @@ class Consumer:
         self.consume()
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional['TracebackType']
+    ) -> None:
         if self.channel and self.channel.connection:
             conn_errors = self.channel.connection.client.connection_errors
             if not isinstance(exc_val, conn_errors):
diff --git a/kombu/resource.py b/kombu/resource.py
index e3617dc4..004eb7a9 100644
--- a/kombu/resource.py
+++ b/kombu/resource.py
@@ -4,11 +4,15 @@ import os
 from collections import deque
 from queue import Empty
 from queue import LifoQueue as _LifoQueue
+from typing import TYPE_CHECKING
 
 from . import exceptions
 from .utils.compat import register_after_fork
 from .utils.functional import lazy
 
+if TYPE_CHECKING:
+    from types import TracebackType
+
 
 def _after_fork_cleanup_resource(resource):
     try:
@@ -191,7 +195,12 @@ class Resource:
             def __enter__(self):
                 pass
 
-            def __exit__(self, type, value, traceback):
+            def __exit__(
+                self,
+                exc_type: type,
+                exc_val: Exception,
+                exc_tb: 'TracebackType'
+            ) -> None:
                 pass
 
         resource = self._resource
diff --git a/kombu/simple.py b/kombu/simple.py
index eee037be..5925ceec 100644
--- a/kombu/simple.py
+++ b/kombu/simple.py
@@ -4,10 +4,14 @@ import socket
 from collections import deque
 from queue import Empty
 from time import monotonic
+from typing import TYPE_CHECKING, Optional, Type
 
 from . import entity, messaging
 from .connection import maybe_channel
 
+if TYPE_CHECKING:
+    from types import TracebackType
+
 __all__ = ('SimpleQueue', 'SimpleBuffer')
 
 
@@ -18,7 +22,12 @@ class SimpleBase:
     def __enter__(self):
         return self
 
-    def __exit__(self, *exc_info):
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional['TracebackType']
+    ) -> None:
         self.close()
 
     def __init__(self, channel, producer, consumer, no_ack=False):
diff --git a/kombu/transport/base.py b/kombu/transport/base.py
index 3083acf4..62a84fda 100644
--- a/kombu/transport/base.py
+++ b/kombu/transport/base.py
@@ -4,6 +4,7 @@
 
 import errno
 import socket
+from typing import TYPE_CHECKING, Optional, Type
 
 from amqp.exceptions import RecoverableConnectionError
 
@@ -13,6 +14,9 @@ from kombu.utils.functional import dictfilter
 from kombu.utils.objects import cached_property
 from kombu.utils.time import maybe_s_to_ms
 
+if TYPE_CHECKING:
+    from types import TracebackType
+
 __all__ = ('Message', 'StdChannel', 'Management', 'Transport')
 
 RABBITMQ_QUEUE_ARGUMENTS = {
@@ -100,7 +104,12 @@ class StdChannel:
     def __enter__(self):
         return self
 
-    def __exit__(self, *exc_info):
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional['TracebackType']
+    ) -> None:
         self.close()
 
 
diff --git a/kombu/transport/virtual/base.py b/kombu/transport/virtual/base.py
index 1e18caa6..6aef17da 100644
--- a/kombu/transport/virtual/base.py
+++ b/kombu/transport/virtual/base.py
@@ -13,6 +13,7 @@ from itertools import count
 from multiprocessing.util import Finalize
 from queue import Empty
 from time import monotonic, sleep
+from typing import TYPE_CHECKING, Optional, Type
 
 from amqp.protocol import queue_declare_ok_t
 
@@ -26,6 +27,9 @@ from kombu.utils.uuid import uuid
 
 from .exchange import STANDARD_EXCHANGE_TYPES
 
+if TYPE_CHECKING:
+    from types import TracebackType
+
 ARRAY_TYPE_H = 'H'
 
 UNDELIVERABLE_FMT = """\
@@ -722,7 +726,8 @@ class Channel(AbstractChannel, base.StdChannel):
         message = message.serializable()
         message['redelivered'] = True
         for queue in self._lookup(
-                delivery_info['exchange'], delivery_info['routing_key']):
+            delivery_info['exchange'],
+                delivery_info['routing_key']):
             self._put(queue, message)
 
     def _restore_at_beginning(self, message):
@@ -799,7 +804,12 @@ class Channel(AbstractChannel, base.StdChannel):
     def __enter__(self):
         return self
 
-    def __exit__(self, *exc_info):
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional['TracebackType']
+    ) -> None:
         self.close()
 
     @property
@@ -947,7 +957,7 @@ class Transport(base.Transport):
         # this channel is then used as the next requested channel.
         # (returned by ``create_channel``).
         self._avail_channels.append(self.create_channel(self))
-        return self     # for drain events
+        return self  # for drain events
 
     def close_connection(self, connection):
         self.cycle.close()