writer_binary_raw.py

# Copyright 2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# A copy of the License is located at:
#
#    http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS
# OF ANY KIND, either express or implied. See the License for the
# specific language governing permissions and limitations under the
# License.

"""Writer for raw binary Ion values, without symbol table management."""

from datetime import datetime
from decimal import Decimal, localcontext
from enum import IntEnum
from functools import partial

import struct

from amazon.ion.equivalence import _is_float_negative_zero
from amazon.ion.symbols import SymbolToken
from .core import IonEventType, IonType, DataEvent, Transition, TimestampPrecision, TIMESTAMP_FRACTION_PRECISION_FIELD, \
    MICROSECOND_PRECISION, TIMESTAMP_PRECISION_FIELD, TIMESTAMP_FRACTIONAL_SECONDS_FIELD, Timestamp
from .util import coroutine, total_seconds
from .writer import NOOP_WRITER_EVENT, WriteEventType, \
                    writer_trampoline, partial_transition, serialize_scalar, \
                    validate_scalar_value, illegal_state_null
from .writer_binary_raw_fields import _write_varuint, _write_uint, _write_varint, _write_int


class _TypeIds(IntEnum):
    NULL = 0x00
    BOOL_FALSE = 0x10
    BOOL_TRUE = 0x11
    POS_INT = 0x20
    NEG_INT = 0x30
    FLOAT = 0x40
    DECIMAL = 0x50
    TIMESTAMP = 0x60
    SYMBOL = 0x70
    STRING = 0x80
    CLOB = 0x90
    BLOB = 0xA0
    LIST = 0xB0
    SEXP = 0xC0
    STRUCT = 0xD0
    ANNOTATION_WRAPPER = 0xE0


class _Zeros(IntEnum):
    """Single-octet encodings, represented by the type ID and a length nibble of zero.

    Notes:
        Blob, clob, list, and sexp also have single-octet encodings, but the current
        implementation handles these implicitly.
    """
    INT = _TypeIds.POS_INT  # Int zero is always encoded using the positive type ID.
    FLOAT = _TypeIds.FLOAT  # Represents zero.
    DECIMAL = _TypeIds.DECIMAL  # Represents 0d0.
    STRING = _TypeIds.STRING  # Represents the zero-length (empty) string.
    SYMBOL = _TypeIds.SYMBOL  # Represents symbol zero.
    STRUCT = _TypeIds.STRUCT  # Represents a struct with zero fields.

_VARINT_NEG_ZERO = 0xC0  # This refers to the variable-length signed integer subfield.
_INT_NEG_ZERO = 0x80  # This refers to the fixed-length signed integer subfield.

_LENGTH_FLOAT_64 = 0x08
_LENGTH_FIELD_THRESHOLD = 14
_LENGTH_FIELD_INDICATOR = 0x0E
_NULL_INDICATOR = 0x0F


def _null(tid):
    return bytearray([tid | _NULL_INDICATOR])


_NULLS = [
    _null(_TypeIds.NULL),
    _null(_TypeIds.BOOL_FALSE),
    _null(_TypeIds.POS_INT),
    _null(_TypeIds.FLOAT),
    _null(_TypeIds.DECIMAL),
    _null(_TypeIds.TIMESTAMP),
    _null(_TypeIds.SYMBOL),
    _null(_TypeIds.STRING),
    _null(_TypeIds.CLOB),
    _null(_TypeIds.BLOB),
    _null(_TypeIds.LIST),
    _null(_TypeIds.SEXP),
    _null(_TypeIds.STRUCT)
]

_BOOL_TRUE = bytearray([_TypeIds.BOOL_TRUE])
_BOOL_FALSE = bytearray([_TypeIds.BOOL_FALSE])


def _serialize_bool(ion_event):
    if ion_event.value:
        return _BOOL_TRUE
    else:
        return _BOOL_FALSE


def _write_length(buf, length, tid):
    if length < _LENGTH_FIELD_THRESHOLD:
        buf.append(tid | length)
    else:
        buf.append(tid | _LENGTH_FIELD_INDICATOR)
        _write_varuint(buf, length)


def _write_int_value(buf, tid, value):
    value_buf = bytearray()
    length = _write_uint(value_buf, value)
    _write_length(buf, length, tid)
    buf.extend(value_buf)


def _serialize_int(ion_event):
    buf = bytearray()
    value = ion_event.value
    validate_scalar_value(value, int)
    if value == 0:
        buf.append(_Zeros.INT)
    else:
        if value < 0:
            value = -value
            tid = _TypeIds.NEG_INT
        else:
            tid = _TypeIds.POS_INT
        _write_int_value(buf, tid, value)
    return buf


