Add utility methods to DefectThermodynamics (len, getattr, `get…

…item`, `delitem`, `setitem`, `contains`, `iter`), and test

doped/analysis.py

@@ -129,7 +129,7 @@ def check_and_set_defect_entry_name(
     site in the bulk cell).
 
     Args:
-        defect_entry (DefectEntry): DefectEntry object.
+        defect_entry (DefectEntry): ``DefectEntry`` object.
         possible_defect_name (str):
             Possible defect name (usually the folder name) to check if
             recognised by ``doped``, otherwise defect name is re-determined.

doped/core.py

@@ -1242,14 +1242,18 @@ def __repr__(self):
         if bulk_supercell is not None:
             formula = bulk_supercell.composition.get_reduced_formula_and_factor(iupac_ordering=True)[0]
         else:
-            formula = self.defect.structure.composition.get_reduced_formula_and_factor(
-                iupac_ordering=True
-            )[0]
+            try:
+                defect_name = self.defect.name
+                formula = self.defect.structure.composition.get_reduced_formula_and_factor(
+                    iupac_ordering=True
+                )[0]
+            except AttributeError:
+                defect_name = "unknown"
+                formula = "unknown"
         properties, methods = _doped_obj_properties_methods(self)
         return (
-            f"doped DefectEntry: {self.name}, with bulk composition: {formula} and defect: "
-            f"{self.defect.name}. Available attributes:\n{properties}\n\n"
-            f"Available methods:\n{methods}"
+            f"doped DefectEntry: {self.name}, with bulk composition: {formula} and defect: {defect_name}. "
+            f"Available attributes:\n{properties}\n\nAvailable methods:\n{methods}"
         )
 
     def __eq__(self, other) -> bool:

doped/generation.py

@@ -2046,7 +2046,7 @@ def from_json(cls, filename: PathLike):
 
         Args:
             filename (PathLike):
-                Filename of json file to load DefectsGenerator
+                Filename of json file to load ``DefectsGenerator``
                 object from.
 
         Returns:
@@ -2056,35 +2056,28 @@ def from_json(cls, filename: PathLike):
 
     def __getattr__(self, attr):
         """
-        Redirects an unknown attribute/method call to the defect_entries
+        Redirects an unknown attribute/method call to the ``defect_entries``
         dictionary attribute, if the attribute doesn't exist in
-        DefectsGenerator.
+        ``DefectsGenerator``.
         """
-        # Return the attribute if it exists in self.__dict__
-        if attr in self.__dict__:
-            return self.__dict__[attr]
-
-        # If trying to access defect_entries and it doesn't exist, raise an error
-        if attr == "defect_entries" or "defect_entries" not in self.__dict__:
-            raise AttributeError(f"'{type(self).__name__}' object has no attribute '{attr}'")
-
-        # Check if the attribute exists in defect_entries
-        if hasattr(self.defect_entries, attr):
+        try:
+            super().__getattribute__(attr)
+        except AttributeError as exc:
+            if attr == "defect_entries":
+                raise exc
             return getattr(self.defect_entries, attr)
 
-        # If all else fails, raise an AttributeError
-        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{attr}'")
-
     def __getitem__(self, key):
         """
-        Makes DefectsGenerator object subscriptable, so that it can be indexed
-        like a dictionary, using the defect_entries dictionary attribute.
+        Makes ``DefectsGenerator`` object subscriptable, so that it can be
+        indexed like a dictionary, using the ``defect_entries`` dictionary
+        attribute.
         """
         return self.defect_entries[key]
 
     def __setitem__(self, key, value):
         """
-        Set the value of a specific key (defect name) in the defect_entries
+        Set the value of a specific key (defect name) in the ``defect_entries``
         dictionary.
 
         Also adds the corresponding defect to the self.defects dictionary, if
@@ -2148,7 +2141,8 @@ def __setitem__(self, key, value):
 
     def __delitem__(self, key):
         """
-        Deletes the specified defect entry from the defect_entries dictionary.
+        Deletes the specified defect entry from the ``defect_entries``
+        dictionary.
 
         Doesn't remove the defect from the defects dictionary attribute, as
         there may be other charge states of the same defect still present.
