1ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh"""A generally useful event scheduler class.
2ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
3ffab958fd8d42ed7227d83007350e61555a1fa36Andrew HsiehEach instance of this class manages its own queue.
4ffab958fd8d42ed7227d83007350e61555a1fa36Andrew HsiehNo multi-threading is implied; you are supposed to hack that
5ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehyourself, or use a single instance per application.
6ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
7ffab958fd8d42ed7227d83007350e61555a1fa36Andrew HsiehEach instance is parametrized with two functions, one that is
8ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehsupposed to return the current time, one that is supposed to
9ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehimplement a delay.  You can implement real-time scheduling by
10ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehsubstituting time and sleep from built-in module time, or you can
11ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehimplement simulated time by writing your own functions.  This can
12ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehalso be used to integrate scheduling with STDWIN events; the delay
13ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehfunction is allowed to modify the queue.  Time can be expressed as
14ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehintegers or floating point numbers, as long as it is consistent.
15ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
16ffab958fd8d42ed7227d83007350e61555a1fa36Andrew HsiehEvents are specified by tuples (time, priority, action, argument).
17ffab958fd8d42ed7227d83007350e61555a1fa36Andrew HsiehAs in UNIX, lower priority numbers mean higher priority; in this
18ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehway the queue can be maintained as a priority queue.  Execution of the
19ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehevent means calling the action function, passing it the argument
20ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehsequence in "argument" (remember that in Python, multiple function
21ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieharguments are be packed in a sequence).
22ffab958fd8d42ed7227d83007350e61555a1fa36Andrew HsiehThe action function may be an instance method so it
23ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehhas another way to reference private data (besides global variables).
24ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh"""
25ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
26ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh# XXX The timefunc and delayfunc should have been defined as methods
27ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh# XXX so you can define new kinds of schedulers using subclassing
28ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh# XXX instead of having to define a module or class just to hold
29ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh# XXX the global state of your particular time and delay functions.
30ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
31ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehimport heapq
32ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehfrom collections import namedtuple
33ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
34ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh__all__ = ["scheduler"]
35ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
36ffab958fd8d42ed7227d83007350e61555a1fa36Andrew HsiehEvent = namedtuple('Event', 'time, priority, action, argument')
37ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
38ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsiehclass scheduler:
39ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    def __init__(self, timefunc, delayfunc):
40ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """Initialize a new instance, passing the time and delay
41ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        functions"""
42ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        self._queue = []
43ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        self.timefunc = timefunc
44ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        self.delayfunc = delayfunc
45ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
46ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    def enterabs(self, time, priority, action, argument):
47ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """Enter a new event in the queue at an absolute time.
48ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
49ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        Returns an ID for the event which can be used to remove it,
50ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        if necessary.
51ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
52ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """
53ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        event = Event(time, priority, action, argument)
54ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        heapq.heappush(self._queue, event)
55ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        return event # The ID
56ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
57ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    def enter(self, delay, priority, action, argument):
58ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """A variant that specifies the time as a relative time.
59ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
60ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        This is actually the more commonly used interface.
61ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
62ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """
63ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        time = self.timefunc() + delay
64ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        return self.enterabs(time, priority, action, argument)
65ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
66ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    def cancel(self, event):
67ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """Remove an event from the queue.
68ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
69ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        This must be presented the ID as returned by enter().
70ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        If the event is not in the queue, this raises ValueError.
71ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
72ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """
73ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        self._queue.remove(event)
74ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        heapq.heapify(self._queue)
75ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
76ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    def empty(self):
77ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """Check whether the queue is empty."""
78ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        return not self._queue
79ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
80ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    def run(self):
81ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """Execute events until the queue is empty.
82ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
83ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        When there is a positive delay until the first event, the
84ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        delay function is called and the event is left in the queue;
85ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        otherwise, the event is removed from the queue and executed
86ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        (its action function is called, passing it the argument).  If
87ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        the delay function returns prematurely, it is simply
88ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        restarted.
89ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
90ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        It is legal for both the delay function and the action
91ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        function to modify the queue or to raise an exception;
92ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        exceptions are not caught but the scheduler's state remains
93ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        well-defined so run() may be called again.
94ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
95ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        A questionable hack is added to allow other threads to run:
96ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        just after an event is executed, a delay of 0 is executed, to
97ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        avoid monopolizing the CPU when other threads are also
98ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        runnable.
99ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
100ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """
101ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        # localize variable access to minimize overhead
102ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        # and to improve thread safety
103ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        q = self._queue
104ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        delayfunc = self.delayfunc
105ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        timefunc = self.timefunc
106ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        pop = heapq.heappop
107ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        while q:
108ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            time, priority, action, argument = checked_event = q[0]
109ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            now = timefunc()
110ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            if now < time:
111ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                delayfunc(time - now)
112ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            else:
113ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                event = pop(q)
114ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                # Verify that the event was not removed or altered
115ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                # by another thread after we last looked at q[0].
116ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                if event is checked_event:
117ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    action(*argument)
118ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    delayfunc(0)   # Let other threads run
119ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                else:
120ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh                    heapq.heappush(q, event)
121ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
122ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    @property
123ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh    def queue(self):
124ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """An ordered list of upcoming events.
125ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
126ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        Events are named tuples with fields for:
127ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh            time, priority, action, arguments
128ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh
129ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        """
130ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        # Use heapq to sort the queue rather than using 'sorted(self._queue)'.
131ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        # With heapq, two events scheduled at the same time will show in
132ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        # the actual order they would be retrieved.
133ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        events = self._queue[:]
134ffab958fd8d42ed7227d83007350e61555a1fa36Andrew Hsieh        return map(heapq.heappop, [events]*len(events))
135