Add multi file download

Author: rodriguez-facundo

netpyne_ui/api.py

@@ -20,6 +20,32 @@ def allowed_file(filename, allowed_extensions=ALLOWED_EXTENSIONS):
            filename.rsplit('.', 1)[1].lower() in allowed_extensions
 
 
+def send_files(handler, file_path, filename):
+    with open(file_path, "rb") as f:
+        handler.set_header('Content-Type', 'application/force-download')
+        handler.set_header('Content-Disposition', f"attachment; filename={filename}")
+
+        try:
+            while True:
+                _buffer = f.read(4096)
+                if _buffer:
+                    handler.write(_buffer)
+                else:
+                    return 
+        except:
+            handler.set_status(500, f"Error sending files")
+
+
+def get_file_paths(handler):
+    file_paths = False
+    if 'uri' in handler.request.arguments:
+        file_paths = []
+        tmp_file_paths = [path.decode('utf-8') for path in handler.request.arguments['uri']]
+        for path in tmp_file_paths:
+            if os.path.exists(path):
+                file_paths.append(path)
+    
+    return file_paths
 
 class NetPyNEController:  # pytest: no cover
 
@@ -57,39 +83,26 @@ def uploads(handler: IPythonHandler):
     @get('/downloads')
     def downloads(handler: IPythonHandler):
 
-        if 'uri' in handler.request.arguments:
-            file_paths = []
-            tmp_file_paths = [path.decode('utf-8') for path in handler.request.arguments['uri']]
-            for path in tmp_file_paths:
-                if os.path.exists(path):
-                    file_paths.append(path)
+        file_paths = get_file_paths(handler)
+        
+        if file_paths:
 
             if len(file_paths) == 0:
                 handler.set_status(400, f"Files not found.")
                 handler.finish()
                 return
 
-            with TemporaryDirectory() as dir_path:
-                tar_gz_file_name = f'{str(uuid.uuid4())}.tar.gz'
-                tar_gz_file_path = os.path.join(dir_path, tar_gz_file_name)
-                with tarfile.open(tar_gz_file_path, mode='w:gz') as tar:
-                    for file_path in file_paths:
-                        tar.add(file_path, os.path.join('download',file_path.split('/')[-1]))
-
-                handler.set_header('Content-Type', 'application/force-download')
-                handler.set_header('Content-Disposition', f'attachment; filename={tar_gz_file_name}')
-
-                with open(tar_gz_file_path, "rb") as f:
-                    try:
-                        while True:
-                            _buffer = f.read(4096)
-                            if _buffer:
-                                handler.write(_buffer)
-                            else:
-                                f.close()
-                                handler.finish()
-                                return
-                    except:
-                        handler.set_status(500, f"Error sending files")
+            if len(file_paths) == 1:
+                send_files(handler, file_paths[0], file_paths[0].split('/')[-1])
+                
+            else :
+                with TemporaryDirectory() as dir_path:
+                    tar_gz_file_name = f'{str(uuid.uuid4())}.tar.gz'
+                    tar_gz_file_path = os.path.join(dir_path, tar_gz_file_name)
+                    with tarfile.open(tar_gz_file_path, mode='w:gz') as tar:
+                        for file_path in file_paths:
+                            tar.add(file_path, os.path.join('download', file_path.split('/')[-1]))
+
+                    send_files(handler, tar_gz_file_path, tar_gz_file_name)
 
         handler.finish()