From 2b8a24c303060f52260d3061785cb29efa7f6978 Mon Sep 17 00:00:00 2001
From: andrewqian2001datadog <andrew.qian@datadoghq.com>
Date: Tue, 2 Jul 2024 11:32:25 -0400
Subject: [PATCH] Add Enums

---
 datadog/dogstatsd/aggregator.py | 34 ++++++++++++++++++---------------
 1 file changed, 19 insertions(+), 15 deletions(-)

diff --git a/datadog/dogstatsd/aggregator.py b/datadog/dogstatsd/aggregator.py
index 8b8fde8a..2c75e731 100644
--- a/datadog/dogstatsd/aggregator.py
+++ b/datadog/dogstatsd/aggregator.py
@@ -1,21 +1,27 @@
 import threading
 import time
+from enum import Enum
 
 from datadog.dogstatsd.metrics import CountMetric, GaugeMetric, SetMetric
 
+class MetricType(Enum):
+    COUNT = 'counts'
+    GAUGE = 'gauges'
+    SET = 'sets'
+
 class Aggregator(object):
     def __init__(self, client):
         self.client = client
 
         self.metrics_map = {
-            'counts': {},
-            'gauges': {},
-            'sets': {}
+            MetricType.COUNT: {},
+            MetricType.GAUGE: {},
+            MetricType.SET: {}
         }
         self.locks = {
-            'counts': threading.RLock(),
-            'gauges': threading.RLock(),
-            'sets': threading.RLock()
+            MetricType.COUNT: threading.RLock(),
+            MetricType.GAUGE: threading.RLock(),
+            MetricType.SET: threading.RLock()
         }
 
         self.closed = threading.Event()
@@ -34,7 +40,7 @@ def tick(self):
 
     def send_metrics(self):
         for metric in self.flush_metrics():
-            # TODO: change the _report function in base.py to handle new data type
+            # TODO: change the _report function in base.py to handle new data
             self.client.send(metric)
 
     def stop(self):
@@ -46,7 +52,7 @@ def stop(self):
     def flush_metrics(self):
         metrics = []
 
-        for metric_type in self.metrics_maps.keys():
+        for metric_type in self.metrics_map.keys():
             with self.locks[metric_type]:
                 current_metrics = self.metrics_map[metric_type]
                 self.metrics_map[metric_type] = {}
@@ -60,15 +66,14 @@ def flush_metrics(self):
     def get_context(self, name, tags):
         return "{}:{}".format(name, ",".join(tags))
 
-    # This function may not be necessary, can just call add_metric directly
     def count(self, name, value, tags, rate, timestamp=0):
-        return self.add_metric('counts', CountMetric, name, value, tags, rate, timestamp)
-    # This function may not be necessary, can just call add_metric directly
+        return self.add_metric(MetricType.COUNT, CountMetric, name, value, tags, rate, timestamp)
+
     def gauge(self, name, value, tags, rate, timestamp=0):
-        return self.add_metric('gauges', GaugeMetric, name, value, tags, rate, timestamp)
-    # This function may not be necessary, can just call add_metric directly
+        return self.add_metric(MetricType.GAUGE, GaugeMetric, name, value, tags, rate, timestamp)
+
     def set(self, name, value, tags, rate, timestamp=0):
-        return self.add_metric('sets', SetMetric, name, value, tags, rate, timestamp)
+        return self.add_metric(MetricType.SET, SetMetric, name, value, tags, rate, timestamp)
 
     def add_metric(self, metric_type, metric_class, name, value, tags, rate, timestamp=0):
         context = self.get_context(name, tags)
@@ -77,4 +82,3 @@ def add_metric(self, metric_type, metric_class, name, value, tags, rate, timesta
                 self.metrics_map[metric_type][context].aggregate(value)
             else:
                 self.metrics_map[metric_type][context] = metric_class(name, value, tags, rate, timestamp)
-        return None