Rewrite the zip tree building algorithm

- Handles the case when the generated ``filelist`` is not in order.
- The outer loop: iterate through each file in the list and break the
  full path into segments.
- The inner loop: iterate through each segment, add a new node if not
  found in the child nodes of its parent and update the node with the
  file's details if the segment is the last one.
- The full tree is complete at the end of only a single pass.

Author: cslzchen

mfr/extensions/zip/render.py

@@ -1,6 +1,6 @@
 import os
-from typing import List
-from zipfile import ZipFile
+from typing import List, Union
+from zipfile import ZipFile, ZipInfo
 
 from mako.lookup import TemplateLookup
 
@@ -27,82 +27,108 @@ def render(self):
 
         zip_file = ZipFile(self.file_path, 'r')
 
-        file_list = self.sanitize_file_list(zip_file.filelist)
-        file_tree = self.file_list_to_tree(file_list)
+        # ``ZipFile.filelist`` contains both files and folder.  Using ``obj`` for better clarity.
+        obj_list = self.sanitize_obj_list(zip_file.filelist)
+        obj_tree = self.obj_list_to_tree(obj_list)
 
-        return self.TEMPLATE.render(data=file_tree, base=self.assets_url)
+        return self.TEMPLATE.render(data=obj_tree, base=self.assets_url)
 
-    def file_list_to_tree(self, file_list: list) -> List[dict]:
-        """Build the file tree and return a "tree".
+    def obj_list_to_tree(self, obj_list: list) -> List[dict]:
+        """Build the object tree from the object list.  Each node is represented using a dictionary,
+        where non-leaf nodes represent folders and leaves represent files.  Return a list which
+        contains only one element: the root node.
 
-        TODO: Fix this algorithm
-        This algorithm only works when the ``file_list`` are in strict alphabetical order. Here is
-        an example file A.zip where list 1 fails while list 2 succeed.
-
-        A.zip
-        --- A/
-            --- A/aa.png
-            --- B/ab.png
-
-        File list 1: [ A/, A/B/, A/A/, A/A/aa.png, A/B/ab.png, ]
-
-        File list 2: [ A/, A/A/, A/A/aa.png, A/B/, A/B/ab.png, ]
-
-        :param file_list: the sanitized file list
+        :param obj_list: the object list
         :rtype: ``List[dict]``
-        :return: a "tree" in form of a list which contains one dictionary as the root node
+        :return: a list which contains only one element: the root node.
         """
 
-        icons_url = self.assets_url + '/img'
-
-        # Build the root of the file tree
-        tree_root = [{
+        # Build the root node of the tree
+        tree_root = {
             'text': self.metadata.name + self.metadata.ext,
-            'icon': icons_url + '/file-ext-zip.png',
+            'icon': self.assets_url + '/img/file-ext-zip.png',
             'children': []
-        }]
-
-        # Iteratively build the file tree for each file and folder.egments.
-        for file in file_list:
-
-            node_path = tree_root[0]
-
-            # Split the full path into segments, add each path segment to the tree if the segment
-            # doesn't already exist.  The segments can be either a folder or a file.
-            paths = [path for path in file.filename.split('/') if path]
-            for path in paths:
-
-                # Add a child to the node
-                if not len(node_path['children']) or node_path['children'][-1]['text'] != path:
-
-                    new_node = {'text': path, 'children': []}
-
-                    date = '%d-%02d-%02d %02d:%02d:%02d' % file.date_time[:6]
-                    size = sizeof_fmt(int(file.file_size)) if file.file_size else ''
-                    new_node['data'] = {'date': date, 'size': size}
+        }
+
+        for obj in obj_list:
+
+            # For each object, always start from the root of the tree
+            parent = tree_root
+            path_from_root = obj.filename
+            is_folder = path_from_root[-1] == '/'
+            path_segments = [segment for segment in path_from_root.split('/') if segment]
+            last_index = len(path_segments) - 1
+
+            # Iterate through the path segments list.  Add the segment to tree if not already there
+            # and update the details with the current object if it is the last one along the path.
+            for index, segment in enumerate(path_segments):
+
+                # Check if the segment has already been added
+                siblings = parent.get('children', [])
+                current_node = self.find_node_among_siblings(segment, siblings)
+
+                # Found
+                if current_node:
+                    if index == last_index:
+                        # If it is the last segment, this node must be a folder and represents the
+                        # current object.  Update it with the objects' info and break.
+                        assert is_folder
+                        self.update_node_with_attributes(current_node, obj, is_folder=is_folder)
+                        break
+                    # Otherwise, jump to the next segment with the current node as the new parent
+                    parent = current_node
+                    continue
+
+                # Not found
+                new_node = {
+                    'text': segment,
+                    'children': [],
+                }
+                if index == last_index:
+                    # If it is the last segment, the node represents the current object.  Update the
+                    # it with the objects' info, add it to the siblings and break.
+                    self.update_node_with_attributes(new_node, obj, is_folder=is_folder)
+                    siblings.append(new_node)
+                    break
+
+                # Otherwise, append the new node to tree, jump to the next segment with the current
+                # node as the new parent
+                siblings.append(new_node)
+                parent = new_node
+                continue
 
-                    if file.filename[-1] == '/':
-                        new_node['icon'] = icons_url + '/folder.png'
-                    else:
-                        ext = os.path.splitext(file.filename)[1].lstrip('.')
-                        if ext:
-                            ext = ext.lower()
-                        if self.icon_exists_for_type(ext):
-                            new_node['icon'] = '{}/file-ext-{}.png'.format(icons_url, ext)
-                        else:
-                            new_node['icon'] = '{}/file-ext-generic.png'.format(icons_url)
+        return [tree_root, ]
 
-                    node_path['children'].append(new_node)
+    def update_node_with_attributes(self, node: dict, obj: ZipInfo, is_folder: bool) -> None:
+        """Update details (date, size, icon, etc.) of the node with the given object.
 
