Merge branch 'master' into remove-leftover-load-profile

aiidalab · Aug 17, 2023 · ae605df · ae605df
2 parents a4faac1 + 415fe7e
commit ae605df
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 8 deletions.
diff --git a/aiidalab_widgets_base/viewers.py b/aiidalab_widgets_base/viewers.py
@@ -452,7 +452,14 @@ def _download_tab(self):
 
         # 1. Choose download file format.
         self.file_format = ipw.Dropdown(
-            options=["xyz", "cif"],
+            label="Extended xyz",
+            # File extension and format may be differen. Therefore, we define both.
+            options=(
+                ("xyz", {"extension": "xyz", "format": "xyz"}),
+                ("cif", {"extension": "cif", "format": "cif"}),
+                ("Extended xyz", {"extension": "xyz", "format": "extxyz"}),
+                ("xsf", {"extension": "xsf", "format": "xsf"}),
+            ),
             layout={"width": "200px"},
             description="File format:",
         )
@@ -710,10 +717,13 @@ def apply_displayed_selection(self, _=None):
 
     def download(self, change=None):  # pylint: disable=unused-argument
         """Prepare a structure for downloading."""
-        suffix = f"-pk-{self.pk}" if self.pk else ""
+        payload = self._prepare_payload()
+        if payload is None:
+            return
+        suffix = f"pk-{self.pk}" if self.pk else "not-stored"
         self._download(
             payload=self._prepare_payload(),
-            filename=f"structure{suffix}.{self.file_format.value}",
+            filename=f"""structure-{suffix}.{self.file_format.value["extension"]}""",
         )
 
     @staticmethod
@@ -737,7 +747,10 @@ def _prepare_payload(self, file_format=None):
         """Prepare binary information."""
         from tempfile import NamedTemporaryFile
 
-        file_format = file_format if file_format else self.file_format.value
+        if not self.structure:
+            return None
+
+        file_format = file_format if file_format else self.file_format.value["format"]
         tmp = NamedTemporaryFile()
         self.structure.write(tmp.name, format=file_format)  # pylint: disable=no-member
         with open(tmp.name, "rb") as raw:

diff --git a/tests/test_viewers.py b/tests/test_viewers.py
@@ -65,22 +65,42 @@ def test_structure_data_viwer(structure_data_object):
 
     # Selection of atoms.
 
-    # Direct selection
+    # Direct selection.
     v._selected_atoms.value = "1..2"
     v.apply_displayed_selection()
     assert v.selection == [0, 1]
 
-    # x coordinate lower than 0.5
+    # The x coordinate lower than 0.5.
     v._selected_atoms.value = "x<0.5"
     v.apply_displayed_selection()
     assert v.selection == [0]
 
-    # id of the second atom
+    # The id of the second atom
     v._selected_atoms.value = "id > 1"
     v.apply_displayed_selection()
     assert v.selection == [1]
 
-    # or of the two selections
+    # or of the two selections.
     v._selected_atoms.value = "x>0.5 or x<0.5"
     v.apply_displayed_selection()
     assert v.selection == [0, 1]
+
+    # Check the `_prepare_payload` function used for downloading.
+    format_cases = [
+        (
+            "Extended xyz",
+            """MgpMYXR0aWNlPSIzLjg0NzM3IDAuMCAwLjAgMS45MjM2ODUgMy4zMzE5MiAwLjAgMS45MjM2ODUgMS4xMTA2NCAzLjE0MTM2NCIgUHJvcGVydGllcz1zcGVjaWVzOlM6MTpwb3M6UjozOm1hc3NlczpSOjEgcGJjPSJUIFQgVCIKU2kgICAgICAgMC4wMDAwMDAwMCAgICAgICAwLjAwMDAwMDAwICAgICAgIDAuMDAwMDAwMDAgICAgICAyOC4wODU1MDAwMApTaSAgICAgICAxLjkyMzY4NTAwICAgICAgIDEuMTEwNjQwMDAgICAgICAgMC43ODUzNDEwMCAgICAgIDI4LjA4NTUwMDAwCg==""",
+        ),
+        (
+            "xsf",
+            """Q1JZU1RBTApQUklNVkVDCiAzLjg0NzM3MDAwMDAwMDAwIDAuMDAwMDAwMDAwMDAwMDAgMC4wMDAwMDAwMDAwMDAwMAogMS45MjM2ODUwMDAwMDAwMCAzLjMzMTkyMDAwMDAwMDAwIDAuMDAwMDAwMDAwMDAwMDAKIDEuOTIzNjg1MDAwMDAwMDAgMS4xMTA2NDAwMDAwMDAwMCAzLjE0MTM2NDAwMDAwMDAwClBSSU1DT09SRAogMiAxCiAxNCAgICAgMC4wMDAwMDAwMDAwMDAwMCAgICAgMC4wMDAwMDAwMDAwMDAwMCAgICAgMC4wMDAwMDAwMDAwMDAwMAogMTQgICAgIDEuOTIzNjg1MDAwMDAwMDAgICAgIDEuMTEwNjQwMDAwMDAwMDAgICAgIDAuNzg1MzQxMDAwMDAwMDAK""",
+        ),
+        (
+            "cif",
+            """ZGF0YV9pbWFnZTAKX2NoZW1pY2FsX2Zvcm11bGFfc3RydWN0dXJhbCAgICAgICBTaTIKX2NoZW1pY2FsX2Zvcm11bGFfc3VtICAgICAgICAgICAgICAiU2kyIgpfY2VsbF9sZW5ndGhfYSAgICAgICAzLjg0NzM3Cl9jZWxsX2xlbmd0aF9iICAgICAgIDMuODQ3MzcKX2NlbGxfbGVuZ3RoX2MgICAgICAgMy44NDczNwpfY2VsbF9hbmdsZV9hbHBoYSAgICA2MApfY2VsbF9hbmdsZV9iZXRhICAgICA2MApfY2VsbF9hbmdsZV9nYW1tYSAgICA2MAoKX3NwYWNlX2dyb3VwX25hbWVfSC1NX2FsdCAgICAiUCAxIgpfc3BhY2VfZ3JvdXBfSVRfbnVtYmVyICAgICAgIDEKCmxvb3BfCiAgX3NwYWNlX2dyb3VwX3N5bW9wX29wZXJhdGlvbl94eXoKICAneCwgeSwgeicKCmxvb3BfCiAgX2F0b21fc2l0ZV90eXBlX3N5bWJvbAogIF9hdG9tX3NpdGVfbGFiZWwKICBfYXRvbV9zaXRlX3N5bW1ldHJ5X211bHRpcGxpY2l0eQogIF9hdG9tX3NpdGVfZnJhY3RfeAogIF9hdG9tX3NpdGVfZnJhY3RfeQogIF9hdG9tX3NpdGVfZnJhY3RfegogIF9hdG9tX3NpdGVfb2NjdXBhbmN5CiAgU2kgIFNpMSAgICAgICAxLjAgIDAuMDAwMDAgIDAuMDAwMDAgIDAuMDAwMDAgIDEuMDAwMAogIFNpICBTaTIgICAgICAgMS4wICAwLjI1MDAwICAwLjI1MDAwICAwLjI1MDAwICAxLjAwMDAK""",
+        ),
+    ]
+
+    for fmt, out in format_cases:
+        v.file_format.label = fmt
+        assert v._prepare_payload() == out