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