Merge pull request #37261 from ajjackson/abins-pydantic

Try using pydantic to cut down on type validation boilerplate

Author: thomashampson

conda/recipes/mantid/meta.yaml

@@ -74,6 +74,8 @@ requirements:
     - joblib
     - orsopy {{ orsopy }}
     - quasielasticbayes
+    - pydantic
+
   run_constrained:
     - matplotlib {{ matplotlib }}
 

mantid-developer-linux.yml

@@ -51,6 +51,7 @@ dependencies:
   - versioningit>=2.1
   - joblib
   - orsopy==1.2.1 # Fix the version to avoid updates being pulled in automatically, which might change the Reflectometry ORSO file content or layout and cause tests to fail.
+  - pydantic
 
   # Not Windows, OpenGL implementation:
   - mesa-libgl-devel-cos7-x86_64>=18.3.4

mantid-developer-osx.yml

@@ -49,6 +49,7 @@ dependencies:
   - versioningit>=2.1
   - joblib
   - orsopy==1.2.1 # Fix the version to avoid updates being pulled in automatically, which might change the Reflectometry ORSO file content or layout and cause tests to fail.
+  - pydantic
 
   # Needed only for development
   - black  # may be out of sync with pre-commit

mantid-developer-win.yml

@@ -48,7 +48,10 @@ dependencies:
   - versioningit>=2.1
   - joblib
   - orsopy==1.2.1 # Fix the version to avoid updates being pulled in automatically, which might change the Reflectometry ORSO file content or layout and cause tests to fail.
+  - pydantic
+
   # Needed only for development
   - black  # may be out of sync with pre-commit
   - cppcheck==2.14.2
   - pre-commit>=2.12.0
+

scripts/abins/abinsdata.py

@@ -6,6 +6,8 @@
 # SPDX - License - Identifier: GPL - 3.0 +
 from typing import Any, Dict, Type, TypedDict, TypeVar
 
+from pydantic import validate_call
+
 import mantid
 from abins.kpointsdata import KpointsData
 from abins.atomsdata import AtomsData