def _serialize_float(ion_event):
    buf = bytearray()
    float_value = ion_event.value
    validate_scalar_value(float_value, float)
    # TODO Assess whether abbreviated encoding of zero is beneficial; it's allowed by spec.
    if float_value.is_integer() and float_value == 0.0 and not _is_float_negative_zero(float_value):
        buf.append(_Zeros.FLOAT)
    else:
        # TODO Add an option for 32-bit representation (length=4) per the spec.
        buf.append(_TypeIds.FLOAT | _LENGTH_FLOAT_64)
        encoded = struct.pack('>d', float_value)
        buf.extend(encoded)
    return buf


def _write_decimal_value(buf, exponent, coefficient, sign=0):
    length = _write_varint(buf, exponent)
    if coefficient:
        # The coefficient is non-zero, so the coefficient field is required.
        length += _write_int(buf, coefficient)
    elif sign:
        # The coefficient is negative zero.
        buf.append(_INT_NEG_ZERO)
        length += 1
    # Else the coefficient is positive zero and the field is omitted.
    return length


def _write_timestamp_fractional_seconds(buf, value):
    sign, digits, exponent = value.as_tuple()
    coefficient = int(value.scaleb(-exponent).to_integral_value())
    if coefficient == 0 and exponent >= 0:
        length = 0
    else:
        length = _write_decimal_value(buf, exponent, coefficient, sign)
    return length


def _serialize_decimal(ion_event):
    buf = bytearray()
    value = ion_event.value
    validate_scalar_value(value, Decimal)
    sign, digits, exponent = value.as_tuple()
    with localcontext() as context:
        # Adjusting precision for taking into account arbitrarily large/small
        # numbers
        context.prec = len(digits)
        coefficient = int(value.scaleb(-exponent).to_integral_value())
    if not sign and not exponent and not coefficient:
        # The value is 0d0; other forms of zero will fall through.
        buf.append(_Zeros.DECIMAL)
    else:
        value_buf = bytearray()
        length = _write_decimal_value(value_buf, exponent, coefficient, sign)
        _write_length(buf, length, _TypeIds.DECIMAL)
        buf.extend(value_buf)
    return buf


def _serialize_string(ion_event):
    buf = bytearray()
    value = ion_event.value
    validate_scalar_value(value, str)
    if not value:
        buf.append(_Zeros.STRING)
    else:
        value_buf = value.encode('utf-8')
        _write_length(buf, len(value_buf), _TypeIds.STRING)
        buf.extend(value_buf)
    return buf


def _serialize_symbol(ion_event):
    buf = bytearray()
    token = ion_event.value
    validate_scalar_value(token, SymbolToken)
    sid = token.sid
    if sid == 0:
        buf.append(_Zeros.SYMBOL)
    else:
        _write_int_value(buf, _TypeIds.SYMBOL, sid)
    return buf


def _serialize_lob_value(event, tid):
    buf = bytearray()
    value = event.value
    _write_length(buf, len(value), tid)
    buf.extend(value)
    return buf


_serialize_blob = partial(_serialize_lob_value, tid=_TypeIds.BLOB)
_serialize_clob = partial(_serialize_lob_value, tid=_TypeIds.CLOB)


