diff --git a/doit/loader.py b/doit/loader.py index b36db4ad..707dd11f 100644 --- a/doit/loader.py +++ b/doit/loader.py @@ -284,12 +284,19 @@ def generate_tasks(func_name, gen_result, gen_doc=None): if isinstance(gen_result, dict): return [_generate_task_from_return(func_name, gen_result, gen_doc)] + if isinstance(gen_result, Task): + return (gen_result,) + # a generator if inspect.isgenerator(gen_result): tasks = OrderedDict() # task_name: task # the generator return subtasks as dictionaries for task_dict, x_doc in flat_generator(gen_result, gen_doc): - _generate_task_from_yield(tasks, func_name, task_dict, x_doc) + if isinstance(task_dict, Task): + tasks[task_dict.name] = task_dict + else: + _generate_task_from_yield(tasks, func_name, task_dict, x_doc) + if tasks: return list(six.itervalues(tasks))