Skip to content

Commit

Permalink
Add option to remove event handler from luigi.Task
Browse files Browse the repository at this point in the history
  • Loading branch information
starhel committed Apr 9, 2024
1 parent 64d6c48 commit 411cacf
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 0 deletions.
7 changes: 7 additions & 0 deletions luigi/task.py
Original file line number Diff line number Diff line change
Expand Up @@ -264,6 +264,13 @@ def wrapped(callback):
return callback
return wrapped

@classmethod
def remove_event_handler(cls, event, callback):
"""
Function to remove the event handler registered previously by the cls.event_handler decorator.
"""
cls._event_callbacks[cls][event].remove(callback)

def trigger_event(self, event, *args, **kwargs):
"""
Trigger that calls all of the specified events associated with this class.
Expand Down
15 changes: 15 additions & 0 deletions test/event_callbacks_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -162,6 +162,21 @@ def test_processing_time_handler_failure(self):
t, result = self._run_processing_time_handler(True)
self.assertEqual(result, [])

def test_remove_event_handler(self):
run_cnt = 0

@EmptyTask.event_handler(luigi.Event.START)
def handler(task):
nonlocal run_cnt
run_cnt += 1

task = EmptyTask()
build([task], local_scheduler=True)
assert run_cnt == 1
EmptyTask.remove_event_handler(luigi.Event.START, handler)
build([task], local_scheduler=True)
assert run_cnt == 1


# A
# / \
Expand Down

0 comments on commit 411cacf

Please sign in to comment.