lib/luci/prpc/encoding.py

# Copyright 2018 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""This file is heavily based off of LUCI encoding.py."""

from chromite.third_party.google.protobuf import json_format
from chromite.third_party.google.protobuf import text_format


class Encoding:
    """Encoding enum-ish class."""

    BINARY = (0, "application/prpc; encoding=binary")
    JSON = (1, "application/json")
    TEXT = (2, "application/prpc; encoding=text")

    @staticmethod
    def media_type(encoding):
        return encoding[1]


def get_decoder(encoding):
    """Returns the appropriate decoder for content type.

    Args:
        encoding: A value from the Encoding enum.

    Returns:
        A callable which takes an encoded string and an empty protobuf message,
        and populates the given protobuf with data from the string. Each decoder
        may raise exceptions of its own based on incorrectly formatted data.
    """
    if encoding == Encoding.BINARY:
        return lambda string, proto: proto.ParseFromString(string)
    elif encoding == Encoding.JSON:
        return json_format.Parse
    elif encoding == Encoding.TEXT:
        return text_format.Merge
    else:
        assert (
            False
        ), "Argument |encoding| was not a value of the Encoding enum."


def get_encoder(encoding):
    """Returns the appropriate encoder for the Accept content type.

    Args:
        encoding: A value from the Encoding enum.

    Returns:
        A callable which takes an initialized protobuf message, and returns a
        string representing its data. Each encoder may raise exceptions of its
        own.
    """
    if encoding == Encoding.BINARY:
        return lambda proto: proto.SerializeToString()
    elif encoding == Encoding.JSON:
        return lambda proto: ")]}'\n" + json_format.MessageToJson(proto)
    elif encoding == Encoding.TEXT:
        return lambda proto: text_format.MessageToString(proto, as_utf8=True)
    else:
        assert (
            False
        ), "Argument |encoding| was not a value of the Encoding enum."