Merge pull request #217 from MolSSI/loriab-patch-2

fix schema for Array classes

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

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):