efro.call

 1# Released under the MIT License. See LICENSE for details.
 2#
 3"""Call related functionality shared between all efro components."""
 4
 5from __future__ import annotations
 6
 7import weakref
 8import threading
 9from typing import TYPE_CHECKING, TypeVar, Generic
10
11T = TypeVar('T')
12
13if TYPE_CHECKING:
14    from typing import Callable
15
16
17class CallbackSet(Generic[T]):
18    """A simple way to manage a set of callbacks.
19
20    Any number of calls can be registered with a callback set. Each
21    registration results in a Registration object that can be used to
22    deregister the call from the set later. Callbacks are also
23    implicitly deregistered when an entry is deallocated, so make sure
24    to hold on to the return value when adding.
25
26    CallbackSet instances should be used from a single thread only
27    (this will be checked in debug mode).
28    """
29
30    def __init__(self) -> None:
31        self._entries: list[weakref.ref[CallbackRegistration[T]]] = []
32        self.thread: threading.Thread
33        if __debug__:
34            self.thread = threading.current_thread()
35
36    def register(self, call: T) -> CallbackRegistration[T]:
37        """Register a new callback."""
38        assert threading.current_thread() == self.thread
39
40        self._prune()
41
42        entry = CallbackRegistration(call, self)
43        self._entries.append(weakref.ref(entry))
44        return entry
45
46    def getcalls(self) -> list[T]:
47        """Return the current set of registered calls.
48
49        Note that this returns a flattened list of calls; generally this
50        should protect against calls which themselves add or remove
51        callbacks.
52        """
53        assert threading.current_thread() == self.thread
54
55        self._prune()
56
57        # Ignore calls that have been deallocated or explicitly
58        # deregistered.
59        entries = [e() for e in self._entries]
60        return [e.call for e in entries if e is not None and e.call is not None]
61
62    def _prune(self) -> None:
63
64        # Quick-out if all our entries are intact.
65        needs_prune = False
66        for entry in self._entries:
67            entrytarget = entry()
68            if entrytarget is None or entrytarget.call is None:
69                needs_prune = True
70                break
71        if not needs_prune:
72            return
73
74        # Ok; something needs pruning. Rebuild the entries list.
75        newentries: list[weakref.ref[CallbackRegistration[T]]] = []
76        for entry in self._entries:
77            entrytarget = entry()
78            if entrytarget is not None and entrytarget.call is not None:
79                newentries.append(entry)
80        self._entries = newentries
81
82
83class CallbackRegistration(Generic[T]):
84    """An entry for a callback set."""
85
86    def __init__(self, call: T, callbackset: CallbackSet[T]) -> None:
87        self.call: T | None = call
88        self.callbackset: CallbackSet[T] | None = callbackset
89
90    def deregister(self) -> None:
91        """Explicitly remove a callback from a CallbackSet."""
92        assert (
93            self.callbackset is None
94            or threading.current_thread() == self.callbackset.thread
95        )
96        # Simply clear the call to mark us as dead.
97        self.call = None