|
25 | 25 | from enum import IntEnum |
26 | 26 | from heapq import heapify, heappop, heappush, nsmallest |
27 | 27 | from types import MethodType |
28 | | -from typing import Any |
| 28 | +from typing import TYPE_CHECKING, Any |
29 | 29 | from weakref import WeakMethod, ref |
30 | 30 |
|
| 31 | +if TYPE_CHECKING: |
| 32 | + from mesa import Model |
| 33 | + |
31 | 34 |
|
32 | 35 | class Priority(IntEnum): |
33 | 36 | """Enumeration of priority levels.""" |
@@ -123,6 +126,181 @@ def __lt__(self, other): # noqa |
123 | 126 | ) |
124 | 127 |
|
125 | 128 |
|
| 129 | +class EventGenerator: |
| 130 | + """A generator that creates recurring events at specified intervals. |
| 131 | +
|
| 132 | + EventGenerator represents a pattern for when things should happen repeatedly. |
| 133 | + Unlike a single SimulationEvent, an EventGenerator is persistent and can be |
| 134 | + stopped or configured with stop conditions. |
| 135 | +
|
| 136 | + Attributes: |
| 137 | + model: The model this generator belongs to |
| 138 | + function: The callable to execute for each generated event |
| 139 | + interval: Time between events (fixed value or callable returning value) |
| 140 | + priority: Priority level for generated events |
| 141 | +
|
| 142 | + """ |
| 143 | + |
| 144 | + def __init__( |
| 145 | + self, |
| 146 | + model: Model, |
| 147 | + function: Callable, |
| 148 | + interval: float | int | Callable[[Model], float | int], |
| 149 | + priority: Priority = Priority.DEFAULT, |
| 150 | + ) -> None: |
| 151 | + """Initialize an EventGenerator. |
| 152 | +
|
| 153 | + Args: |
| 154 | + model: The model this generator belongs to |
| 155 | + function: The callable to execute for each generated event. |
| 156 | + Use functools.partial to bind arguments. |
| 157 | + interval: Time between events. Can be a fixed value or a callable |
| 158 | + that takes the model and returns the interval. |
| 159 | + priority: Priority level for generated events |
| 160 | +
|
| 161 | + """ |
| 162 | + self.model = model |
| 163 | + self.function = function |
| 164 | + self.interval = interval |
| 165 | + self.priority = priority |
| 166 | + |
| 167 | + self._active: bool = False |
| 168 | + self._current_event: SimulationEvent | None = None |
| 169 | + self._execution_count: int = 0 |
| 170 | + |
| 171 | + # Stop conditions (mutually exclusive) |
| 172 | + self._max_count: int | None = None |
| 173 | + self._end_time: float | None = None |
| 174 | + |
| 175 | + @property |
| 176 | + def is_active(self) -> bool: |
| 177 | + """Return whether the generator is currently active.""" |
| 178 | + return self._active |
| 179 | + |
| 180 | + @property |
| 181 | + def execution_count(self) -> int: |
| 182 | + """Return the number of times this generator has executed.""" |
| 183 | + return self._execution_count |
| 184 | + |
| 185 | + def _get_interval(self) -> float | int: |
| 186 | + """Get the next interval value.""" |
| 187 | + if callable(self.interval): |
| 188 | + return self.interval(self.model) |
| 189 | + return self.interval |
| 190 | + |
| 191 | + def _should_stop(self, next_time: float) -> bool: |
| 192 | + """Check if the generator should stop before scheduling the next event.""" |
| 193 | + return ( |
| 194 | + self._max_count is not None and self._execution_count >= self._max_count |
| 195 | + ) or (self._end_time is not None and next_time > self._end_time) |
| 196 | + |
| 197 | + def _execute_and_reschedule(self) -> None: |
| 198 | + """Execute the function and schedule the next event.""" |
| 199 | + if not self._active: |
| 200 | + return |
| 201 | + |
| 202 | + self.function() |
| 203 | + self._execution_count += 1 |
| 204 | + |
| 205 | + # Schedule next event if we shouldn't stop |
| 206 | + next_time = self.model.time + self._get_interval() |
| 207 | + if not self._should_stop(next_time): |
| 208 | + self._schedule_next(next_time) |
| 209 | + else: |
| 210 | + self._active = False |
| 211 | + self._current_event = None |
| 212 | + |
| 213 | + def _schedule_next(self, time: float) -> None: |
| 214 | + """Schedule the next event at the given time.""" |
| 215 | + self._current_event = SimulationEvent( |
| 216 | + time, |
| 217 | + self._execute_and_reschedule, |
| 218 | + priority=self.priority, |
| 219 | + ) |
| 220 | + self.model._simulator.event_list.add_event(self._current_event) |
| 221 | + |
| 222 | + def start( |
| 223 | + self, |
| 224 | + at: float | None = None, |
| 225 | + after: float | None = None, |
| 226 | + ) -> EventGenerator: |
| 227 | + """Start the event generator. |
| 228 | +
|
| 229 | + Args: |
| 230 | + at: Absolute time to start generating events |
| 231 | + after: Relative time from now to start generating events |
| 232 | +
|
| 233 | + Returns: |
| 234 | + Self for method chaining |
| 235 | +
|
| 236 | + Raises: |
| 237 | + ValueError: If both `at` and `after` are specified |
| 238 | +
|
| 239 | + """ |
| 240 | + if self._active: |
| 241 | + return self |
| 242 | + |
| 243 | + if at is not None and after is not None: |
| 244 | + raise ValueError("Cannot specify both 'at' and 'after'") |
| 245 | + |
| 246 | + if at is None and after is None: |
| 247 | + # Default: start at next interval from now |
| 248 | + start_time = self.model.time + self._get_interval() |
| 249 | + elif at is not None: |
| 250 | + if at < self.model.time: |
| 251 | + raise ValueError(f"Cannot start in the past: {at} < {self.model.time}") |
| 252 | + start_time = at |
| 253 | + else: # after is not None |
| 254 | + if after < 0: |
| 255 | + raise ValueError(f"Cannot start in the past: after={after}") |
| 256 | + start_time = self.model.time + after |
| 257 | + |
| 258 | + self._active = True |
| 259 | + self._schedule_next(start_time) |
| 260 | + return self |
| 261 | + |
| 262 | + def stop( |
| 263 | + self, |
| 264 | + at: float | None = None, |
| 265 | + after: float | None = None, |
| 266 | + count: int | None = None, |
| 267 | + ) -> EventGenerator: |
| 268 | + """Stop the event generator. |
| 269 | +
|
| 270 | + Args: |
| 271 | + at: Absolute time to stop generating events |
| 272 | + after: Relative time from now to stop generating events |
| 273 | + count: Number of additional executions before stopping |
| 274 | +
|
| 275 | + Returns: |
| 276 | + Self for method chaining |
| 277 | +
|
| 278 | + Raises: |
| 279 | + ValueError: If more than one stop condition is specified |
| 280 | +
|
| 281 | + """ |
| 282 | + conditions = sum(x is not None for x in [at, after, count]) |
| 283 | + if conditions > 1: |
| 284 | + raise ValueError("Can only specify one of 'at', 'after', or 'count'") |
| 285 | + |
| 286 | + if conditions == 0: |
| 287 | + # Immediate stop |
| 288 | + self._active = False |
| 289 | + if self._current_event is not None: |
| 290 | + self._current_event.cancel() |
| 291 | + self._current_event = None |
| 292 | + return self |
| 293 | + |
| 294 | + if at is not None: |
| 295 | + self._end_time = at |
| 296 | + elif after is not None: |
| 297 | + self._end_time = self.model.time + after |
| 298 | + elif count is not None: |
| 299 | + self._max_count = self._execution_count + count |
| 300 | + |
| 301 | + return self |
| 302 | + |
| 303 | + |
126 | 304 | class EventList: |
127 | 305 | """An event list. |
128 | 306 |
|
|
0 commit comments