diff --git a/qcelemental/models/align.py b/qcelemental/models/align.py index 00db470a..aa3a8358 100644 --- a/qcelemental/models/align.py +++ b/qcelemental/models/align.py @@ -1,5 +1,7 @@ +from typing import Optional + import numpy as np -from pydantic import validator +from pydantic import Field, validator from ..util import blockwise_contract, blockwise_expand from .basemodels import ProtoModel @@ -19,10 +21,18 @@ class AlignmentMill(ProtoModel): """ - shift: Array[float] # type: ignore - rotation: Array[float] # type: ignore - atommap: Array[int] # type: ignore - mirror: bool = False + shift: Optional[Array[float]] = Field( # type: ignore + None, description="Translation array (3,) for coordinates.", + ) + rotation: Optional[Array[float]] = Field( # type: ignore + None, description="Rotation array (3, 3) for coordinates.", + ) + atommap: Optional[Array[int]] = Field( # type: ignore + None, description="Atom exchange map (nat,) for coordinates.", + ) + mirror: bool = Field( + False, description="Do mirror invert coordinates?", + ) class Config: force_skip_defaults = True diff --git a/qcelemental/models/types.py b/qcelemental/models/types.py index 8457efbb..ee961535 100644 --- a/qcelemental/models/types.py +++ b/qcelemental/models/types.py @@ -1,3 +1,5 @@ +from typing import Any, Dict + import numpy as np @@ -15,6 +17,19 @@ def validate(cls, v): return v + @classmethod + def __modify_schema__(cls, field_schema: Dict[str, Any]) -> None: + dt = cls._dtype + if dt is int or np.issubdtype(dt, np.integer): + items = {"type": "number", "multipleOf": 1.0} + elif dt is float or np.issubdtype(dt, np.floating): + items = {"type": "number"} + elif dt is str or np.issubdtype(dt, np.string_): + items = {"type": "string"} + elif dt is bool or np.issubdtype(dt, np.bool_): + items = {"type": "boolean"} + field_schema.update(type="array", items=items) + class ArrayMeta(type): def __getitem__(self, dtype):