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

pymatgen/io/vasp/outputs.py

@@ -2010,6 +2010,8 @@ def __init__(self, filename):
                 r"\n{3}-{104}\n{3}",
                 r".+plane waves:\s+(\*{6,}|\d+)",
                 r"maximum and minimum number of plane-waves",
+                last_one_only=False,
+                first_one_only=True,
             )
         ]
         self.data["nplwvs_at_kpoints"] = [None for n in nplwvs_at_kpoints]
@@ -2165,6 +2167,7 @@ def read_table_pattern(
         postprocess=str,
         attribute_name=None,
         last_one_only=True,
+        first_one_only=False,
     ):
         r"""
         Parse table-like data. A table composes of three parts: header,
@@ -2191,14 +2194,21 @@ def read_table_pattern(
             last_one_only (bool): All the tables will be parsed, if this option
                 is set to True, only the last table will be returned. The
                 enclosing list will be removed. i.e. Only a single table will
-                be returned. Default to be True.
+                be returned. Default to be True. Incompatible with first_one_only.
+            first_one_only (bool): Only the first occurence of the table will be
+                parsed and the parsing procedure will stop. The enclosing list
+                will be removed. i.e. Only a single table will be returned.
+                Incompatible with last_one_only.
 
         Returns:
             List of tables. 1) A table is a list of rows. 2) A row if either a list of
             attribute values in case the the capturing group is defined without name in
             row_pattern, or a dict in case that named capturing groups are defined by
             row_pattern.
         """
+        if last_one_only and first_one_only:
+            raise ValueError("last_one_only and first_one_only options are incompatible")
+
         with zopen(self.filename, "rt") as f:
             text = f.read()
         table_pattern_text = header_pattern + r"\s*^(?P<table_body>(?:\s+" + row_pattern + r")+)\s+" + footer_pattern
@@ -2220,7 +2230,9 @@ def read_table_pattern(
                     processed_line = [postprocess(v) for v in ml.groups()]
                 table_contents.append(processed_line)
             tables.append(table_contents)
-        if last_one_only:
+            if first_one_only:
+                break
+        if last_one_only or first_one_only:
             retained_data = tables[-1]
         else:
             retained_data = tables

pymatgen/io/vasp/tests/test_outputs.py

@@ -1543,6 +1543,34 @@ def test_energies(self):
         self.assertAlmostEqual(o.final_energy_wo_entrp, -15.83863167)
         self.assertAlmostEqual(o.final_fr_energy, -15.92115453)
 
+    def test_read_table_pattern(self):
+        outcar = Outcar(self.TEST_FILES_DIR / "OUTCAR")
+
+        header_pattern = r"\(the norm of the test charge is\s+[\.\-\d]+\)"
+        table_pattern = r"((?:\s+\d+\s*[\.\-\d]+)+)"
+        footer_pattern = r"\s+E-fermi :"
+
+        pots = outcar.read_table_pattern(header_pattern, table_pattern, footer_pattern, last_one_only=True)
+        ref_last = [
+            ["       1 -26.0704       2 -45.5046       3 -45.5046       4 -72.9539       5 -73.0621"],
+            ["       6 -72.9539       7 -73.0621"],
+        ]
+        self.assertEqual(pots, ref_last)
+
+        pots = outcar.read_table_pattern(
+            header_pattern, table_pattern, footer_pattern, last_one_only=False, first_one_only=True
+        )
+        ref_first = [
+            ["       1 -26.1149       2 -45.5359       3 -45.5359       4 -72.9831       5 -73.1068"],
+            ["       6 -72.9831       7 -73.1068"],
+        ]
+        self.assertEqual(pots, ref_first)
+
+        with self.assertRaises(ValueError):
+            outcar.read_table_pattern(
+                header_pattern, table_pattern, footer_pattern, last_one_only=True, first_one_only=True
+            )
+
 
 class BSVasprunTest(PymatgenTest):
     _multiprocess_shared_ = True