-                    node_path = new_node
-                # Go one level deeper
-                else:
-                    node_path = node_path['children'][-1]
+        :param node: the node to update
+        :param obj: the object that the node represents
+        :param is_folder: the folder flag
+        """
 
-        return tree_root
+        date = '%d-%02d-%02d %02d:%02d:%02d' % obj.date_time[:6]
+        size = sizeof_fmt(int(obj.file_size)) if obj.file_size else ''
+
+        if is_folder:
+            icon_path = self.assets_url + '/img/folder.png'
+        else:
+            ext = (os.path.splitext(obj.filename)[1].lstrip('.')).lower()
+            if self.icon_exists(ext):
+                icon_path = '{}/img/file-ext-{}.png'.format(self.assets_url, ext)
+            else:
+                icon_path = '{}/img/file-ext-generic.png'.format(self.assets_url)
+
+        node.update({
+            'icon': icon_path,
+            'data': {
+                'date': date,
+                'size': size,
+            },
+        })
 
     @staticmethod
-    def icon_exists_for_type(ext: str) -> bool:
+    def icon_exists(ext: str) -> bool:
         """Check if an icon exists for the given file type.  The extension string is converted to
         lower case.
 
@@ -119,28 +145,44 @@ def icon_exists_for_type(ext: str) -> bool:
         ))
 
     @staticmethod
-    def sanitize_file_list(file_list: list) -> list:
+    def sanitize_obj_list(obj_list: list) -> list:
         """Remove macOS system and temporary files.  Current implementation only removes '__MACOSX/'
         and '.DS_Store'.  If necessary, extend the sanitizer to exclude more file types.
 
-        :param file_list: the list of the path for each file and folder in the zip
+        :param obj_list: a list of full paths for each file and folder in the zip
         :rtype: ``list``
         :return: a sanitized list
         """
 
-        sanitized_file_list = []
+        sanitized_obj_list = []
 
-        for file in file_list:
+        for obj in obj_list:
 
-            file_path = file.filename
+            obj_path = obj.filename
             # Ignore macOS '__MACOSX' folder for zip file
-            if file_path.startswith('__MACOSX/'):
+            if obj_path.startswith('__MACOSX/'):
                 continue
-
             # Ignore macOS '.DS_STORE' file
-            if file_path == '.DS_Store' or file_path.endswith('/.DS_Store'):
+            if obj_path == '.DS_Store' or obj_path.endswith('/.DS_Store'):
                 continue
 
-            sanitized_file_list.append(file)
+            sanitized_obj_list.append(obj)
+
+        return sanitized_obj_list
+
+    @staticmethod
+    def find_node_among_siblings(segment: str, siblings: list) -> Union[dict, None]:
+        """Find if the folder or file represented by the path segment has already been added.
+
+        :param segment: the path segment
+        :param siblings: the list containing all added sibling nodes
+        :rtype: ``Union[dict, None]``
+        :return: the node if found or ``None`` otherwise
+        """
+
+        for sibling in siblings:
+
+            if sibling.get('text', '') == segment:
+                return sibling
 
-        return sanitized_file_list
+        return None