test return type and reduce function calls in test_plot_periodic_heatmap()

Author: janosh

pymatgen/io/abinit/pseudos.py

@@ -1693,7 +1693,7 @@ def as_dict(self, **kwargs):
             # k, count = p.element, 1
             # Handle multiple-pseudos with the same name!
             while k in dct:
-                k += k.split("#")[0] + "#" + str(count)
+                k += f"{k.split('#')[0]}#{count}"
                 count += 1
             dct.update({k: p.as_dict()})
         dct["@module"] = type(self).__module__

pymatgen/io/abinit/tests/test_abiobjects.py

@@ -136,16 +136,16 @@ def test_base(self):
 
         assert fd1ev
 
-        same_fd = Smearing.as_smearing("fermi_dirac:" + str(1.0 / Ha_to_eV))
+        same_fd = Smearing.as_smearing(f"fermi_dirac:{1.0 / Ha_to_eV}")
 
         assert same_fd == fd1ev
 
-        nosmear = Smearing.nosmearing()
-        assert nosmear == Smearing.as_smearing("nosmearing")
+        no_smear = Smearing.nosmearing()
+        assert no_smear == Smearing.as_smearing("nosmearing")
 
-        assert not nosmear
-        assert nosmear != fd1ev
-        self.assert_msonable(nosmear)
+        assert not no_smear
+        assert no_smear != fd1ev
+        self.assert_msonable(no_smear)
 
         new_fd1ev = Smearing.from_dict(fd1ev.as_dict())
         assert new_fd1ev == fd1ev
@@ -227,7 +227,7 @@ def test_base(self):
         godby.to_abivars()
         assert godby
 
-        same_godby = PPModel.as_ppmodel("godby:" + str(12.0 / Ha_to_eV))
+        same_godby = PPModel.as_ppmodel(f"godby:{12.0 / Ha_to_eV}")
         assert same_godby == godby
 
         no_pp_model = PPModel.get_noppmodel()

pymatgen/io/abinit/tests/test_netcdf.py

@@ -81,7 +81,7 @@ def test_read_Si2(self):
 
             data.print_tree()
             for group in data.walk_tree():
-                print("group: " + str(group))
+                print(f"{group=}")
 
             # Initialize pymatgen structure from GSR.
             structure = data.read_structure()

pymatgen/io/feff/outputs.py

@@ -200,25 +200,25 @@ def charge_transfer_from_file(feff_inp_file, ldos_file):
             dicts = Potential.pot_dict_from_string(pot_string)
             pot_dict = dicts[1]
 
-        for i in range(0, len(dicts[0]) + 1):
-            if len(str(i)) == 1:
-                with zopen(f"{ldos_file}0{i}.dat", "rt") as fobject:
-                    f = fobject.readlines()
-                    s = float(f[3].split()[2])
-                    p = float(f[4].split()[2])
-                    d = float(f[5].split()[2])
-                    f1 = float(f[6].split()[2])
-                    tot = float(f[1].split()[4])
-                    cht[str(i)] = {pot_dict[i]: {"s": s, "p": p, "d": d, "f": f1, "tot": tot}}
+        for idx in range(len(dicts[0]) + 1):
+            if len(str(idx)) == 1:
+                with zopen(f"{ldos_file}0{idx}.dat", "rt") as file:
+                    lines = file.readlines()
+                    s = float(lines[3].split()[2])
+                    p = float(lines[4].split()[2])
+                    d = float(lines[5].split()[2])
+                    f1 = float(lines[6].split()[2])
+                    tot = float(lines[1].split()[4])
+                    cht[str(idx)] = {pot_dict[idx]: {"s": s, "p": p, "d": d, "f": f1, "tot": tot}}
             else:
