1 files changed, 57 insertions, 2 deletions
diff --git a/ur_upgrade.py b/ur_upgrade.py
index 3b149fa..9718137 100755
--- a/ur_upgrade.py
+++ b/ur_upgrade.py
@@ -41,6 +41,45 @@ from itertools import chain
 DEFAULT_MIRROR = "http://ssb.stsci.edu/ureka"
 DEFAULT_PUBLIC_RELEASE = "{}/public_releases.txt".format(DEFAULT_MIRROR)
 
+def retype(t):
+    for cast in (int, float):
+        try:
+            return cast(t)
+        except ValueError:
+            pass
+    return t
+
+def get_vfs_data(path):
+    path = os.path.abspath(path)
+    if not os.path.exists(path):
+        return False
+
+    headers = ['device', 'type', 'size', 'used', 'available', 'percent', 'mountpoint']
+    command = 'df -Tk {}'.format(path).split()
+    proc = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=open(os.devnull, 'w'))
+    stdout, _ = proc.communicate()
+
+    data = stdout.split('\n')[1].split()
+    for index, record in enumerate(data):
+        data[index] = retype(record)
+
+    vfs_data = zip(headers, data)
+    return dict(vfs_data)
+
+def ur_storage_strategy(required_size, storage_map):
+    vfs = []
+    for d in storage_map:
+        data = get_vfs_data(d)
+        data['path'] = d
+        if (data['available'] - required_size) > 0:
+            vfs.append(data)
+
+    try:
+        best = max([dev['size'] for dev in vfs])
+        if best in [dev['size'] for dev in vfs]:
+            return dev['path']
+    except ValueError:
+        return None
 
 def ur_getenv(ur_dir):
     ''' Evaluates environment variables produced by ur-setup-real
@@ -142,7 +181,11 @@ class Upgrade(object):
         self.ureka = Ureka(ur_dir)
         self.ureka_next = None
         self.to_version = to_version
-        self.tmp = tempfile.mkdtemp(prefix="upgrade")
+        _tmp = ur_storage_strategy(storage_requires, storage_areas)
+        self.tmp = tempfile.mkdtemp(dir=_tmp, prefix="upgrade")
+        if _tmp is None:
+            print("++ failed to determine best upgrade storage strategy!")
+            print("++ defaulting to {}".format(self.tmp))
         self.tmp_dist = os.path.join(self.tmp, self.to_version)
         self.force = False
         self.backup = True
@@ -302,9 +345,21 @@ if __name__ == "__main__":
     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')
+
+    if 'UR_DIR' in os.environ:
+        args = parser.parse_args()
+        args.UR_DIR = os.environ['UR_DIR']
+    else:
+        parser.add_argument('UR_DIR',
+                            action='store',
+                            help='Absolute path to Ureka installation')
     args = parser.parse_args()
 
+    storage_requires = (1024 ** 2) * 7  # 7GB
+    storage_areas = [tempfile.gettempdir(),
+                     os.path.abspath(os.curdir),
+                     os.environ['HOME']]
+
     mirror = DEFAULT_MIRROR
     mirror_releases = DEFAULT_PUBLIC_RELEASE