fix return type, fix analysis

zincware · Dec 13, 2023 · c7d327e · c7d327e
1 parent a401b1d
commit c7d327e
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/zndraw/zndraw.py b/zndraw/zndraw.py
@@ -142,7 +142,7 @@ def extend(
                 value = Frame.from_atoms(value)
             self[size + idx] = value
 
-    def __getitem__(self, index) -> Frame:
+    def __getitem__(self, index) -> t.Union[ase.Atoms, list[ase.Atoms]]:
         length = len(self)
         is_scalar = isinstance(index, int)
         is_sclice = isinstance(index, slice)
@@ -161,7 +161,7 @@ def __getitem__(self, index) -> Frame:
         atoms_list = []
 
         for val in downloaded_data.values():
-            atoms_list.append(Frame.from_dict(val))
+            atoms_list.append(Frame.from_dict(val).to_atoms())
 
         data = atoms_list[0] if is_scalar else atoms_list
         if data == [] and not is_sclice:
@@ -300,7 +300,7 @@ def initialize_webclient(self, sid):
         with self._set_sid(sid):
             for idx, frames in enumerate(self.read_data()):
                 if idx == 0:
-                    # self.analysis_schema(atoms)
+                    self.analysis_schema(frames)
                     self.selection_schema()
                     self.draw_schema()
                 self[idx] = frames