Add type hints to all public APIs

.github/workflows/ci.yml

@@ -32,6 +32,8 @@ jobs:
       with:
         python-version: '3.11'
 
-    - run: pip install flake8
+    - run: pip install flake8 pyright .
 
-    - run: flake8 statsd
+    - run: flake8 statsd tests
+
+    - run: pyright --verifytypes statsd

MANIFEST.in

@@ -1,3 +1,4 @@
 include AUTHORS CHANGES LICENSE MANIFEST.in README.rst
 include setup.py
+include statsd/py.typed
 recursive-include docs *

docs/contributing.rst

@@ -36,8 +36,8 @@ You can also run the tests with tox::
 
     $ tox
 
-Tox will run the tests in Pythons 2.5, 2.6, 2.7, 3.2, 3.3, 3.4, and
-PyPy, if they're available.
+Tox will run the tests in Pythons 3.7, 3.8, 3.9, 3.10, 3.11, and
+PyPy3, if they're available.
 
 
 Writing Tests

statsd/client/__init__.py

@@ -1,2 +1,9 @@
-from .stream import TCPStatsClient, UnixSocketStatsClient  # noqa
-from .udp import StatsClient  # noqa
+from .stream import TCPStatsClient, UnixSocketStatsClient
+from .udp import StatsClient
+
+
+__all__ = [
+    'TCPStatsClient',
+    'UnixSocketStatsClient',
+    'StatsClient',
+]

statsd/client/base.py

@@ -1,27 +1,33 @@
 import random
 from collections import deque
 from datetime import timedelta
+from types import TracebackType
+from typing import Any, Optional, Type, TypeVar, Union
 
 from .timer import Timer
 
 
+_PipelineBaseT = TypeVar('_PipelineBaseT', bound='PipelineBase')
+
+
 class StatsClientBase:
     """A Base class for various statsd clients."""
 
-    def close(self):
+    def close(self) -> None:
         """Used to close and clean up any underlying resources."""
         raise NotImplementedError()
 
     def _send(self):
         raise NotImplementedError()
 
-    def pipeline(self):
+    def pipeline(self) -> 'PipelineBase':
         raise NotImplementedError()
 
-    def timer(self, stat, rate=1):
+    def timer(self, stat: str, rate: float = 1) -> Timer:
         return Timer(self, stat, rate)
 
