Improve & generalize zipfile code.

Moved some heavy lifting out of the flask album.zip endpoint
and into helpers.py.
Renamed some things for clarity.

etiquette/helpers.py

@@ -10,6 +10,7 @@ import mimetypes
 import os
 import PIL.Image
 import unicodedata
+import zipstream
 
 from . import constants
 from . import exceptions
@@ -17,11 +18,13 @@ from . import exceptions
 from voussoirkit import bytestring
 from voussoirkit import pathclass
 
-def album_zip_directories(album, recursive=True):
+def album_as_directory_map(album, recursive=True):
     '''
     Given an album, produce a dictionary mapping Album objects to directory
     names as they will appear inside the zip archive.
     Sub-albums become subfolders.
+
+    If an album is a child of multiple albums, only one instance is used.
     '''
     directories = {}
     if album.title:
@@ -33,30 +36,28 @@ def album_zip_directories(album, recursive=True):
     directories[album] = root_folder
     if recursive:
         for child_album in album.get_children():
-            child_directories = album_zip_directories(child_album, recursive=True)
+            child_directories = album_as_directory_map(child_album, recursive=True)
             for (child_album, child_directory) in child_directories.items():
                 child_directory = os.path.join(root_folder, child_directory)
                 directories[child_album] = child_directory
+
     return directories
 
-def album_zip_filenames(album, recursive=True):
+def album_photos_as_filename_map(album, recursive=True):
     '''
-    Given an album, produce a dictionary mapping local filepaths to the
+    Given an album, produce a dictionary mapping Photo objects to the
     filenames that will appear inside the zip archive.
     This includes creating subfolders for sub albums.
 
-    If a photo appears in multiple albums, only the first is used.
+    If a photo appears in multiple albums, only one instance is used.
     '''
     arcnames = {}
-    directories = album_zip_directories(album, recursive=recursive)
+    directories = album_as_directory_map(album, recursive=recursive)
     for (album, directory) in directories.items():
         photos = album.get_photos()
         for photo in photos:
-            filepath = photo.real_path.absolute_path
-            if filepath in arcnames:
-                continue
             photo_name = f'{photo.id} - {photo.basename}'
-            arcnames[filepath] = os.path.join(directory, photo_name)
+            arcnames[photo] = os.path.join(directory, photo_name)
 
     return arcnames
 
@@ -495,6 +496,57 @@ def truthystring(s):
         return None
     return False
 
+def zip_album(album, recursive=True):
+    '''
+    Given an album, return a zipstream zipfile that contains the album's
+    photos (recursive = include childen's photos) organized into folders
+    for each album. Each album folder also gets a text file containing
+    the album's name and description if applicable.
+
+    If an album is a child of multiple albums, only one instance is used.
+    '''
+    zipfile = zipstream.ZipFile()
+
+    # Add the photos.
+    arcnames = album_photos_as_filename_map(album, recursive=recursive)
+    for (photo, arcname) in arcnames.items():
+        zipfile.write(filename=photo.real_path.absolute_path, arcname=arcname)
+
+    # Add the album metadata as an {id}.txt file within each directory.
+    directories = album_as_directory_map(album, recursive=recursive)
+    for (inner_album, directory) in directories.items():
+        metafile_text = []
+        if inner_album.title:
+            metafile_text.append(f'Title: {inner_album.title}')
+
+        if inner_album.description:
+            metafile_text.append(f'Description: {inner_album.description}')
+
+        if not metafile_text:
+            continue
+
+        metafile_text = '\r\n\r\n'.join(metafile_text)
+        metafile_text = metafile_text.encode('utf-8')
+        metafile_name = f'album {inner_album.id}.txt'
+        metafile_name = os.path.join(directory, metafile_name)
+        zipfile.writestr(
+            arcname=metafile_name,
+            data=metafile_text,
+        )
+
+    return zipfile
+
+def zip_photos(photos):
+    '''
+    Given some photos, return a zipstream zipfile that contains the files.
+    '''
+    zipfile = zipstream.ZipFile()
+
+    for photo in photos:
+        arcname = os.path.join('photos', f'{photo.id} - {photo.basename}')
+        zipfile.write(filename=photo.real_path.absolute_path, arcname=arcname)
+
+    return zipfile
 
 _numerical_characters = set('0123456789.')
 def _unitconvert(value):

frontends/etiquette_flask/etiquette_flask/endpoints/album_endpoints.py

@@ -45,40 +45,20 @@ def get_album_zip(album_id):
     recursive = request.args.get('recursive', True)
     recursive = etiquette.helpers.truthystring(recursive)
 
-    arcnames = etiquette.helpers.album_zip_filenames(album, recursive=recursive)
-
-    streamed_zip = zipstream.ZipFile()
-    for (real_filepath, arcname) in arcnames.items():
-        streamed_zip.write(real_filepath, arcname=arcname)
-
-    # Add the album metadata as an {id}.txt file within each directory.
-    directories = etiquette.helpers.album_zip_directories(album, recursive=recursive)
-    for (inner_album, directory) in directories.items():
-        text = []
-        if inner_album.title:
-            text.append('Title: ' + inner_album.title)
-        if inner_album.description:
-            text.append('Description: ' + inner_album.description)
-        if not text:
-            continue
-        text = '\r\n\r\n'.join(text)
-        streamed_zip.writestr(
-            arcname=os.path.join(directory, 'album %s.txt' % inner_album.id),
-            data=text.encode('utf-8'),
-        )
+    streamed_zip = etiquette.helpers.zip_album(album, recursive=recursive)
 
     if album.title:
-        download_as = 'album %s - %s.zip' % (album.id, album.title)
+        download_as = f'album {album.id} - {album.title}.zip'
     else:
-        download_as = 'album %s.zip' % album.id
+        download_as = f'album {album.id}.zip'
 
     download_as = etiquette.helpers.remove_path_badchars(download_as)
     download_as = urllib.parse.quote(download_as)
     outgoing_headers = {
         'Content-Type': 'application/octet-stream',
-        'Content-Disposition': 'attachment; filename*=UTF-8\'\'%s' % download_as,
-
+        'Content-Disposition': f'attachment; filename*=UTF-8\'\'{download_as}',
     }
+
     return flask.Response(streamed_zip, headers=outgoing_headers)
 
 # Album photo operations ###########################################################################