Linting + py2.5 support.

alisaifee · Jun 10, 2015 · 64df320 · 64df320
1 parent 585781f
commit 64df320
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 26 deletions.
diff --git a/sifr/__init__.py b/sifr/__init__.py
@@ -19,25 +19,25 @@ def incr(self, key, amount=1, resolutions=None):
 
     def incr_unique(self, key, identity, resolutions=None):
         self.client.call(
-                'incr_unique', key,
-                resolutions or self.resolutions, identity
-            )
+            'incr_unique', key,
+            resolutions or self.resolutions, identity
+        )
 
     def track(self, key, identity, resolutions=None):
         self.client.call(
-                'track', key,
-                resolutions or self.resolutions, identity
-            )
+            'track', key,
+            resolutions or self.resolutions, identity
+        )
 
     def count(self, key, at, resolution):
         return self.client.call(
-                'count', key, time.mktime(at.timetuple()), resolution
-            )
+            'count', key, time.mktime(at.timetuple()), resolution
+        )
 
     def cardinality(self, key, at, resolution):
         return self.client.call(
-                'cardinality', key, time.mktime(at.timetuple()), resolution
-            )
+            'cardinality', key, time.mktime(at.timetuple()), resolution
+        )
 
     def uniques(self, key, at, resolution):
         return set(

diff --git a/sifr/storage.py b/sifr/storage.py
@@ -205,13 +205,18 @@ def cardinality(self, span):
         return int(value) if value is not None else 0
 
 
-
 class RiakStorage(Storage):
     def __init__(self, riak):
         self.riak = riak
-        self.counter_bucket = self.riak.bucket_type("maps").bucket("sifr_counter")
-        self.unique_counters_bucket = self.riak.bucket_type("maps").bucket("sifr_unique_counter")
-        self.uniques_bucket = self.riak.bucket_type("maps").bucket("sifr_uniques")
+        self.counter_bucket = self.riak.bucket_type(
+            "maps"
+        ).bucket("sifr_counter")
+        self.unique_counters_bucket = self.riak.bucket_type(
+            "maps"
+        ).bucket("sifr_unique_counter")
+        self.uniques_bucket = self.riak.bucket_type(
+            "maps"
+        ).bucket("sifr_uniques")
 
     def count(self, span):
         map = self.counter_bucket.get(span.namespace)
@@ -224,22 +229,26 @@ def incr(self, span, amount=1):
         counter.increment()
         map.store()
 
+    def get_maps(self, bucket, spans, create=False):
+        maps = {}
+        namespaces = set(span.namespace for span in spans)
+        for namespace in namespaces:
+            if create:
+                maps[namespace] = bucket.new(namespace)
+            else:
+                maps[namespace] = bucket.get(namespace)
+        return maps
+
     def track_multi(self, spans, identifier):
-        maps = {
-            namespace:self.uniques_bucket.new(namespace)
-            for namespace in set(span.namespace for span in spans)
-            }
+        maps = self.get_maps(self.uniques_bucket, spans, True)
         for span in spans:
             riak_set = maps[span.namespace].sets.get(span.timestamp)
             riak_set.add(str(identifier))
         for map in maps.values():
             map.store()
 
     def incr_unique_multi(self, spans, identifier):
-        maps = {
-            namespace:self.unique_counters_bucket.new(namespace)
-            for namespace in set(span.namespace for span in spans)
-        }
+        maps = self.get_maps(self.unique_counters_bucket, spans, True)
         for span in spans:
             counter = maps[span.namespace].sets.get(span.timestamp)
             counter.add(str(identifier))
@@ -259,10 +268,7 @@ def incr_unique(self, span, identifier):
         map.store()
 
     def incr_multi(self, spans, amount=1):
-        maps = {
-            namespace:self.counter_bucket.new(namespace)
-            for namespace in set(span.namespace for span in spans)
-            }
+        maps = self.get_maps(self.counter_bucket, spans, True)
         for span in spans:
             counter = maps[span.namespace].counters.get(span.timestamp)
             counter.increment(amount)