Merge pull request #4 from jimustafa/develop

separate raw parsers and add schema validation

Author: jimustafa

requirements/dev-requirements.in

@@ -2,4 +2,5 @@ mkdocs-material
 mkdocstrings[python]
 pip-tools
 pre-commit
+pydantic
 pytest

requirements/dev-requirements.txt

@@ -96,6 +96,8 @@ pre-commit==2.17.0
     # via -r dev-requirements.in
 py==1.11.0
     # via pytest
+pydantic==1.9.1
+    # via -r dev-requirements.in
 pygments==2.12.0
     # via mkdocs-material
 pymdown-extensions==9.4
@@ -131,6 +133,7 @@ tomli==2.0.1
 typing-extensions==4.1.1
     # via
     #   importlib-metadata
+    #   pydantic
     #   pytkdocs
 virtualenv==20.13.1
     # via pre-commit

src/w90io/__main__.py

@@ -23,7 +23,7 @@ def parse_win(args):
             'blocks': blocks,
         })
     else:
-        parsed_win = w90io.parse_win(contents)
+        parsed_win = w90io.parse_win_raw(contents)
         if args.parameters:
             pp.pprint({
                 parameter: parsed_win['parameters'][parameter]
@@ -71,7 +71,7 @@ def parse_nnkp(args):
             'blocks': blocks,
         })
     else:
