Refactor and download with ".zip" suffix when downloading multiple files

diff --git a/src/pgwui_sql/views/sql.py b/src/pgwui_sql/views/sql.py
index fa8d336cdcd38dc7b79153df15617ec4e3ab6786..2d326e86132186aa6a75270b25e262d79d6ec9d5 100644 (file)
--- a/src/pgwui_sql/views/sql.py
+++ b/src/pgwui_sql/views/sql.py
@@ -40,7 +40,7 @@ from pgwui_core.constants import (
     CSV,
     CHECKED,
     UNCHECKED)
-from pgwui_sql.constants import MANY_FILES
+from pgwui_sql.constants import MANY_FILES_VALUE
 
 log = logging.getLogger(__name__)
 
@@ -380,6 +380,28 @@ def log_response(response):
              f'By user {response["user"]}']))
 
 
+def generate_suffix(uf):
+    '''Return the suffix to use for a downloaded file
+    '''
+    if uf['download_as'] == MANY_FILES_VALUE:
+        return 'zip'
+    elif uf['download_fmt'] == CSV:
+        return 'csv'
+    else:
+        return 'txt'  # Tab delimited
+
+
+def generate_content_type(uf):
+    '''Return the content-type to use for a downloaded file
+    '''
+    if uf['download_as'] == MANY_FILES_VALUE:
+        return 'application/zip'
+    elif uf['download_fmt'] == CSV:
+        return 'text/csv'
+    else:
+        return 'text/plain'  # Tab delimited
+
+
 @view_config(route_name='pgwui_sql',
              renderer='pgwui_sql:templates/sql.mak')
 @auth_base_view
@@ -391,18 +413,10 @@ def sql_view(request):
     if uh.uf['download'] and not response['errors']:
         pmd_response = pyramid.response.Response()
         pmd_response.cache_control = 'public; max-age=0'
-        if uh.uf['download_as'] == MANY_FILES:
-            pmd_response.content_type = 'application/zip'
-        elif uh.uf['download_fmt'] == CSV:
-            pmd_response.content_type = 'text/csv'
-        else:
-            pmd_response.content_type = 'text/plain'
-        if uh.uf['download_fmt'] == CSV:
-            suffix = 'csv'
-        else:
-            suffix = 'txt'  # Tab delimited
-        pmd_response.content_disposition = \
-            f'attachment; filename={response["dl_filename"]}.{suffix}'
+        pmd_response.content_type = generate_content_type(uh.uf)
+        pmd_response.content_disposition = (
+            'attachment;'
+            f' filename={response["dl_filename"]}.{generate_suffix(uh.uf)}')
 
         pmd_response.app_iter = codecs.iterencode(uh.tfile, 'utf_8')