Refactor common parts from mass_migrate.

Signed-off-by: Edward Z. Yang <ezyang@mit.edu>

diff --git a/wizard/command/mass_migrate.py b/wizard/command/mass_migrate.py
index b3198dcbef2cb9a9f89395bd94f70e5746462453..15fde37ef6ca0b85e4cffde9174a3cb4dfe18669 100644 (file)
--- a/wizard/command/mass_migrate.py
+++ b/wizard/command/mass_migrate.py
@@ -10,14 +10,13 @@ import itertools
 
 import wizard
 from wizard import deploy, util, shell, sset, command
-from wizard.command import migrate
 
 def main(argv, baton):
     options, args = parse_args(argv, baton)
     app = args[0]
     base_args = calculate_base_args(options)
-    sh = make_shell(options)
-    seen = make_serialized_set(options)
+    sh = shell.ParallelShell.make(options.no_parallelize, options.max)
+    seen = sset.make(options)
     is_root = not os.getuid()
     warnings_log, errors_log = open_aggregate_logs(options)
     # loop stuff
@@ -89,7 +88,7 @@ untrusted repositories."""
             default=False, help="Turn off parallelization")
     parser.add_option("--dry-run", dest="dry_run", action="store_true",
             default=False, help="Print commands that would be run. Implies --no-parallelize")
-    parser.add_option("--max", dest="max",
+    parser.add_option("--max", dest="max", type="int",
             default=10, help="Maximum subprocesses to run concurrently")
     parser.add_option("--force", dest="force", action="store_true",
             default=False, help="Force migrations to occur even if .scripts or .git exists.")
@@ -148,14 +147,3 @@ def calculate_base_args(options):
 def calculate_log_name(i, dir):
     return "%04d" % i + dir.replace('/', '-') + ".log"
 
-def make_shell(options):
-    if options.no_parallelize:
-        return shell.DummyParallelShell()
-    else:
-        return shell.ParallelShell(max=int(options.max))
-
-def make_serialized_set(options):
-    if options.seen:
-        return sset.SerializedSet(options.seen)
-    else:
-        return sset.DummySerializedSet()

diff --git a/wizard/command/migrate.py b/wizard/command/migrate.py
index b159427dbed5b2972f378a1867cd1205a94b6d38..9053037f3c03c42299a636bd4b58e72864f32b31 100644 (file)
--- a/wizard/command/migrate.py
+++ b/wizard/command/migrate.py
@@ -33,9 +33,7 @@ def main(argv, baton):
     except deploy.NotMigratedError:
         pass
     except (deploy.CorruptedAutoinstallError, AlreadyMigratedError):
-        if options.force:
-            perform_force(options)
-        else:
+        if not options.force:
             raise
 
     deployment.verifyTag(options.srv_path)
@@ -58,6 +56,7 @@ def main(argv, baton):
             elif e.errno == errno.EACCES:
                 raise command.PermissionsError(dir)
             raise
+        if options.force: perform_force(options)
         make_repository(sh, options, repo, tag)
         check_variables(deployment, options)
     finally:

diff --git a/wizard/shell.py b/wizard/shell.py
index c66147a0f9d876fe9fe10b2249d4cfd7f6771a49..79b2a71fa3e93b49def7d7c8d703fc07dc91005c 100644 (file)
--- a/wizard/shell.py
+++ b/wizard/shell.py
@@ -244,6 +244,13 @@ class ParallelShell(Shell):
         super(ParallelShell, self).__init__(dry=dry)
         self.running = {}
         self.max = max # maximum of commands to run in parallel
+    @staticmethod
+    def make(self, no_parallelize, max):
+        """Convenience method oriented towards command modules."""
+        if no_parallelize:
+            return DummyParallelShell()
+        else:
+            return ParallelShell(max=max)
     def _async(self, proc, args, python, on_success, on_error, **kwargs):
         """
         Gets handed a :class:`subprocess.Proc` object from our deferred

diff --git a/wizard/sset.py b/wizard/sset.py
index a5c37d4af7faca940c2c6a49c1214f0be080ed0c..22cc971858c54124c38d60804db0749d4dd2b1cd 100644 (file)
--- a/wizard/sset.py
+++ b/wizard/sset.py
@@ -1,5 +1,11 @@
 import os.path
 
+def make(seen_file):
+    if seen_file:
+        return SerializedSet(seen_file)
+    else:
+        return DummySerializedSet(seen_file)
+
 class ISerializedSet(object):
     def put(self, name):
         raise NotImplementedError