ResultSet: added Keys() and items() functions

aviau · aviau · commit 922a0a58c550 · 2015-04-08T16:28:14.000-04:00
diff --git a/influxdb/resultset.py b/influxdb/resultset.py
@@ -35,40 +35,23 @@ def __getitem__(self, key):
         if not isinstance(name, (str, type(None))):
             raise TypeError('serie_name must be an str or None')
 
-        for result in self.raw['results']:
-            for serie in result['series']:
-                serie_name = serie.get('name', None)
-                if serie_name is None:
-                    # this is a "system" query or a query which
-                    # doesn't return a name attribute.
-                    # like 'show retention policies' ..
-                    if key is None:
-                        for point in serie['values']:
-                            yield Point(None, serie['columns'], point)
+        for serie in self._get_series():
+            serie_name = serie.get('name', None)
+            if serie_name is None:
+                # this is a "system" query or a query which
+                # doesn't return a name attribute.
+                # like 'show retention policies' ..
+                if key is None:
+                    for point in serie['values']:
+                        yield Point(None, serie['columns'], point)
 
-                elif name in (None, serie_name):
-                    # by default if no tags was provided then
-                    # we will matches every returned serie
-                    serie_matches = True
-                    serie_tags = serie.get('tags', {})
-                    if tags:
-                        serie_matches = False
-                        # if there are some tags requested,
-                        # let's check them:
-                        for tag_name, tag_value in tags.items():
-                            # using _sentinel as I'm not sure that "None"
-                            # could be used, because it could be a valid
-                            # serie_tags value : when a serie has no such tag
-                            # then I think it's set to /null/None/.. TBC..
-                            serie_tag_value = serie_tags.get(tag_name, _sentinel)
-                            if serie_tag_value != tag_value:
-                                break
-                        else:
-                            serie_matches = True
-
-                    if serie_matches:
-                        for point in serie['values']:
-                            yield Point(serie_name, serie['columns'], point, serie_tags)
+            elif name in (None, serie_name):
+                # by default if no tags was provided then
+                # we will matches every returned serie
+                serie_tags = serie.get('tags', {})
+                if tags is None or self._tag_matches(serie_tags, tags):
+                    for point in serie['values']:
+                        yield Point(serie_name, serie['columns'], point, serie_tags)
 
     def __repr__(self):
         return str(self.raw)
@@ -80,5 +63,42 @@ def __iter__(self):
             for serie in results['series']:
                 yield serie
 
-    #def __len__(self):
-    #    return len(self.raw)
+    def _tag_matches(self, tags, filter):
+        """Checks if all key/values in filter match in tags"""
+        for tag_name, tag_value in filter.items():
+            # using _sentinel as I'm not sure that "None"
+            # could be used, because it could be a valid
+            # serie_tags value : when a serie has no such tag
+            # then I think it's set to /null/None/.. TBC..
+            serie_tag_value = tags.get(tag_name, _sentinel)
+            if serie_tag_value != tag_value:
+                return False
+        return True
+
+    def _get_series(self):
+        """Returns all series"""
+        series = []
+        try:
+            for result in self.raw['results']:
+                series.extend(result['series'])
+        except KeyError:
+            pass
+        return series
+
+    def __len__(self):
+        return len(self.keys())
+
+    def keys(self):
+        keys = []
+        for serie in self._get_series():
+            keys.append((serie['name'], serie['tags']))
+        return keys
+
+    def items(self):
+        items = []
+        for serie in self._get_series():
+            serie_key = (serie['name'], serie['tags'])
+            items.append(
+                (serie_key, self.__getitem__(serie_key))
+            )
+        return items
diff --git a/tests/influxdb/resultset_test.py b/tests/influxdb/resultset_test.py
@@ -72,3 +72,47 @@ def test_filter_by_tags(self):
                       tags={"host": "server02", "region": "us-west"}),
             ]
         )
+
+    def test_keys(self):
+        self.assertItemsEqual(
+            self.rs.keys(),
+            [
+                ('cpu_load_short', {'host': 'server01', 'region': 'us-west'}),
+                ('cpu_load_short', {'host': 'server02', 'region': 'us-west'}),
+                ('other_serie', {'host': 'server01', 'region': 'us-west'})
+            ]
+        )
+
+    def test_len(self):
+        self.assertEqual(
+            len(self.rs),
+            3
+        )
+
+    def test_items(self):
+        items = list(self.rs.items())
+        items_lists = [(item[0], list(item[1])) for item in items]
+
+        self.assertEqual(
+            items_lists,
+            [
+                (
+                    ('cpu_load_short', {'host': 'server01', 'region': 'us-west'}),
+                    [Point("cpu_load_short", ["time", "value"],
+                           ["2015-01-29T21:51:28.968422294Z", 0.64],
+                           tags={"host": "server01", "region": "us-west"})]
+                ),
+                (
+                    ('cpu_load_short', {'host': 'server02', 'region': 'us-west'}),
+                    [Point("cpu_load_short", ["time", "value"],
+                           ["2015-01-29T21:51:28.968422294Z", 0.64],
+                           tags={"host": "server02", "region": "us-west"})]
+                ),
+                (
+                    ('other_serie', {'host': 'server01', 'region': 'us-west'}),
+                    [Point("other_serie", ["time", "value"],
+                           ["2015-01-29T21:51:28.968422294Z", 0.64],
+                           tags={"host": "server01", "region": "us-west"})]
+                )
+            ]
+        )
