Simplify tasklist by passing a set to the modules

base/main.py

@@ -26,14 +26,10 @@ def get_args():
 def run(args):
 	from manifest import Manifest
 	manifest = Manifest(args.manifest)
-	provider = manifest.modules['provider']
-	plugins = manifest.modules['plugins']
 
 	from tasklist import TaskList
 	tasklist = TaskList()
-	provider.resolve_tasks(tasklist, manifest)
-	for plugin in plugins:
-		plugin.resolve_tasks(tasklist, manifest)
+	tasklist.load('resolve_tasks', manifest)
 
 	from bootstrapinfo import BootstrapInformation
 	bootstrap_info = BootstrapInformation(manifest=manifest, debug=args.debug)
@@ -46,11 +42,13 @@ def run(args):
 		if args.pause_on_error:
 			raw_input("Press Enter to commence rollback")
 		log.error('Rolling back')
+
 		rollback_tasklist = TaskList()
-		provider.resolve_rollback_tasks(rollback_tasklist, tasklist.tasks_completed, manifest)
-		for plugin in plugins:
-			resolve_rollback_tasks = getattr(plugin, 'resolve_rollback_tasks', None)
-			if callable(resolve_rollback_tasks):
-				resolve_rollback_tasks(rollback_tasklist, tasklist.tasks_completed, manifest)
+
+		def counter_task(task, counter):
+			if task in tasklist.tasks_completed and counter not in tasklist.tasks_completed:
+				rollback_tasklist.tasks.add(counter)
+		rollback_tasklist.load('resolve_rollback_tasks', manifest, counter_task)
+
 		rollback_tasklist.run(info=bootstrap_info, dry_run=args.dry_run)
 		log.info('Successfully completed rollback')

base/tasklist.py

@@ -9,15 +9,15 @@ class TaskList(object):
 		self.tasks = set()
 		self.tasks_completed = []
 
-	def add(self, *args):
-		self.tasks.update(args)
-
-	def remove(self, *args):
-		for task in args:
-			self.tasks.discard(task)
+	def load(self, function, manifest, *args):
+		getattr(manifest.modules['provider'], function)(self.tasks, manifest, *args)
+		for plugin in manifest.modules['plugins']:
+			fn = getattr(plugin, function, None)
+			if callable(fn):
+				fn(self.tasks, manifest, *args)
 
 	def run(self, info={}, dry_run=False):
-		task_list = self.create_list(self.tasks)
+		task_list = self.create_list()
 		log.debug('Tasklist:\n\t{list}'.format(list='\n\t'.join(repr(task) for task in task_list)))
 
 		for task_type in task_list:
@@ -28,19 +28,19 @@ class TaskList(object):
 				log.info('Running {task}'.format(task=task))
 			if not dry_run:
 				task.run(info)
-			self.tasks_completed.append(task)
+			self.tasks_completed.append(task_type)
 
-	def create_list(self, tasks):
+	def create_list(self):
 		from common.phases import order
 		graph = {}
-		for task in tasks:
+		for task in self.tasks:
 			self.check_ordering(task)
 			successors = set()
 			successors.update(task.successors)
-			successors.update(filter(lambda succ: task in succ.predecessors, tasks))
+			successors.update(filter(lambda succ: task in succ.predecessors, self.tasks))
 			succeeding_phases = order[order.index(task.phase) + 1:]
-			successors.update(filter(lambda succ: succ.phase in succeeding_phases, tasks))
-			graph[task] = filter(lambda succ: succ in tasks, successors)
+			successors.update(filter(lambda succ: succ.phase in succeeding_phases, self.tasks))
+			graph[task] = filter(lambda succ: succ in self.tasks, successors)
 
 		components = self.strongly_connected_components(graph)
 		cycles_found = 0

common/exceptions.py

@@ -9,8 +9,9 @@ class ManifestError(Exception):
 	def __str__(self):
 		if self.json_path is not None:
 			path = '.'.join(self.json_path)
