bedrock.external.decorators.websocket

 1import functools  # pragma: unit
 2
 3from bedrock.external.websocket import broadcast_websocket_message  # pragma: unit
 4
 5
 6class _WebsocketBroadcast:  # pragma: unit
 7    """
 8    The intended behaviour here is for the endpoint to receive a new attribute, __uses_websockets__, which is detected
 9    in our build scripts to determine whether to spin up websocket infrastructure.
10
11    When another decorator wraps a method before this decorator executes, it loses ownership of the endpoint class,
12    and hence cannot modify/add the __uses_websockets__ attribute.
13    """
14
15    def __init__(self, fn, topics: list[str]):
16        self.fn = fn
17        self.topics = topics
18        functools.update_wrapper(self, fn)
19
20    def __set_name__(self, owner, name):
21        owner.__uses_websockets__ = True
22        owner.__websocket_topics__ = list(set(owner.__websocket_topics__) | set(self.topics)) if hasattr(owner,
23                                                                                                         '__websocket_topics__') \
24            else self.topics
25
26    def __call__(self, *args, **kwargs):
27        return self.fn(*args, **kwargs)
28
29    def __get__(self, instance, owner=None):
30        if instance is None:
31            return self
32
33        @functools.wraps(self.fn)
34        def wrapper(*args, **kwargs):
35            status, content = self.fn(instance, *args, **kwargs)
36
37            if 200 <= int(status) < 300 and self.topics:
38                related_model = getattr(instance, "related_model", None)
39                broadcast_websocket_message(content, self.topics, related_model)
40
41            return status, content
42
43        return wrapper
44
45
46def websocket_broadcast(topics: str | list[str]):  # pragma: unit
47    """
48    WARNING: When multiple decorators are used to wrap an endpoint method, you need this decorator to be executed last.
49    To do this you must place it visually at the top of the list of decorator you use for a method.
50
51    For example:
52    ```python
53    @websocket_broadcast("entity") <- Place decorator at the top here (executed last)
54    @protected()
55    def method_a():
56        pass
57    ```
58
59    This behaviour is being verified in `tests/unit/external/decorators/test_websocket.py`
60    """
61    topic_list = [] if not topics else [topics] if isinstance(topics, str) else topics
62
63    def decorator(fn):
64        return _WebsocketBroadcast(fn, topic_list)
65
66    return decorator