Retrieving an existing gauge, and possibly replacing it.

metrology/registry.py

@@ -1,6 +1,7 @@
 import inspect
 
 from threading import RLock
+from collections import defaultdict
 
 from metrology.exceptions import RegistryException
 from metrology.instruments import (
@@ -17,6 +18,8 @@ class Registry(object):
     def __init__(self):
         self.lock = RLock()
         self.metrics = {}
+        self.metrics_by_tag = defaultdict(lambda: [])
+        self.tags_by_metric = {}
 
     def clear(self):
         with self.lock:
@@ -52,31 +55,86 @@ def derive(self, name):
         return self.add_or_get(name, Derive)
 
     def get(self, name):
+        name = safe_key(name)
         with self.lock:
             return self.metrics[name]
 
     def add(self, name, metric):
+        key = safe_key(name)
         with self.lock:
-            if name in self.metrics:
+            if key in self.metrics:
                 raise RegistryException("{0} already present "
                                         "in the registry.".format(name))
             else:
-                self.metrics[name] = metric
+                self.metrics[key] = metric
+                self._index(name, metric)
 
     def add_or_get(self, name, klass):
+        """Creates an instance of `klass`, and registered it with `name`.
+        `klass` may also be an instance of a metric, and will be registered
+        directly.
+
+        If `name` is already registered:
+
+        - If a klass is passed, return the existing instance.
+        - If an instance is passed, replace the existing metric with the new
+          instance.
+        """
+        key = safe_key(name)
         with self.lock:
-            metric = self.metrics.get(name)
+            metric = self.metrics.get(key)
+
             if metric is not None:
-                if not isinstance(metric, klass):
-                    raise RegistryException("{0} is not of "
-                                            "type {1}.".format(name, klass))
+                # If a klass was given, return the existing metric.
+                if inspect.isclass(klass):
+                    if not isinstance(metric, klass):
+                        raise RegistryException("{0} is not of "
+                                                "type {1}.".format(name, klass))
+                    return metric
+
+                # If a metric object was passed, make sure the registry has
+                # the most recent version.
+                else:
+                    self.metrics[key] = klass
+                    self._index(name, klass)
+
             else:
                 if inspect.isclass(klass):
                     metric = klass()
                 else:
                     metric = klass
-                self.metrics[name] = metric
-            return metric
+                self.metrics[key] = metric
+                self._index(name, metric)
+                return metric
+
+    def _index(self, name, metric):
+        if not isinstance(name, dict):
+            return
+
+        for key, value in name.items():
+            self.metrics_by_tag[(key, value)].append(metric)
+        self.tags_by_metric[metric] = name
+
+    def filter_metrics(self, filters):
+        """
+        Find all metrics matching the tags given in `filters`. For each
+        metric, remain a 2-tuple (metric, other_tags).
+        """
+        result = None
+
+        for filter_tag, filter_value in filters.items():
+            local_match = set(self.metrics_by_tag[(filter_tag, filter_value)])
+            if result is None:
+                result = local_match
+            else:
+                result = result.intersection(local_match)
+
+        for metric in result:
+            tags = self.tags_by_metric[metric].copy()
+            # Do not include any tags the caller has queried for
+            for tag in filters:
+                del tags[tag]
+            yield metric, tags
 
     def stop(self):
         self.clear()
@@ -87,4 +145,10 @@ def __iter__(self):
                 yield name, metric
 
 
+def safe_key(name):
+    if isinstance(name, dict):
+        return tuple(name.items())
+    return name
+
+
 registry = Registry()