From ca3f065f0cf3455e12b852a2d77e3e8373affff1 Mon Sep 17 00:00:00 2001
From: Emanuele Bosoni <bosonie@tcd.ie>
Date: Mon, 17 May 2021 16:15:53 +0200
Subject: [PATCH] `PsmlData`: Add parsing of Z-valence from file

---
 src/aiida_pseudo/data/pseudo/psml.py | 51 +++++++++++++++++++++++++++-
 tests/data/pseudo/test_psml.py       |  3 ++
 tests/fixtures/pseudos/psml/Ar.psml  |  2 +-
 tests/fixtures/pseudos/psml/He.psml  |  2 +-
 tests/fixtures/pseudos/psml/Ne.psml  |  2 +-
 5 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/src/aiida_pseudo/data/pseudo/psml.py b/src/aiida_pseudo/data/pseudo/psml.py
index db52c02..623bccb 100644
--- a/src/aiida_pseudo/data/pseudo/psml.py
+++ b/src/aiida_pseudo/data/pseudo/psml.py
@@ -11,7 +11,7 @@
 REGEX_ELEMENT = re.compile(r"""\s*(?P<element>[a-zA-Z]{1}[a-z]?)\s+.*""")
 
 
-def parse_element(stream: typing.BinaryIO):
+def parse_element(stream: typing.BinaryIO) -> str:
     """Parse the content of the PSML file to determine the element.
 
     :param stream: a filelike object with the binary content of the file.
@@ -28,9 +28,36 @@ def parse_element(stream: typing.BinaryIO):
     return element.capitalize()
 
 
+def parse_z_valence(stream: typing.BinaryIO) -> int:
+    """Parse the content of the PSML file to determine the Z valence.
+
+    :param stream: a filelike object with the binary content of the file.
+    :return: the Z valence.
+    """
+    from xml.dom.minidom import parse
+
+    try:
+        xml = parse(stream)
+        z_valence = xml.getElementsByTagName('pseudo-atom-spec')[0].attributes['z-pseudo'].value
+    except (AttributeError, IndexError, KeyError) as exception:
+        raise ValueError(f'could not parse the Z valence from the PSML content: {exception}') from exception
+
+    try:
+        z_valence = float(z_valence)
+    except ValueError as exception:
+        raise ValueError(f'parsed value for the Z valence `{z_valence}` is not a valid number.') from exception
+
+    if int(z_valence) != z_valence:
+        raise ValueError(f'parsed value for the Z valence `{z_valence}` is not an integer.')
+
+    return int(z_valence)
+
+
 class PsmlData(PseudoPotentialData):
     """Data plugin to represent a pseudo potential in PSML format."""
 
+    _key_z_valence = 'z_valence'
+
     def set_file(self, source: typing.Union[str, pathlib.Path, typing.BinaryIO], filename: str = None, **kwargs):  # pylint: disable=arguments-differ
         """Set the file content and parse other optional attributes from the content.
 
@@ -54,3 +81,25 @@ def set_file(self, source: typing.Union[str, pathlib.Path, typing.BinaryIO], fil
         super().set_file(source, filename, **kwargs)
         source.seek(0)
         self.element = parse_element(source)
+        source.seek(0)
+        self.z_valence = parse_z_valence(source)
+
+    @property
+    def z_valence(self) -> typing.Optional[int]:
+        """Return the Z valence.
+
+        :return: the Z valence.
+        """
+        return self.base.attributes.get(self._key_z_valence, None)
+
+    @z_valence.setter
+    def z_valence(self, value: int):
+        """Set the Z valence.
+
+        :param value: the Z valence.
+        :raises ValueError: if the value is not a positive integer.
+        """
+        if not isinstance(value, int) or value < 0:
+            raise ValueError(f'`{value}` is not a positive integer')
+
+        self.base.attributes.set(self._key_z_valence, value)
diff --git a/tests/data/pseudo/test_psml.py b/tests/data/pseudo/test_psml.py
index 6b6c884..9e2a113 100644
--- a/tests/data/pseudo/test_psml.py
+++ b/tests/data/pseudo/test_psml.py
@@ -51,13 +51,16 @@ def test_set_file(filepath_pseudos, get_pseudo_potential_data):
     """
     pseudo = get_pseudo_potential_data(element='Ar', entry_point='psml')
     assert pseudo.element == 'Ar'
+    assert pseudo.z_valence == 8
 
     with (filepath_pseudos('psml') / 'He.psml').open('rb') as handle:
         pseudo.set_file(handle)
         assert pseudo.element == 'He'
+        assert pseudo.z_valence == 2
 
         pseudo.store()
         assert pseudo.element == 'He'
+        assert pseudo.z_valence == 2
 
         with pytest.raises(ModificationNotAllowed):
             pseudo.set_file(handle)
diff --git a/tests/fixtures/pseudos/psml/Ar.psml b/tests/fixtures/pseudos/psml/Ar.psml
index b30c7c0..dfe6921 100644
--- a/tests/fixtures/pseudos/psml/Ar.psml
+++ b/tests/fixtures/pseudos/psml/Ar.psml
@@ -1,4 +1,4 @@
 <?xml version="1.0" encoding="UTF-8" ?>
 <psml version="1.1" energy_unit="hartree" length_unit="bohr">
-<pseudo-atom-spec atomic-label="Ar" atomic-number="18"></pseudo-atom-spec>
+<pseudo-atom-spec atomic-label="Ar" atomic-number="18" z-pseudo="8"></pseudo-atom-spec>
 </psml>
diff --git a/tests/fixtures/pseudos/psml/He.psml b/tests/fixtures/pseudos/psml/He.psml
index 434a6b7..0117fac 100644
--- a/tests/fixtures/pseudos/psml/He.psml
+++ b/tests/fixtures/pseudos/psml/He.psml
@@ -1,4 +1,4 @@
 <?xml version="1.0" encoding="UTF-8" ?>
 <psml version="1.1" energy_unit="hartree" length_unit="bohr">
-<pseudo-atom-spec atomic-label="He" atomic-number="2"></pseudo-atom-spec>
+<pseudo-atom-spec atomic-label="He" atomic-number="2"  z-pseudo="2"></pseudo-atom-spec>
 </psml>
diff --git a/tests/fixtures/pseudos/psml/Ne.psml b/tests/fixtures/pseudos/psml/Ne.psml
index 243240e..7d02383 100644
--- a/tests/fixtures/pseudos/psml/Ne.psml
+++ b/tests/fixtures/pseudos/psml/Ne.psml
@@ -1,4 +1,4 @@
 <?xml version="1.0" encoding="UTF-8" ?>
 <psml version="1.1" energy_unit="hartree" length_unit="bohr">
-<pseudo-atom-spec atomic-label="Ne" atomic-number="10"></pseudo-atom-spec>
+<pseudo-atom-spec atomic-label="Ne" atomic-number="10" z-pseudo="8"></pseudo-atom-spec>
 </psml>