Include album info as txt in zip; fix normalize_filepath bugs

etiquette.py

@@ -279,11 +279,20 @@ def get_album_zip(albumid):
     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'),
-    #    )
+    directories = helpers.album_zip_directories(album, recursive=recursive)
+    for (inner_album, directory) in directories.items():
+        text = []
+        if inner_album.title:
+            text.append(inner_album.title)
+        if inner_album.description:
+            text.append(inner_album.description)
+        if not text:
+            continue
+        text = '\r\n\r\n'.join(text)
+        streamed_zip.writestr(
+            arcname=os.path.join(directory, '%s.txt' % inner_album.id),
+            data=text.encode('utf-8'),
+        )
 
     if album.title:
         download_as = '%s - %s.zip' % (album.id, album.title)

helpers.py

@@ -9,6 +9,27 @@ import exceptions
 
 from voussoirkit import bytestring
 
+def album_zip_directories(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.
+    '''
+    directories = {}
+    if album.title:
+        root_folder = '%s - %s' % (album.id, normalize_filepath(album.title))
+    else:
+        root_folder = '%s' % album.id
+
+    directories[album] = root_folder
+    if recursive:
+        for child_album in album.children():
+            child_directories = album_zip_directories(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):
     '''
     Given an album, produce a dictionary mapping local filepaths to the filenames
@@ -17,25 +38,16 @@ def album_zip_filenames(album, recursive=True):
 
     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)
+    directories = album_zip_directories(album, recursive=recursive)
+    for (album, directory) in directories.items():
+        photos = album.photos()
+        for photo in photos:
+            if photo.real_filepath in arcnames:
+                continue
+            photo_name = '%s - %s' % (photo.id, photo.basename)
+            arcnames[photo.real_filepath] = os.path.join(directory, 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):
@@ -165,13 +177,8 @@ 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)
+    badchars = remove_characters(constants.FILENAME_BADCHARS, allowed)
+    filepath = remove_characters(filepath, badchars)
 
     filepath = filepath.replace('/', os.sep)
     filepath = filepath.replace('\\', os.sep)
@@ -206,13 +213,18 @@ def read_filebytes(filepath, range_min, range_max, chunk_size=2 ** 20):
             yield chunk
             sent_amount += len(chunk)
 
+def remove_characters(text, characters):
+    translator = {ord(c): None for c in characters}
+    text = text.translate(translator)
+    return text
+
 def remove_control_characters(text):
     '''
     Thanks SilentGhost
     http://stackoverflow.com/a/4324823
     '''
-    kill = dict.fromkeys(range(32))
-    text = text.translate(kill)
+    translator = dict.fromkeys(range(32))
+    text = text.translate(translator)
     return text
 
 def seconds_to_hms(seconds):

objects.py

@@ -175,6 +175,9 @@ class Album(ObjectBase, GroupableMixin):
         self.name = 'Album %s' % self.id
         self.group_getter = self.photodb.get_album
 
+    def __hash__(self):
+        return hash(self.id)
+
     def __repr__(self):
         return 'Album:{id}'.format(id=self.id)
 
@@ -280,7 +283,7 @@ class Photo(ObjectBase):
 
         self.id = row_tuple['id']
         self.real_filepath = row_tuple['filepath']
-        self.real_filepath = helpers.normalize_filepath(self.real_filepath)
+        self.real_filepath = helpers.normalize_filepath(self.real_filepath, allowed=':\\/')
         self.real_path = pathclass.Path(self.real_filepath)
         self.filepath = row_tuple['override_filename'] or self.real_filepath
         self.basename = row_tuple['override_filename'] or os.path.basename(self.real_filepath)
@@ -572,7 +575,7 @@ class Photo(ObjectBase):
         old_path = self.real_path
         old_path.correct_case()
 
-        new_filename = helpers.normalize_filepath(new_filename)
+        new_filename = helpers.normalize_filepath(new_filename, allowed=':\\/')
         if os.path.dirname(new_filename) == '':
             new_path = old_path.parent.with_child(new_filename)
         else:

templates/album.html

@@ -47,7 +47,11 @@ p
         {% endfor %}
     </ul>
     {% endif %}
-    <span><a href="/album/{{album.id}}.zip">(download .zip)</a></span>
+    <span>
+        Download:
+        <a href="/album/{{album.id}}.zip?recursive=no">These files</a>
+        {% if sub_albums %} | <a href="/album/{{album.id}}.zip?recursive=yes">Include children</a>{% endif %}
+    </span>
     {% set photos = album.photos() %}
     {% if photos %}
     <h3>Photos</h3>