from collections import defaultdict

class PartitionMonitor:
    def __init__(self):
        self.assignments = defaultdict(set)

    def on_partitions_assigned(self, consumer, partitions):
        self.assignments.clear()
        for tp in partitions:
            self.assignments[tp.topic].add(tp.partition)