-			return "{2}\n\tFile: {0}\n\tJSON path: {1}".format(self.manifest_path, path, self.message)
-		return "{0}: {1}".format(self.manifest_path, self.message)
+			return ('{msg}\n\tFile: {file}\n\tJSON path: {jsonpath}'
+			        .format(msg=self.message, file=self.manifest_path, jsonpath=path))
+		return '{file}: {msg}'.format(msg=self.message, file=self.manifest_path)
 
 
 class TaskListError(Exception):

plugins/admin_user/__init__.py

@@ -6,13 +6,14 @@ def validate_manifest(data, validator, error):
 	validator(data, schema_path)
 
 
-def resolve_tasks(tasklist, manifest):
+def resolve_tasks(taskset, manifest):
 	import tasks
 	from providers.ec2.tasks import initd
-	if initd.AddEC2InitScripts in tasklist.tasks:
-		tasklist.add(tasks.AdminUserCredentials)
+	if initd.AddEC2InitScripts in taskset:
+		taskset.add(tasks.AdminUserCredentials)
 
-	tasklist.add(tasks.AddSudoPackage,
+	taskset.update([tasks.AddSudoPackage,
 	                tasks.CreateAdminUser,
 	                tasks.PasswordlessSudo,
-	             tasks.DisableRootLogin)
+	                tasks.DisableRootLogin,
+	                ])

plugins/build_metadata/__init__.py

@@ -1,5 +1,5 @@
 
 
-def resolve_tasks(tasklist, manifest):
+def resolve_tasks(taskset, manifest):
 	from tasks import WriteMetadata
-	tasklist.add(WriteMetadata)
+	taskset.add(WriteMetadata)

plugins/cloud_init/__init__.py

@@ -6,24 +6,25 @@ def validate_manifest(data, validator, error):
 	validator(data, schema_path)
 
 
-def resolve_tasks(tasklist, manifest):
+def resolve_tasks(taskset, manifest):
 	import tasks
 	import providers.ec2.tasks.initd as initd_ec2
 	from common.tasks import initd
 
 	if manifest.system['release'] in ['wheezy', 'stable']:
-		tasklist.add(tasks.AddBackports)
+		taskset.add(tasks.AddBackports)
 
-	tasklist.add(tasks.AddCloudInitPackages,
-	             tasks.SetMetadataSource)
+	taskset.update([tasks.AddCloudInitPackages,
+	                tasks.SetMetadataSource,
+	                ])
 
 	options = manifest.plugins['cloud_init']
 	if 'username' in options:
-		tasklist.add(tasks.SetUsername)
+		taskset.add(tasks.SetUsername)
 	if 'disable_modules' in options:
-		tasklist.add(tasks.DisableModules)
+		taskset.add(tasks.DisableModules)
 
-	tasklist.remove(initd_ec2.AddEC2InitScripts,
-	                initd.AddExpandRoot,
-	                initd.AdjustExpandRootScript,
-	                initd.AddSSHKeyGeneration)
+	taskset.discard(initd_ec2.AddEC2InitScripts)
+	taskset.discard(initd.AddExpandRoot)
+	taskset.discard(initd.AdjustExpandRootScript)
+	taskset.discard(initd.AddSSHKeyGeneration)

plugins/image_commands/__init__.py

@@ -6,6 +6,6 @@ def validate_manifest(data, validator, error):
 	validator(data, schema_path)
 
 
-def resolve_tasks(tasklist, manifest):
+def resolve_tasks(taskset, manifest):
 	from tasks import ImageExecuteCommand
-	tasklist.add(ImageExecuteCommand)
+	taskset.add(ImageExecuteCommand)

plugins/minimize_size/__init__.py

@@ -1,16 +1,11 @@
 import tasks
 
 
-def resolve_tasks(tasklist, manifest):
-	tasklist.add(tasks.AddFolderMounts,
-	             tasks.RemoveFolderMounts)
+def resolve_tasks(taskset, manifest):
+	taskset.update([tasks.AddFolderMounts,
+	                tasks.RemoveFolderMounts,
+	                ])
 
 
-def resolve_rollback_tasks(tasklist, tasks_completed, manifest):
-	completed = [type(task) for task in tasks_completed]
-
-	def counter_task(task, counter):
-		if task in completed and counter not in completed:
-			tasklist.add(counter)
-
+def resolve_rollback_tasks(taskset, manifest, counter_task):
 	counter_task(tasks.AddFolderMounts, tasks.RemoveFolderMounts)