@@ -23,13 +25,9 @@ class AbinsData:
 
     """
 
+    @validate_call(config=dict(arbitrary_types_allowed=True, strict=True))
     def __init__(self, *, k_points_data: KpointsData, atoms_data: AtomsData) -> None:
-        if not isinstance(k_points_data, KpointsData):
-            raise TypeError("Invalid type of k-points data.: {}".format(type(k_points_data)))
         self._k_points_data = k_points_data
-
-        if not isinstance(atoms_data, AtomsData):
-            raise TypeError("Invalid type of atoms data.")
         self._atoms_data = atoms_data
         self._check_consistent_dimensions()
 

scripts/abins/atomsdata.py

@@ -6,8 +6,9 @@
 # SPDX - License - Identifier: GPL - 3.0 +
 import collections.abc
 import numbers
-from typing import Dict, List, Optional, overload, Union, TypedDict
+from typing import Dict, List, Optional, overload, TypedDict, Union
 import re
+
 import numpy as np
 
 import abins
@@ -146,12 +147,10 @@ def __len__(self) -> int:
         return len(self._data)
 
     @overload
-    def __getitem__(self, item: int) -> _AtomData:
-        ...
+    def __getitem__(self, item: int) -> _AtomData: ...
 
     @overload
-    def __getitem__(self, item: slice) -> List[_AtomData]:
-        ...
+    def __getitem__(self, item: slice) -> List[_AtomData]: ...
 
     def __getitem__(self, item):
         return self._data[item]

scripts/abins/constants.py

@@ -5,6 +5,7 @@
 #   Institut Laue - Langevin & CSNS, Institute of High Energy Physics, CAS
 # SPDX - License - Identifier: GPL - 3.0 +
 import math
+from typing import Literal
 import warnings
 
 import numpy as np
@@ -164,6 +165,7 @@
 
 # ALL_SAMPLE_FORMS = ["SingleCrystal", "Powder"]  # valid forms of samples
 ALL_SAMPLE_FORMS = ["Powder"]  # valid forms of samples
+ALL_SAMPLE_FORMS_TYPE = Literal["Powder"]
 
 # keywords which define data structure of KpointsData
 ALL_KEYWORDS_K_DATA = ["weights", "k_vectors", "frequencies", "atomic_displacements", "unit_cell"]
@@ -191,6 +193,9 @@
 INT_ID = np.dtype(np.uint32).num
 INT_TYPE = np.dtype(np.uint32)
 
+# Valid types for hdf5 attr read/write
+HDF5_ATTR_TYPE = np.int64 | int | np.float64 | float | str | bytes | bool
+
 HIGHER_ORDER_QUANTUM_EVENTS = 3  # number of quantum order effects taken into account
 HIGHER_ORDER_QUANTUM_EVENTS_DIM = HIGHER_ORDER_QUANTUM_EVENTS
 

scripts/abins/io.py

@@ -10,46 +10,41 @@
 import os
 import subprocess
 import shutil
+from typing import List, Optional
+
 import h5py
+
 import numpy as np
+from pydantic import BaseModel, ConfigDict, Field, validate_call
 
 import abins
-from abins.constants import AB_INITIO_FILE_EXTENSIONS, BUF
+from abins.constants import AB_INITIO_FILE_EXTENSIONS, BUF, HDF5_ATTR_TYPE
 from mantid.kernel import logger, ConfigService
 
 
-class IO(object):
+class IO(BaseModel):
     """
     Class for Abins I/O HDF file operations.
     """
 
-    def __init__(self, input_filename=None, group_name=None, setting="", autoconvolution: bool = False, temperature: float = None):
-        self._setting = setting
-        self._autoconvolution = autoconvolution
-        self._temperature = temperature
-
-        if isinstance(input_filename, str):
-            self._input_filename = input_filename
-            try:
-                self._hash_input_filename = self.calculate_ab_initio_file_hash()
-            except IOError as err:
-                logger.error(str(err))
-            except ValueError as err:
-                logger.error(str(err))
+    model_config = ConfigDict(strict=True)
 
-            # extract name of file from the full path in the platform independent way
-            filename = os.path.basename(self._input_filename)
+    input_filename: str = Field(min_length=1)
+    group_name: str = Field(min_length=1)
+    setting: str = ""
+    autoconvolution: int = 10
+    temperature: float = None
 
-            if filename.strip() == "":
-                raise ValueError("Name of the file cannot be an empty string.")
+    def model_post_init(self, __context):
+        try:
+            self._hash_input_filename = self.calculate_ab_initio_file_hash()
+        except IOError as err:
+            logger.error(str(err))
+        except ValueError as err:
+            logger.error(str(err))
 
-        else:
-            raise ValueError("Invalid name of input file. String was expected.")
-
-        if isinstance(group_name, str):
-            self._group_name = group_name
-        else:
-            raise ValueError("Invalid name of the group. String was expected.")
+        # extract name of file from the full path in the platform independent way
+        filename = os.path.basename(self.input_filename)
 
         if filename.split(".")[-1] in AB_INITIO_FILE_EXTENSIONS:
             core_name = filename[0 : filename.rfind(".")]  # e.g. NaCl.phonon -> NaCl (core_name) -> NaCl.hdf5
@@ -85,15 +80,15 @@ def _valid_setting(self):
         :returns: True if consistent, otherwise False.
         """
         saved_setting = self.load(list_of_attributes=["setting"])
-        return self._setting == saved_setting["attributes"]["setting"]
+        return self.setting == saved_setting["attributes"]["setting"]
 
     def _valid_autoconvolution(self):
         """
         Check if autoconvolution setting matches content of HDF file
         :returns: True if consistent, otherwise False
         """
         saved_autoconvolution = self.load(list_of_attributes=["autoconvolution"])
-        return self._autoconvolution == saved_autoconvolution["attributes"]["autoconvolution"]
+        return self.autoconvolution == saved_autoconvolution["attributes"]["autoconvolution"]
 
     def _valid_temperature(self):
         """
@@ -104,7 +99,7 @@ def _valid_temperature(self):
 
         :returns: True if consistent or temperature not set for Clerk, otherwise False
         """
-        if self._temperature is None:
+        if self.temperature is None:
             return True
 
         else:
