Support generic type of download file

Author: DanSava

CHANGELOG.md

@@ -14,6 +14,8 @@ now has search functionality (using [`docsify` full text search](https://docsify
 enforced by `webviz-config` (inline script hashes are added automatically).
 
 ### Changed
+- [#294](https://github.com/equinor/webviz-config/pull/294) - Plugin authors can now define file type to download
+(including specifying MIME type). Before only `.zip` archives were supported.
 - [#281](https://github.com/equinor/webviz-config/pull/281) - Now uses `importlib` instead of `pkg_resources` for
 detecting plugin entry points and package versions.
 

CONTRIBUTING.md

@@ -83,7 +83,7 @@ a button toolbar. The default buttons to appear is stored in the class constant
 `WebvizPluginABC.TOOLBAR_BUTTONS`. If you want to override which buttons should
 appear, redefine this class constant in your subclass. To remove all buttons,
 simply define it as an empty list. See [this section](#data-download-callback)
-for more information regarding the `data_download` button.
+for more information regarding downloading plugin data.
 
 ### Callbacks
 
@@ -150,17 +150,40 @@ There are three fundamental additions to the minimal example without callbacks:
 
 There is a [data download button](#override-plugin-toolbar) provided by
 the `WebvizPluginABC` class. However, it will only appear if the corresponding
-callback is set. A typical data download callback will look like
+callback is set. A typical compressed data download callback will look like
 
 ```python
 @app.callback(self.plugin_data_output,
-              [self.plugin_data_requested])
+              self.plugin_data_requested)
 def _user_download_data(data_requested):
-    return WebvizPluginABC.plugin_data_compress(
-        [{'filename': 'some_file.txt',
-          'content': 'Some download data'}]
-    ) if data_requested else ''
+    return (
+        WebvizPluginABC.plugin_compressed_data(
+            filename="webviz-data.zip",
+            content=[{"filename": "some_file.txt", "content": "Some download data"}],
+        )
+        if data_requested
+        else None
+    )
 ```
+
+A typical CSV data download from e.g. a `pandas.DataFrame` will look like:
+```python
+@app.callback(self.plugin_data_output,
+              self.plugin_data_requested)
+def _user_download_data(data_requested):
+    return (
+        {
+            "filename": "some-file.csv",
+            "content": base64.b64encode(
+                some_pandas_dataframe.to_csv().encode()
+            ).decode("ascii"),
+            "mime_type": "text/csv",
+        }
+        if data_requested
+        else None
+    )     
+```
+
 By letting the plugin define the callback, the plugin author is able
 to utilize the whole callback machinery, including e.g. state of the individual
 components in the plugin. This way the data downloaded can e.g. depend on
@@ -170,8 +193,8 @@ The attributes `self.plugin_data_output` and `self.plugin_data_requested`
 are Dash `Output` and `Input` instances respectively, and are provided by
 the base class `WebvizPluginABC` (i.e. include them as shown here).
 
-The function `WebvizPluginABC.plugin_data_compress` is a utility function
-which takes a list of dictionaries, giving filenames and corresponding data,
+The function `WebvizPluginABC.plugin_compressed_data` is a utility function
+which takes a file name and a list of dictionaries, containing file names and corresponding data,
 and compresses them to a zip archive which is then downloaded by the user.
 
 ### User provided arguments

setup.py

@@ -48,7 +48,7 @@
         "tqdm>=4.8",
         "importlib-metadata>=1.7; python_version<'3.8'",
         "typing-extensions>=3.7; python_version<'3.8'",
-        "webviz-core-components>=0.0.19",
+        "webviz-core-components>=0.1.0",
     ],
     tests_require=TESTS_REQUIRES,
     extras_require={"tests": TESTS_REQUIRES},

webviz_config/__init__.py

@@ -8,7 +8,7 @@
 from ._theme_class import WebvizConfigTheme
 from ._localhost_token import LocalhostToken
 from ._is_reload_process import is_reload_process
-from ._plugin_abc import WebvizPluginABC
+from ._plugin_abc import WebvizPluginABC, EncodedFile, ZipFileMember
 from ._shared_settings_subscriptions import SHARED_SETTINGS_SUBSCRIPTIONS
 
 try:

webviz_config/_plugin_abc.py

@@ -2,15 +2,39 @@
 import abc
 import base64
 import zipfile
+import warnings
+
 from uuid import uuid4
 from typing import List, Optional, Type, Union
 
+try:
+    # Python 3.8+
+    # pylint: disable=ungrouped-imports
+    from typing import TypedDict  # type: ignore
+except (ImportError, ModuleNotFoundError):
+    # Python < 3.8
+    from typing_extensions import TypedDict  # type: ignore
+
+
 import bleach
 from dash.development.base_component import Component
 from dash.dependencies import Input, Output
 import webviz_core_components as wcc
 
 
+class ZipFileMember(TypedDict):
+    filename: str  # Filename to be given within the archive
+    content: str  # String of file to be added
+
+
+class EncodedFile(ZipFileMember):
+    # Same keys as in ZipFileMember, with the following changes:
+    # - filename is now name of the actual downloaded file.
+    # - content is now a base64 encoded ASCII string.
+    # - mime_type needs to be added as well.
+    mime_type: str
+
+
 class WebvizPluginABC(abc.ABC):
     """All webviz plugins need to subclass this abstract base class,
     e.g.
@@ -40,7 +64,6 @@ def layout(self):
     TOOLBAR_BUTTONS = [
         "screenshot",
         "expand",
-        "download_zip",
         "contact_person",
         "guided_tour",
     ]
@@ -62,6 +85,7 @@ def __init__(self, screenshot_filename: str = "webviz-screenshot.png") -> None:
 
         self._plugin_uuid = uuid4()
         self._screenshot_filename = screenshot_filename
+        self._add_download_button = False
 
     def uuid(self, element: str) -> str:
         """Typically used to get a unique ID for some given element/component in
@@ -94,10 +118,8 @@ def _plugin_wrapper_id(self) -> str:
 
     @property
     def plugin_data_output(self) -> Output:
-        # pylint: disable=attribute-defined-outside-init
-        # We do not have a __init__ method in this abstract base class
         self._add_download_button = True
-        return Output(self._plugin_wrapper_id, "zip_base64")
+        return Output(self._plugin_wrapper_id, "download")
 
     @property
     def plugin_data_requested(self) -> Input:
@@ -110,16 +132,28 @@ def _reformat_tour_steps(steps: List[dict]) -> List[dict]:
         ]
 
     @staticmethod
-    def plugin_data_compress(content: List[dict]) -> str:
-        byte_io = io.BytesIO()
-
-        with zipfile.ZipFile(byte_io, "w") as zipped_data:
-            for data in content:
-                zipped_data.writestr(data["filename"], data["content"])
-
-        byte_io.seek(0)
+    def plugin_compressed_data(
+        filename: str, content: List[ZipFileMember]
+    ) -> EncodedFile:
+        with io.BytesIO() as bytes_io:
+            with zipfile.ZipFile(bytes_io, "w") as zipped_data:
+                for data in content:
+                    zipped_data.writestr(data["filename"], data["content"])
+
+            bytes_io.seek(0)
+            return EncodedFile(
+                filename=filename,
+                content=base64.b64encode(bytes_io.read()).decode("ascii"),
+                mime_type="application/zip",
+            )
 
-        return base64.b64encode(byte_io.read()).decode("ascii")
+    @staticmethod
+    def plugin_data_compress(content: List[ZipFileMember]) -> EncodedFile:
+        warnings.warn(
+            "Use 'plugin_compressed_data' instead of 'plugin_data_compress'",
+            DeprecationWarning,
+        )
+        return WebvizPluginABC.plugin_compressed_data("webviz-data.zip", content)
 
     def plugin_layout(
         self, contact_person: Optional[dict] = None
@@ -145,8 +179,8 @@ def plugin_layout(
             for key in contact_person:
                 contact_person[key] = bleach.clean(str(contact_person[key]))
 
-        if "download_zip" in buttons and not hasattr(self, "_add_download_button"):
-            buttons.remove("download_zip")
+        if self._add_download_button:
+            buttons.append("download")
 
         if buttons:
             # pylint: disable=no-member

webviz_config/plugins/_example_data_download.py

@@ -1,7 +1,9 @@
+from typing import Optional
+
 import dash_html_components as html
 from dash import Dash
 
-from .. import WebvizPluginABC
+from .. import WebvizPluginABC, EncodedFile
 
 
 class ExampleDataDownload(WebvizPluginABC):
@@ -16,12 +18,15 @@ def layout(self) -> html.H1:
         return html.H1(self.title)
 
     def set_callbacks(self, app: Dash) -> None:
-        @app.callback(self.plugin_data_output, [self.plugin_data_requested])
-        def _user_download_data(data_requested: bool) -> str:
+        @app.callback(self.plugin_data_output, self.plugin_data_requested)
+        def _user_download_data(data_requested: bool) -> Optional[EncodedFile]:
             return (
-                WebvizPluginABC.plugin_data_compress(
-                    [{"filename": "some_file.txt", "content": "Some download data"}]
+                WebvizPluginABC.plugin_compressed_data(
+                    filename="webviz-data.zip",
+                    content=[
+                        {"filename": "some_file.txt", "content": "Some download data"}
+                    ],
                 )
                 if data_requested
-                else ""
+                else None
             )

webviz_config/plugins/_table_plotter.py

@@ -1,7 +1,8 @@
+import base64
+import inspect
 from pathlib import Path
 from collections import OrderedDict
 from typing import Optional, List, Dict, Any
-import inspect
 
 import numpy as np
 import pandas as pd
@@ -12,7 +13,7 @@
 from dash import Dash
 import webviz_core_components as wcc
 
-from .. import WebvizPluginABC
+from .. import WebvizPluginABC, EncodedFile
 from ..webviz_store import webvizstore
 from ..common_cache import CACHE
 
@@ -356,19 +357,18 @@ def plot_input_callbacks(self) -> List[Input]:
         return inputs
 
     def set_callbacks(self, app: Dash) -> None:
-        @app.callback(self.plugin_data_output, [self.plugin_data_requested])
-        def _user_download_data(data_requested: Optional[int]) -> str:
+        @app.callback(self.plugin_data_output, self.plugin_data_requested)
+        def _user_download_data(data_requested: Optional[int]) -> Optional[EncodedFile]:
             return (
-                WebvizPluginABC.plugin_data_compress(
-                    [
-                        {
-                            "filename": "table_plotter.csv",
-                            "content": get_data(self.csv_file).to_csv(),
-                        }
-                    ]
-                )
+                {
+                    "filename": "table-plotter.csv",
+                    "content": base64.b64encode(
+                        get_data(self.csv_file).to_csv().encode()
+                    ).decode("ascii"),
+                    "mime_type": "text/csv",
+                }
                 if data_requested
-                else ""
+                else None
             )
 
         @app.callback(self.plot_output_callbacks, self.plot_input_callbacks)