Add zip support to python.to_lorry ext

Change-Id: I04ed11272c7ddfb22829c71607772b86f216d457

2 files changed, 37 insertions, 18 deletions

diff --git a/baserockimport/exts/python.to_lorry b/baserockimport/exts/python.to_lorry
index 30f38e7..2a0cebb 100755
--- a/baserockimport/exts/python.to_lorry
+++ b/baserockimport/exts/python.to_lorry
@@ -37,15 +37,12 @@ from importer_python_common import *
 from utils import warn, error
 import utils
 
-def filter_urls(urls):
-    allowed_extensions = ['tar.gz', 'tgz', 'tar.Z', 'tar.bz2', 'tbz2',
-                          'tar.lzma', 'tar.xz', 'tlz', 'txz', 'tar']
+def filter_urls(urls, extensions):
+    def match(url):
+        return ('.'.join(url['url'].split('.')[-2:]) in extensions
+                or url['url'].split('.')[-1] in extensions)
 
-    def allowed_extension(url):
-        return ('.'.join(url['url'].split('.')[-2:]) in allowed_extensions
-            or url['url'].split('.')[-1:] in allowed_extensions)
-
-    return filter(allowed_extension, urls)
+    return filter(match, urls)
 
 def get_releases(client, package_name):
     try:
@@ -55,7 +52,7 @@ def get_releases(client, package_name):
 
     return releases
 
-def generate_tarball_lorry(lorry_prefix, client, package_name, version=None):
+def generate_archive_lorry(lorry_prefix, client, package_name, version=None):
     releases = get_releases(client, package_name)
 
     if len(releases) == 0:
@@ -79,12 +76,20 @@ def generate_tarball_lorry(lorry_prefix, client, package_name, version=None):
     except Exception as e:
         error("Couldn't fetch release urls:", e)
 
-    tarball_urls = filter_urls(urls)
+    tar_extensions = ('tar.gz', 'tgz', 'tar.Z', 'tar.bz2', 'tbz2',
+                      'tar.lzma', 'tar.xz', 'tlz', 'txz', 'tar')
+    tarball_urls = filter_urls(urls, tar_extensions)
+    zip_urls = filter_urls(urls, ('zip',))
+    archive_type = None
 
     if len(tarball_urls) > 0:
         urls = tarball_urls
+        archive_type = 'tarball'
+    elif len(zip_urls) > 0:
+        urls = zip_urls
+        archive_type = 'zip'
     elif len(urls) > 0:
-        warn("None of these urls look like tarballs:")
+        warn("None of these urls look like archives:")
         for url in urls:
             warn("\t%s" % url['url'])
         error("Cannot proceed")
@@ -94,7 +99,8 @@ def generate_tarball_lorry(lorry_prefix, client, package_name, version=None):
     url = urls[0]['url']
 
     # TODO: shouldn't be hardcoding ext name here
-    return utils.str_tarball_lorry('python', lorry_prefix, package_name, url)
+    return utils.str_archive_lorry(archive_type,
+                                   'python', lorry_prefix, package_name, url)
 
 
 class PythonLorryExtension(ImportExtension):
@@ -133,7 +139,7 @@ class PythonLorryExtension(ImportExtension):
         logging.debug('Treating %s as %s' % (package_name, new_proj_name))
         package_name = new_proj_name
 
-        print(generate_tarball_lorry(lorry_prefix, client,
+        print(generate_archive_lorry(lorry_prefix, client,
                                      package_name, version))
 
 if __name__ == '__main__':
diff --git a/baserockimport/exts/utils.py b/baserockimport/exts/utils.py
index aadefdf..2e2d741 100644
--- a/baserockimport/exts/utils.py
+++ b/baserockimport/exts/utils.py
@@ -44,13 +44,19 @@ def get_compression(url):
 
     return None
 
-# Assumption: url passed to this function must have a 'standard' tar extension
-def str_tarball_lorry(extname, lorry_prefix, package_name, url):
-    ''' Create a json lorry file '''
+def str_archive_lorry(archive_type, extname, lorry_prefix, package_name, url):
+    ''' Create a json lorry file from a url
+
+        Assumption: url passed to this function must have either:
+            * a 'standard' tar extension
+            * a zip extension
+    '''
+
+    assert archive_type in ('tarball', 'zip')
 
     name = '%s/%s' % (lorry_prefix, package_name)
 
-    lorry = {'type': 'tarball',
+    lorry = {'type': archive_type,
              'url': url,
              'x-products-%s' % extname: [package_name]}
 
@@ -58,7 +64,14 @@ def str_tarball_lorry(extname, lorry_prefix, package_name, url):
     if compression:
         lorry['compression'] = compression
 
-    return json.dumps({name + "-tarball": lorry}, indent=4, sort_keys=True)
+    return json.dumps({name + "-%s" % archive_type: lorry},
+                      indent=4, sort_keys=True)
+
+def str_tarball_lorry(extname, lorry_prefix, package_name, url):
+    return str_archive_lorry('tarball', extname, lorry_prefix, package_name, url)
+
+def str_zip_lorry(extname, lorry_prefix, package_name, url):
+    return str_archive_lorry('zip', extname, lorry_prefix, package_name, url)
 
 def str_repo_lorry(extname, lorry_prefix, package_name, repo_type, url):
     name = '%s/%s' % (lorry_prefix, package_name)