Minor reordering of functions and classes. Implement --list-available, --latest, --no-backup, --backup-dir, and --force arguments. Add additional error checking. Fix incorrect upgrade ordering.

1 files changed, 135 insertions, 63 deletions

diff --git a/ur_upgrade.py b/ur_upgrade.py
index ff23a10..c4a9857 100755
--- a/ur_upgrade.py
+++ b/ur_upgrade.py
@@ -31,6 +31,7 @@ import tempfile
 import urllib2
 import signal
 import tarfile
+import time
 from distutils.version import StrictVersion
 from distutils.spawn import find_executable as which
 from collections import namedtuple
@@ -38,42 +39,8 @@ from string import Template
 from itertools import chain
 
 DEFAULT_MIRROR = "http://ssb.stsci.edu/ureka"
+DEFAULT_PUBLIC_RELEASE = "{}/public_releases.txt".format(DEFAULT_MIRROR)
 
-class Ureka(object):
-    '''
-    '''
-    def __init__(self, basepath):
-        self.path = basepath
-        self.path_data = os.path.abspath(os.path.join(self.path, 'misc'))
-        self._files = ['os', 'bits', 'version', 'name']
-        if os.path.exists(self.path_data):
-            self.info = self._info()
-        else:
-            self.info = {}
-
-    def _info(self):
-        data = []
-        for key in self._files:
-            path = os.path.join(self.path_data, key)
-            item = None
-            try:
-                item = open(path, 'r').readline().strip()
-            except:
-                item = ''
-            data.append(item)
-
-        ureka_info = namedtuple('ureka_info', self._files)
-        return ureka_info._make(data)
-
-    def __getitem__(self, key):
-        info = self.info._asdict()
-        if key not in info:
-            return None
-        return info[key]
-
-    def __iter__(self):
-        for item in self.info._asdict().iteritems():
-            yield item
 
 def ur_getenv(ur_dir):
     ''' Evaluates environment variables produced by ur-setup-real
@@ -112,13 +79,62 @@ def ur_getenv(ur_dir):
 
     return output_env
 
+
 def ur_check_version(urobj, vers):
     if StrictVersion(urobj['version']) > StrictVersion(vers):
         return False
+    elif StrictVersion(urobj['version']) == StrictVersion(vers):
+        return False
     return True
 
-class VersionError(Exception):
-    pass
+def ur_get_public_releases(m):
+    data = []
+    try:
+        req_data = urllib2.urlopen(m)
+        data = req_data.readlines()
+        data = [ x.strip() for x in data ]
+    except:
+        return []
+    return data
+
+
+class Ureka(object):
+    def __init__(self, basepath):
+        self.path = os.path.abspath(basepath)
+        self.path_data = os.path.join(self.path, 'misc')
+        self._files = ['os', 'bits', 'version', 'name']
+
+        if not os.path.exists(self.path):
+            print('{} does not exist.'.format(self.path))
+            exit(1)
+
+        try:
+            self.info = self._info()
+        except:
+            print('{} does not contain a valid Ureka installation.'.format(self.path))
+            exit(1)
+
+    def _info(self):
+        data = []
+        for key in self._files:
+            path = os.path.join(self.path_data, key)
+            item = None
+            item = open(path, 'r').readline().strip()
+            data.append(item)
+
+        ureka_info = namedtuple('ureka_info', self._files)
+        return ureka_info._make(data)
+
+    def __getitem__(self, key):
+        info = self.info._asdict()
+        if key not in info:
+            return None
+        return info[key]
+
+    def __iter__(self):
+        for item in self.info._asdict().iteritems():
+            yield item
+
 
 class Upgrade(object):
     def __init__(self, ur_dir, to_version, mirror=DEFAULT_MIRROR, **kwargs):
@@ -128,6 +144,9 @@ class Upgrade(object):
         self.to_version = to_version
         self.tmp = tempfile.mkdtemp(prefix="upgrade")
         self.tmp_dist = os.path.join(self.tmp, self.to_version)
+        self.force = False
+        self.backup = True
+        self.backup_path = os.path.abspath(os.curdir)
         self.archive_ext = '.tar.gz'
 
         if 'archive_ext' in kwargs:
@@ -142,11 +161,6 @@ class Upgrade(object):
                                              self.archive)
         self.archive_path = os.path.abspath(os.path.join(self.tmp, self.archive))
 
-        if not ur_check_version(self.ureka, self.to_version):
-            self._cleanup()
-            raise VersionError("Refusing to downgrade from {} to {}".format(self.ureka['version'], self.to_version))
-            exit(1)
-
         signal.signal(signal.SIGINT, self._cleanup_on_signal)
         signal.signal(signal.SIGTERM, self._cleanup_on_signal)
 
@@ -155,8 +169,19 @@ class Upgrade(object):
             exit(1)
 
     def run(self):
+        if not self.force:
+            if not ur_check_version(self.ureka, self.to_version):
+                self._cleanup()
+                print("Refusing upgrade from {} to {}. Use --force to override.".format(self.ureka['version'], self.to_version))
+                exit(1)
+
         self._get_archive()
+
         ureka_next_path = self._unpack_archive()
+
+        if self.backup:
+            self._backup()
+
         self._pre()
 
         # Populate temporary Ureka upgrade object
@@ -200,11 +225,11 @@ class Upgrade(object):
         total = 0
         with open(self.archive_path, 'wb') as installer:
             for chunk in iter(lambda: req_data.read(16 * 1024), ''):
-                total += float(len(chunk))
-                installer.write(chunk)
                 print("\r{:.2f}% [{:.2f}/{:.2f} MB]".format((total / remote_size * 100),
                                                             (total / 1024 ** 2),
                                                             (remote_size / 1024 ** 2))),
+                total += float(len(chunk))
+                installer.write(chunk)
         print("")
 
     def _unpack_archive(self):
@@ -222,6 +247,15 @@ class Upgrade(object):
 
         return os.path.join(self.tmp_dist, 'Ureka')
 
+    def _backup(self):
+        filename = '{}-{}-{}-backup.tar'.format(self.ureka['name'],
+                                            self.ureka['version'],
+                                            str(time.time()))
+        backup_path = os.path.join(self.backup_path, filename)
+        print('+ Generating backup... {}'.format(filename))
+        with tarfile.open(backup_path, 'w') as tar:
+            tar.add(self.ureka.path, arcname=os.path.basename(self.ureka.path))
+
     def _pre(self):
         print("+ Executing pre-upgrade tasks...")
         misc = os.path.join(self.tmp_dist, 'Ureka', 'misc', 'name')
@@ -229,10 +263,10 @@ class Upgrade(object):
             fp.write(self.ureka['name'] + os.linesep)
 
     def _upgrade(self, src, dest):
-        print("+ Upgrade in progress ({} to {})...".format(src['version'],
-                                                               dest['version']))
-        command = 'rsync -a -u {} {}'.format(dest.path,
-                                             os.path.dirname(src.path)).split()
+        print("+ Upgrade in progress ({} to {})...".format(dest['version'],
+                                                               src['version']))
+        command = 'rsync -a -u {} {}'.format(src.path,
+                                             os.path.dirname(dest.path)).split()
         proc = subprocess.Popen(command)
         proc.wait()
 
@@ -255,31 +289,69 @@ class Upgrade(object):
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description='Ureka Upgrade Utility')
-    parser.add_argument('UR_DIR', action='store', type=str, help='Absolute path to Ureka installation')
-    parser.add_argument('--latest', action='store_true', default=False, help='Upgrade to the latest (stable) version')
+    parser.add_argument('--list-available', action='store_true', help='List available releases')
+    parser.add_argument('--latest', action='store_true', help='Upgrade to the latest (stable) version')
     parser.add_argument('--request', type=str, help='Upgrade to a specific version')
     parser.add_argument('--mirror', type=str, help='Use a Ureka download mirror')
+    parser.add_argument('--no-backup', action='store_true', help='Do not backup existing installation')
+    parser.add_argument('--backup-dir', action='store_true', help='Alternative backup storage location')
+    parser.add_argument('--force', action='store_true', help='Ignore version checking')
+    parser.add_argument('UR_DIR', action='store', help='Absolute path to Ureka installation')
     args = parser.parse_args()
 
+    mirror = DEFAULT_MIRROR
+    mirror_releases = DEFAULT_PUBLIC_RELEASE
+
+    if args.mirror:
+        mirror = args.mirror
+        mirror_releases = "{}/public_releases.txt".format(mirror)
+
+    if args.list_available:
+        ureka = Ureka(args.UR_DIR)
+        for release in ur_get_public_releases(mirror_releases):
+            if StrictVersion(ureka['version']) < StrictVersion(release):
+                flag = 'Available for upgrade'
+            elif StrictVersion(ureka['version']) == StrictVersion(release):
+                flag = 'Currently installed'
+            else:
+                flag = 'Older release'
+            print('{:>6s} - {}'.format(release, flag))
+        exit(0)
+
     if args.latest and args.request:
-        print("--latest and --request are mutually exclusive options.")
+        print('--latest and --request are mutually exclusive options.')
         exit(1)
 
-    if args.latest:
-        # Need to work out a deal with SSB.
-        # "LATEST-STABLE" and "LATEST-DEVEL" links could be useful
-        print("--latest is not implemented yet.")
-        exit(1)
+    if args.latest and not args.request:
+        try:
+            request = ur_get_public_releases(DEFAULT_PUBLIC_RELEASE)[-1:][-1]
+        except:
+            request = ''
 
-    if not args.request:
-        print("No version requested. Use --request (e.g. 1.4.1)")
-        exit(1)
+        if not request:
+            print('Unable to retrieve the latest public release from {}.'.format(DEFAULT_PUBLIC_RELEASE))
+            exit(1)
 
-    mirror = DEFAULT_MIRROR
-    if args.mirror:
-        mirror = args.mirror
+    if not args.latest:
+        if not args.request:
+            print("No version requested. Use --request (e.g. 1.4.1)")
+            exit(1)
+
+    upgrade = Upgrade(args.UR_DIR, request, mirror)
+
+    if args.force:
+        upgrade.force = True
+
+    if args.backup_dir:
+        if not os.path.exists(args.backup_dir):
+            print("Backup directory {} does not exist.".format(args.backup_dir))
+            exit(1)
+
+        upgrade.backup_path = os.path.abspath(args.backup_dir)
+
+    if args.no_backup:
+        upgrade.backup = False
 
-    upgrade = Upgrade(args.UR_DIR, args.request, mirror)
     exit_code = upgrade.run()
 
     exit(exit_code)