From 9054c729f55d5e45375393e4a75080bc0735dac3 Mon Sep 17 00:00:00 2001
From: Dave Shawley <daves@aweber.com>
Date: Tue, 20 Jul 2021 17:32:51 -0400
Subject: [PATCH] Add some missing typing.

This doesn't completely appease pyright & mypy since I insist on calling
the sprockets.http.Application methods when guarded by isinstance()
checks.
---
 sprockets_statsd/statsd.py  | 16 ++++++++++------
 sprockets_statsd/tornado.py | 10 +++++-----
 2 files changed, 15 insertions(+), 11 deletions(-)

diff --git a/sprockets_statsd/statsd.py b/sprockets_statsd/statsd.py
index 73d0393..9571c2c 100644
--- a/sprockets_statsd/statsd.py
+++ b/sprockets_statsd/statsd.py
@@ -67,6 +67,9 @@ class Timer:
     twice in a row will result in a :exc:`RuntimeError` as well.
 
     """
+    _start_time: typing.Union[None, float]
+    _finish_time: typing.Union[None, float]
+
     def __init__(self, connector: 'AbstractConnector', path: str):
         self._connector = connector
         self._path = path
@@ -94,7 +97,7 @@ class Timer:
             self.send()
         return self
 
-    def send(self):
+    def send(self) -> None:
         """Send the recorded timing to the connector.
 
         This method will raise a :exc:`RuntimeError` if a timing has
@@ -112,13 +115,14 @@ class Timer:
             max(self._finish_time, self._start_time) - self._start_time)
         self._start_time, self._finish_time = None, None
 
-    def __enter__(self):
+    def __enter__(self) -> 'Timer':
         self.start()
         return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(self, exc_type: typing.Union[typing.Type[Exception], None],
+                 exc_val: typing.Union[Exception, None],
+                 exc_tb: typing.Union[typing.Tuple, None]) -> None:
         self.stop()
-        return False
 
 
 class AbstractConnector:
@@ -187,7 +191,7 @@ class AbstractConnector:
             seconds = seconds.total_seconds()
         self.inject_metric(f'timers.{path}', str(seconds * 1000.0), 'ms')
 
-    def timer(self, path) -> Timer:
+    def timer(self, path: str) -> Timer:
         """Send a timer metric using a context manager.
 
         :param path: timer to append the measured time to
@@ -482,7 +486,7 @@ class Processor:
 
     logger: logging.Logger
     protocol: typing.Optional[StatsdProtocol]
-    queue: asyncio.Queue
+    queue: asyncio.Queue[bytes]
     _create_transport: typing.Callable[[], typing.Coroutine[
         typing.Any, typing.Any, typing.Tuple[asyncio.BaseTransport,
                                              StatsdProtocol]]]
diff --git a/sprockets_statsd/tornado.py b/sprockets_statsd/tornado.py
index 316dd7d..78c1d1d 100644
--- a/sprockets_statsd/tornado.py
+++ b/sprockets_statsd/tornado.py
@@ -121,7 +121,7 @@ class Application(web.Application):
         except AttributeError:
             pass
 
-    async def start_statsd(self, *_) -> None:
+    async def start_statsd(self, *_: typing.Any) -> None:
         """Start the connector during startup.
 
         Call this method during application startup to enable the statsd
@@ -157,7 +157,7 @@ class Application(web.Application):
 
             await self.statsd_connector.start()
 
-    async def stop_statsd(self, *_) -> None:
+    async def stop_statsd(self, *_: typing.Any) -> None:
         """Stop the connector during shutdown.
 
         If the connector was started, then this method will gracefully
@@ -171,7 +171,7 @@ class Application(web.Application):
 
     def __handle_fatal_error(self,
                              message: str,
-                             exc: typing.Optional[Exception] = None):
+                             exc: typing.Optional[Exception] = None) -> None:
         logger = self.__get_logger()
         if exc is not None:
             logger.exception('%s', message)
@@ -184,7 +184,7 @@ class Application(web.Application):
 
     def __get_logger(self) -> logging.Logger:
         try:
-            return getattr(self, 'logger')
+            return typing.cast(logging.Logger, getattr(self, 'logger'))
         except AttributeError:
             return logging.getLogger(__package__).getChild(
                 'tornado.Application')
@@ -192,7 +192,7 @@ class Application(web.Application):
 
 class RequestHandler(web.RequestHandler):
     """Mix this into your handler to send metrics to a statsd server."""
-    statsd_connector: typing.Optional[statsd.Connector]
+    statsd_connector: typing.Optional[statsd.AbstractConnector]
 
     def initialize(self, **kwargs: typing.Any) -> None:
         super().initialize(**kwargs)