UpfData: automatically parse the Z valence from the file

aiida_pseudo/data/pseudo/upf.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 """Module for data plugin to represent a pseudo potential in UPF format."""
 import re
-from typing import BinaryIO
+from typing import BinaryIO, Union
 
 from .pseudo import PseudoPotentialData
 
@@ -10,30 +10,58 @@
 REGEX_ELEMENT_V1 = re.compile(r"""(?P<element>[a-zA-Z]{1,2})\s+Element""")
 REGEX_ELEMENT_V2 = re.compile(r"""\s*element\s*=\s*['"]\s*(?P<element>[a-zA-Z]{1,2})\s*['"].*""")
 
+PATTERN_FLOAT = r'[+-]?(\d+(\.\d*)?|\.\d+)([eE][+-]?\d+)?'
+REGEX_Z_VALENCE_V1 = re.compile(r"""(?P<z_valence>""" + PATTERN_FLOAT + r""")\s+Z valence""")
+REGEX_Z_VALENCE_V2 = re.compile(r"""\s*z_valence\s*=\s*['"]\s*(?P<z_valence>""" + PATTERN_FLOAT + r""")\s*['"].*""")
 
-def parse_element(stream: BinaryIO):
+
+def parse_element(content: str):
     """Parse the content of the UPF file to determine the element.
 
     :param stream: a filelike object with the binary content of the file.
     :return: the symbol of the element following the IUPAC naming standard.
     """
-    lines = stream.read().decode('utf-8')
-    match = REGEX_ELEMENT_V2.search(lines)
+    for regex in [REGEX_ELEMENT_V2, REGEX_ELEMENT_V1]:
+
+        match = regex.search(content)
+
+        if match:
+            return match.group('element')
+
+    raise ValueError(f'could not parse the element from the UPF content: {content}')
+
+
+def parse_z_valence(content: str) -> int:
+    """Parse the content of the UPF file to determine the Z valence.
+
+    :param stream: a filelike object with the binary content of the file.
+    :return: the Z valence.
+    """
+    for regex in [REGEX_Z_VALENCE_V2, REGEX_Z_VALENCE_V1]:
+
+        match = regex.search(content)
+
+        if match:
+            z_valence = match.group('z_valence')
 
-    if match:
-        return match.group('element')
+            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
 
-    match = REGEX_ELEMENT_V1.search(lines)
+            if int(z_valence) != z_valence:
+                raise ValueError(f'parsed value for the Z valence `{z_valence}` is not an integer.')
 
-    if match:
-        return match.group('element')
+            return int(z_valence)
 
-    raise ValueError('could not parse the element from the UPF content.')
+    raise ValueError(f'could not parse the Z valence from the UPF content: {content}')
 
 
 class UpfData(PseudoPotentialData):
     """Data plugin to represent a pseudo potential in UPF format."""
 
+    _key_z_valence = 'z_valence'
+
     def set_file(self, stream: BinaryIO, filename: str = None, **kwargs):  # pylint: disable=arguments-differ
         """Set the file content.
 
@@ -42,6 +70,30 @@ def set_file(self, stream: BinaryIO, filename: str = None, **kwargs):  # pylint:
         :raises ValueError: if the element symbol is invalid.
         """
         stream.seek(0)
-        self.element = parse_element(stream)
+
+        content = stream.read().decode('utf-8')
+        self.element = parse_element(content)
+        self.z_valence = parse_z_valence(content)
+
         stream.seek(0)
         super().set_file(stream, filename, **kwargs)
+
+    @property
+    def z_valence(self) -> Union[int, None]:
+        """Return the Z valence.
+
+        :return: the Z valence.
+        """
+        return self.get_attribute(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.set_attribute(self._key_z_valence, value)

tests/data/pseudo/test_upf.py

@@ -6,6 +6,7 @@
 
 from aiida.common.exceptions import ModificationNotAllowed
 from aiida_pseudo.data.pseudo import UpfData
+from aiida_pseudo.data.pseudo.upf import parse_z_valence
 
 
 @pytest.mark.usefixtures('clear_db')
@@ -38,3 +39,21 @@ def test_set_file(filepath_pseudos, get_pseudo_potential_data):
 
         with pytest.raises(ModificationNotAllowed):
             pseudo.set_file(handle)
+
+
+@pytest.mark.parametrize(
+    'content', (
+        'z_valence="1"',
+        'z_valence="1.0"',
+        'z_valence="1.000"',
+        'z_valence="1.00E+01"',
+        'z_valence="1500."',
+        "z_valence='1.0'",
+        'z_valence="    1"',
+        'z_valence="1    "',
+        '1.0     Z valence',
+    )
+)
+def test_parse_z_valence(content):
+    """Test the ``parse_z_valence`` method."""
+    assert parse_z_valence(content)

tests/fixtures/pseudos/upf/Ar.upf

@@ -8,6 +8,7 @@
         author="sphuber"
         date="200411"
         element="Ar"
+        z_valence="1.0"
         pseudo_type="NC"
     />
 </UPF>

tests/fixtures/pseudos/upf/He.upf

@@ -8,6 +8,7 @@
         author="sphuber"
         date="200411"
         element="He"
+        z_valence="1.0"
         pseudo_type="NC"
     />
 </UPF>

tests/fixtures/pseudos/upf/Kr.upf

@@ -8,6 +8,7 @@
         author="sphuber"
         date="200411"
         element="Kr"
+        z_valence="1.0"
         pseudo_type="NC"
     />
 </UPF>

tests/fixtures/pseudos/upf/Ne.upf

@@ -8,6 +8,7 @@
         author="sphuber"
         date="200411"
         element="Ne"
+        z_valence="1.0"
         pseudo_type="NC"
     />
 </UPF>

tests/fixtures/pseudos/upf/Rn.upf

@@ -8,6 +8,7 @@
         author="sphuber"
         date="200411"
         element="Rn"
+        z_valence="1.0"
         pseudo_type="NC"
     />
 </UPF>