-    def timing(self, stat, delta, rate=1):
+    def timing(self, stat: str, delta: Union[float, timedelta],
+               rate: float = 1) -> None:
         """
         Send new timing information.
 
@@ -32,15 +38,16 @@ def timing(self, stat, delta, rate=1):
             delta = delta.total_seconds() * 1000.
         self._send_stat(stat, '%0.6f|ms' % delta, rate)
 
-    def incr(self, stat, count=1, rate=1):
+    def incr(self, stat: str, count: int = 1, rate: float = 1) -> None:
         """Increment a stat by `count`."""
         self._send_stat(stat, '%s|c' % count, rate)
 
-    def decr(self, stat, count=1, rate=1):
+    def decr(self, stat: str, count: int = 1, rate: float = 1) -> None:
         """Decrement a stat by `count`."""
         self.incr(stat, -count, rate)
 
-    def gauge(self, stat, value, rate=1, delta=False):
+    def gauge(self, stat: str, value: float, rate: float = 1,
+              delta: bool = False) -> None:
         """Set a gauge value."""
         if value < 0 and not delta:
             if rate < 1:
@@ -53,7 +60,7 @@ def gauge(self, stat, value, rate=1, delta=False):
             prefix = '+' if delta and value >= 0 else ''
             self._send_stat(stat, '{}{}|g'.format(prefix, value), rate)
 
-    def set(self, stat, value, rate=1):
+    def set(self, stat: str, value: Any, rate: float = 1) -> None:
         """Set a set value."""
         self._send_stat(stat, '%s|s' % value, rate)
 
@@ -78,7 +85,7 @@ def _after(self, data):
 
 class PipelineBase(StatsClientBase):
 
-    def __init__(self, client):
+    def __init__(self, client: StatsClientBase) -> None:
         self._client = client
         self._prefix = client._prefix
         self._stats = deque()
@@ -90,16 +97,21 @@ def _after(self, data):
         if data is not None:
             self._stats.append(data)
 
-    def __enter__(self):
+    def __enter__(self: _PipelineBaseT) -> _PipelineBaseT:
         return self
 
-    def __exit__(self, typ, value, tb):
+    def __exit__(
+        self,
+        typ: Optional[Type[BaseException]],
+        value: Optional[BaseException],
+        tb: Optional[TracebackType],
+    ) -> None:
         self.send()
 
-    def send(self):
+    def send(self) -> None:
         if not self._stats:
             return
         self._send()
 
-    def pipeline(self):
+    def pipeline(self) -> 'PipelineBase':
         return self.__class__(self)

statsd/client/stream.py

@@ -1,4 +1,5 @@
 import socket
+from typing import Optional
 
 from .base import StatsClientBase, PipelineBase
 
@@ -10,19 +11,19 @@ def _send(self):
 
 
 class StreamClientBase(StatsClientBase):
-    def connect(self):
+    def connect(self) -> None:
         raise NotImplementedError()
 
-    def close(self):
+    def close(self) -> None:
         if self._sock and hasattr(self._sock, 'close'):
             self._sock.close()
         self._sock = None
 
-    def reconnect(self):
+    def reconnect(self) -> None:
         self.close()
         self.connect()
 
-    def pipeline(self):
+    def pipeline(self) -> StreamPipeline:
         return StreamPipeline(self)
 
     def _send(self, data):
@@ -38,8 +39,12 @@ def _do_send(self, data):
 class TCPStatsClient(StreamClientBase):
     """TCP version of StatsClient."""
 
-    def __init__(self, host='localhost', port=8125, prefix=None,
-                 timeout=None, ipv6=False):
+    def __init__(self,
+                 host: str = 'localhost',
+                 port: int = 8125,
+                 prefix: Optional[str] = None,
+                 timeout: Optional[float] = None,
+                 ipv6: bool = False) -> None:
         """Create a new client."""
         self._host = host
         self._port = port
@@ -48,7 +53,7 @@ def __init__(self, host='localhost', port=8125, prefix=None,
         self._prefix = prefix
         self._sock = None
 
-    def connect(self):
+    def connect(self) -> bool:
         fam = socket.AF_INET6 if self._ipv6 else socket.AF_INET
         family, _, _, _, addr = socket.getaddrinfo(
             self._host, self._port, fam, socket.SOCK_STREAM)[0]
@@ -60,14 +65,15 @@ def connect(self):
 class UnixSocketStatsClient(StreamClientBase):
     """Unix domain socket version of StatsClient."""
 
-    def __init__(self, socket_path, prefix=None, timeout=None):
+    def __init__(self, socket_path: str, prefix: Optional[str] = None,
+                 timeout: Optional[float] = None):
         """Create a new client."""
         self._socket_path = socket_path
         self._timeout = timeout
         self._prefix = prefix
         self._sock = None
 
-    def connect(self):
+    def connect(self) -> None:
         self._sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
         self._sock.settimeout(self._timeout)
         self._sock.connect(self._socket_path)

statsd/client/timer.py

@@ -1,9 +1,19 @@
 import functools
 from inspect import iscoroutinefunction
 from time import perf_counter as time_now
+from types import TracebackType
+from typing import TYPE_CHECKING, Any, Callable, Optional, Type, TypeVar
 
 
-def safe_wraps(wrapper, *args, **kwargs):
+if TYPE_CHECKING:
+    from statsd.client.base import StatsClientBase
+
+
+_F = TypeVar('_F', bound=Callable[..., Any])
+_TimerT = TypeVar('_TimerT', bound='Timer')
+
+
+def _safe_wraps(wrapper, *args, **kwargs):
     """Safely wraps partial functions."""
     while isinstance(wrapper, functools.partial):
         wrapper = wrapper.func
@@ -13,18 +23,24 @@ def safe_wraps(wrapper, *args, **kwargs):
 class Timer:
     """A context manager/decorator for statsd.timing()."""
 
-    def __init__(self, client, stat, rate=1):
+    client: 'StatsClientBase'
+    stat: str
+    rate: float
+    ms: Optional[float]
+
+    def __init__(self, client: 'StatsClientBase', stat: str,
+                 rate: float = 1) -> None:
         self.client = client
         self.stat = stat
         self.rate = rate
         self.ms = None
         self._sent = False
         self._start_time = None
 
-    def __call__(self, f):
+    def __call__(self, f: _F) -> _F:
         """Thread-safe timing function decorator."""
         if iscoroutinefunction(f):
-            @safe_wraps(f)
+            @_safe_wraps(f)
             async def _async_wrapped(*args, **kwargs):
                 start_time = time_now()
                 try:
@@ -34,7 +50,7 @@ async def _async_wrapped(*args, **kwargs):
                     self.client.timing(self.stat, elapsed_time_ms, self.rate)
             return _async_wrapped
 
-        @safe_wraps(f)
+        @_safe_wraps(f)
         def _wrapped(*args, **kwargs):
             start_time = time_now()
             try:
@@ -44,19 +60,24 @@ def _wrapped(*args, **kwargs):
                 self.client.timing(self.stat, elapsed_time_ms, self.rate)
         return _wrapped
 
-    def __enter__(self):
+    def __enter__(self: _TimerT) -> _TimerT:
         return self.start()
 
-    def __exit__(self, typ, value, tb):
+    def __exit__(
+        self,
+        typ: Optional[Type[BaseException]],
+        value: Optional[BaseException],
+        tb: Optional[TracebackType],
+    ) -> None:
         self.stop()
 
-    def start(self):
+    def start(self: _TimerT) -> _TimerT:
         self.ms = None
         self._sent = False
         self._start_time = time_now()
         return self
 
-    def stop(self, send=True):
+    def stop(self: _TimerT, send: bool = True) -> _TimerT:
         if self._start_time is None:
             raise RuntimeError('Timer has not started.')
         dt = time_now() - self._start_time
@@ -65,7 +86,7 @@ def stop(self, send=True):
             self.send()
         return self
 
-    def send(self):
+    def send(self) -> None:
         if self.ms is None:
             raise RuntimeError('No data recorded.')
         if self._sent:

statsd/client/udp.py

@@ -1,11 +1,12 @@
 import socket
+from typing import Optional
 
 from .base import StatsClientBase, PipelineBase
 
 
 class Pipeline(PipelineBase):
 
-    def __init__(self, client):
+    def __init__(self, client: StatsClientBase) -> None:
         super().__init__(client)
         self._maxudpsize = client._maxudpsize
 
@@ -25,8 +26,9 @@ def _send(self):
 class StatsClient(StatsClientBase):
     """A client for statsd."""
 
-    def __init__(self, host='localhost', port=8125, prefix=None,
-                 maxudpsize=512, ipv6=False):
+    def __init__(self, host: str = 'localhost', port: int = 8125,
+                 prefix: Optional[str] = None,
+                 maxudpsize: int = 512, ipv6: bool = False) -> None:
         """Create a new client."""
         fam = socket.AF_INET6 if ipv6 else socket.AF_INET
         family, _, _, _, addr = socket.getaddrinfo(
@@ -44,10 +46,10 @@ def _send(self, data):
             # No time for love, Dr. Jones!
             pass
 
-    def close(self):
+    def close(self) -> None:
         if self._sock and hasattr(self._sock, 'close'):
             self._sock.close()
         self._sock = None
 
-    def pipeline(self):
+    def pipeline(self) -> Pipeline:
         return Pipeline(self)

statsd/defaults/__init__.py

@@ -1,5 +1,8 @@
-HOST = 'localhost'
-PORT = 8125
-IPV6 = False
-PREFIX = None
-MAXUDPSIZE = 512
+from typing import Optional
+
+
+HOST: str = 'localhost'
+PORT: int = 8125
+IPV6: bool = False
+PREFIX: Optional[str] = None
+MAXUDPSIZE: int = 512

statsd/defaults/django.py

@@ -1,16 +1,14 @@
+from typing import Optional
 from django.conf import settings
 
 from statsd import defaults
 from statsd.client import StatsClient
 
 
-statsd = None
-
-if statsd is None:
-    host = getattr(settings, 'STATSD_HOST', defaults.HOST)
-    port = getattr(settings, 'STATSD_PORT', defaults.PORT)
-    prefix = getattr(settings, 'STATSD_PREFIX', defaults.PREFIX)
-    maxudpsize = getattr(settings, 'STATSD_MAXUDPSIZE', defaults.MAXUDPSIZE)
-    ipv6 = getattr(settings, 'STATSD_IPV6', defaults.IPV6)
-    statsd = StatsClient(host=host, port=port, prefix=prefix,
-                         maxudpsize=maxudpsize, ipv6=ipv6)
+host: str = getattr(settings, 'STATSD_HOST', defaults.HOST)
+port: int = getattr(settings, 'STATSD_PORT', defaults.PORT)
+prefix: Optional[str] = getattr(settings, 'STATSD_PREFIX', defaults.PREFIX)
+maxudpsize: int = getattr(settings, 'STATSD_MAXUDPSIZE', defaults.MAXUDPSIZE)
+ipv6: bool = getattr(settings, 'STATSD_IPV6', defaults.IPV6)
+statsd: StatsClient = StatsClient(host=host, port=port, prefix=prefix,
+                                  maxudpsize=maxudpsize, ipv6=ipv6)