Fix download of binary files in FolderDataViewer (#579)

aiidalab_widgets_base/viewers.py

@@ -1521,20 +1521,25 @@ def __init__(self, folder, downloadable=True, **kwargs):
         super().__init__(children, **kwargs)
 
     def change_file_view(self, _=None):
-        with self._folder.base.repository.open(self.files.value) as fobj:
-            self.text.value = fobj.read()
+        try:
+            with self._folder.base.repository.open(self.files.value) as fobj:
+                self.text.value = fobj.read()
+        except UnicodeDecodeError:
+            self.text.value = "[Binary file, preview not available]"
 
     def download(self, _=None):
-        """Prepare for downloading."""
+        """Download selected file."""
         from IPython.display import Javascript
 
-        payload = base64.b64encode(
-            self._folder.get_object_content(self.files.value).encode()
-        ).decode()
+        # TODO: Preparing large files for download might take a while.
+        # Can we do a streaming solution?
+        raw_bytes = self._folder.get_object_content(self.files.value, "rb")
+        base64_payload = base64.b64encode(raw_bytes).decode()
+
         javas = Javascript(
             f"""
             var link = document.createElement('a');
-            link.href = "data:;base64,{payload}"
+            link.href = "data:;base64,{base64_payload}"
             link.download = "{self.files.value}"
             document.body.appendChild(link);
             link.click();

tests/conftest.py

@@ -264,12 +264,15 @@ def folder_data_object():
     """Return a `FolderData` object."""
     FolderData = plugins.DataFactory("core.folder")  # noqa: N806
     folder_data = FolderData()
-    with io.StringIO("content of test1 filelike") as fobj:
+    with io.StringIO("content of test1.txt") as fobj:
         folder_data.put_object_from_filelike(fobj, path="test1.txt")
-    with io.StringIO("content of test2 filelike") as fobj:
+    with io.StringIO("content of test2.txt") as fobj:
         folder_data.put_object_from_filelike(fobj, path="test2.txt")
-    with io.StringIO("content of test_long file" * 1000) as fobj:
+    with io.StringIO("content of test_long.txt" * 1000) as fobj:
         folder_data.put_object_from_filelike(fobj, path="test_long.txt")
+    # NOTE: The byte-sequence is chosen so that it is not valid UTF-8
+    with io.BytesIO(b"\xf8\x01") as fobj:
+        folder_data.put_object_from_filelike(fobj, path="test.bin")
 
     return folder_data
 

tests/test_viewers.py

@@ -22,9 +22,7 @@ def test_pbc_structure_data_viewer(structure_data_object):
 
 
 @pytest.mark.usefixtures("aiida_profile_clean")
-def test_several_data_viewers(
-    bands_data_object, folder_data_object, generate_calc_job_node
-):
+def test_several_data_viewers(bands_data_object, generate_calc_job_node):
     v = viewers.viewer(orm.Int(1))
 
     # No viewer for Int, so it should return the input
@@ -38,10 +36,6 @@ def test_several_data_viewers(
     v = viewers.viewer(bands_data_object)
     assert isinstance(v, viewers.BandsDataViewer)
 
-    # FolderDataViewer
-    v = viewers.viewer(folder_data_object)
-    assert isinstance(v, viewers.FolderDataViewer)
-
     # ProcessNodeViewer
     process = generate_calc_job_node(
         inputs={
@@ -55,6 +49,27 @@ def test_several_data_viewers(
     assert isinstance(v, viewers.ProcessNodeViewerWidget)
 
 
+@pytest.mark.usefixtures("aiida_profile_clean")
+def test_folder_data_viewer(folder_data_object):
+    v = viewers.viewer(folder_data_object)
+    assert isinstance(v, viewers.FolderDataViewer)
+
+    v.files.value = "test1.txt"
+    assert v.text.value == "content of test1.txt"
+
+    v.files.value = "test2.txt"
+    assert v.text.value == "content of test2.txt"
+    v.download_btn.click()
+    # NOTE: We're testing the download() method directly as well,
+    # since triggering it via self.download_btn.click() callback
+    # seems to swallow all exceptions.
+    v.download()
+
+    v.files.value = "test.bin"
+    assert v.text.value == "[Binary file, preview not available]"
+    v.download()
+
+
 @pytest.mark.usefixtures("aiida_profile_clean")
 def test_structure_data_viewer_storage(structure_data_object):
     v = viewers.viewer(structure_data_object)