Handle both args cases

gianlucaborello · Nov 23, 2018 · 81de9a5 · 81de9a5
1 parent f4965d6
commit 81de9a5
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/cassandradump.py b/cassandradump.py
@@ -74,7 +74,11 @@ def make_value_encoders(tableval):
         return dict((to_utf8(k), make_value_encoder(cql_type(v))) for k, v in six.iteritems(tableval.columns))
 
     def make_row_encoder():
-        def type_selector(k, v):
+        def type_selector(*args):
+            if len(args) == 1:
+                (k, v) = args[0]
+            else:
+                (k, v) = args
             return cql_type(v) == 'counter'
 
         partitions = dict(