@@ -115,7 +110,7 @@ def _valid_temperature(self):
                 return False
 
             else:
-                return np.abs(self._temperature - saved_temperature["attributes"]["temperature"]) < T_THRESHOLD
+                return np.abs(self.temperature - saved_temperature["attributes"]["temperature"]) < T_THRESHOLD
 
     @classmethod
     def _close_enough(cls, previous, new):
@@ -215,7 +210,8 @@ def erase_hdf_file(self):
         with h5py.File(self._hdf_filename, "w") as hdf_file:
             hdf_file.close()
 
-    def add_attribute(self, name=None, value=None):
+    @validate_call(config=ConfigDict(arbitrary_types_allowed=True, strict=True))
+    def add_attribute(self, name: str, value: HDF5_ATTR_TYPE | None) -> None:
         """
         Adds attribute to the dictionary with other attributes.
         :param name: name of the attribute
@@ -228,10 +224,10 @@ def add_file_attributes(self):
         Add attributes for input data filename, hash of file, advanced parameters to data for HDF5 file
         """
         self.add_attribute("hash", self._hash_input_filename)
-        self.add_attribute("setting", self._setting)
-        self.add_attribute("autoconvolution", self._autoconvolution)
-        self.add_attribute("temperature", self._temperature)
-        self.add_attribute("filename", self._input_filename)
+        self.add_attribute("setting", self.setting)
+        self.add_attribute("autoconvolution", self.autoconvolution)
+        self.add_attribute("temperature", self.temperature)
+        self.add_attribute("filename", self.input_filename)
         self.add_attribute("advanced_parameters", json.dumps(abins.parameters.non_performance_parameters))
 
     def add_data(self, name=None, value=None):
@@ -250,15 +246,10 @@ def _save_attributes(self, group=None):
         :param group: group to which attributes should be saved.
         """
         for name in self._attributes:
-            if isinstance(self._attributes[name], (np.int64, int, np.float64, float, str, bytes, bool)):
-                group.attrs[name] = self._attributes[name]
-            elif self._attributes[name] is None:
+            if self._attributes[name] is None:
                 group.attrs[name] = "None"
             else:
-                raise ValueError(
-                    "Invalid value of attribute. String, "
-                    "int, bool or bytes was expected! " + name + "= (invalid type : %s) " % type(self._attributes[name])
-                )
+                group.attrs[name] = self._attributes[name]
 
     def _recursively_save_structured_data_to_group(self, hdf_file=None, path=None, dic=None):
         """
@@ -309,17 +300,17 @@ def _save_data(self, hdf_file=None, group=None):
             elif isinstance(self._data[item], dict):
                 self._recursively_save_structured_data_to_group(hdf_file=hdf_file, path=group.name + "/" + item + "/", dic=self._data[item])
             else:
-                raise ValueError("Invalid structured dataset. Cannot save %s type" % type(item))
+                raise TypeError("Invalid structured dataset. Cannot save %s type" % type(item))
 
     def save(self):
         """
         Saves datasets and attributes to an hdf file.
         """
 
         with h5py.File(self._hdf_filename, "a") as hdf_file:
-            if self._group_name not in hdf_file:
-                hdf_file.create_group(self._group_name)
-            group = hdf_file[self._group_name]
+            if self.group_name not in hdf_file:
+                hdf_file.create_group(self.group_name)
+            group = hdf_file[self.group_name]
 
             if len(self._attributes.keys()) > 0:
                 self._save_attributes(group=group)
@@ -342,7 +333,8 @@ def save(self):
             pass
 
     @staticmethod