-                with zopen(ldos_file + str(i) + ".dat", "rt") as fid:
-                    f = fid.readlines()
-                    s = float(f[3].split()[2])
-                    p = float(f[4].split()[2])
-                    d = float(f[5].split()[2])
-                    f1 = float(f[6].split()[2])
-                    tot = float(f[1].split()[4])
-                    cht[str(i)] = {pot_dict[i]: {"s": s, "p": p, "d": d, "f": f1, "tot": tot}}
+                with zopen(f"{ldos_file}{idx}.dat", "rt") as file:
+                    lines = file.readlines()
+                    s = float(lines[3].split()[2])
+                    p = float(lines[4].split()[2])
+                    d = float(lines[5].split()[2])
+                    f1 = float(lines[6].split()[2])
+                    tot = float(lines[1].split()[4])
+                    cht[str(idx)] = {pot_dict[idx]: {"s": s, "p": p, "d": d, "f": f1, "tot": tot}}
 
         return cht
 

pymatgen/io/lobster/inputs.py

@@ -267,7 +267,7 @@ def write_lobsterin(self, path="lobsterin", overwritedict=None):
             for key in Lobsterin.AVAILABLEKEYWORDS:
                 if key.lower() in [element.lower() for element in self]:
                     if key.lower() in [element.lower() for element in Lobsterin.FLOAT_KEYWORDS]:
-                        f.write(key + " " + str(self.get(key)) + "\n")
+                        f.write(f"{key} {self.get(key)}\n")
                     elif key.lower() in [element.lower() for element in Lobsterin.BOOLEAN_KEYWORDS]:
                         # checks if entry is True or False
                         for key_here in self:
@@ -276,10 +276,10 @@ def write_lobsterin(self, path="lobsterin", overwritedict=None):
                         if self.get(new_key):
                             f.write(key + "\n")
                     elif key.lower() in [element.lower() for element in Lobsterin.STRING_KEYWORDS]:
-                        f.write(key + " " + str(self.get(key) + "\n"))
+                        f.write(f"{key} {self.get(key)}\n")
                     elif key.lower() in [element.lower() for element in Lobsterin.LISTKEYWORDS]:
                         for entry in self.get(key):
-                            f.write(key + " " + str(entry) + "\n")
+                            f.write(f"{key} {entry}\n")
 
     def as_dict(self):
         """:return: MSONable dict"""

pymatgen/io/pwscf.py

@@ -97,7 +97,7 @@ def __str__(self):
                         name = k
 
                 if name is None:
-                    name = site.specie.symbol + str(c)
+                    name = f"{site.specie.symbol}{c}"
                     site_descriptions[name] = site.properties
                     c += 1
 

pymatgen/transformations/standard_transformations.py

@@ -310,9 +310,7 @@ def apply_transformation(self, structure: Structure) -> Structure:
         return struct
 
     def __str__(self):
-        return "Substitution Transformation :" + ", ".join(
-            [str(k) + "->" + str(v) for k, v in self._species_map.items()]
-        )
+        return "Substitution Transformation :" + ", ".join([f"{k}->{v}" for k, v in self._species_map.items()])
 
     def __repr__(self):
         return str(self)

pymatgen/util/convergence.py

@@ -441,9 +441,9 @@ def multi_reciprocal_extra(xs, ys, noise=False):
 def print_plot_line(function, popt, xs, ys, name, tol: float = 0.05, extra=""):
     """Print the gnuplot command line to plot the x, y data with the fitted function using the popt parameters."""
     idp = id_generator()
-    with open("convdat." + str(idp), mode="w") as f:
+    with open(f"convdat.{idp}", mode="w") as f:
         for n in range(0, len(ys), 1):
-            f.write(str(xs[n]) + " " + str(ys[n]) + "\n")
+            f.write(f"{xs[n]} {ys[n]}\n")
     tol = abs(tol)
     line = f"plot 'convdat.{idp}' pointsize 4 lt 0, "
     line += f"{popt[0]} lt 3, {popt[0] - tol} lt 4, {popt[0] + tol} lt 4, "

pymatgen/util/string.py

