Greatly improve zip endpoint with python-zipstream

constants.py

@@ -11,6 +11,8 @@ except converter.ffmpeg.FFMpegError:
     traceback.print_exc()
     ffmpeg = None
 
+FILENAME_BADCHARS = '\\/:*?<>|"'
+
 ALLOWED_ORDERBY_COLUMNS = [
     'extension',
     'width',

etiquette.py

@@ -5,6 +5,7 @@ import mimetypes
 import os
 import random
 import warnings
+import zipstream
 
 import constants
 import decorators
@@ -14,10 +15,6 @@ import jsonify
 import phototagger
 import sessions
 
-# pip install
-# https://raw.githubusercontent.com/voussoir/else/master/_voussoirkit/voussoirkit.zip
-from voussoirkit import webstreamzip
-
 site = flask.Flask(__name__)
 site.config.update(
     SEND_FILE_MAX_AGE_DEFAULT=180,
@@ -270,14 +267,34 @@ def get_album_json(albumid):
     return jsonify.make_json_response(album)
 
 
-@site.route('/album/<albumid>.tar')
-def get_album_tar(albumid):
+@site.route('/album/<albumid>.zip')
+def get_album_zip(albumid):
     album = P_album(albumid)
-    photos = list(album.walk_photos())
-    zipname_map = {p.real_filepath: '%s - %s' % (p.id, p.basename) for p in photos}
-    streamed_zip = webstreamzip.stream_tar(zipname_map)
-    #content_length = sum(p.bytes for p in photos)
-    outgoing_headers = {'Content-Type': 'application/octet-stream'}
+
+    recursive = request.args.get('recursive', True)
+    recursive = helpers.truthystring(recursive)
+    arcnames = 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)
+
+    #if album.description:
+    #    streamed_zip.writestr(
+    #        arcname='%s.txt' % album.id,
+    #        data=album.description.encode('utf-8'),
+    #    )
+
+    if album.title:
+        download_as = '%s - %s.zip' % (album.id, album.title)
+    else:
+        download_as = '%s.zip' % album.id
+    download_as = download_as.replace('"', '\\"')
+    outgoing_headers = {
+        'Content-Type': 'application/octet-stream',
+        'Content-Disposition': 'attachment; filename=%s' % download_as,
+
+    }
     return flask.Response(streamed_zip, headers=outgoing_headers)
 
 
@@ -325,8 +342,6 @@ def get_file(photoid):
         else:
             download_as = photo.id + '.' + photo.extension
 
-        ## Sorry, but otherwise the attachment filename gets terminated
-        #download_as = download_as.replace(';', '-')
         download_as = download_as.replace('"', '\\"')
         response = flask.make_response(send_file(photo.real_filepath))
         response.headers['Content-Disposition'] = 'attachment; filename="%s"' % download_as

helpers.py

@@ -9,6 +9,35 @@ import exceptions
 
 from voussoirkit import bytestring
 
+def album_zip_filenames(album, recursive=True):
+    '''
+    Given an album, produce a dictionary mapping local filepaths 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 album.title:
+        root_folder = '%s - %s' % (album.id, normalize_filepath(album.title))
+    else:
+        root_folder = '%s' % album.id
+
+    photos = album.photos()
+    arcnames = {}
+    for photo in photos:
+        photo_name = '%s - %s' % (photo.id, photo.basename)
+        arcnames[photo.real_filepath] = os.path.join(root_folder, photo_name)
+
+    if recursive:
+        for child_album in album.children():
+            child_arcnames = album_zip_filenames(child_album)
+            for (filepath, arcname) in child_arcnames.items():
+                if filepath in arcnames:
+                    continue
+                arcname = os.path.join(root_folder, arcname)
+                arcnames[filepath] = arcname
+    return arcnames
+
 def chunk_sequence(sequence, chunk_length, allow_incomplete=True):
     '''
     Given a sequence, divide it into sequences of length `chunk_length`.
@@ -132,14 +161,20 @@ def is_xor(*args):
     '''
     return [bool(a) for a in args].count(True) == 1
 
-def normalize_filepath(filepath):
+def normalize_filepath(filepath, allowed=''):
     '''
     Remove some bad characters.
     '''
+    badchars = constants.FILENAME_BADCHARS
+    for character in allowed:
+        badchars = badchars.replace(allowed, '')
+
+    filepath = remove_control_characters(filepath)
+    badchars = dict.fromkeys(badchars)
+    filepath = filepath.translate(badchars)
+
     filepath = filepath.replace('/', os.sep)
     filepath = filepath.replace('\\', os.sep)
-    filepath = filepath.replace('<', '')
-    filepath = filepath.replace('>', '')
     return filepath
 
 def now(timestamp=True):
@@ -171,6 +206,15 @@ def read_filebytes(filepath, range_min, range_max, chunk_size=2 ** 20):
             yield chunk
             sent_amount += len(chunk)
 
+def remove_control_characters(text):
+    '''
+    Thanks SilentGhost
+    http://stackoverflow.com/a/4324823
+    '''
+    kill = dict.fromkeys(range(32))
+    text = text.translate(kill)
+    return text
+
 def seconds_to_hms(seconds):
     '''
     Convert integer number of seconds to an hh:mm:ss string.

phototagger.py

@@ -992,7 +992,7 @@ class PhotoDB(PDBAlbumMixin, PDBPhotoMixin, PDBTagMixin, PDBUserMixin):
             data_directory = constants.DEFAULT_DATADIR
 
         # DATA DIR PREP
-        data_directory = helpers.normalize_filepath(data_directory)
+        data_directory = helpers.normalize_filepath(data_directory, allowed='/\\')
         self.data_directory = os.path.abspath(data_directory)
         os.makedirs(self.data_directory, exist_ok=True)
 

requirements.txt

@@ -2,5 +2,6 @@ bcrypt
 flask
 gevent
 pillow
+zipstream
 https://raw.githubusercontent.com/voussoir/else/master/_voussoirkit/voussoirkit.zip
 git+https://github.com/senko/python-video-converter.git

templates/album.html

@@ -47,7 +47,7 @@ p
         {% endfor %}
     </ul>
     {% endif %}
-    <span><a href="/album/{{album.id}}.tar">(download .tar)</a></span>
+    <span><a href="/album/{{album.id}}.zip">(download .zip)</a></span>
     {% set photos = album.photos() %}
     {% if photos %}
     <h3>Photos</h3>