Merge branch 'master' of github.com:materialsproject/pymatgen

Author: shyuep

src/pymatgen/analysis/phase_diagram.py

@@ -1624,7 +1624,7 @@ class PatchedPhaseDiagram(PhaseDiagram):
 
     def __init__(
         self,
-        entries: Sequence[PDEntry] | set[PDEntry],
+        entries: Sequence[Entry] | set[Entry],
         elements: Sequence[Element] | None = None,
         keep_all_spaces: bool = False,
         verbose: bool = False,

src/pymatgen/entries/compatibility.py

@@ -9,7 +9,7 @@
 import os
 import warnings
 from collections import defaultdict
-from typing import TYPE_CHECKING, TypeAlias, cast
+from typing import TYPE_CHECKING, TypeAlias, TypeVar, cast
 
 import numpy as np
 from joblib import Parallel, delayed
@@ -77,6 +77,7 @@
     raise RuntimeError("MP2020Compatibility.yaml expected to have the same Hubbard U corrections for O and F")
 
 AnyComputedEntry: TypeAlias = ComputedEntry | ComputedStructureEntry
+TypeVarAnyEntry = TypeVar("TypeVarAnyEntry", bound=AnyComputedEntry)
 
 
 class CompatibilityError(Exception):
@@ -577,10 +578,10 @@ def get_adjustments(self, entry: AnyComputedEntry) -> list[EnergyAdjustment]:
 
     def process_entry(
         self,
-        entry: ComputedEntry,
+        entry: TypeVarAnyEntry,
         inplace: bool = True,
         **kwargs,
-    ) -> ComputedEntry | None:
+    ) -> TypeVarAnyEntry | None:
         """Process a single entry with the chosen Corrections.
         Note that this method may change the original entry.
 
@@ -595,29 +596,29 @@ def process_entry(
         if not inplace:
             entry = copy.deepcopy(entry)
 
-        _entry: tuple[ComputedEntry, bool] | None = self._process_entry_inplace(entry, **kwargs)
+        _entry: tuple[TypeVarAnyEntry, bool] | None = self._process_entry_inplace(entry, **kwargs)
 
         return _entry[0] if _entry is not None else None
 
     def _process_entry_inplace(
         self,
-        entry: AnyComputedEntry,
+        entry: TypeVarAnyEntry,
         clean: bool = True,
         on_error: Literal["ignore", "warn", "raise"] = "ignore",
-    ) -> tuple[ComputedEntry, bool] | None:
+    ) -> tuple[TypeVarAnyEntry, bool] | None:
         """Process a single entry with the chosen Corrections.
         Note that this method will change the original entry.
 
         Args:
-            entry (AnyComputedEntry): An AnyComputedEntry object.
+            entry (TypeVarAnyEntry): A TypeVarAnyEntry object.
             clean (bool): Whether to remove any previously-applied energy adjustments.
                 If True, all EnergyAdjustment are removed prior to processing the Entry.
                 Defaults to True.
             on_error ("ignore" | "warn" | "raise"): What to do when get_adjustments(entry)
                 raises CompatibilityError. Defaults to "ignore".
 
         Returns:
-            tuple[AnyComputedEntry, ignore_entry (bool)] if entry is compatible, else None.
+            tuple[TypeVarAnyEntry, ignore_entry (bool)] if entry is compatible, else None.
         """
         ignore_entry: bool = False
         # If clean, remove all previous adjustments from the entry
@@ -662,13 +663,13 @@ def _process_entry_inplace(
 
     def process_entries(
         self,
-        entries: AnyComputedEntry | list[AnyComputedEntry],
+        entries: TypeVarAnyEntry | list[TypeVarAnyEntry],
         clean: bool = True,
         verbose: bool = False,
         inplace: bool = True,
         n_workers: int = 1,
         on_error: Literal["ignore", "warn", "raise"] = "ignore",
-    ) -> list[AnyComputedEntry]:
+    ) -> list[TypeVarAnyEntry]:
         """Process a sequence of entries with the chosen Compatibility scheme.
 
         Warning: This method changes entries in place! All changes can be undone and original entries
@@ -695,7 +696,7 @@ def process_entries(
         if isinstance(entries, ComputedEntry):  # True for ComputedStructureEntry too
             entries = [entries]
 
-        processed_entry_list: list[AnyComputedEntry] = []
+        processed_entry_list: list[TypeVarAnyEntry] = []
 
         # if inplace = False, process entries on a copy
         if not inplace:
@@ -1570,13 +1571,13 @@ def get_adjustments(self, entry: ComputedEntry) -> list[EnergyAdjustment]:
 
     def process_entries(
         self,
-        entries: list[AnyComputedEntry],
+        entries: list[TypeVarAnyEntry],
         clean: bool = False,
         verbose: bool = False,
         inplace: bool = True,
         n_workers: int = 1,
         on_error: Literal["ignore", "warn", "raise"] = "ignore",
-    ) -> list[AnyComputedEntry]:
+    ) -> list[TypeVarAnyEntry]:
         """Process a sequence of entries with the chosen Compatibility scheme.
 
         Args:

src/pymatgen/entries/mixing_scheme.py

@@ -16,10 +16,10 @@
 from pymatgen.analysis.phase_diagram import PhaseDiagram
 from pymatgen.analysis.structure_matcher import StructureMatcher
 from pymatgen.entries.compatibility import (
-    AnyComputedEntry,
     Compatibility,
     CompatibilityError,
     MaterialsProject2020Compatibility,
+    TypeVarAnyEntry,
 )
 from pymatgen.entries.computed_entries import ComputedStructureEntry, ConstantEnergyAdjustment
 from pymatgen.entries.entry_tools import EntrySet
@@ -121,12 +121,12 @@ def __init__(
 
     def process_entries(
         self,
-        entries: AnyComputedEntry | list[AnyComputedEntry],
+        entries: TypeVarAnyEntry | list[TypeVarAnyEntry],
         clean: bool = True,
         verbose: bool = False,
         inplace: bool = True,
         mixing_state_data=None,
-    ) -> list[AnyComputedEntry]:
+    ) -> list[TypeVarAnyEntry]:
         """Process a sequence of entries with the DFT mixing scheme. Note
         that this method will change the data of the original entries.