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