From a41cae06d3db9ca3d6664b55ef174f02f7f067cc Mon Sep 17 00:00:00 2001 From: Timofey Kukushkin Date: Tue, 11 Jul 2023 14:30:30 +0400 Subject: [PATCH 1/2] Add type hints to all public APIs --- MANIFEST.in | 1 + statsd/client/__init__.py | 11 ++++++++-- statsd/client/base.py | 38 ++++++++++++++++++++++------------ statsd/client/stream.py | 24 ++++++++++++++-------- statsd/client/timer.py | 41 ++++++++++++++++++++++++++++--------- statsd/client/udp.py | 12 ++++++----- statsd/defaults/__init__.py | 13 +++++++----- statsd/defaults/django.py | 18 ++++++++-------- statsd/defaults/env.py | 18 ++++++++-------- statsd/py.typed | 0 10 files changed, 112 insertions(+), 64 deletions(-) create mode 100644 statsd/py.typed diff --git a/MANIFEST.in b/MANIFEST.in index 72ce170db..9c7227b08 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,3 +1,4 @@ include AUTHORS CHANGES LICENSE MANIFEST.in README.rst include setup.py +include statsd/py.typed recursive-include docs * diff --git a/statsd/client/__init__.py b/statsd/client/__init__.py index bfacdae54..a0b0091e4 100644 --- a/statsd/client/__init__.py +++ b/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', +] diff --git a/statsd/client/base.py b/statsd/client/base.py index 4fbbe5cb4..2e49b662e 100644 --- a/statsd/client/base.py +++ b/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) diff --git a/statsd/client/stream.py b/statsd/client/stream.py index 79951d255..5ccb99cdc 100644 --- a/statsd/client/stream.py +++ b/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) diff --git a/statsd/client/timer.py b/statsd/client/timer.py index 5354a47d5..6d9114802 100644 --- a/statsd/client/timer.py +++ b/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,7 +23,13 @@ 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 @@ -21,10 +37,10 @@ def __init__(self, client, stat, rate=1): 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: diff --git a/statsd/client/udp.py b/statsd/client/udp.py index ec10fc737..77dce5bbc 100644 --- a/statsd/client/udp.py +++ b/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) diff --git a/statsd/defaults/__init__.py b/statsd/defaults/__init__.py index 21851e3a9..0797e09e2 100644 --- a/statsd/defaults/__init__.py +++ b/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 diff --git a/statsd/defaults/django.py b/statsd/defaults/django.py index 664127e8a..8a18b5ba9 100644 --- a/statsd/defaults/django.py +++ b/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) diff --git a/statsd/defaults/env.py b/statsd/defaults/env.py index 51fc7101a..e72c9ce5e 100644 --- a/statsd/defaults/env.py +++ b/statsd/defaults/env.py @@ -1,16 +1,14 @@ import os +from typing import Optional from statsd import defaults from statsd.client import StatsClient -statsd = None - -if statsd is None: - host = os.getenv('STATSD_HOST', defaults.HOST) - port = int(os.getenv('STATSD_PORT', defaults.PORT)) - prefix = os.getenv('STATSD_PREFIX', defaults.PREFIX) - maxudpsize = int(os.getenv('STATSD_MAXUDPSIZE', defaults.MAXUDPSIZE)) - ipv6 = bool(int(os.getenv('STATSD_IPV6', defaults.IPV6))) - statsd = StatsClient(host=host, port=port, prefix=prefix, - maxudpsize=maxudpsize, ipv6=ipv6) +host: str = os.getenv('STATSD_HOST', defaults.HOST) +port: int = int(os.getenv('STATSD_PORT', defaults.PORT)) +prefix: Optional[str] = os.getenv('STATSD_PREFIX', defaults.PREFIX) +maxudpsize: int = int(os.getenv('STATSD_MAXUDPSIZE', defaults.MAXUDPSIZE)) +ipv6: bool = bool(int(os.getenv('STATSD_IPV6', defaults.IPV6))) +statsd: StatsClient = StatsClient(host=host, port=port, prefix=prefix, + maxudpsize=maxudpsize, ipv6=ipv6) diff --git a/statsd/py.typed b/statsd/py.typed new file mode 100644 index 000000000..e69de29bb From 6141bd14f15f3b388ed9b7ea4405396e57efe4c9 Mon Sep 17 00:00:00 2001 From: Timofey Kukushkin Date: Tue, 11 Jul 2023 14:45:03 +0400 Subject: [PATCH 2/2] Move tests to a separate package --- .github/workflows/ci.yml | 6 ++++-- docs/contributing.rst | 4 ++-- tests/__init__.py | 0 {statsd => tests}/tests.py | 0 tox.ini | 2 +- 5 files changed, 7 insertions(+), 5 deletions(-) create mode 100644 tests/__init__.py rename {statsd => tests}/tests.py (100%) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ea1e4fe82..9acf8a553 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -31,6 +31,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 diff --git a/docs/contributing.rst b/docs/contributing.rst index c16769b95..8f71bae65 100644 --- a/docs/contributing.rst +++ b/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 diff --git a/tests/__init__.py b/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/statsd/tests.py b/tests/tests.py similarity index 100% rename from statsd/tests.py rename to tests/tests.py diff --git a/tox.ini b/tox.ini index 42c4c687a..5291c5355 100644 --- a/tox.ini +++ b/tox.ini @@ -8,4 +8,4 @@ deps= coverage commands= - nose2 statsd --with-coverage + nose2 tests --with-coverage