def _serialize_timestamp(ion_event):
    buf = bytearray()
    dt = ion_event.value
    precision = getattr(dt, TIMESTAMP_PRECISION_FIELD, TimestampPrecision.SECOND)
    if precision is None:  # TODO should this defaulting be pushed into Timestamp itself?
        precision = TimestampPrecision.SECOND
    validate_scalar_value(dt, datetime)
    value_buf = bytearray()
    if dt.tzinfo is None:
        value_buf.append(_VARINT_NEG_ZERO)  # This signifies an unknown local offset.
        length = 1
    else:
        # Normalize to UTC and write the offset field.
        offset = dt.utcoffset()
        dt -= offset
        length = _write_varint(value_buf, int(total_seconds(offset) // 60))
    length += _write_varuint(value_buf, dt.year)
    if precision.includes_month:
        length += _write_varuint(value_buf, dt.month)
    if precision.includes_day:
        length += _write_varuint(value_buf, dt.day)
    if precision.includes_minute:
        length += _write_varuint(value_buf, dt.hour)
        length += _write_varuint(value_buf, dt.minute)
    if precision.includes_second:
        length += _write_varuint(value_buf, dt.second)
        if isinstance(ion_event.value, Timestamp):
            fractional_seconds = getattr(ion_event.value, TIMESTAMP_FRACTIONAL_SECONDS_FIELD, None)
            if fractional_seconds is not None:
                length += _write_timestamp_fractional_seconds(value_buf, fractional_seconds)
        else:
            # This must be a normal datetime, which always has a range-validated microsecond value.
            length += _write_decimal_value(value_buf, -MICROSECOND_PRECISION, dt.microsecond)

    _write_length(buf, length, _TypeIds.TIMESTAMP)
    buf.extend(value_buf)
    return buf


_SERIALIZE_SCALAR_JUMP_TABLE = {
    IonType.NULL: illegal_state_null,
    IonType.BOOL: _serialize_bool,
    IonType.INT: _serialize_int,
    IonType.FLOAT: _serialize_float,
    IonType.DECIMAL: _serialize_decimal,
    IonType.TIMESTAMP: _serialize_timestamp,
    IonType.SYMBOL: _serialize_symbol,
    IonType.STRING: _serialize_string,
    IonType.CLOB: _serialize_clob,
    IonType.BLOB: _serialize_blob,
}


_serialize_scalar = partial(
    serialize_scalar, jump_table=_SERIALIZE_SCALAR_JUMP_TABLE, null_table=_NULLS
)


def _serialize_annotation_wrapper(output_buf, annotations):
    value_length = output_buf.current_container_length
    annot_length_buf = bytearray()
    annot_length = 0
    for annotation in annotations:
        annot_length += _write_varuint(annot_length_buf, annotation.sid)
    header = bytearray()
    length_buf = bytearray()
    length = _write_varuint(length_buf, annot_length) + annot_length + value_length
    _write_length(header, length, _TypeIds.ANNOTATION_WRAPPER)
    header.extend(length_buf)
    header.extend(annot_length_buf)
    output_buf.end_container(header)


def _serialize_container(output_buf, ion_event):
    ion_type = ion_event.ion_type
    length = output_buf.current_container_length
    header = bytearray()
    if ion_type is IonType.STRUCT:
        if length == 0:
            header.append(_Zeros.STRUCT)
        else:
            # TODO Support sorted field name symbols, per the spec.
            header.append(_TypeIds.STRUCT | _LENGTH_FIELD_INDICATOR)
            _write_varuint(header, length)
    else:
        tid = _TypeIds.LIST
        if ion_type is IonType.SEXP:
            tid = _TypeIds.SEXP
        _write_length(header, length, tid)
    output_buf.end_container(header)


_WRITER_EVENT_NEEDS_INPUT_EMPTY = DataEvent(WriteEventType.NEEDS_INPUT, b'')


@coroutine
def _raw_writer_coroutine(writer_buffer, depth=0, container_event=None,
                          whence=None, pending_annotations=None):

    def fail():
        raise TypeError('Invalid event: %s at depth %d' % (ion_event, depth))

    write_result = None
    while True:
        ion_event, self = (yield write_result)
        delegate = self
        curr_annotations = ion_event.annotations
        writer_event = _WRITER_EVENT_NEEDS_INPUT_EMPTY
        if depth > 0 and container_event.ion_type is IonType.STRUCT \
                and ion_event.event_type.begins_value:
            # A field name symbol ID is required at this position.
            sid_buffer = bytearray()
            _write_varuint(sid_buffer, ion_event.field_name.sid)  # Write the field name's symbol ID.
            writer_buffer.add_scalar_value(sid_buffer)
        if ion_event.event_type.begins_value and curr_annotations:
            writer_buffer.start_container()
        if ion_event.event_type is IonEventType.SCALAR:
            scalar_buffer = _serialize_scalar(ion_event)
            writer_buffer.add_scalar_value(scalar_buffer)
            if curr_annotations:
                _serialize_annotation_wrapper(writer_buffer, curr_annotations)
        elif ion_event.event_type is IonEventType.STREAM_END:
            if depth != 0:
                fail()
            for partial_value in writer_buffer.drain():
                yield partial_transition(partial_value, self)
            writer_event = NOOP_WRITER_EVENT
        elif ion_event.event_type is IonEventType.CONTAINER_START:
            if not ion_event.ion_type.is_container:
                raise TypeError('Expected container type')
            writer_buffer.start_container()
            delegate = _raw_writer_coroutine(writer_buffer, depth + 1,
                                             ion_event, self, curr_annotations)
        elif ion_event.event_type is IonEventType.CONTAINER_END:
            if depth < 1:
                fail()
            _serialize_container(writer_buffer, container_event)
            if pending_annotations:
                _serialize_annotation_wrapper(writer_buffer, pending_annotations)
            pending_annotations = None
            delegate = whence
        else:
            fail()
        write_result = Transition(writer_event, delegate)


def _raw_binary_writer(writer_buffer):
    """Returns a raw binary writer co-routine.

    Yields:
        DataEvent: serialization events to write out

        Receives :class:`amazon.ion.core.IonEvent`.
    """
    return writer_trampoline(_raw_writer_coroutine(writer_buffer))