Suggestions after review

equinor · Sep 20, 2020 · 5860483 · 5860483
1 parent 342f640
commit 5860483
Show file tree

Hide file tree

Showing 7 changed files with 63 additions and 52 deletions.
diff --git a/CHANGELOG.md b/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.
 

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -156,28 +156,32 @@ callback is set. A typical compressed data download callback will look like
 @app.callback(self.plugin_data_output,
               self.plugin_data_requested)
 def _user_download_data(data_requested):
-    return WebvizPluginABC.plugin_compressed_data(
-        file_name="webviz-data.zip",
-        content=[{'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 raw CSV data download will look like:
+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):
-    if data_requested:
-        byte_io = io.BytesIO()
-        byte_io.write(get_data(self.csv_file).to_csv().encode('ascii'))
-        byte_io.seek(0)
-        return {
-                    "file_name": "file_name.csv",
-                    "content": base64.b64encode(byte_io.read()).decode("ascii"), 
-                    "mime_type": "text/csv"
-                }
-    return {}      
+    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

diff --git a/setup.py b/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.0.19",  # bump before merge
     ],
     tests_require=TESTS_REQUIRES,
     extras_require={"tests": TESTS_REQUIRES},

diff --git a/webviz_config/__init__.py b/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:

diff --git a/webviz_config/_plugin_abc.py b/webviz_config/_plugin_abc.py
@@ -22,9 +22,12 @@
 import webviz_core_components as wcc
 
 
-class FileDefinition(TypedDict, total=False):
-    file_name: str
+class ZipFileMember(TypedDict):
+    filename: str
     content: str
+
+
+class EncodedFile(ZipFileMember):
     mime_type: str
 
 
@@ -126,23 +129,22 @@ def _reformat_tour_steps(steps: List[dict]) -> List[dict]:
 
     @staticmethod
     def plugin_compressed_data(
-        file_name: str, content: List[FileDefinition]
-    ) -> FileDefinition:
-        byte_io = io.BytesIO()
-
-        with zipfile.ZipFile(byte_io, "w") as zipped_data:
-            for data in content:
-                zipped_data.writestr(data["file_name"], data["content"])
-
-        byte_io.seek(0)
-        return {
-            "file_name": file_name,
-            "content": base64.b64encode(byte_io.read()).decode("ascii"),
-            "mime_type": "application/zip",
-        }
+        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",
+            )
 
     @staticmethod
-    def plugin_data_compress(content: List[dict]) -> FileDefinition:
+    def plugin_data_compress(content: List[ZipFileMember]) -> EncodedFile:
         warnings.warn(
             "Use 'plugin_compressed_data' instead of 'plugin_data_compress'",
             DeprecationWarning,

diff --git a/webviz_config/plugins/_example_data_download.py b/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, ZipFileMember
 
 
 class ExampleDataDownload(WebvizPluginABC):
@@ -17,14 +19,16 @@ def layout(self) -> html.H1:
 
     def set_callbacks(self, app: Dash) -> None:
         @app.callback(self.plugin_data_output, self.plugin_data_requested)
-        def _user_download_data(data_requested: bool) -> dict:
+        def _user_download_data(data_requested: bool) -> Optional[EncodedFile]:
             return (
                 WebvizPluginABC.plugin_compressed_data(
-                    file_name="webviz-data.zip",
+                    filename="webviz-data.zip",
                     content=[
-                        {"filename": "some_file.txt", "content": "Some download data"}
+                        ZipFileMember(
+                            filename="some_file.txt", content="Some download data"
+                        )
                     ],
                 )
                 if data_requested
-                else {}
+                else None
             )
diff --git a/webviz_config/plugins/_table_plotter.py b/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
 
@@ -357,19 +358,17 @@ def plot_input_callbacks(self) -> List[Input]:
 
     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]) -> dict:
+        def _user_download_data(data_requested: Optional[int]) -> Optional[EncodedFile]:
             return (
-                WebvizPluginABC.plugin_compressed_data(
-                    file_name="webviz-data.zip",
-                    content=[
-                        {
-                            "filename": "table_plotter.csv",
-                            "content": get_data(self.csv_file).to_csv(),
-                        }
-                    ],
+                EncodedFile(
+                    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)