-    def _list_of_str(list_str=None):
+    @validate_call
+    def _list_of_str(list_str: Optional[List[str]]) -> bool:
         """
         Checks if all elements of the list are strings.
         :param list_str: list to check
@@ -351,9 +343,6 @@ def _list_of_str(list_str=None):
         if list_str is None:
             return False
 
-        if not (isinstance(list_str, list) and all([isinstance(list_str[item], str) for item in range(len(list_str))])):
-            raise ValueError("Invalid list of items to load!")
-
         return True
 
     def _load_attributes(self, list_of_attributes=None, group=None):
@@ -422,39 +411,34 @@ def _encode_utf8_if_text(item):
         else:
             return item
 
-    def _load_dataset(self, hdf_file=None, name=None, group=None):
+    @validate_call(config=ConfigDict(arbitrary_types_allowed=True))
+    def _load_dataset(self, *, hdf_file: h5py.File, name: str, group: h5py.Group):
         """
         Loads one structured dataset.
         :param hdf_file:  hdf file object from which structured dataset should be loaded.
         :param name:  name of dataset
         :param group: name of the main group
         :returns: loaded dataset
         """
-        if not isinstance(name, str):
-            raise ValueError("Invalid name of the dataset.")
-
         if name in group:
             hdf_group = group[name]
         else:
-            raise ValueError("Invalid name of the dataset.")
+            raise ValueError(f"Dataset {name} not found in group {group.name}.")
 
         # noinspection PyUnresolvedReferences,PyProtectedMember
-        if isinstance(hdf_group, h5py._hl.dataset.Dataset):
+        if isinstance(hdf_group, h5py.Dataset):
             return hdf_group[()]
         elif all([self._get_subgrp_name(hdf_group[el].name).isdigit() for el in hdf_group.keys()]):
-            structured_dataset_list = []
-            # here we make an assumption about keys which have a numeric values; we assume that always : 1, 2, 3... Max
-            num_keys = len(hdf_group.keys())
-            for item in range(num_keys):
-                structured_dataset_list.append(
-                    self._recursively_load_dict_contents_from_group(hdf_file=hdf_file, path=hdf_group.name + "/%s" % item)
-                )
+            structured_dataset_list = [
+                self._recursively_load_dict_contents_from_group(hdf_file=hdf_file, path=f"{hdf_group.name}/{key}")
+                for key in sorted(hdf_group.keys(), key=int)
+            ]
             return structured_dataset_list
         else:
             return self._recursively_load_dict_contents_from_group(hdf_file=hdf_file, path=hdf_group.name + "/")
 
     @classmethod
-    def _recursively_load_dict_contents_from_group(cls, hdf_file=None, path=None):
+    def _recursively_load_dict_contents_from_group(cls, *, hdf_file: h5py.File, path: str):
         """
         Loads structure dataset which has form of Python dictionary.
         :param hdf_file:  hdf file object from which dataset is loaded
@@ -468,7 +452,7 @@ def _recursively_load_dict_contents_from_group(cls, hdf_file=None, path=None):
             if isinstance(item, h5py._hl.dataset.Dataset):
                 ans[key] = item[()]
             elif isinstance(item, h5py._hl.group.Group):
-                ans[key] = cls._recursively_load_dict_contents_from_group(hdf_file, path + key + "/")
+                ans[key] = cls._recursively_load_dict_contents_from_group(hdf_file=hdf_file, path=f"{path}{key}/")
         return ans
 
     def load(self, list_of_attributes=None, list_of_datasets=None):
@@ -484,10 +468,10 @@ def load(self, list_of_attributes=None, list_of_datasets=None):
 
         results = {}
         with h5py.File(self._hdf_filename, "r") as hdf_file:
-            if self._group_name not in hdf_file:
-                raise ValueError("No group %s in hdf file." % self._group_name)
+            if self.group_name not in hdf_file:
+                raise ValueError("No group %s in hdf file." % self.group_name)
 
-            group = hdf_file[self._group_name]
+            group = hdf_file[self.group_name]
 
             if self._list_of_str(list_str=list_of_attributes):
                 results["attributes"] = self._load_attributes(list_of_attributes=list_of_attributes, group=group)
@@ -521,7 +505,7 @@ def _calculate_hash(filename=None):
         return hash_calculator.hexdigest()
 
     def get_input_filename(self):
-        return self._input_filename
+        return self.input_filename
 
     def calculate_ab_initio_file_hash(self):
         """
@@ -530,4 +514,4 @@ def calculate_ab_initio_file_hash(self):
         :returns: string representation of hash for file with vibrational data which contains only hexadecimal digits
         """
 
-        return self._calculate_hash(filename=self._input_filename)
+        return self._calculate_hash(filename=self.input_filename)