add ZnTrack option (#249)

* add ZnTrack option

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

zndraw/cli.py

@@ -11,7 +11,8 @@
 @cli.command()
 def main(
     filename: Optional[str] = typer.Argument(
-        None, help="Path to the file which should be visualized in ZnDraw."
+        None,
+        help="Path to the file which should be visualized in ZnDraw. Can also be the name and attribute of a ZnTrack Node like 'MyNode.atoms' if at least '--remote .' is provided. ",
     ),
     webview: bool = typer.Option(
         True,
@@ -55,6 +56,14 @@ def main(
         False,
         help="Use a token to authenticate the ZnDraw server.  This is useful if you are running ZnDraw as a server application.",
     ),
+    remote: str = typer.Option(
+        None,
+        help="URL to a ZnTrack repository to stream data from.",
+    ),
+    rev: str = typer.Option(
+        None,
+        help="Revision of the ZnTrack repository to stream data from.",
+    ),
 ):
     """Start the ZnDraw server.
 
@@ -75,4 +84,6 @@ def main(
         compute_bonds=compute_bonds,
         upgrade_insecure_requests=upgrade_insecure_requests,
         use_token=use_token,
+        remote=remote,
+        rev=rev,
     )

zndraw/view.py

@@ -50,14 +50,16 @@ def view(
     compute_bonds: bool,
     use_token: bool = False,
     upgrade_insecure_requests: bool = False,
+    remote: str = None,
+    rev: str = None,
 ):
     if not use_token:
         app.config["token"] = "notoken"
     app.config["upgrade_insecure_requests"] = upgrade_insecure_requests
     app.config["compute_bonds"] = compute_bonds
     url = f"http://127.0.0.1:{port}"
 
-    file_io = FileIO(filename, start, stop, step)
+    file_io = FileIO(filename, start, stop, step, remote, rev)
 
     proc = mp.Process(
         target=ZnDrawDefault,

zndraw/zndraw.py

@@ -31,6 +31,8 @@ class FileIO:
     start: int = 0
     stop: int = None
     step: int = 1
+    remote: str = None
+    rev: str = None
 
 
 @dataclasses.dataclass
@@ -222,7 +224,20 @@ def read_data(self):
         if self.file_io.name is None:
             return
 
-        if pathlib.Path(self.file_io.name).suffix == ".h5":
+        if self.file_io.remote is not None:
+            node_name, attribute = self.file_io.name.split(".", 1)
+            try:
+                import zntrack
+
+                node = zntrack.from_rev(
+                    node_name, remote=self.file_io.remote, rev=self.file_io.rev
+                )
+                generator = getattr(node, attribute)
+            except ImportError as err:
+                raise ImportError(
+                    "You need to install ZnTrack to use the remote feature"
+                ) from err
+        elif pathlib.Path(self.file_io.name).suffix == ".h5":
             reader = znh5md.ASEH5MD(self.file_io.name)
             generator = reader.get_atoms_list()
         else: