Rewrite the zip tree building algorithm

- Short version: now it handles the case when the
  ``filelist`` by ``ZipFile`` are not in order.
- Long version:
  - Outer loop: iterate through each file in the
    list and break the full path into segments
  - 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 one 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