JsonColumn for working with SQLAlchemy JSON and JSONB columns

jsoncolumn.py

class JsonColumn(object):
    """
    Descriptor class for accessing values from `JSONB` columns.

    **In Expressions**:

    As class attribute the descriptor resolves the class property
    pointing to prepared `JSONB` element using the ``'column'`` value's
    :attr:`key` attribute. It then uses the element's item accessor with the
    ``*keys`` and returns it :attr:`astext`, if not ``'astext'`` is True.

    For example given the (simplified) class

    .. code-block:: python

        Klass(ModelBase):

            json = Column(JSONB)
            some_value = JsonColumn(json, 'path', 'to', 'value')

    these two are equivalent:

    .. code-block:: python

        Klass.some_value
        Klass.json['path', 'to', 'value'].astext

    **On Instances**:

    Similar to the machinery on classes, but accesses the actual `JSON`
    data or returns `None`, if the key does not exist.

    :param column:
        SQLAlchemy :class:`sqlalchemy.dialects.postgresql.JSON` column.

    :param *keys:
        Keys to `JSONB` structure for accessing the value.

    :param astext:
        If True, apply :attr:`astext` to
        :class:`sqlalchemy.dialects.postgresql.json.JSONElement`
        expression.

    """

    def __init__(self, column, *keys, astext=True):
        self._keys = keys
        self._column = column
        self._astext = astext

    @staticmethod
    def _extract(obj, key):
        if isinstance(obj, dict):
            return obj[key]

    def __get__(self, instance, owner):
        prop = getattr(instance or owner, self._column.key)

        if instance is None:
            element = prop[self._keys]
            if self._astext:
                return element.astext

            return element

        try:
            return reduce(self._extract, self._keys, prop)

        except KeyError:
            pass