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

This commit is contained in:
Dana Powers
2015-02-06 11:57:49 -08:00
parent bb731eb89c
commit fb03c69544
3 changed files with 8 additions and 7 deletions

View File

@@ -12,14 +12,13 @@ class Partitioner(object):
"""
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')

View File

@@ -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

View File

@@ -15,9 +15,9 @@ class RoundRobinPartitioner(Partitioner):
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)