Committing internal C.C.P. changes to this module. These changes extend the scheduling interface to allow easy use of a pre-emptive approach.

Author: rmtew

libraries/stacklesslib/main.py

@@ -167,21 +167,21 @@ def wakeup_tasklets(self, time):
     def run_tasklets(self, run_for=0):
         """ Run tasklets for as long as necessary """
         try:
-            stackless.run(run_for)
+            return stackless.run(run_for)
         except Exception:
             self.handle_run_error(sys.exc_info())
 
     def handle_run_error(self, ei):
         traceback.print_exception(*ei)
 
-    def pump(self):
+    def pump(self, run_for=0):
         t = time.time()
         wait_time = self.get_wait_time(t)
         if wait_time:
             self.wait(wait_time)
             t = elapsed_time()
         self.wakeup_tasklets(t + 0.001) #fuzz
-        self.run_tasklets()
+        return self.run_tasklets(run_for=run_for)
 
     def run(self):
         while self.running:

libraries/stacklesslib/test/testmainloop.py

@@ -0,0 +1,64 @@
+import stackless
+import unittest
+import stacklesslib.main
+
+class TestMainLoop(unittest.TestCase):
+    def setUp(self):
+        pass
+ 
+    def tearDown(self):
+        pass
+
+    def checkLeftThingsClean(self):
+        self.assertEqual(len(stacklesslib.main.event_queue.queue_a), 0) 
+        self.assertEqual(len(stacklesslib.main.event_queue.queue_b), 0) 
+        return True
+
+    def testPreemptiveRun(self):
+        """
+        Create a tasklet and run it pre-emptively, ensuring that we
+        get the tasklet returned from 'run_tasklets' when it is
+        interrupted.
+        """
+   
+        t = stackless.tasklet(ArbitraryFunc)()
+        t.run()
+
+        while t.alive:
+            stacklesslib.main.mainloop.wakeup_tasklets(0)
+            ret = stacklesslib.main.mainloop.run_tasklets(100)
+            if ret is None and t.alive:
+                continue
+            self.assertEqual(ret, t)
+            break
+        else:
+            self.fail("Tasklet was not interrupted")
+
+        self.checkLeftThingsClean() # Boilerplate check. 
+
+    def testCooperativeRun(self):
+        """
+        Create a tasklet and run it cooperatively, ensuring we never
+        get it interrupted and returned from run_tasklets.
+        """
+    
+        t = stackless.tasklet(ArbitraryFunc)()
+        t.run()
+        
+        while t.alive:
+            stacklesslib.main.mainloop.wakeup_tasklets(0)
+            ret = stacklesslib.main.mainloop.run_tasklets()
+            self.assertFalse(ret)
+    
+        self.checkLeftThingsClean() # Boilerplate check. 
+
+
+def ArbitraryFunc():
+    sum = 0
+    for i in range(1000):
+        for j in range(1000):
+            sum += 10
+        stacklesslib.main.sleep(0)
+
+if __name__ == '__main__':
+    unittest.main()

libraries/stacklesslib/util.py

@@ -38,6 +38,20 @@ def block_trap(trap=True):
     finally:
         c.block_trap = old
 
+@contextlib.contextmanager
+def ignore_nesting(flag=True):
+    """
+    A context manager which allows the current tasklet to engage the
+    ignoring of nesting levels.  By default pre-emptive switching can
+    only happen at the top nesting level, setting this allows it to
+    happen at all nesting levels.  Defaults to setting it to True.
+    """
+    c = stackless.getcurrent()
+    old = c.set_ignore_nesting(flag)
+    try:
+        yield
+    finally:
+        c.set_ignore_nesting(old)
 
 class local(object):
     """Tasklet local storage.  Similar to threading.local"""