PintField implementation for uMongo

pintfield.py

import collections

from marshmallow import ValidationError
from umongo.fields import BaseField
import numpy as np
from json import JSONEncoder

from ..._units import units as units_quantity


class PintField(BaseField):
    """Custom uMongo Field for serializing/de-serializing Pint values."""

    def _serialize(self, value, attr, obj, **kwargs):
        if value is None:
            return []
        mag = value.m
        units = value.u
        if isinstance(mag, np.ndarray):
            mag = mag.tolist()
        return (mag, str(units))

    def _serialize_to_mongo(self, obj):
        return self._serialize(obj, None, None)

    def _deserialize(self, value, attr, data, **kwargs):
        if isinstance(value, units_quantity.Quantity):
            return value
        if not isinstance(value, collections.Iterable) or len(value) != 2:
            raise ValidationError('Pint fields must be lists with exactly 2 items (magnitude and units).')
        mag = value[0]
        units = value[1]
        if isinstance(mag, collections.Iterable):
            mag = np.array(mag)
        return units_quantity.Quantity(mag, units)

    def _deserialize_from_mongo(self, value):
        return self._deserialize(value, None, None)


class EnhancedJSONEncoder(JSONEncoder):
    def default(self, obj):
        if isinstance(obj, units_quantity.Quantity):
            return obj.to_tuple()
        if isinstance(obj, np.ndarray):
            return obj.tolist()
        return super().default(obj)