@@ -2157,26 +2151,26 @@ def __delitem__(self, key):
 
     def __contains__(self, key):
         """
-        Returns True if the defect_entries dictionary contains the specified
-        defect name.
+        Returns ``True`` if the ``defect_entries`` dictionary contains the
+        specified defect name.
         """
         return key in self.defect_entries
 
     def __len__(self):
         """
-        Returns the number of entries in the defect_entries dictionary.
+        Returns the number of entries in the ``defect_entries`` dictionary.
         """
         return len(self.defect_entries)
 
     def __iter__(self):
         """
-        Returns an iterator over the defect_entries dictionary.
+        Returns an iterator over the ``defect_entries`` dictionary.
         """
         return iter(self.defect_entries)
 
     def __str__(self):
         """
-        Returns a string representation of the DefectsGenerator object.
+        Returns a string representation of the ``DefectsGenerator`` object.
         """
         formula = self.primitive_structure.composition.get_reduced_formula_and_factor(iupac_ordering=True)[
             0

doped/thermodynamics.py

@@ -402,7 +402,6 @@ def group_defects_by_name(entry_list: list[DefectEntry]) -> dict[str, list[Defec
     return grouped_entries
 
 
-# TODO: Make entries sub-selectable using dict indexing like DefectsGenerator
 class DefectThermodynamics(MSONable):
     """
     Class for analysing the calculated thermodynamics of defects in solids.
@@ -3387,6 +3386,64 @@ def __repr__(self):
             f"{properties}\n\nAvailable methods:\n{methods}"
         )
 
+    def __getattr__(self, attr):
+        """
+        Redirects an unknown attribute/method call to the ``defect_entries``
+        dictionary attribute, if the attribute doesn't exist in
+        ``DefectThermodynamics``.
+        """
+        try:
+            super().__getattribute__(attr)
+        except AttributeError as exc:
+            if attr == "_defect_entries":
+                raise exc
+            return getattr(self._defect_entries, attr)
+
+    def __getitem__(self, key):
+        """
+        Makes ``DefectThermodynamics`` object subscriptable, so that it can be
+        indexed like a dictionary, using the ``defect_entries`` dictionary
+        attribute.
+        """
+        return self.defect_entries[key]
+
+    def __setitem__(self, key, value):
+        """
+        Set the value of a specific key (defect name) in the ``defect_entries``
+        dictionary, using ``add_entries`` (to check compatibility).
+        """
+        self.add_entries(
+            [
+                value,
+            ]
+        )
+
+    def __delitem__(self, key):
+        """
+        Deletes the specified defect entry from the ``defect_entries``
+        dictionary.
+        """
+        del self.defect_entries[key]
+
+    def __contains__(self, key):
+        """
+        Returns ``True`` if the ``defect_entries`` dictionary contains the
+        specified defect name.
+        """
+        return key in self.defect_entries
+
+    def __len__(self):
+        """
+        Returns the number of entries in the ``defect_entries`` dictionary.
+        """
+        return len(self.defect_entries)
+
+    def __iter__(self):
+        """
+        Returns an iterator over the ``defect_entries`` dictionary.
+        """
+        return iter(self.defect_entries)
+
 
 def _group_defect_charge_state_concentrations(
     conc_df: pd.DataFrame, per_site: bool = False, skip_formatting: bool = False

tests/test_thermodynamics.py

@@ -427,6 +427,29 @@ def _check_defect_thermo(
         first_entry = next(iter(defect_thermo.defect_entries.values()))
         assert np.mean(guessed_def_pos_deviations) < np.max(first_entry.bulk_supercell.lattice.abc) * 0.2
 
+        print("Checking dict attributes passed to defect_entries successfully")
+        assert len(defect_thermo) == len(defect_thermo.defect_entries)  # __len__()
+        assert dict(defect_thermo.items()) == defect_thermo.defect_entries  # __iter__()
+        assert all(
+            defect_entry_name in defect_thermo
+            for defect_entry_name in defect_thermo.defect_entries  # __contains__()
+        )
+
+        random_name, random_defect_entry = random.choice(list(defect_thermo.defect_entries.items()))
+        print(f"Checking editing DefectsGenerator, using {random_defect_entry.name}")
+        assert (
+            defect_thermo[random_defect_entry.name]
+            == defect_thermo.defect_entries[random_defect_entry.name]
+        )  # __getitem__()
+        assert (
+            defect_thermo.get(random_defect_entry.name)
+            == defect_thermo.defect_entries[random_defect_entry.name]
+        )  # get()
+        random_defect_entry = defect_thermo.defect_entries[random_defect_entry.name]
+        del defect_thermo[random_defect_entry.name]  # __delitem__()
+        assert random_defect_entry.name not in defect_thermo
+        defect_thermo[random_defect_entry.name] = random_defect_entry  # __setitem__()
+
     def _check_CdTe_example_dist_tol(self, defect_thermo, num_grouped_defects):
         print(f"Testing CdTe updated dist_tol: {defect_thermo.dist_tol}")
         tl_df = defect_thermo.get_transition_levels()
@@ -1735,7 +1758,11 @@ def test_parse_chempots_CdTe(self):
             print(i)
             assert i in str(exc.value)
 
-    def test_add_entries(self):
+    @pytest.mark.parametrize("using_dict", [False, True])
+    def test_add_entries(self, using_dict=False):
+        """
+        If ``using_dict``, test this through the ``__setitem__`` method.
+        """
         partial_defect_thermo = DefectThermodynamics(list(self.CdTe_defect_dict.values())[:4])
         assert not partial_defect_thermo.get_formation_energies().equals(
             self.CdTe_defect_thermo.get_formation_energies()
@@ -1744,7 +1771,11 @@ def test_add_entries(self):
             self.CdTe_defect_thermo.get_symmetries_and_degeneracies()
         )
 
-        partial_defect_thermo.add_entries(list(self.CdTe_defect_dict.values())[4:])
+        if using_dict:
+            for name, entry in list(self.CdTe_defect_dict.items())[4:]:
+                partial_defect_thermo[name] = entry
+        else:
+            partial_defect_thermo.add_entries(list(self.CdTe_defect_dict.values())[4:])
         assert partial_defect_thermo.get_formation_energies().equals(
             self.CdTe_defect_thermo.get_formation_energies()
         )