From 04eb5fecec293bfa2ac010aeba54b32b4911c86d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?S=C3=A9amus=20=C3=93=20Ceanainn?=
 <seamusoceanainn@gmail.com>
Date: Tue, 5 Nov 2024 17:56:07 +0000
Subject: [PATCH] Add typing to timed function and TimedContextManagerDecorator
 class

---
 datadog/dogstatsd/base.py    | 10 +++++++--
 datadog/dogstatsd/context.py | 43 +++++++++++++++++++++++++++---------
 2 files changed, 41 insertions(+), 12 deletions(-)

diff --git a/datadog/dogstatsd/base.py b/datadog/dogstatsd/base.py
index 28a62d896..da6cdee35 100644
--- a/datadog/dogstatsd/base.py
+++ b/datadog/dogstatsd/base.py
@@ -413,7 +413,7 @@ def __init__(
         if namespace is not None:
             namespace = text(namespace)
         self.namespace = namespace
-        self.use_ms = use_ms
+        self.use_ms = use_ms  # type: bool
         self.default_sample_rate = default_sample_rate
 
         # Origin detection
@@ -988,7 +988,13 @@ def timing(
         """
         self._report(metric, "ms", value, tags, sample_rate)
 
-    def timed(self, metric=None, tags=None, sample_rate=None, use_ms=None):
+    def timed(
+        self,
+        metric=None,  # type: Optional[Text]
+        tags=None,  # type: Optional[List[str]]
+        sample_rate=None,  # type: Optional[float]
+        use_ms=None,  # type: Optional[bool]
+    ):  # type(...) -> TimedContextManagerDecorator
         """
         A decorator or context manager that will measure the distribution of a
         function's/context's run time. Optionally specify a list of tags or a
diff --git a/datadog/dogstatsd/context.py b/datadog/dogstatsd/context.py
index 90e9ce90e..abc9cb5b5 100644
--- a/datadog/dogstatsd/context.py
+++ b/datadog/dogstatsd/context.py
@@ -3,6 +3,7 @@
 # Copyright 2015-Present Datadog, Inc
 # stdlib
 from functools import wraps
+from typing import Any, Callable, List, Optional, Text, TYPE_CHECKING, Union
 
 try:
     from time import monotonic  # type: ignore[attr-defined]
@@ -13,6 +14,9 @@
 from datadog.dogstatsd.context_async import _get_wrapped_co
 from datadog.util.compat import iscoroutinefunction
 
+if TYPE_CHECKING:
+    from datadog.dogstatsd.base import DogStatsd
+
 
 class TimedContextManagerDecorator(object):
     """
@@ -20,16 +24,25 @@ class TimedContextManagerDecorator(object):
     the context OR in a function call.
     """
 
-    def __init__(self, statsd, metric=None, tags=None, sample_rate=1, use_ms=None):
+    def __init__(
+        self,
+        statsd,  # type: DogStatsd
+        metric=None,  # type: Optional[Text]
+        tags=None,  # type: Optional[List[str]]
+        sample_rate=1,  # type: Optional[float]
+        use_ms=None,  # type: Optional[bool]
+    ):  # type(...) -> None
         self.statsd = statsd
         self.timing_func = statsd.timing
         self.metric = metric
         self.tags = tags
         self.sample_rate = sample_rate
         self.use_ms = use_ms
-        self.elapsed = None
+        self.elapsed = None  # type: Optional[Union[float, int]]
 
-    def __call__(self, func):
+    def __call__(
+        self, func  # type: Callable[..., Any]
+    ):  # type(...) -> Callable[..., Any]
         """
         Decorator which returns the elapsed time of the function call.
 
@@ -53,27 +66,30 @@ def wrapped(*args, **kwargs):
 
         return wrapped
 
-    def __enter__(self):
+    def __enter__(self):  # type(...) -> TimedContextManagerDecorator
         if not self.metric:
             raise TypeError("Cannot used timed without a metric!")
         self._start = monotonic()
         return self
 
-    def __exit__(self, type, value, traceback):
+    def __exit__(self, type, value, traceback):  # type(...) -> None
         # Report the elapsed time of the context manager.
         self._send(self._start)
 
-    def _send(self, start):
+    def _send(
+        self,
+        start,  # type: float
+    ):  # type(...) -> None
         elapsed = monotonic() - start
         use_ms = self.use_ms if self.use_ms is not None else self.statsd.use_ms
         elapsed = int(round(1000 * elapsed)) if use_ms else elapsed
-        self.timing_func(self.metric, elapsed, self.tags, self.sample_rate)
+        self.timing_func(self.metric, elapsed, self.tags, self.sample_rate)  # type: ignore
         self.elapsed = elapsed
 
-    def start(self):
+    def start(self):  # type(...) -> None
         self.__enter__()
 
-    def stop(self):
+    def stop(self):  # type(...) -> None
         self.__exit__(None, None, None)
 
 
@@ -83,6 +99,13 @@ class DistributedContextManagerDecorator(TimedContextManagerDecorator):
     the context OR in a function call using the custom distribution metric.
     """
 
-    def __init__(self, statsd, metric=None, tags=None, sample_rate=1, use_ms=None):
+    def __init__(
+        self,
+        statsd,  # type: DogStatsd
+        metric=None,  # type: Optional[Text]
+        tags=None,  # type: Optional[List[str]]
+        sample_rate=1,  # type: Optional[float]
+        use_ms=None,  # type: Optional[bool]
+    ):  # type(...) -> None
         super(DistributedContextManagerDecorator, self).__init__(statsd, metric, tags, sample_rate, use_ms)
         self.timing_func = statsd.distribution