fix test_entry_like() not including all subclasses

also add GibbsComputedStructureEntry to EntryLike Union
materialsproject · Nov 8, 2022 · 32068db · 32068db
1 parent 205ea3e
commit 32068db
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 4 deletions.
diff --git a/pymatgen/util/tests/test_typing.py b/pymatgen/util/tests/test_typing.py
@@ -24,10 +24,20 @@ def _type_str(some_type: Any) -> str:
 
 
 def test_entry_like():
-    assert (
-        _type_str(EntryLike) == "Union[Dict[str, Any], ForwardRef('Entry'), "
-        "ForwardRef('PDEntry'), ForwardRef('ComputedEntry'), ForwardRef('ComputedStructureEntry')]"
+    entries = (
+        "Entry",
+        "ComputedEntry",
+        "ComputedStructureEntry",
+        "PDEntry",
+        "ExpEntry",
+        "TransformedPDEntry",
+        "GrandPotPDEntry",
+        "CostEntry",
+        "GibbsComputedStructureEntry",
     )
+    type_str = _type_str(EntryLike)
+    for entry in entries:
+        assert entry in type_str
     assert Entry.__name__ in str(EntryLike)
 
 

diff --git a/pymatgen/util/typing.py b/pymatgen/util/typing.py
@@ -30,7 +30,11 @@
         TransformedPDEntry,
     )
     from pymatgen.entries import Entry
-    from pymatgen.entries.computed_entries import ComputedEntry, ComputedStructureEntry
+    from pymatgen.entries.computed_entries import (
+        ComputedEntry,
+        ComputedStructureEntry,
+        GibbsComputedStructureEntry,
+    )
     from pymatgen.entries.exp_entries import ExpEntry
 
 VectorLike = Union[Sequence[float], np.ndarray]
@@ -55,4 +59,5 @@
     "TransformedPDEntry",
     "GrandPotPDEntry",
     "CostEntry",
+    "GibbsComputedStructureEntry",
 ]