Support multiple packages in single upload, refactor unpack script

master_unpack.py

@@ -21,14 +21,22 @@
 
 # <pep8 compliant>
 
+import argparse
 import os
 import shutil
 import sys
 import zipfile
 
+# Parse package filename to extract branch and platform
+class Package:
+    def __init__(self, zipname):
+        self.zipname = zipname
+        self.filename = os.path.basename(zipname)
+        self.platform = self._get_platform(self.filename)
+        self.branch = self._get_branch(self.filename)
 
     # extension stripping
-def strip_extension(filename):
+    def _strip_extension(self, filename):
         extensions = '.zip', '.tar', '.bz2', '.gz', '.tgz', '.tbz', '.exe'
 
         for ext in extensions:
@@ -37,14 +45,13 @@ def strip_extension(filename):
 
         return filename
 
-
     # extract platform from package name
-def get_platform(filename):
+    def _get_platform(self, filename):
         # name is blender-version-platform.extension. we want to get the
         # platform out, but there may be some variations, so we fiddle a
         # bit to handle current and hopefully future names
-    filename = strip_extension(filename)
+        filename = self._strip_extension(filename)
-    filename = strip_extension(filename)
+        filename = self._strip_extension(filename)
 
         tokens = filename.split("-")
         platforms = ('osx', 'mac', 'bsd',
@@ -66,7 +73,7 @@ def get_platform(filename):
         return '-'.join(platform_tokens)
 
 
-def get_branch(filename):
+    def _get_branch(self, filename):
         tokens = filename.split("-")
         branch = ""
 
@@ -81,41 +88,8 @@ def get_branch(filename):
 
         return ""
 
-# get filename
-if len(sys.argv) < 2:
-    sys.stderr.write("Not enough arguments, expecting file to unpack\n")
-    sys.exit(1)
 
-filename = sys.argv[1]
+def get_download_dir(branch):
-
-# open zip file
-if not os.path.exists(filename):
-    sys.stderr.write("File %r not found.\n" % filename)
-    sys.exit(1)
-
-try:
-    z = zipfile.ZipFile(filename, "r")
-except Exception as ex:
-    sys.stderr.write('Failed to open zip file: %s\n' % str(ex))
-    sys.exit(1)
-
-if len(z.namelist()) != 1:
-    sys.stderr.write("Expected one file in %r." % filename)
-    sys.exit(1)
-
-package = z.namelist()[0]
-packagename = os.path.basename(package)
-
-# detect platform and branch
-platform = get_platform(packagename)
-branch = get_branch(packagename)
-
-if platform == '':
-    sys.stderr.write('Failed to detect platform ' +
-        'from package: %r\n' % packagename)
-    sys.exit(1)
-
-# extract
     download_prefix = "/data/www/vhosts/builder.blender.org/webroot/download/"
     if not branch or branch == 'master':
         directory = download_prefix
@@ -125,26 +99,91 @@ else:
         directory = os.path.join(download_prefix, branch)
         os.makedirs(directory, exist_ok=True)
 
+    return directory
+
+def open_zipfile(filename):
+    # Open zip file
+    if not os.path.exists(filename):
+        sys.stderr.write("File %r not found.\n" % filename)
+        sys.exit(1)
+
     try:
-    filename = os.path.join(directory, packagename)
+        return zipfile.ZipFile(filename, "r")
-    zf = z.open(package)
+    except Exception as ex:
-    f = open(filename, "wb")
+        sys.stderr.write('Failed to open zip file: %s\n' % str(ex))
+        sys.exit(1)
+
+def get_zipfile_packages(zipfile):
+    # List packages in zip file
+    packages = [Package(zipname) for zipname in zipfile.namelist()]
+    if len(packages) == 0:
+        sys.stderr.write('Empty zip file\n')
+        sys.exit(1)
+    return packages
+
+def get_branch_platform(packages):
+    # Extract branch and platform names
+    branch = packages[0].branch
+    platform = packages[0].platform
+
+    if platform == '':
+        sys.stderr.write('Failed to detect platform ' +
+            'from package: %r\n' % packages[0].filename)
+        sys.exit(1)
+
+    for package in packages:
+        if package.branch != branch or package.platform != platform:
+            sys.stderr.write('All packages in the zip file must have the same branch and platform\n')
+            sys.exit(1)
+
+    return branch, platform
+
+def extract_zipfile_packages(zipfile, packages, branch):
+    # Extract packages from zip file
+    directory = get_download_dir(branch)
+
+    for package in packages:
+        filepath = os.path.join(directory, package.filename)
+
+        try:
+            zf = zipfile.open(package.zipname)
+            f = open(filepath, "wb")
 
             shutil.copyfileobj(zf, f)
-    os.chmod(filename, 0o644)
+            os.chmod(filepath, 0o644)
 
             zf.close()
-    z.close()
         except Exception as ex:
             sys.stderr.write('Failed to unzip package: %s\n' % str(ex))
             sys.exit(1)
 
-# remove other files from the same platform and branch
+def remove_older_packages(branch, platform, new_packages):
+    # Remove other files from the same platform and branch
+    directory = get_download_dir(branch)
+
+    for filename in os.listdir(directory):
+        package = Package(filename)
+        if package.platform == platform and package.branch == branch:
+            is_new_package = False
+            for new_package in new_packages:
+                if package.filename == new_package.filename:
+                    is_new_package = True
+
+            if not is_new_package:
                 try:
-    for f in os.listdir(directory):
+                    os.remove(os.path.join(directory, filename))
-        if get_platform(f) == platform and get_branch(f) == branch:
-            if f != packagename:
-                os.remove(os.path.join(directory, f))
                 except Exception as ex:
                     sys.stderr.write('Failed to remove old packages: %s\n' % str(ex))
-    sys.exit(1)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument('filename')
+    args = parser.parse_args()
+
+    with open_zipfile(args.filename) as zipfile:
+        packages = get_zipfile_packages(zipfile)
+        branch, platform = get_branch_platform(packages)
+        extract_zipfile_packages(zipfile, packages, branch)
+
+    remove_older_packages(branch, platform, packages)