-        parsed_nnkp = w90io.parse_nnkp(contents)
+        parsed_nnkp = w90io.parse_nnkp_raw(contents)
         if args.parameters:
             pp.pprint({
                 parameter: parsed_nnkp['parameters'][parameter]

src/w90io/_nnkp.py

@@ -6,7 +6,7 @@
 from . import _core
 
 
-__all__ = ['parse_nnkp']
+__all__ = ['parse_nnkp_raw']
 
 
 patterns = {
@@ -38,13 +38,12 @@ def parse_lattice(string: str) -> dict:
     match = patterns['lattice_vectors'].search(string)
 
     if match is not None:
-        v1 = np.fromstring(match.group('v1'), sep=' ')
-        v2 = np.fromstring(match.group('v2'), sep=' ')
-        v3 = np.fromstring(match.group('v3'), sep=' ')
+        v1 = [float(x) for x in match.group('v1').split()]
+        v2 = [float(x) for x in match.group('v2').split()]
+        v3 = [float(x) for x in match.group('v3').split()]
 
         return {
             'v1': v1, 'v2': v2, 'v3': v3,
-            'lattice_vectors': np.array([v1, v2, v3]),
         }
     else:
         return None
@@ -55,7 +54,6 @@ def parse_direct_lattice(string: str) -> dict:
 
     return {
         'a1': lattice['v1'], 'a2': lattice['v2'], 'a3': lattice['v3'],
-        'lattice_vectors': lattice['lattice_vectors']
     }
 
 
@@ -64,15 +62,16 @@ def parse_reciprocal_lattice(string: str) -> dict:
 
     return {
         'b1': lattice['v1'], 'b2': lattice['v2'], 'b3': lattice['v3'],
-        'lattice_vectors': lattice['lattice_vectors']
     }
 
 
 def parse_kpoints(string: str) -> dict:
     match = patterns['kpoints'].search(string)
 
     return {
-        'kpoints': np.fromstring(match.group('kpoints'), sep='\n').reshape((len(match.group('kpoints').splitlines()), -1))[:, :3]
+        'kpoints': [
+            [float(x) for x in line.split()] for line in match.group('kpoints').splitlines()
+        ]
     }
 
 
@@ -119,19 +118,17 @@ def parse_spinor_projections(string: str) -> dict:
 def parse_exclude_bands(string: str) -> dict:
     match = patterns['exclude_bands'].search(string)
 
-    exclude_bands = np.fromstring(match.group('exclude_bands'), sep='\n', dtype=int)
-
-    if exclude_bands.size > 0:
+    if match is not None:
         return {
-            'exclude_bands': exclude_bands
+            'exclude_bands': [int(line) for line in match.group('exclude_bands').splitlines()]
         }
     else:
         return {
             'exclude_bands': None
         }
 
 
-def parse_nnkp(string: str) -> dict:
+def parse_nnkp_raw(string: str) -> dict:
     """
     Parse NNKP
 

src/w90io/_schema.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+import typing
+
+import pydantic
+
+
+class UnitCell(pydantic.BaseModel):
+    units: typing.Optional[str]
+    a1: list
+    a2: list
+    a3: list
+
+
+class DirectLattice(pydantic.BaseModel):
+    a1: list
+    a2: list
+    a3: list
+
+
+class ReciprocalLattice(pydantic.BaseModel):
+    b1: list
+    b2: list
+    b3: list
+
+
+class Atom(pydantic.BaseModel):
+    species: str
+    basis_vector: list
+
+
+class Atoms(pydantic.BaseModel):
+    units: typing.Optional[str]
+    atoms: typing.List[Atoms]
+
+
+class Projections(pydantic.BaseModel):
+    units: typing.Optional[str]
+    projections: typing.List[str]
+
+
+class Kpoints(pydantic.BaseModel):
+    kpoints: typing.List[list]
+
+
+class ExcludeBands(pydantic.BaseModel):
+    exclude_bands: typing.Optional[typing.List[int]]
+
+
+class WIN(pydantic.BaseModel):
+    comments: typing.List[str]
+    parameters: dict
+    blocks: dict
+    unit_cell_cart: typing.Optional[UnitCell]
+    # atoms_frac: Atoms
+    projections: Projections
+    kpoints: Kpoints
+
+
+class NNKP(pydantic.BaseModel):
+    comments: typing.List[str]
+    parameters: dict
+    blocks: dict
+    direct_lattice: DirectLattice
+    reciprocal_lattice: ReciprocalLattice
+    kpoints: Kpoints
+    exclude_bands: ExcludeBands

src/w90io/_win.py

@@ -3,10 +3,8 @@
 
 from . import _core
 
-import numpy as np
 
-
-__all__ = ['parse_win']
+__all__ = ['parse_win_raw']
 
 
 patterns = {
@@ -34,14 +32,13 @@ def parse_unit_cell(string: str) -> dict:
     match = patterns['unit_cell'].search(string)
 
     if match is not None:
-        a1 = np.fromstring(match.group('a1'), sep=' ')
-        a2 = np.fromstring(match.group('a2'), sep=' ')
-        a3 = np.fromstring(match.group('a3'), sep=' ')
+        a1 = [float(x) for x in match.group('a1').split()]
+        a2 = [float(x) for x in match.group('a2').split()]
+        a3 = [float(x) for x in match.group('a3').split()]
 
         return {
             'units': match.group('units'),
             'a1': a1, 'a2': a2, 'a3': a3,
-            'lattice_vectors': np.array([a1, a2, a3]),
         }
     else:
         return None
@@ -56,7 +53,7 @@ def parse_atoms(string: str) -> dict:
             'atoms': [
                 {
                     'species': line.split()[0],
-                    'basis_vector': np.fromiter(map(float, line.split()[1:]), dtype=float),
+                    'basis_vector': [float(x) for x in line.split()[1:]],
                 }
                 for line in match.group('atoms').splitlines()
             ]
@@ -79,11 +76,13 @@ def parse_projections(string: str) -> dict:
 
 def parse_kpoints(string: str) -> dict:
     return {
-        'kpoints': np.fromstring(string, sep='\n').reshape((len(string.splitlines()), -1))[:, :3]
+        'kpoints': [
+            [float(x) for x in line.split()] for line in string.splitlines()
+        ]
     }
 
 
-def parse_win(string: str) -> dict:
+def parse_win_raw(string: str) -> dict:
     """
     Parse WIN
 

tests/test_nnkp.py

@@ -2,7 +2,7 @@
 
 import pytest
 
-import w90io
+import w90io._schema
 
 
 @pytest.mark.parametrize('example', [f'example{i:02d}' for i in [1, 2, 3, 4, 5, 6, 7, 9, 10, 11, 13, 17, 18, 19, 20]])
@@ -11,6 +11,11 @@ def test_parse_nnkp(wannier90, example):
         contents = fh.read()
 
     try:
-        w90io.parse_nnkp(contents)
+        parsed_nnkp = w90io.parse_nnkp_raw(contents)
+    except Exception:
+        assert False
+
+    try:
+        w90io._schema.NNKP.parse_obj(parsed_nnkp)
     except Exception:
         assert False

tests/test_win.py

@@ -11,6 +11,11 @@ def test_parse_win(wannier90, example):
         contents = fh.read()
 
     try:
-        w90io.parse_win(contents)
+        parsed_win = w90io.parse_win_raw(contents)
+    except Exception:
+        assert False
+
+    try:
+        w90io._schema.WIN.parse_obj(parsed_win)
     except Exception:
         assert False