Support non-FloatQuantity and non-ArrayQuantity types

[mikemhenry]

I expected this to work:

from openff.models.models import DefaultModel
from openff.models.types import ArrayQuantity, FloatQuantity
from openff.units import unit

class Atom(DefaultModel):
    mass: FloatQuantity["atomic_mass_constant"]
    charge: FloatQuantity["elementary_charge"]
    some_array: ArrayQuantity["nanometer"]
    baz: int

atom = Atom(
    mass=12.011 * unit.atomic_mass_constant,
    charge=0.0 * unit.elementary_charge,
    some_array=unit.Quantity([4, -1, 0], unit.nanometer),
    baz = 2,
)

Atom.parse_raw(atom.json())

But if fails on the round trip since it can't parse baz properly:

{'baz': 2,
 'charge': '{"val": 0.0, "unit": "elementary_charge"}',
 'mass': '{"val": 12.011, "unit": "atomic_mass_constant"}',
 'some_array': '{"val": [4, -1, 0], "unit": "nanometer"}'}

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
File ~/miniconda3/envs/gufe/lib/python3.10/site-packages/pydantic/main.py:534, in pydantic.main.BaseModel.parse_raw()

File ~/miniconda3/envs/gufe/lib/python3.10/site-packages/pydantic/parse.py:37, in pydantic.parse.load_str_bytes()

File ~/miniconda3/envs/gufe/lib/python3.10/site-packages/openff/models/types.py:135, in json_loader(data)
    132 try:
    133     # Directly look for an encoded FloatQuantity/ArrayQuantity,
    134     # which is itself a dict
--> 135     v = json.loads(val)
    136 except json.JSONDecodeError:
    137     # Handles some cases of the val being a primitive type

File ~/miniconda3/envs/gufe/lib/python3.10/json/__init__.py:339, in loads(s, cls, object_hook, parse_float, parse_int, parse_constant, object_pairs_hook, **kw)
    338 if not isinstance(s, (bytes, bytearray)):
--> 339     raise TypeError(f'the JSON object must be str, bytes or bytearray, '
    340                     f'not {s.__class__.__name__}')
    341 s = s.decode(detect_encoding(s), 'surrogatepass')

TypeError: the JSON object must be str, bytes or bytearray, not int

During handling of the above exception, another exception occurred:

ValidationError                           Traceback (most recent call last)
Input In [4], in <cell line: 20>()
     10     baz: int
     13 atom = Atom(
     14     mass=12.011 * unit.atomic_mass_constant,
     15     charge=0.0 * unit.elementary_charge,
     16     some_array=unit.Quantity([4, -1, 0], unit.nanometer),
     17     baz = 2,
     18 )
---> 20 Atom.parse_raw(atom.json())

File ~/miniconda3/envs/gufe/lib/python3.10/site-packages/pydantic/main.py:543, in pydantic.main.BaseModel.parse_raw()

ValidationError: 1 validation error for Atom
__root__
  the JSON object must be str, bytes or bytearray, not int (type=type_error)

This bit https://github.com/mattwthompson/openff-models/blob/61260ec1d2746a0431671498dfe1540b006354db/openff/models/types.py#L137 doesn't actually seem to handle the case where the type isn't a float or array quantity.

[mattwthompson]

Weird - that's definitely not good. I'll have a look!

[mattwthompson]

Slight clarification: I'm pretty sure the issue is isolated to the JSON parser forgetting that primitives exist.

[mattwthompson]

Closed #4 - thanks Mike!