From 754a18429a26550069fee596109952400409a038 Mon Sep 17 00:00:00 2001
From: andrewtarzia <andrew.tarzia@gmail.com>
Date: Tue, 24 Sep 2024 23:25:57 +0200
Subject: [PATCH] Shift handling to stk functions.

---
 python/chemiscope/input.py               |  3 +--
 python/chemiscope/structures/__init__.py | 13 +++++++------
 python/chemiscope/structures/_stk.py     | 17 +++++++++++++++--
 3 files changed, 23 insertions(+), 10 deletions(-)

diff --git a/python/chemiscope/input.py b/python/chemiscope/input.py
index 4c9490b7a..3ac232d45 100644
--- a/python/chemiscope/input.py
+++ b/python/chemiscope/input.py
@@ -380,8 +380,7 @@ def create_input(
     # Check to tell the user they might have forgotten some properties coming
     # from the frames (that chemiscope used to automatically extract). This code
     # should be removed in version 0.6 of chemiscope.
-    # stk cannot have atom properties or structure properties, so skipping.
-    if frames is not None and not from_stk:
+    if frames is not None:
         found_one_from_frame = False
 
         atom_properties = _list_atom_properties(frames)
diff --git a/python/chemiscope/structures/__init__.py b/python/chemiscope/structures/__init__.py
index 28124eebb..9aa8c6506 100644
--- a/python/chemiscope/structures/__init__.py
+++ b/python/chemiscope/structures/__init__.py
@@ -26,6 +26,8 @@
     convert_stk_bonds_as_shapes,
     _stk_all_atomic_environments,
     _stk_composition_properties,
+    _stk_list_atom_properties,
+    _stk_list_structure_properties,
 )
 
 __all__ = [
@@ -83,9 +85,9 @@ def _list_atom_properties(frames):
     if adapter == "ASE":
         return _ase_list_atom_properties(frames)
     elif adapter == "stk":
-        # Do not check, because stk does not contain its own properties
-        # and a dictionary must be added.
-        pass
+        # stk does not contain properties inside the structure objects.
+        return _stk_list_atom_properties(frames)
+
     else:
         raise Exception("reached unreachable code")
 
@@ -101,9 +103,8 @@ def _list_structure_properties(frames):
     if adapter == "ASE":
         return _ase_list_structure_properties(frames)
     elif adapter == "stk":
-        # Do not check, because stk does not contain its own properties
-        # and a dictionary must be added.
-        pass
+        # stk does not contain properties inside the structure objects.
+        return _stk_list_structure_properties(frames)
     else:
         raise Exception("reached unreachable code")
 
diff --git a/python/chemiscope/structures/_stk.py b/python/chemiscope/structures/_stk.py
index 6c6d6b80e..4ffb34b60 100644
--- a/python/chemiscope/structures/_stk.py
+++ b/python/chemiscope/structures/_stk.py
@@ -41,8 +41,11 @@ def _stk_to_json(molecule: stk.Molecule) -> dict[str : int | list]:
     return data
 
 
-def _stk_all_atomic_environments(frames, cutoff):
-    "Extract all atomic environments out of a set of ASE Atoms objects"
+def _stk_all_atomic_environments(
+    frames: list[stk.Molecule],
+    cutoff: float,
+) -> list[tuple[int, int, float]]:
+    "Extract all atomic environments out of a set of stk Molecule objects"
     environments = []
     for structure_i, frame in enumerate(frames):
         for atom in frame.get_atoms():
@@ -196,3 +199,13 @@ def convert_stk_bonds_as_shapes(
                 shape_dict[bname]["parameters"]["structure"].append(fake_bond)
 
     return shape_dict
+
+
+def _stk_list_atom_properties(frames: list[stk.Molecule]) -> list:
+    # stk cannot have atom properties or structure properties, so skipping.
+    return []
+
+
+def _stk_list_structure_properties(frames: list[stk.Molecule]) -> list:
+    # stk cannot have atom properties or structure properties, so skipping.
+    return []