Source code for wolframclient.serializers.wxfencoder.wxfnumpyencoder

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

from __future__ import absolute_import, print_function, unicode_literals

from wolframclient.serializers.wxfencoder.constants import (
    ARRAY_TYPES, WXF_CONSTANTS, WXF_HEADER_COMPRESS, WXF_HEADER_SEPARATOR,
    WXF_VERSION)
from wolframclient.serializers.wxfencoder.wxfencoder import WXFEncoder
from wolframclient.serializers.wxfencoder.wxfexpr import (WXFExprNumericArray,
                                                          WXFExprPackedArray)
from wolframclient.utils.api import numpy

__all__ = [
    'NumPyWXFEncoder',
]


class NumPyWXFEncoder(WXFEncoder):
    '''
    NumPy array encoder. Encode numpy array as instances of packed array and / or raw array.
    By default only packed arrays are generated. Unsigned integer data are cast to a type that
    can fit the maximum value.

    It's possible to add support for raw arrays only for unsigned data, in which case both
    `packed_array_support` and `numeric_array_support` must be true.

    >>> NumPyWXFEncoder(packed_array_support=True, numeric_array_support=True)

    Finally it's possible to only output raw arrays with:

    >>> NumPyWXFEncoder(packed_array_support=False, numeric_array_support=True)

    '''

    def __init__(self, packed_array_support=True, numeric_array_support=False):
        if not packed_array_support and not numeric_array_support:
            raise ValueError(
                'At least one of the two parameters packed_array_support or numeric_array_support must be True.'
            )
        self.packed_array_support = packed_array_support
        self.numeric_array_support = numeric_array_support

    def encode(self, python_expr):
        if isinstance(python_expr, numpy.ndarray):
            if self.packed_array_support:
                array_class = WXFExprPackedArray
            else:
                array_class = WXFExprNumericArray

            if python_expr.dtype == numpy.int8:
                value_type = ARRAY_TYPES.Integer8
                data = python_expr.astype('<i1')
            elif python_expr.dtype == numpy.int16:
                data = python_expr.astype('<i2')
                value_type = ARRAY_TYPES.Integer16
            elif python_expr.dtype == numpy.int32:
                data = python_expr.astype('<i4')
                value_type = ARRAY_TYPES.Integer32
            elif python_expr.dtype == numpy.int64:
                data = python_expr.astype('<i8')
                value_type = ARRAY_TYPES.Integer64
            elif python_expr.dtype == numpy.uint8:
                if self.numeric_array_support:
                    value_type = ARRAY_TYPES.UnsignedInteger8
                    data = python_expr.astype('<u1')
                    array_class = WXFExprNumericArray
                else:
                    value_type = ARRAY_TYPES.Integer16
                    data = python_expr.astype('<i2')
            elif python_expr.dtype == numpy.uint16:
                if self.numeric_array_support:
                    value_type = ARRAY_TYPES.UnsignedInteger16
                    data = python_expr.astype('<u2')
                    array_class = WXFExprNumericArray
                else:
                    value_type = ARRAY_TYPES.Integer32
                    data = python_expr.astype('<i4')
            elif python_expr.dtype == numpy.uint32:
                if self.numeric_array_support:
                    value_type = ARRAY_TYPES.UnsignedInteger32
                    data = python_expr.astype('<u4')
                    array_class = WXFExprNumericArray
                else:
                    value_type = ARRAY_TYPES.Integer64
                    data = python_expr.astype('<i8')
            # no one to one mapping to signed values, even if most of the time
            # the values would fit
            elif python_expr.dtype == numpy.uint64:
                if self.numeric_array_support:
                    value_type = ARRAY_TYPES.UnsignedInteger64
                    data = python_expr.astype('<u8')
                    array_class = WXFExprNumericArray
                else:
                    TypeError(
                        'Cannot represent data of type uint64 as signed int64')
            elif python_expr.dtype == numpy.float32:
                value_type = ARRAY_TYPES.Real32
                data = python_expr
            elif python_expr.dtype == numpy.float64:
                value_type = ARRAY_TYPES.Real64
                data = python_expr
            elif python_expr.dtype == numpy.complex64:
                value_type = ARRAY_TYPES.ComplexReal32
                data = python_expr
            elif python_expr.dtype == numpy.complex128:
                value_type = ARRAY_TYPES.ComplexReal64
                data = python_expr
            else:
                raise NotImplementedError(
                    'NumPy serialization not implemented for %s' % repr(
                        python_expr.dtype))

            if hasattr(data, 'tobytes'):
                #Numpy 1.9+ support array.tobytes, but previous versions don't and use tostring instead.
                yield array_class(python_expr.shape, value_type,
                                  data.tobytes())
            else:
                yield array_class(python_expr.shape, value_type,
                                  data.tostring())