plugins/opennebula/__init__.py

@@ -1,5 +1,5 @@
 
 
-def resolve_tasks(tasklist, manifest):
+def resolve_tasks(taskset, manifest):
 	import tasks
-	tasklist.add(tasks.OpenNebulaContext)
+	taskset.add(tasks.OpenNebulaContext)

plugins/prebootstrapped/__init__.py

@@ -18,7 +18,7 @@ def validate_manifest(data, validator, error):
 	validator(data, schema_path)
 
 
-def resolve_tasks(tasklist, manifest):
+def resolve_tasks(taskset, manifest):
 	settings = manifest.plugins['prebootstrapped']
 	skip_tasks = [ebs.Create,
 	              loopback.Create,
@@ -35,25 +35,19 @@ def resolve_tasks(tasklist, manifest):
 	              bootstrap.Bootstrap]
 	if manifest.volume['backing'] == 'ebs':
 		if 'snapshot' in settings and settings['snapshot'] is not None:
-			tasklist.add(CreateFromSnapshot)
-			tasklist.remove(*skip_tasks)
+			taskset.add(CreateFromSnapshot)
+			[taskset.discard(task) for task in skip_tasks]
 		else:
-			tasklist.add(Snapshot)
+			taskset.add(Snapshot)
 	else:
 		if 'image' in settings and settings['image'] is not None:
-			tasklist.add(CreateFromImage)
-			tasklist.remove(*skip_tasks)
+			taskset.add(CreateFromImage)
+			[taskset.discard(task) for task in skip_tasks]
 		else:
-			tasklist.add(CopyImage)
+			taskset.add(CopyImage)
 
 
-def resolve_rollback_tasks(tasklist, tasks_completed, manifest):
-	completed = [type(task) for task in tasks_completed]
-
-	def counter_task(task, counter):
-		if task in completed and counter not in completed:
-			tasklist.add(counter)
-
+def resolve_rollback_tasks(taskset, manifest, counter_task):
 	if manifest.volume['backing'] == 'ebs':
 		counter_task(CreateFromSnapshot, volume.Delete)
 	else:

plugins/root_password/__init__.py

@@ -6,8 +6,8 @@ def validate_manifest(data, validator, error):
 	validator(data, schema_path)
 
 
-def resolve_tasks(tasklist, manifest):
+def resolve_tasks(taskset, manifest):
 	from common.tasks.security import DisableSSHPasswordAuthentication
 	from tasks import SetRootPassword
-	tasklist.remove(DisableSSHPasswordAuthentication)
-	tasklist.add(SetRootPassword)
+	taskset.discard(DisableSSHPasswordAuthentication)
+	taskset.add(SetRootPassword)

plugins/unattended_upgrades/__init__.py

@@ -6,7 +6,7 @@ def validate_manifest(data, validator, error):
 	validator(data, schema_path)
 
 
-def resolve_tasks(tasklist, manifest):
+def resolve_tasks(taskset, manifest):
 	import tasks
-	tasklist.add(tasks.AddUnattendedUpgradesPackage,
-	             tasks.EnablePeriodicUpgrades)
+	taskset.add(tasks.AddUnattendedUpgradesPackage)
+	taskset.add(tasks.EnablePeriodicUpgrades)

plugins/vagrant/__init__.py

@@ -7,27 +7,21 @@ def validate_manifest(data, validator, error):
 	validator(data, schema_path)
 
 
-def resolve_tasks(tasklist, manifest):
+def resolve_tasks(taskset, manifest):
 	from common.tasks import security
 	from common.tasks import loopback
-	tasklist.remove(security.DisableSSHPasswordAuthentication,
-	                loopback.MoveImage,
-	                )
+	taskset.discard(security.DisableSSHPasswordAuthentication)
+	taskset.discard(loopback.MoveImage)
+
 	from common.tasks import volume
