Source code for iodata.attrutils

# IODATA is an input and output module for quantum chemistry.
# Copyright (C) 2011-2019 The IODATA Development Team
# This file is part of IODATA.
# IODATA is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 3
# of the License, or (at your option) any later version.
# IODATA is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# GNU General Public License for more details.
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <>
# --
"""Utilities for building attr classes."""

from typing import Callable

import numpy as np

__all__ = ["convert_array_to", "validate_shape"]

[docs] def convert_array_to(dtype): """Return a function to convert arrays to the given type.""" def converter(array): if array is None: return None return np.asarray(array, dtype=dtype) return converter
[docs] def validate_shape(*shape_requirements: tuple) -> Callable: """Return a validator for the shape of an array or the length of an iterable. Parameters ---------- shape_requirements Specifications for the required shape. Every item of the tuple describes the required size of the corresponding axis of an array. Also the number of items should match the dimensionality of the array. When the validator is used for general iterables, this tuple should contain just one element. Possible values for each item are explained in the "Notes" section below. Returns ------- A validator function for the attr library. Notes ----- Every element of ``shape_requirements`` defines the expected size of an array along the corresponding axis. An item in this tuple at position (or index) ``i`` can be one of the following: 1. An integer, which is taken as the expected size along axis ``i``. 2. None. In this case, the size of the array along axis ``i`` is not checked. 3. A string, which should be the name of another integer attribute with the expected size along axis ``i``. The other attribute is always an attribute of the same object as the attribute being checked. 4. A 2-tuple containing a name and an integer. In this case, the name refers to another attribute which is an array or an iterable. When the integer is 0, just the length of the other attribute is used. When the integer is non-zero, the other attribute must be an array and the integer selects an axis. The size of the other array along the selected axis is then used as the expected size of the array being checked along axis ``i``. """ def validator(obj, attribute, value): # Build the expected shape, with the rules from the docstring. expected_shape = [] for item in shape_requirements: if isinstance(item, int) or item is None: expected_shape.append(item) elif isinstance(item, str): expected_shape.append(getattr(obj, item)) elif isinstance(item, tuple) and len(item) == 2: other_name, other_axis = item other = getattr(obj, other_name) if other is None: raise TypeError(f"Other attribute '{other_name}' is not set.") if other_axis == 0: expected_shape.append(len(other)) else: if other_axis >= other.ndim or other_axis < 0: raise TypeError( "Cannot get length along axis " f"{other_axis} of attribute {other_name} with ndim {other.ndim}." ) expected_shape.append(other.shape[other_axis]) else: raise ValueError(f"Cannot interpret item in shape_requirements: {item}") expected_shape = tuple(expected_shape) # Get the actual shape observed_shape = value.shape if isinstance(value, np.ndarray) else (len(value),) # Compare match = True if len(expected_shape) != len(observed_shape): match = False if match: for es, os in zip(expected_shape, observed_shape): if es is None: continue if es != os: match = False break # Raise TypeError if needed. if not match: raise TypeError( f"Expecting shape {expected_shape} for attribute {}, " f"got {observed_shape}" ) return validator