[tor-commits] [compass/master] Refactor RelayStats._get_group_function

[karsten at torproject.org]

commit a1bbdbd5e7eada06db1ff4070f0ce1be69916a25
Author: Lunar <lunar at torproject.org>
Date:   Sun May 26 20:26:41 2013 +0200

    Refactor RelayStats._get_group_function
    
    We need it to scale for more than two grouping options
---
 compass.py |   17 +++++++++--------
 1 file changed, 9 insertions(+), 8 deletions(-)

diff --git a/compass.py b/compass.py
index c55fe98..283424e 100755
--- a/compass.py
+++ b/compass.py
@@ -237,14 +237,15 @@ class RelayStats(object):
         return filters
 
     def _get_group_function(self, options):
-        if options.by_country and options.by_as:
-            return lambda relay: (relay.get('country', None), relay.get('as_number', None))
-        elif options.by_country:
-            return lambda relay: relay.get('country', None)
-        elif options.by_as:
-            return lambda relay: relay.get('as_number', None)
-        else:
-            return lambda relay: relay.get('fingerprint')
+        funcs = []
+        if options.by_country:
+            funcs.append(lambda relay: relay.get('country', None))
+        if options.by_as:
+            funcs.append(lambda relay: relay.get('as_number', None))
+        # Default on grouping by fingerprint
+        if len(funcs) == 0:
+            funcs.append(lambda relay: relay.get('fingerprint'))
+        return lambda relay: tuple([func(relay) for func in funcs])
 
     def add_relay(self, relay):
         key = self._get_group(relay)