-	tasklist.add(tasks.CreateVagrantBoxDir,
+	taskset.update([tasks.CreateVagrantBoxDir,
 	                tasks.AddPackages,
 	                tasks.AddInsecurePublicKey,
 	                tasks.PackageBox,
 	                tasks.RemoveVagrantBoxDir,
 	                volume.Delete,
-	             )
+	                ])
 
 
-def resolve_rollback_tasks(tasklist, tasks_completed, manifest):
-	completed = [type(task) for task in tasks_completed]
-
-	def counter_task(task, counter):
-		if task in completed and counter not in completed:
-			tasklist.add(counter)
-
+def resolve_rollback_tasks(taskset, manifest, counter_task):
 	counter_task(tasks.CreateVagrantBoxDir, tasks.RemoveVagrantBoxDir)

providers/ec2/__init__.py

@@ -48,23 +48,18 @@ def validate_manifest(data, validator, error):
 			error('HVM AMIs only support extlinux as a bootloader', ['system', 'bootloader'])
 
 
-def resolve_tasks(tasklist, manifest):
-	from common.task_sets import base_set
-	from common.task_sets import mounting_set
-	from common.task_sets import apt_set
-	from common.task_sets import locale_set
-	from common.task_sets import ssh_set
-	tasklist.add(*base_set)
-	tasklist.add(*mounting_set)
-	tasklist.add(*apt_set)
-	tasklist.add(*locale_set)
-	tasklist.add(*ssh_set)
+def resolve_tasks(taskset, manifest):
+	import common.task_sets
+	taskset.update(common.task_sets.base_set)
+	taskset.update(common.task_sets.mounting_set)
+	taskset.update(common.task_sets.apt_set)
+	taskset.update(common.task_sets.locale_set)
+	taskset.update(common.task_sets.ssh_set)
 
 	if manifest.volume['partitions']['type'] != 'none':
-		from common.task_sets import partitioning_set
-		tasklist.add(*partitioning_set)
+		taskset.update(common.task_sets.partitioning_set)
 
-	tasklist.add(tasks.host.HostDependencies,
+	taskset.update([tasks.host.HostDependencies,
 	                tasks.packages.DefaultPackages,
 	                tasks.connection.GetCredentials,
 	                tasks.host.GetInfo,
@@ -86,13 +81,14 @@ def resolve_tasks(tasklist, manifest):
 	                cleanup.ClearMOTD,
 	                cleanup.CleanTMP,
 
-	             tasks.ami.RegisterAMI)
+	                tasks.ami.RegisterAMI,
+	                ])
 
 	if manifest.system['bootloader'] == 'pvgrub':
-		tasklist.add(boot.AddGrubPackage, tasks.boot.ConfigurePVGrub)
+		taskset.add(boot.AddGrubPackage)
+		taskset.add(tasks.boot.ConfigurePVGrub)
 	else:
-		from common.task_sets import bootloader_set
-		tasklist.add(*bootloader_set.get(manifest.system['bootloader']))
+		taskset.update(common.task_sets.bootloader_set.get(manifest.system['bootloader']))
 
 	backing_specific_tasks = {'ebs': [tasks.ebs.Create,
 	                                  tasks.ebs.Attach,
@@ -104,29 +100,22 @@ def resolve_tasks(tasklist, manifest):
 	                                 tasks.ami.BundleImage,
 	                                 tasks.ami.UploadImage,
 	                                 tasks.ami.RemoveBundle]}
-	tasklist.add(*backing_specific_tasks.get(manifest.volume['backing'].lower()))
-	tasklist.add(filesystem.Format,
+	taskset.update(backing_specific_tasks.get(manifest.volume['backing'].lower()))
+	taskset.update([filesystem.Format,
 	                volume.Detach,
-	             volume.Delete)
+	                volume.Delete,
+	                ])
 
 	if manifest.bootstrapper.get('tarball', False):
-		tasklist.add(bootstrap.MakeTarball)
+		taskset.add(bootstrap.MakeTarball)
 
