/
events.py
232 lines (193 loc) · 8.63 KB
/
events.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
"""Event handler.
The :class:`EventHandler` just holds a list of functions
which can get called once a certain "event" happens.
Examples are given in the class doc-string.
"""
# Copyright (C) TeNPy Developers, GNU GPLv3
from collections import namedtuple
import warnings
from .hdf5_io import find_global
__all__ = ['Listener', 'EventHandler']
Listener = namedtuple('Listener', "listener_id, callback, priority, extra_kwargs")
class EventHandler:
"""Handler for an event represented by an instance of this class.
All in all, events provide a flexible extension mechanism for classes to define "checkpoints"
in the code where the user of a class might want to run something else,
for example doing some measurements or saving intermediate results.
Parameters
----------
arg_descr : str
An informative description how the callback function is called.
An empty string indicates no arguments.
Attributes
----------
arg_descr : str
An informative description how the callback function is called.
listeners : list of (int, function, int)
Entries are tuples ``(listener_id, callback, priority)``.
Examples
--------
Instances of this class typically get defined during class initialization and define an event.
The event "happens" each time :meth:`emit` or :meth:`emit_until_result` is called,
typically inside a method of the class defining the event. Example:
.. testsetup :: EventHandler
from tenpy.tools.events import EventHandler
.. doctest :: EventHandler
>>> class MyAlgorithm:
... def __init__(self):
... self.checkpoint = EventHandler("algorithm, iteration")
... self.data = 0
... def run(self):
... for i in range(4):
... self.data += i # do some complicated stuff
... self.checkpoint.emit(self, i)
Other code with access to the event can then connect a `listener` to the event, i.e.,
give a function to the event that should be called each time the event is :meth:`emit`-ed.
.. doctest :: EventHandler
>>> my_alg = MyAlgorithm()
>>> def my_listener(algorithm, iteration):
... print("my_listener called: iteration", iteration, "with data", algorithm.data)
>>> my_alg.checkpoint.connect(my_listener) # doctest: +ELLIPSIS
<function my_listener at 0x...>
>>> my_alg.run()
my_listener called: iteration 0 with data 0
my_listener called: iteration 1 with data 1
my_listener called: iteration 2 with data 3
my_listener called: iteration 3 with data 6
As you can see, the function `my_listener` has been called during the ``MyAlgorithm.run()``
and had full access to the current status of the algorithm class.
This is convenient to e.g. perform measurements of the state so far, print a status message
of the progress or save intermediate results.
If the EventHandler is already initialized when you define the function, you can also
use :meth:`connect` as a function property like this:
.. doctest :: EventHandler
>>> @my_alg.checkpoint.connect
... def another_one(algorithm, iteration):
... print("another_one called: iteration", iteration)
>>> @my_alg.checkpoint.connect(priority=5)
... def high_priority(algorithm, iteration):
... print("high_priority call: iteration", iteration)
>>> my_alg.run()
high_priority call: iteration 0
my_listener called: iteration 0 with data 6
another_one called: iteration 0
high_priority call: iteration 1
my_listener called: iteration 1 with data 7
another_one called: iteration 1
high_priority call: iteration 2
my_listener called: iteration 2 with data 9
another_one called: iteration 2
high_priority call: iteration 3
my_listener called: iteration 3 with data 12
another_one called: iteration 3
"""
def __init__(self, arg_descr=None):
self.arg_descr = arg_descr
self.listeners = []
self._id_counter = 0
def copy(self):
"""Make a (shallow) copy."""
cp = EventHandler(self.arg_descr)
cp.listeners = self.listeners[:]
cp._id_counter = self._id_counter
return cp
@property
def id_of_last_connected(self):
if self._id_counter == 0:
raise ValueError("connect() hasn't been called yet!'")
return self._id_counter - 1
def connect(self, callback=None, priority=0, extra_kwargs=None):
"""Register a `callback` function as a listener to the event.
You can either call this function directly or use it as a function decorator,
see the example in :class:`EventHandler`.
If you ever plan to :meth:`disconnect` again, you can read it out with
:attr:`id_of_last_connected` right after connecting, i.e., right after calling this method.
Parameters
----------
callback : callable
A function to be called during each :meth:`emit` of the event.
priority : int
Higher priority indicates that the callback function should be called before other
possibly registered callback functions.
extra_kwargs : None | dict
Optional extra keyword arguments given directly to the function.
``None`` is equivalent to ``{}``.
Returns
-------
callback : callable
The callback function exactly as given.
"""
if callback is None:
# handle the case that we got called as property like this::
# @ev_handler.connect(priority=2)
# def my_function():
# pass
def property(callback):
self.connect(callback, priority)
return callback
return property
listener_id = self._id_counter
self._id_counter += 1
if extra_kwargs is None:
extra_kwargs = {}
self.listeners.append(Listener(listener_id, callback, priority, extra_kwargs))
return callback
def connect_by_name(self, module_name, func_name, extra_kwargs=None, priority=0):
"""Connect to a function given by the name in a module, optionally inserting arguments.
Parameters
----------
module_name : str
The name of the module containing the function to be used. Gets imported.
func_name : str
The (qualified) name of the function inside the module.
extra_kwargs : dict
Optional extra keyword-arguments to be given to the function.
``None`` is equivalent to ``{}``.
priority : int
Higher priority indicates that the callback function should be called before other
possibly registered callback functions.
"""
func = find_global(module_name, func_name)
self.connect(func, priority, extra_kwargs)
def disconnect(self, listener_id):
"""De-register a listener.
Parameters
----------
listener_id : int
The id of the listener, as given by :attr:`id_of_last_connected`
right after calling :meth:`connect`.
"""
for i, listener in enumerate(self.listeners):
if listener.listener_id == 0:
del self.listeners[i]
return
warnings.warn("No listener with listener_id {id_:d} found".format(id_=listener_id))
def emit(self, *args, **kwargs):
"""Call the `callback` functions of all listeners.
Parameters
----------
*args, **kwargs:
(Keyword) arguments to give to the callback functions.
Recall that there can also be additional `extra_kwargs` for each callback function.
Returns
-------
results : list
List of results returned by the individual callback functions.
"""
self._prepare_emit()
results = []
for _, callback, _, extra_kwargs in self.listeners:
res = callback(*args, **kwargs, **extra_kwargs)
results.append(res)
return results
def emit_until_result(self, *args, **kwargs):
"""Call the listeners `callback` until one returns not `None`."""
self._prepare_emit()
for _, callback, _, extra_kwargs in self.listeners:
res = callback(*args, **kwargs, **extra_kwargs)
if res is not None:
return res
return None
def _prepare_emit(self):
# sort listeners: highest priority first
self.listeners = sorted(self.listeners, key=lambda listener: -listener.priority)