from base64 import b64decode
from tempfile import NamedTemporaryFile
from dash import dcc, html
from dash.dependencies import Input, Output
from dash.exceptions import PreventUpdate
from monty.serialization import loadfn
from pymatgen.core.structure import Molecule, Structure
from pymatgen.io.vasp.outputs import Chgcar
from crystal_toolkit.core.mpcomponent import MPComponent
from crystal_toolkit.helpers.layouts import (
Icon,
MessageBody,
MessageContainer,
MessageHeader,
)
[docs]class StructureMoleculeUploadComponent(MPComponent):
@property
def _sub_layouts(self):
# this is a very custom component based on Bulma css styles
upload_layout = html.Div(
html.Label(
[
html.Span(
[
Icon(kind="upload"),
html.Span(
"Choose a file to upload or drag and drop",
className="file-label",
),
],
className="file-cta",
),
# TODO: CSS fix style and un-hide file name
html.Span(
id=self.id("upload_label"),
className="file-name",
style={"display": "none"},
),
],
className="file-label",
),
className="file is-boxed",
# TODO: CSS set sensible max-width, don't hard-code
style={"max-width": "312px"},
)
upload = html.Div(
[
html.Label("Load from your computer: ", className="mpc-label"),
dcc.Upload(upload_layout, id=self.id("upload_data"), multiple=False),
html.Div(id=self.id("error_message_container")),
]
)
return {"upload": upload}
[docs] def generate_callbacks(self, app, cache):
@app.callback(
Output(self.id("upload_label"), "children"),
[Input(self.id("upload_data"), "filename")],
)
def show_filename_on_upload(filename):
if not filename:
raise PreventUpdate
return filename
@app.callback(
Output(self.id("error_message_container"), "children"),
[Input(self.id(), "data")],
)
def update_error_message(data):
if not data:
raise PreventUpdate
if not data["error"]:
return html.Div()
else:
return html.Div(
[
html.Br(),
MessageContainer(
[MessageHeader("Error"), MessageBody([data["error"]])],
kind="danger",
size="small",
),
]
)
@app.callback(
Output(self.id(), "data"),
[
Input(self.id("upload_data"), "contents"),
Input(self.id("upload_data"), "filename"),
Input(self.id("upload_data"), "last_modified"),
],
)
def callback_update_structure(contents, filename, last_modified):
if not contents:
raise PreventUpdate
# assume we only want the first input for now
content_type, content_string = contents.split(",")
decoded_contents = b64decode(content_string)
error = None
struct_or_mol = None
# necessary to write to file so pymatgen's filetype detection can work
with NamedTemporaryFile(suffix=filename) as tmp:
tmp.write(decoded_contents)
tmp.flush()
try:
struct_or_mol = Structure.from_file(tmp.name)
except Exception:
try:
struct_or_mol = Molecule.from_file(tmp.name)
except Exception:
try:
struct_or_mol = Chgcar.from_file(tmp.name)
except Exception:
# TODO: fix these horrible try/excepts, loadfn may be dangerous
try:
struct_or_mol = loadfn(tmp.name)
except Exception:
error = (
"Could not parse uploaded file. "
"If this seems like a bug, please report it. "
"Crystal Toolkit understands all crystal "
"structure file types and molecule file types "
"supported by pymatgen."
)
return {"data": struct_or_mol, "error": error}