@@ -300,27 +300,27 @@ def transformation_to_string(matrix, translation_vec=(0, 0, 0), components=("x",
     :return: xyz string.
     """
     parts = []
-    for i in range(3):
-        s = ""
-        m = matrix[i]
-        t = translation_vec[i]
+    for idx in range(3):
+        string = ""
+        m = matrix[idx]
+        offset = translation_vec[idx]
         for j, dim in enumerate(components):
             if m[j] != 0:
                 f = Fraction(m[j]).limit_denominator()
-                if s != "" and f >= 0:
-                    s += "+"
+                if string != "" and f >= 0:
+                    string += "+"
                 if abs(f.numerator) != 1:
-                    s += str(f.numerator)
+                    string += str(f.numerator)
                 elif f < 0:
-                    s += "-"
-                s += c + dim
+                    string += "-"
+                string += c + dim
                 if f.denominator != 1:
-                    s += "/" + str(f.denominator)
-        if t != 0:
-            s += ("+" if (t > 0 and s != "") else "") + str(Fraction(t).limit_denominator())
-        if s == "":
-            s += "0"
-        parts.append(s)
+                    string += "/" + str(f.denominator)
+        if offset != 0:
+            string += ("+" if (offset > 0 and string != "") else "") + str(Fraction(offset).limit_denominator())
+        if string == "":
+            string += "0"
+        parts.append(string)
     return delim.join(parts)
 
 

pymatgen/util/tests/test_plotting.py

@@ -1,32 +1,35 @@
 from __future__ import annotations
 
+import matplotlib.pyplot as plt
+
 from pymatgen.util.plotting import periodic_table_heatmap, van_arkel_triangle
 from pymatgen.util.testing import PymatgenTest
 
 
 class FuncTestCase(PymatgenTest):
     def test_plot_periodic_heatmap(self):
-        random_data = {
-            "Te": 0.11083818874391202,
-            "Au": 0.7575629917425387,
-            "Th": 1.2475885304040335,
-            "Ni": -2.0354391922547705,
-        }
-        _ = periodic_table_heatmap(random_data)
-        _ = periodic_table_heatmap(random_data, cmap="plasma")
-        _ = periodic_table_heatmap(random_data, max_row=7)
-        _ = periodic_table_heatmap(random_data, max_row=10)
-        _ = periodic_table_heatmap(random_data, cbar_label_size=18)
-        _ = periodic_table_heatmap(random_data, cmap_range=[0, 1])
-        _ = periodic_table_heatmap(random_data, cbar_label="Hello World")
-        _ = periodic_table_heatmap(random_data, blank_color="white")
-        _ = periodic_table_heatmap(random_data, value_format="%.4f")
-        _ = periodic_table_heatmap(random_data, edge_color="black")
-        _ = periodic_table_heatmap(random_data, value_fontsize=12)
-        _ = periodic_table_heatmap(random_data, symbol_fontsize=18)
-        _ = periodic_table_heatmap(random_data, readable_fontcolor=True)
+        random_data = {"Te": 0.11083, "Au": 0.75756, "Th": 1.24758, "Ni": -2.0354}
+        ret_val = periodic_table_heatmap(random_data)
+        assert ret_val is plt
+
+        # Test all keywords
+        periodic_table_heatmap(
+            random_data,
+            cmap="plasma",
+            max_row=10,
+            cbar_label_size=18,
+            cmap_range=[0, 1],
+            cbar_label="Hello World",
+            blank_color="white",
+            value_format="%.4f",
+            edge_color="black",
+            value_fontsize=12,
+            symbol_fontsize=18,
+            readable_fontcolor=True,
+        )
 
     def test_van_arkel_triangle(self):
         random_list = [("Fe", "C"), ("Ni", "F")]
-        _ = van_arkel_triangle(random_list)
-        _ = van_arkel_triangle(random_list, annotate=True)
+        ret_val = van_arkel_triangle(random_list)
+        assert ret_val is plt
+        van_arkel_triangle(random_list, annotate=True)