1from contextlib import ContextDecorator 2from operator import itemgetter 3 4from wagtail.utils.apps import get_app_submodules 5 6 7_hooks = {} 8 9 10def register(hook_name, fn=None, order=0): 11 """ 12 Register hook for ``hook_name``. Can be used as a decorator:: 13 14 @register('hook_name') 15 def my_hook(...): 16 pass 17 18 or as a function call:: 19 20 def my_hook(...): 21 pass 22 register('hook_name', my_hook) 23 """ 24 25 # Pretend to be a decorator if fn is not supplied 26 if fn is None: 27 def decorator(fn): 28 register(hook_name, fn, order=order) 29 return fn 30 return decorator 31 32 if hook_name not in _hooks: 33 _hooks[hook_name] = [] 34 _hooks[hook_name].append((fn, order)) 35 36 37class TemporaryHook(ContextDecorator): 38 def __init__(self, hooks, order): 39 self.hooks = hooks 40 self.order = order 41 42 def __enter__(self): 43 for hook_name, fn in self.hooks: 44 if hook_name not in _hooks: 45 _hooks[hook_name] = [] 46 _hooks[hook_name].append((fn, self.order)) 47 48 def __exit__(self, exc_type, exc_value, traceback): 49 for hook_name, fn in self.hooks: 50 _hooks[hook_name].remove((fn, self.order)) 51 52 53def register_temporarily(hook_name_or_hooks, fn=None, *, order=0): 54 """ 55 Register hook for ``hook_name`` temporarily. This is useful for testing hooks. 56 57 Can be used as a decorator:: 58 59 def my_hook(...): 60 pass 61 62 class TestMyHook(Testcase): 63 @hooks.register_temporarily('hook_name', my_hook) 64 def test_my_hook(self): 65 pass 66 67 or as a context manager:: 68 69 def my_hook(...): 70 pass 71 72 with hooks.register_temporarily('hook_name', my_hook): 73 # Hook is registered here 74 75 # Hook is unregistered here 76 77 To register multiple hooks at the same time, pass in a list of 2-tuples: 78 79 def my_hook(...): 80 pass 81 82 def my_other_hook(...): 83 pass 84 85 with hooks.register_temporarily([ 86 ('hook_name', my_hook), 87 ('hook_name', my_other_hook), 88 ]): 89 # Hooks are registered here 90 """ 91 if not isinstance(hook_name_or_hooks, list) and fn is not None: 92 hooks = [(hook_name_or_hooks, fn)] 93 else: 94 hooks = hook_name_or_hooks 95 96 return TemporaryHook(hooks, order) 97 98 99_searched_for_hooks = False 100 101 102def search_for_hooks(): 103 global _searched_for_hooks 104 if not _searched_for_hooks: 105 list(get_app_submodules('wagtail_hooks')) 106 _searched_for_hooks = True 107 108 109def get_hooks(hook_name): 110 """ Return the hooks function sorted by their order. """ 111 search_for_hooks() 112 hooks = _hooks.get(hook_name, []) 113 hooks = sorted(hooks, key=itemgetter(1)) 114 return [hook[0] for hook in hooks] 115