| # -*- coding: utf-8 -*- |
| # Copyright 2020 The Chromium OS Authors. All rights reserved. |
| # Use of this source code is governed by a BSD-style license that can be |
| # found in the LICENSE file. |
| |
| """Protobuf message utilities. |
| |
| The Serializer classes are adapters to standardize the reading and writing of |
| different protobuf message serialization formats to and from a message. |
| |
| The base MessageHandler class encapsulates the functionality of reading |
| a file containing serialized data into a protobuf message instance, and |
| writing serialized data from a message instance out to a file. |
| """ |
| |
| from __future__ import print_function |
| |
| import os |
| import sys |
| |
| from google.protobuf import json_format |
| |
| from chromite.lib import cros_logging as logging |
| from chromite.lib import osutils |
| |
| |
| assert sys.version_info >= (3, 6), 'This module requires Python 3.6+' |
| |
| |
| FORMAT_BINARY = 1 |
| FORMAT_JSON = 2 |
| VALID_FORMATS = (FORMAT_BINARY, FORMAT_JSON) |
| |
| |
| class Error(Exception): |
| """Base error class for the module.""" |
| |
| |
| class InvalidHandlerError(Error): |
| """Raised when a message handler has no input/output argument when needed.""" |
| |
| |
| class InvalidInputFileError(Error): |
| """Raised when the input file cannot be read.""" |
| |
| |
| class InvalidInputFormatError(Error): |
| """Raised when the passed input protobuf can't be parsed.""" |
| |
| |
| class InvalidOutputFileError(Error): |
| """Raised when the output file cannot be written.""" |
| |
| |
| class UnknownHandlerError(Error): |
| """Raised when a valid type has not been implemented yet. |
| |
| This should only ever be raised when under active development. |
| See: get_message_handler. |
| """ |
| |
| |
| def get_message_handler(path, msg_format): |
| """Get a message handler to handle the given message format.""" |
| assert msg_format in VALID_FORMATS |
| |
| if msg_format == FORMAT_BINARY: |
| return MessageHandler( |
| path=path, |
| serializer=BinarySerializer(), |
| binary=True, |
| input_arg='--input-binary', |
| output_arg='--output-binary', |
| config_arg='--config-binary') |
| elif msg_format == FORMAT_JSON: |
| return MessageHandler( |
| path=path, |
| serializer=JsonSerializer(), |
| binary=False, |
| input_arg='--input-json', |
| output_arg='--output-json', |
| config_arg='--config-json') |
| else: |
| # Unexpected. Your new format type needs a case in this function if |
| # you got this error. |
| raise UnknownHandlerError('Unknown format type.') |
| |
| |
| class Serializer(object): |
| """Base (and null) serializer class.""" |
| |
| def deserialize(self, data, message): |
| """Deserialize the data into the given message. |
| |
| Args: |
| data (str): The message data to deserialize. |
| message (google.protobuf.Message): The message to load the data into. |
| """ |
| |
| # pylint: disable=unused-argument |
| def serialize(self, message): |
| """Serialize the message data. |
| |
| Args: |
| message (google.protobuf.Message): The message to be serialized. |
| |
| Returns: |
| str: The message's serialized data. |
| """ |
| return '' |
| |
| |
| class BinarySerializer(Serializer): |
| """Protobuf binary serializer class.""" |
| |
| def deserialize(self, data, message): |
| """Deserialize the data into the given message. |
| |
| See: Serializer.deserialize |
| """ |
| message.ParseFromString(data) |
| |
| def serialize(self, message): |
| """Serialize the message data. |
| |
| See: Serializer.serialize |
| """ |
| return message.SerializeToString() |
| |
| |
| class JsonSerializer(Serializer): |
| """Protobuf json serializer class.""" |
| |
| def deserialize(self, data, message): |
| """Deserialize the data into the given message. |
| |
| See: Serializer.deserialize |
| """ |
| try: |
| json_format.Parse(data, message, ignore_unknown_fields=True) |
| except json_format.ParseError as e: |
| raise InvalidInputFormatError('Unable to parse the input json: %s' % e) |
| |
| def serialize(self, message): |
| """Serialize the message data. |
| |
| See: Serializer.serialize |
| """ |
| return json_format.MessageToJson( |
| message, sort_keys=True, use_integers_for_enums=True) or '{}' |
| |
| |
| class MessageHandler(object): |
| """Class to handle message (de)serialization to and from files. |
| |
| The class is fairly tightly coupled to the build api, but we currently have |
| no other projected use cases for this, so it's handy. In particular, if we |
| scrap the "maintain the same input/output/config serialization when reexecing |
| inside the chroot" convention, this implementation is much less useful and |
| can be fairly trivially generalized. |
| |
| The instance's path is the primary path the message handler was built for. |
| For the Build API, this means one of the input/output/config arguments. In |
| practice, it's largely a convenience/shortcut so we don't have to either |
| track which input files are what types (which we know from the argument used |
| to pass them in), or create another containing data class for the |
| functionality provided by the handler and serializer classes and the build |
| api data. |
| |
| Examples: |
| message_handler = MessageHandler(path, ...) |
| message = ... |
| # Parse path into message. |
| message_handler.read_into(message) |
| # Write message to a different file. |
| message_handler.write_into(message, path=other_path) |
| """ |
| |
| def __init__(self, path, serializer, binary, input_arg, output_arg, |
| config_arg): |
| """MessageHandler init. |
| |
| Args: |
| path (str): The path to the main file associated with this handler. |
| serializer (Serializer): The serializer to be used for the messages. |
| binary (bool): Whether the serialized content is binary. |
| input_arg (str): The --input-x argument used for this type. Used for |
| reexecution inside the chroot. |
| output_arg (str): The --output-x argument used for this type. Used for |
| reexecution inside the chroot. |
| config_arg (str): The --config-x argument used for this type. Used for |
| reexecution inside the chroot. |
| """ |
| self.path = path |
| self.serializer = serializer |
| self.read_mode = 'rb' if binary else 'r' |
| self.write_mode = 'wb' if binary else 'w' |
| self.input_arg = input_arg |
| self.output_arg = output_arg |
| self.config_arg = config_arg |
| |
| def read_into(self, message, path=None): |
| """Read a file containing serialized data into a message. |
| |
| Args: |
| message (google.protobuf.Message): The message to populate. |
| path (str|None): A path to read. Uses the instance's path when not given. |
| |
| Raises: |
| InvalidInputFileError: When a path has not been given, does not exist, |
| or cannot be read. |
| """ |
| target_path = path or self.path |
| if not target_path: |
| raise InvalidInputFileError('No input file has been specified.') |
| if not os.path.exists(target_path): |
| raise InvalidInputFileError('The input file does not exist.') |
| |
| try: |
| content = osutils.ReadFile(target_path, mode=self.read_mode) |
| except IOError as e: |
| raise InvalidInputFileError('Unable to read input file: %s' % e) |
| |
| if content: |
| self.serializer.deserialize(content, message) |
| else: |
| logging.warning('No content found in %s to deserialize.', target_path) |
| |
| def write_from(self, message, path=None): |
| """Write serialized data from the message to a file. |
| |
| Args: |
| message (google.protobuf.Message): The message to serialize and persist. |
| path (str|None): An optional override of the instance's path. |
| |
| Raises: |
| InvalidOutputFileError: When no path given, or it cannot be written to. |
| """ |
| if not path and not self.path: |
| raise InvalidOutputFileError('No output file has been specified.') |
| |
| try: |
| osutils.WriteFile( |
| path or self.path, |
| self.serializer.serialize(message), |
| mode=self.write_mode) |
| except IOError as e: |
| raise InvalidOutputFileError('Cannot write output file: %s' % e) |