Merge pull request #17 from benshi97/class_format

Class_format

src/quacc/types.py

@@ -101,7 +101,6 @@ class MaxwellBoltzmanDistributionKwargs(TypedDict, total=False):
 
     # ----------- Atoms handling type hints -----------
 
-
     ElementStr = Literal[
         "H",
         "He",
@@ -248,6 +247,11 @@ class FindAdsSitesKwargs(TypedDict, total=False):
 
     # ----------- Atoms (skzcam) handling type hints -----------
 
+    class SKZCAMOutput(TypedDict):
+        adsorbate_slab_embedded_cluster: Atoms
+        quantum_cluster_indices_set: list[list[int]]
+        ecp_region_indices_set: list[list[int]]
+
     class ElementInfo(TypedDict):
         core: int
         basis: str

tests/core/atoms/conftest.py

@@ -9,22 +9,27 @@
 FILE_DIR = Path(__file__).parent
 
 
-def mock_generate_chemshell_cluster(
-    slab, slab_center_idx, atom_oxi_states, filepath, **kwargs
-):
-    with (
-        gzip.open(
-            Path(FILE_DIR, "skzcam_files", "REF_ChemShell_cluster.xyz.gz"), "rb"
-        ) as f_in,
-        Path(filepath, "ChemShell_cluster.xyz").open(mode="wb") as f_out,
-    ):
-        shutil.copyfileobj(f_in, f_out)
+def mock_run_chemshell(*args, filepath=".", write_xyz_file=False, **kwargs):
+    if write_xyz_file:
+        with (
+            gzip.open(
+                Path(FILE_DIR, "skzcam_files", "REF_ChemShell_Cluster.xyz.gz"), "rb"
+            ) as f_in,
+            Path(filepath).with_suffix(".xyz").open(mode="wb") as f_out,
+        ):
+            shutil.copyfileobj(f_in, f_out)
+    else:
+        with (
+            gzip.open(
+                Path(FILE_DIR, "skzcam_files", "ChemShell_Cluster.pun.gz"), "rb"
+            ) as f_in,
+            Path(filepath).with_suffix(".pun").open(mode="wb") as f_out,
+        ):
+            shutil.copyfileobj(f_in, f_out)
 
 
 @pytest.fixture(autouse=True)
-def patch_generate_chemshell_cluster(monkeypatch):
-    from quacc.atoms import skzcam
+def patch_run_chemshell(monkeypatch):
+    from quacc.atoms.skzcam import CreateSKZCAMClusters
 
-    monkeypatch.setattr(
-        skzcam, "generate_chemshell_cluster", mock_generate_chemshell_cluster
-    )
+    monkeypatch.setattr(CreateSKZCAMClusters, "run_chemshell", mock_run_chemshell)