task generators

dtool/src/pystub/pystub.cxx

@@ -30,6 +30,7 @@ extern "C" {
   EXPCL_DTOOLCONFIG int PyDict_Size(...);
   EXPCL_DTOOLCONFIG int PyDict_Type(...);
   EXPCL_DTOOLCONFIG int PyErr_Clear(...);
+  EXPCL_DTOOLCONFIG int PyErr_ExceptionMatches(...);
   EXPCL_DTOOLCONFIG int PyErr_Fetch(...);
   EXPCL_DTOOLCONFIG int PyErr_Format(...);
   EXPCL_DTOOLCONFIG int PyErr_Occurred(...);
@@ -39,11 +40,11 @@ extern "C" {
   EXPCL_DTOOLCONFIG int PyEval_InitThreads(...);
   EXPCL_DTOOLCONFIG int PyEval_RestoreThread(...);
   EXPCL_DTOOLCONFIG int PyEval_SaveThread(...);
-  EXPCL_DTOOLCONFIG int PyExc_TypeError(...);
-  EXPCL_DTOOLCONFIG int PyExc_ValueError(...);
   EXPCL_DTOOLCONFIG int PyFloat_AsDouble(...);
   EXPCL_DTOOLCONFIG int PyFloat_FromDouble(...);
   EXPCL_DTOOLCONFIG int PyFloat_Type(...);
+  EXPCL_DTOOLCONFIG int PyGen_Check(...);
+  EXPCL_DTOOLCONFIG int PyGen_Type(...);
   EXPCL_DTOOLCONFIG int PyGILState_Ensure(...);
   EXPCL_DTOOLCONFIG int PyGILState_Release(...);
   EXPCL_DTOOLCONFIG int PyInt_AsLong(...);
@@ -119,6 +120,9 @@ extern "C" {
   EXPCL_DTOOLCONFIG int _Py_RefTotal(...);
 
   EXPCL_DTOOLCONFIG extern void *PyExc_AssertionError;
+  EXPCL_DTOOLCONFIG extern void *PyExc_StopIteration;
+  EXPCL_DTOOLCONFIG extern void *PyExc_TypeError;
+  EXPCL_DTOOLCONFIG extern void *PyExc_ValueError;
   EXPCL_DTOOLCONFIG extern void *_Py_NoneStruct;
   EXPCL_DTOOLCONFIG extern void *_Py_NotImplementedStruct;
 };
@@ -139,6 +143,7 @@ int PyDict_SetItemString(...) { return 0; };
 int PyDict_Size(...){ return 0; }
 int PyDict_Type(...) { return 0; };
 int PyErr_Clear(...) { return 0; };
+int PyErr_ExceptionMatches(...) { return 0; };
 int PyErr_Fetch(...) { return 0; }
 int PyErr_Format(...) { return 0; };
 int PyErr_Occurred(...) { return 0; }
@@ -148,11 +153,11 @@ int PyErr_SetString(...) { return 0; }
 int PyEval_InitThreads(...) { return 0; }
 int PyEval_RestoreThread(...) { return 0; }
 int PyEval_SaveThread(...) { return 0; }
-int PyExc_TypeError(...) { return 0; }
-int PyExc_ValueError(...) { return 0; }
 int PyFloat_AsDouble(...) { return 0; }
 int PyFloat_FromDouble(...) { return 0; }
 int PyFloat_Type(...) { return 0; }
+int PyGen_Check(...) { return 0; }
+int PyGen_Type(...) { return 0; }
 int PyGILState_Ensure(...) { return 0; }
 int PyGILState_Release(...) { return 0; }
 int PyInt_AsLong(...) { return 0; }
@@ -228,8 +233,10 @@ int _Py_NegativeRefcount(...) { return 0; };
 int _Py_RefTotal(...) { return 0; };
 
 
-
 void *PyExc_AssertionError = (void *)NULL;
+void *PyExc_StopIteration = (void *)NULL;
+void *PyExc_TypeError = (void *)NULL;
+void *PyExc_ValueError = (void *)NULL;
 void *_Py_NoneStruct = (void *)NULL;
 void *_Py_NotImplementedStruct = (void *)NULL;
 

panda/src/event/asyncTask.cxx

@@ -346,6 +346,14 @@ unlock_and_do_task() {
 //               DS_abort: abort the task, and interrupt the whole
 //               AsyncTaskManager.
 //
+//               DS_restart: like DS_cont, but next time call the
+//               function from the beginning.  This only has meaning
+//               to a PythonTask that has already used the yield
+//               expression to return a generator, in which case it
+//               provides a way to abort the generator and create a
+//               new one by calling the function again.  In other
+//               contexts, this behaves exactly the same as DS_cont.
+//
 //               This function is called with the lock *not* held.
 ////////////////////////////////////////////////////////////////////
 AsyncTask::DoneStatus AsyncTask::

panda/src/event/asyncTask.h

@@ -53,6 +53,7 @@ PUBLISHED:
     DS_cont,   // run task again next epoch
     DS_again,  // run task again after get_delay() seconds
     DS_abort,  // abort the task and interrupt the whole task manager
+    DS_restart, // like cont, but next time start the task from the beginning.  Only meaningful for a PythonTask which has yielded a generator.
   };
 
   enum State {

panda/src/event/asyncTaskChain.cxx

@@ -670,6 +670,7 @@ service_one_task(AsyncTaskChain::AsyncTaskChainThread *thread) {
       } else {
         switch (ds) {
         case AsyncTask::DS_cont:
+        case AsyncTask::DS_restart:
           // The task is still alive; put it on the next frame's active
           // queue.
           task->_state = AsyncTask::S_active;

panda/src/event/pythonTask.cxx

@@ -37,6 +37,7 @@ PythonTask(PyObject *function, const string &name) :
   _args = NULL;
   _upon_death = NULL;
   _owner = NULL;
+  _generator = NULL;
 
   set_function(function);
   set_args(Py_None, true);
@@ -62,6 +63,7 @@ PythonTask::
   Py_DECREF(_function);
   Py_DECREF(_args);
   Py_DECREF(_dict);
+  Py_XDECREF(_generator);
 }
 
 ////////////////////////////////////////////////////////////////////
@@ -295,12 +297,6 @@ PyObject *PythonTask::
 __getattr__(const string &attr_name) const {
   if (attr_name == "time") {
     return PyFloat_FromDouble(get_elapsed_time());
-  } else if (attr_name == "done") {
-    return PyInt_FromLong(DS_done);
-  } else if (attr_name == "cont") {
-    return PyInt_FromLong(DS_cont);
-  } else if (attr_name == "again") {
-    return PyInt_FromLong(DS_again);
   } else if (attr_name == "name") {
     return PyString_FromString(get_name().c_str());
   } else if (attr_name == "id") {
@@ -317,10 +313,48 @@ __getattr__(const string &attr_name) const {
 ////////////////////////////////////////////////////////////////////
 AsyncTask::DoneStatus PythonTask::
 do_task() {
-  PyObject *args = get_args();
-  PyObject *result = 
-    Thread::get_current_thread()->call_python_func(_function, args);
-  Py_DECREF(args);
+  PyObject *result = NULL;
+
+  if (_generator == (PyObject *)NULL) {
+    // We are calling the function directly.
+    PyObject *args = get_args();
+    result = 
+      Thread::get_current_thread()->call_python_func(_function, args);
+    Py_DECREF(args);
+
+    if (result != (PyObject *)NULL && PyGen_Check(result)) {
+      // The function has yielded a generator.  We will call into that
+      // henceforth, instead of calling the function from the top
+      // again.
+      if (task_cat.is_debug()) {
+        PyObject *str = PyObject_Repr(_function);
+        task_cat.debug() 
+          << PyString_AsString(str) << " in " << *this
+          << " yielded a generator.\n";
+        Py_DECREF(str);
+      }
+      _generator = result;
+      result = NULL;
+    }
+  }
+
+  if (_generator != (PyObject *)NULL) {
+    // We are calling a generator.
+    PyObject *func = PyObject_GetAttrString(_generator, "next");
+    nassertr(func != (PyObject *)NULL, DS_abort);
+
+    result = PyObject_CallObject(func, NULL);
+    Py_DECREF(func);
+
+    if (result == (PyObject *)NULL && PyErr_Occurred() &&
+        PyErr_ExceptionMatches(PyExc_StopIteration)) {
+      // "Catch" StopIteration and treat it like DS_done.
+      PyErr_Clear();
+      Py_DECREF(_generator);
+      _generator = NULL;
+      return DS_done;
+    }
+  }
 
   if (result == (PyObject *)NULL) {
     task_cat.error()
@@ -336,6 +370,11 @@ do_task() {
   if (PyInt_Check(result)) {
     int retval = PyInt_AS_LONG(result);
     switch (retval) {
+    case DS_restart:
+      Py_XDECREF(_generator);
+      _generator = NULL;
+      // Fall through.
+
     case DS_done:
     case DS_cont:
     case DS_again:

panda/src/event/pythonTask.h

@@ -63,6 +63,8 @@ private:
   PyObject *_owner;
   PyObject *_dict;
 
+  PyObject *_generator;
+
 public:
   static TypeHandle get_class_type() {
     return _type_handle;