-	from common.task_sets import get_fs_specific_set
-	tasklist.add(*get_fs_specific_set(manifest.volume['partitions']))
+	taskset.update(common.task_sets.get_fs_specific_set(manifest.volume['partitions']))
 
 	if 'boot' in manifest.volume['partitions']:
-		from common.task_sets import boot_partition_set
-		tasklist.add(*boot_partition_set)
+		taskset.update(common.task_sets.boot_partition_set)
 
 
-def resolve_rollback_tasks(tasklist, tasks_completed, manifest):
-	completed = [type(task) for task in tasks_completed]
-
-	def counter_task(task, counter):
-		if task in completed and counter not in completed:
-			tasklist.add(counter)
-
+def resolve_rollback_tasks(taskset, manifest, counter_task):
 	counter_task(tasks.ebs.Create, volume.Delete)
 	counter_task(tasks.ebs.Attach, volume.Detach)
 

providers/virtualbox/__init__.py

@@ -22,25 +22,19 @@ def validate_manifest(data, validator, error):
 
 
 def resolve_tasks(tasklist, manifest):
-	from common.task_sets import base_set
-	from common.task_sets import volume_set
-	from common.task_sets import mounting_set
-	from common.task_sets import apt_set
-	from common.task_sets import locale_set
-	tasklist.add(*base_set)
-	tasklist.add(*volume_set)
-	tasklist.add(*mounting_set)
-	tasklist.add(*apt_set)
-	tasklist.add(*locale_set)
+	import common.task_sets
+	tasklist.update(common.task_sets.base_set)
+	tasklist.update(common.task_sets.volume_set)
+	tasklist.update(common.task_sets.mounting_set)
+	tasklist.update(common.task_sets.apt_set)
+	tasklist.update(common.task_sets.locale_set)
 
-	from common.task_sets import bootloader_set
-	tasklist.add(*bootloader_set.get(manifest.system['bootloader']))
+	tasklist.update(common.task_sets.bootloader_set.get(manifest.system['bootloader']))
 
 	if manifest.volume['partitions']['type'] != 'none':
-		from common.task_sets import partitioning_set
-		tasklist.add(*partitioning_set)
+		tasklist.update(common.task_sets.partitioning_set)
 
-	tasklist.add(tasks.packages.DefaultPackages,
+	tasklist.update([tasks.packages.DefaultPackages,
 
 	                 loopback.Create,
 
@@ -53,32 +47,26 @@ def resolve_tasks(tasklist, manifest):
 	                 cleanup.ClearMOTD,
 	                 cleanup.CleanTMP,
 
-	             loopback.MoveImage)
+	                 loopback.MoveImage,
+	                 ])
 
 	if manifest.bootstrapper.get('guest_additions', False):
 		from tasks import guest_additions
-		tasklist.add(guest_additions.CheckGuestAdditionsPath,
+		tasklist.update([guest_additions.CheckGuestAdditionsPath,
 		                 guest_additions.AddGuestAdditionsPackages,
-		             guest_additions.InstallGuestAdditions)
+		                 guest_additions.InstallGuestAdditions,
+		                 ])
 
 	if manifest.bootstrapper.get('tarball', False):
 		tasklist.add(bootstrap.MakeTarball)
 
-	from common.task_sets import get_fs_specific_set
-	tasklist.add(*get_fs_specific_set(manifest.volume['partitions']))
+	tasklist.update(common.task_sets.get_fs_specific_set(manifest.volume['partitions']))
 
 	if 'boot' in manifest.volume['partitions']:
-		from common.task_sets import boot_partition_set
-		tasklist.add(*boot_partition_set)
+		tasklist.update(common.task_sets.boot_partition_set)
 
 
-def resolve_rollback_tasks(tasklist, tasks_completed, manifest):
-	completed = [type(task) for task in tasks_completed]
-
-	def counter_task(task, counter):
-		if task in completed and counter not in completed:
-			tasklist.add(counter)
-
+def resolve_rollback_tasks(tasklist, manifest, counter_task):
 	counter_task(loopback.Create, volume.Delete)
 	counter_task(filesystem.CreateMountDir, filesystem.DeleteMountDir)
 	counter_task(partitioning.MapPartitions, partitioning.UnmapPartitions)