diff --git a/src/tqdm_publisher/publisher.py b/src/tqdm_publisher/publisher.py index 8085f4f..a5830d5 100644 --- a/src/tqdm_publisher/publisher.py +++ b/src/tqdm_publisher/publisher.py @@ -9,8 +9,9 @@ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.callbacks = {} + # Override the update method to call callbacks - def update(self, n=1, always_callback=False): + def update(self, n=1): displayed = super().update(n) for id in list(self.callbacks): callback = self.callbacks.get(id) @@ -19,16 +20,79 @@ def update(self, n=1, always_callback=False): return displayed + def subscribe(self, callback: callable): + """ + Subscribe to updates from the progress bar. + + This method assigns a unique ID to the given callback and stores it in an internal + dictionary. It allows the callback to be referenced and activated in the overridden + update function for TQDM. The unique callback ID is returned, which can be used + for future operations such as deregistering the callback. + + Parameters + ---------- + callback : callable + A callable object (like a function) that will be called back by this object. + The callback should be able to be invoked with a single argument, the progress + bar's format_dict. + + Returns + ------- + str + A unique identifier for the callback. This ID is a UUID string and can be used + to reference the registered callback in future operations. + + Examples + -------- + >>> def my_callback(format_dict): + >>> print("Progress update", format_dict) + >>> + >>> publisher = TQDMPublisher() + >>> callback_id = publisher.subscribe(my_callback) + >>> print(callback_id) # Prints the unique callback ID + """ - # Subscribe to updates - def subscribe(self, callback): callback_id = str(uuid4()) self.callbacks[callback_id] = callback return callback_id - # Unsubscribe from updates - def unsubscribe(self, callback_id): + def unsubscribe(self, callback_id: str): + """ + Unsubscribe a previously registered callback from the progress bar updates. + + This method removes the callback associated with the given unique ID from the internal + dictionary. It is used to deregister callbacks that were previously added via the + `subscribe` method. Once a callback is removed, it will no longer be called during + the progress bar's update events. + + Parameters + ---------- + callback_id : str + The unique identifier of the callback to be unsubscribed. This is the same UUID string + that was returned by the `subscribe` method when the callback was registered. + + Returns + ------- + bool + Returns True if the callback was successfully removed, or False if no callback was + found with the given ID. + + Examples + -------- + >>> publisher = TQDMPublisher() + >>> callback_id = publisher.subscribe(my_callback) + >>> print(callback_id) # Prints the unique callback ID + >>> unsubscribed = publisher.unsubscribe(callback_id) + >>> print(unsubscribed) # True if successful, False otherwise + + Notes + ----- + It's important to manage the lifecycle of callbacks, especially in cases where the + number of subscribers might grow large or change frequently. Unsubscribing callbacks + when they are no longer needed can help prevent memory leaks and other performance issues. + """ if callback_id not in self.callbacks: return False + del self.callbacks[callback_id] return True \ No newline at end of file