Implement startOver, queueOwnTask and queueOwnTasks

Author: mshima

lib/index.js

@@ -783,6 +783,64 @@ class Generator extends EventEmitter {
     });
   }
 
+  /**
+   * @private
+   * Schedule a generator's method on a run queue.
+   *
+   * @param {String}  name: The method name to schedule.
+   * @param {TaskOptions} [taskOptions]: options.
+   */
+  queueOwnTask(name, taskOptions = {}) {
+    const property = Object.getOwnPropertyDescriptor(Object.getPrototypeOf(this), name);
+    const item = property.value ? property.value : property.get.call(this);
+
+    const priority = this._queues[name];
+    taskOptions = {
+      ...priority,
+      cancellable: true,
+      run: false,
+      ...taskOptions
+    };
+
+    // Name points to a function; run it!
+    if (typeof item === 'function') {
+      taskOptions.taskName = name;
+      taskOptions.method = item;
+      this.queueTask(taskOptions);
+      return;
+    }
+
+    // Not a queue hash; stop
+    if (!priority) {
+      return;
+    }
+
+    this.queueTaskGroup(item, taskOptions);
+  }
+
+  /**
+   * @private
+   * Schedule every generator's methods on a run queue.
+   *
+   * @param {TaskOptions} [taskOptions]: options.
+   */
+  queueOwnTasks(taskOptions) {
+    this._running = true;
+    this._taskStatus = { cancelled: false, timestamp: new Date() };
+
+    const methods = Object.getOwnPropertyNames(Object.getPrototypeOf(this));
+    const validMethods = methods.filter(methodIsValid);
+    if (!validMethods.length) {
+      const error = new Error(
+        'This Generator is empty. Add at least one method for it to run.'
+      );
+      this.emit('error', error);
+      throw error;
+    }
+
+    validMethods.forEach(methodName => this.queueOwnTask(methodName, taskOptions));
+  }
+
   /**
    * Schedule tasks on a run queue.
    *
@@ -810,6 +868,7 @@ class Generator extends EventEmitter {
       namespace = self.options.namespace;
     }
 
+    const taskStatus = this._taskStatus;
     debug(
       `Queueing ${namespace}#${methodName} with options %o`,
       _.omit(task, ['method'])
@@ -820,23 +879,18 @@ class Generator extends EventEmitter {
       continueQueue => {
         debug(`Running ${namespace}#${methodName}`);
         self.emit(`method:${methodName}`);
-        const taskCancelled = task.cancellable && !self._running;
+        const taskCancelled = task.cancellable && taskStatus.cancelled;
+        if (taskCancelled) {
+          continueQueue();
+          return;
+        }
 
         runAsync(function() {
-          if (taskCancelled) {
-            return Promise.resolve();
-          }
-
           self.async = () => this.async();
           self.runningState = { namespace, queueName, methodName };
           return method.apply(self, self.args);
         })()
           .then(function() {
-            if (taskCancelled) {
-              continueQueue();
-              return;
-            }
-
             delete self.runningState;
             const eventName = `done$${namespace || 'unknownnamespace'}#${methodName}`;
             debug(`Emiting event ${eventName}`);
@@ -878,6 +932,19 @@ class Generator extends EventEmitter {
    */
   cancelCancellableTasks() {
     this._running = false;
+    this._taskStatus.cancelled = true;
+    delete this._taskStatus;
+  }
+
+  /**
+   * Start the generator again.
+   *
+   * @param {Object} [options]: options.
+   */
+  startOver(options = {}) {
+    this.cancelCancellableTasks();
+    Object.assign(this.options, options);
+    this.queueOwnTasks();
   }
 
   /**
@@ -933,8 +1000,6 @@ class Generator extends EventEmitter {
     }
 
     const promise = new Promise((resolve, reject) => {
-      const self = this;
-      this._running = true;
       this.debug('Generator is starting');
       this.emit('run');
 
@@ -948,52 +1013,15 @@ class Generator extends EventEmitter {
         this.on('error', reject);
       }
 
-      const methods = Object.getOwnPropertyNames(Object.getPrototypeOf(this));
-      const validMethods = methods.filter(methodIsValid);
-      if (!validMethods.length) {
-        return this.emit(
-          'error',
-          new Error('This Generator is empty. Add at least one method for it to run.')
-        );
-      }
-
       this.env.runLoop.once('end', () => {
         this.debug('Generator has ended');
         this.emit('end');
         resolve();
       });
 
-      function addInQueue(name) {
-        const property = Object.getOwnPropertyDescriptor(
-          Object.getPrototypeOf(self),
-          name
-        );
-        const item = property.value ? property.value : property.get.call(self);
-
-        const priority = self._queues[name];
-        let taskOptions = {
-          ...priority,
-          cancellable: true,
-          run: false,
-          generatorReject: usePromise ? undefined : reject
-        };
-
-        // Name points to a function; run it!
-        if (typeof item === 'function') {
-          taskOptions.taskName = name;
-          taskOptions.method = item;
-          return self.queueTask(taskOptions);
-        }
-
-        // Not a queue hash; stop
-        if (!priority) {
-          return;
-        }
-
-        self.queueTaskGroup(item, taskOptions);
-      }
-
-      validMethods.forEach(addInQueue);
+      this.queueOwnTasks({
+        generatorReject: usePromise ? undefined : reject
+      });
 
       const writeFiles = () => {
         this.env.runLoop.add('conflicts', this._writeFiles.bind(this), {
@@ -1016,6 +1044,7 @@ class Generator extends EventEmitter {
       });
 
       this._composedWith.forEach(runGenerator);
+      this._composedWith = [];
     });
 
     // For composed generators, otherwise error will not be catched.

test/base.js

@@ -556,6 +556,49 @@ describe('Base', () => {
 
       this.testGen.run().then(done);
     });
+
+    it('can start over the generator', function(done) {
+      const spy1 = sinon.spy();
+      const spy2 = sinon.spy();
+
+      this.TestGenerator.prototype.cancel = function() {
+        spy1();
+        if (!this.startedOver) {
+          this.startOver({ startedOver: true });
+          this.startedOver = true;
+        }
+      };
+
+      this.TestGenerator.prototype.after = function() {
+        assert(this.options.startedOver);
+        assert(this.startedOver);
+        spy2();
+      };
+
+      this.testGen.run().then(() => {
+        assert(spy1.calledTwice);
+        assert(spy2.calledOnce);
+        done();
+      });
+    });
+
+    it('can queue a method again', function(done) {
+      const spy1 = sinon.spy();
+
+      this.TestGenerator.prototype.cancel = function() {
+        spy1();
+        if (!this.startedOver) {
+          this.queueOwnTask('cancel');
+          this.startOver();
+          this.startedOver = true;
+        }
+      };
+
+      this.testGen.run().then(() => {
+        assert(spy1.calledTwice);
+        done();
+      });
+    });
   });
 
   describe('#argument()', () => {