make partitions arg optional in *Partitioner.partition, defaulting to self.partitions from init

Author: Dana Powers

kafka/partitioner/base.py

@@ -12,14 +12,13 @@ def __init__(self, partitions):
         """
         self.partitions = partitions
 
-    def partition(self, key, partitions):
+    def partition(self, key, partitions=None):
         """
         Takes a string key and num_partitions as argument and returns
         a partition to be used for the message
 
         Arguments:
-            partitions: The list of partitions is passed in every call. This
-                may look like an overhead, but it will be useful
-                (in future) when we handle cases like rebalancing
+            key: the key to use for partitioning
+            partitions: (optional) a list of partitions.
         """
         raise NotImplementedError('partition function has to be implemented')

kafka/partitioner/hashed.py

@@ -5,7 +5,9 @@ class HashedPartitioner(Partitioner):
     Implements a partitioner which selects the target partition based on
     the hash of the key
     """
-    def partition(self, key, partitions):
+    def partition(self, key, partitions=None):
+        if not partitions:
+            partitions = self.partitions
         size = len(partitions)
         idx = hash(key) % size
 

kafka/partitioner/roundrobin.py

@@ -15,9 +15,9 @@ def _set_partitions(self, partitions):
         self.partitions = partitions
         self.iterpart = cycle(partitions)
 
-    def partition(self, key, partitions):
+    def partition(self, key, partitions=None):
         # Refresh the partition list if necessary
-        if self.partitions != partitions:
+        if partitions and self.partitions != partitions:
             self._set_partitions(partitions)
 
         return next(self.iterpart)