Source code for wolframclient.serializers.encoders.pandas

# -*- coding: utf-8 -*-

from __future__ import absolute_import, print_function, unicode_literals

from collections import OrderedDict
from itertools import starmap

from wolframclient.utils.api import pandas
from wolframclient.utils.dispatch import Dispatch

encoder = Dispatch()


def safe_pandas_length(o):
    """ Return the length of a pandas Series and DataFrame as expected for WL serialization.

    - The length of a Series is the only value of the tuple `shape`.
    - The length of a dataframe is the number of columns. It's the second value of `shape`.

    This function is safe, when the shape does not have the expected number of elements, it fails silently and
    returns `None`, the object is later traversed to find out how many elements it contains.
    """
    try:
        return o.shape[-1]
    except (TypeError, IndexError):
        return


def encode_as_dataset(serializer, o, length):
    return serializer.serialize_function(
        serializer.serialize_symbol(b"Dataset"),
        (encode_as_association(serializer, o, length),),
    )


def encode_as_list(serializer, o, length):
    return serializer.serialize_iterable(
        starmap(
            lambda k, v: serializer.serialize_rule(k, v), encoded_kv_tuples(serializer, o)
        ),
        length=length,
    )


def encode_as_association(serializer, o, length):
    return serializer.serialize_association(encoded_kv_tuples(serializer, o), length=length)


def encoded_kv_tuples(serializer, o):
    return ((serializer.encode(k), serializer.encode(v)) for k, v in o.items())


def encode_as_timeseries(serializer, o, length):
    return serializer.serialize_function(
        serializer.serialize_symbol(b"TimeSeries"),
        (
            serializer.serialize_iterable(
                (
                    serializer.serialize_iterable(item, length=2)
                    for item in encoded_kv_tuples(serializer, o)
                ),
                length=length,
            ),
        ),
    )


def _distribute_multikey(o):
    expr_dict = OrderedDict()
    for multikey, value in o.iteritems():
        cur_dict = expr_dict
        for key in multikey[:-1]:
            if key not in cur_dict:
                cur_dict[key] = OrderedDict()
            cur_dict = cur_dict[key]
        cur_dict[multikey[-1]] = value
    return expr_dict


def encode_multiindex_as_assoc(serializer, o, length):
    return serializer.encode(_distribute_multikey(o))


def encode_multiindex_as_dataset(serializer, o, length):
    return serializer.serialize_function(
        serializer.serialize_symbol(b"Dataset"), (serializer.encode(_distribute_multikey(o)),)
    )


PANDAS_PROPERTIES = {
    "pandas_series_head": {"dataset", "list", "association"},
    "pandas_dataframe_head": {"dataset", "association"},
    "timeseries": True,
}

ENCODERS = {
    "default": {
        "dataset": encode_as_dataset,
        "list": encode_as_list,
        "association": encode_as_association,
    },
    "datetimeindex": encode_as_timeseries,
    "multiindex": {
        "association": encode_multiindex_as_assoc,
        "list": encode_multiindex_as_assoc,
        "dataset": encode_multiindex_as_dataset,
    },
}


def get_series_encoder_from_index(index, use_ts, form):
    if use_ts and isinstance(index, pandas.DatetimeIndex):
        return ENCODERS["datetimeindex"]
    elif isinstance(index, pandas.MultiIndex):
        return ENCODERS["multiindex"][form or "dataset"]
    else:
        return ENCODERS["default"][form or "association"]


INVALID_PROPERTY_MSG = "Invalid property %s, expecting %s"


def normalized_prop_timeseries(serializer):
    prop = serializer.get_property("timeseries", d=True)
    if not isinstance(prop, bool):
        raise ValueError(
            "Invalid value for property 'timeseries'. Expecting a boolean, got %s." % prop
        )
    return prop


def normalized_prop_pandas_series_head(serializer):
    """ Check property `pandas_series_head` only if specified (not None). """
    prop = serializer.get_property("pandas_series_head", d=None)
    if prop and prop not in PANDAS_PROPERTIES["pandas_series_head"]:
        raise ValueError(
            "Invalid value for property 'pandas_series_head'. Expecting one of (%s), got %s."
            % (", ".join(PANDAS_PROPERTIES["pandas_series_head"]), prop)
        )
    return prop


@encoder.dispatch(pandas.Series)
def encode_panda_series(serializer, o):
    use_ts = normalized_prop_timeseries(serializer)
    form = normalized_prop_pandas_series_head(serializer)
    encoder = get_series_encoder_from_index(o.index, use_ts, form)
    return encoder(serializer, o, safe_pandas_length(o))


def encode_dataframe_as_assoc(serializer, o, length):
    use_ts = normalized_prop_timeseries(serializer)
    return serializer.serialize_association(
        (
            (
                serializer.encode(k),
                get_series_encoder_from_index(v.index, use_ts, "association")(
                    serializer, v, safe_pandas_length(v)
                ),
            )
            for k, v in o.T.items()
        ),
        length=length,
    )


def encode_dataframe_as_dataset(serializer, o, length):
    return serializer.serialize_function(
        serializer.serialize_symbol(b"Dataset"),
        (encode_dataframe_as_assoc(serializer, o, length),),
    )


@encoder.dispatch(pandas.DataFrame)
def encoder_panda_dataframe(serializer, o):
    head = serializer.get_property("pandas_dataframe_head", d=None)
    if head is None or head == "dataset":
        return encode_dataframe_as_dataset(serializer, o, safe_pandas_length(o.index))
    elif head in PANDAS_PROPERTIES["pandas_dataframe_head"]:
        return encode_dataframe_as_assoc(serializer, o, safe_pandas_length(o.index))
    else:
        raise ValueError(
            "Invalid value for property 'pandas_dataframe_head'. Expecting one of (%s), got %s."
            % (", ".join(PANDAS_PROPERTIES["pandas_dataframe_head"]), head)
        )