typing: add typing for dogstatsd

Author: KowalskiThomas

datadog/__init__.py

@@ -13,8 +13,7 @@
 # stdlib
 import logging
 import os
-import os.path
-from typing import Any, List, Optional
+from typing import Any, List, Optional  # noqa: F401
 
 # datadog
 from datadog import api
@@ -43,7 +42,7 @@ def initialize(
     statsd_use_default_route=False,  # type: bool
     statsd_socket_path=None,  # type: Optional[str]
     statsd_namespace=None,  # type: Optional[str]
-    statsd_max_samples_per_context=0,  # type: Optional[int]
+    statsd_max_samples_per_context=0,  # type: int
     statsd_constant_tags=None,  # type: Optional[List[str]]
     return_raw_response=False,  # type: bool
     hostname_from_config=True,  # type: bool

datadog/dogstatsd/aggregator.py

@@ -1,13 +1,16 @@
 import threading
+from typing import Any, Dict, List, Optional
+
 from datadog.dogstatsd.metrics import (
     CountMetric,
     GaugeMetric,
     SetMetric,
+    MetricAggregator,
 )
 from datadog.dogstatsd.max_sample_metric import (
     HistogramMetric,
     DistributionMetric,
-    TimingMetric
+    TimingMetric,
 )
 from datadog.dogstatsd.metric_types import MetricType
 from datadog.dogstatsd.max_sample_metric_context import MaxSampleMetricContexts
@@ -16,12 +19,13 @@
 
 class Aggregator(object):
     def __init__(self, max_samples_per_context=0, cardinality=None):
+        # type: (int, Optional[str]) -> None
         self.max_samples_per_context = max_samples_per_context
         self.metrics_map = {
             MetricType.COUNT: {},
             MetricType.GAUGE: {},
             MetricType.SET: {},
-        }
+        }  # type: Dict[str, Dict[str, MetricAggregator]]
         self.max_sample_metric_map = {
             MetricType.HISTOGRAM: MaxSampleMetricContexts(HistogramMetric),
             MetricType.DISTRIBUTION: MaxSampleMetricContexts(DistributionMetric),
@@ -35,7 +39,8 @@ def __init__(self, max_samples_per_context=0, cardinality=None):
         self.cardinality = cardinality
 
     def flush_aggregated_metrics(self):
-        metrics = []
+        # type: () -> List[MetricAggregator]
+        metrics = []  # type: List[MetricAggregator]
         for metric_type in self.metrics_map.keys():
             with self._locks[metric_type]:
                 current_metrics = self.metrics_map[metric_type]
@@ -46,38 +51,45 @@ def flush_aggregated_metrics(self):
         return metrics
 
     def set_max_samples_per_context(self, max_samples_per_context=0):
+        # type: (int) -> None
         self.max_samples_per_context = max_samples_per_context
 
     def flush_aggregated_sampled_metrics(self):
-        metrics = []
+        # type: () -> List[MetricAggregator]
+        metrics = []  # type: List[MetricAggregator]
         for metric_type in self.max_sample_metric_map.keys():
             metric_context = self.max_sample_metric_map[metric_type]
             for metricList in metric_context.flush():
                 metrics.extend(metricList)
         return metrics
 
     def get_context(self, name, tags):
+        # type: (str, Optional[List[str]]) -> str
         tags_str = u",".join(tags) if tags is not None else ""
         return u"{}:{}".format(name, tags_str)
 
     def count(self, name, value, tags, rate, timestamp=0, cardinality=None):
+        # type: (str, Any, Optional[List[str]], Optional[float], int, Optional[str]) -> None
         return self.add_metric(
             MetricType.COUNT, CountMetric, name, value, tags, rate, timestamp, cardinality
         )
 
     def gauge(self, name, value, tags, rate, timestamp=0, cardinality=None):
+        # type: (str, Any, Optional[List[str]], Optional[float], int, Optional[str]) -> None
         return self.add_metric(
             MetricType.GAUGE, GaugeMetric, name, value, tags, rate, timestamp, cardinality
         )
 
     def set(self, name, value, tags, rate, timestamp=0, cardinality=None):
+        # type: (str, Any, Optional[List[str]], Optional[float], int, Optional[str]) -> None
         return self.add_metric(
             MetricType.SET, SetMetric, name, value, tags, rate, timestamp, cardinality
         )
 
     def add_metric(
         self, metric_type, metric_class, name, value, tags, rate, timestamp=0, cardinality=None
     ):
+        # type: (str, Any, str, Any, Optional[List[str]], Optional[float], int, Optional[str]) -> None
         context = self.get_context(name, tags)
         with self._locks[metric_type]:
             if context in self.metrics_map[metric_type]:
@@ -91,23 +103,27 @@ def add_metric(
                 )
 
     def histogram(self, name, value, tags, rate, cardinality=None):
+        # type: (str, Any, Optional[List[str]], Optional[float], Optional[str]) -> None
         return self.add_max_sample_metric(
             MetricType.HISTOGRAM, name, value, tags, rate, cardinality
         )
 
     def distribution(self, name, value, tags, rate, cardinality=None):
+        # type: (str, Any, Optional[List[str]], Optional[float], Optional[str]) -> None
         return self.add_max_sample_metric(
             MetricType.DISTRIBUTION, name, value, tags, rate, cardinality
         )
 
     def timing(self, name, value, tags, rate, cardinality=None):
+        # type: (str, Any, Optional[List[str]], Optional[float], Optional[str]) -> None
         return self.add_max_sample_metric(
             MetricType.TIMING, name, value, tags, rate, cardinality
         )
 
     def add_max_sample_metric(
         self, metric_type, name, value, tags, rate, cardinality=None
     ):
+        # type: (str, str, Any, Optional[List[str]], Optional[float], Optional[str]) -> None
         if rate is None:
             rate = 1
         context_key = self.get_context(name, tags)