From 6dfdc2ed6db637fcbebdc6aee558bed4789a36a8 Mon Sep 17 00:00:00 2001 From: Anders Ingemann Date: Mon, 1 Jul 2013 00:11:51 +0200 Subject: [PATCH] make tasklist.remove() behave like discard() --- base/tasklist.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/base/tasklist.py b/base/tasklist.py index 9b61c35..758fe95 100644 --- a/base/tasklist.py +++ b/base/tasklist.py @@ -13,15 +13,17 @@ class TaskList(object): self.tasks.update(args) def remove(self, *args): - for task in args: - self.tasks.discard(self.get(task)) + for task_type in args: + task = self.get(task_type) + if task is not None: + self.tasks.discard(task) def replace(self, task, replacement): self.remove(task) self.add(replacement) def get(self, ref): - return next(task for task in self.tasks if type(task) is ref) + return next((task for task in self.tasks if type(task) is ref), None) def run(self, bootstrap_info): task_list = self.create_list(self.tasks)