def _set_assignment(self, partition_states, randomize=False):
        """Batch partition assignment by topic (self.assignment is OrderedDict)"""
        self.assignment.clear()
        topics = [tp.topic for tp in six.iterkeys(partition_states)]
        if randomize:
            random.shuffle(topics)
        topic_partitions = OrderedDict({topic: [] for topic in topics})
        for tp in six.iterkeys(partition_states):
            topic_partitions[tp.topic].append(tp)
        for topic in six.iterkeys(topic_partitions):
            for tp in topic_partitions[topic]:
                self.assignment[tp] = partition_states[tp]
Batch partition assignment by topic (self.assignment is OrderedDict)
kafka/consumer/subscription_state.py