Copyright (C) 2020-2023 HERE Global B.V. and its affiliate(s).

All rights reserved.

This software and other materials contain proprietary information

controlled by HERE and are protected by applicable copyright legislation.

Any use and utilization of this software and other materials and

disclosure to any third parties is conditional upon having a separate

agreement with HERE for the access, use, utilization or disclosure of this

software. In the absence of such agreement, the use of the software is not

allowed.

"""
Protobuf Parser module
"""

import io
import os
import zipfile
from pathlib import Path
from typing import Any, Dict, List, Optional, Union

from google.protobuf import descriptor_pool
from google.protobuf.descriptor import FieldDescriptor
from google.protobuf.descriptor_pb2 import FileDescriptorSet
from google.protobuf.message import Message
from google.protobuf.message_factory import GetMessageClassesForFiles
from here.platform.exceptions import SchemaException
from here.platform.schema.parser import Parser

_PROTOBUF_EXTENSION = ".proto"

[docs]
class ProtobufParser(Parser):
"""A class for parsing Protobuf data of a catalog layer partition."""

def init(self, schema_pkg_file: Union[str, Path, io.BytesIO]):
"""
Initialize a ProtobufParser instance.

:param schema_pkg_file: the schema package zip-file as a str, Path or file-like object
:raises SchemaException: if the required meta information is missing in schema package
"""
layer_manifest, layer_descriptor_data = Parser.extract_layer_info(schema_pkg_file)
if not layer_manifest:
raise SchemaException(
f"Invalid schema. Layer manifest file {Parser.LAYER_MANIFEST} "
"not found in schema package file" # noqa: E713
)
if not layer_descriptor_data:
raise SchemaException(
f"Invalid schema. Layer descriptor file {Parser.LAYER_DESCRIPTORS} "
"not found in schema package file" # noqa: E713
)

main_partition_class_name = layer_manifest["main"]["message"]

descriptor_set: FileDescriptorSet = FileDescriptorSet.FromString(layer_descriptor_data)
self.message_classes = ProtobufParser._create_message_classes(descriptor_set.file)

if main_partition_class_name not in self.message_classes:
raise SchemaException(
f"Unable to find main class {main_partition_class_name} present under"
f"{Parser.LAYER_MANIFEST} in {Parser.LAYER_DESCRIPTORS} for the schema artifact"
)

self.partition_class = self.message_classes[main_partition_class_name]

@staticmethod
def _extract_protobuf_sources(input_zip: Union[str, Path, io.BytesIO]) -> Dict[str, str]:
"""
Extract Protobuf from an input zipfile.

:param input_zip: input Protobuf schema zipfile
:returns: dictionary of schema file names and content
"""
input_zipfile = zipfile.ZipFile(input_zip)
protobuf_sources =
for name in input_zipfile.namelist():
if name.endswith(_PROTOBUF_EXTENSION):
protobuf_sources[os.path.basename(name)] = input_zipfile.read(name).decode("utf-8")

return protobuf_sources

@staticmethod
def _create_message_classes(file_protos) -> dict:
"""
Create message classes.
This code is taken from google.protobuf.message_factory.GetMessages
and modified to use a local message factory instance.

:param file_protos: A list of protobuf files.
:return: a dictionary mapping proto names to the message classes.
"""
file_by_name = {file_proto.name: file_proto for file_proto in file_protos}
pool = descriptor_pool.DescriptorPool()

def _add_file(file_proto):
for dependency in file_proto.dependency:
if dependency in file_by_name:

Remove from elements to be visited, in order to cut cycles.

_add_file(file_by_name.pop(dependency))
pool.Add(file_proto)

while fileby_name: add_file(file_by_name.popitem()[1])

return GetMessageClassesForFiles([file_proto.name for file_proto in file_protos], pool)

[docs]
def get_partition_class(self):
"""
Returns partition class.

:return: Partition class
"""
return self.partition_class

[docs]
def parse_message(self, message: bytes) -> Message:
"""
Parse a Protobuf message and returns the corresponding partition class instance.

:param message: message received as blob
:return: parsed Protobuf message
"""
msg = self.partition_class.FromString(message)
assert isinstance(msg, Message)
return msg

[docs]
def parse_message_by_class(
self, message: bytes, fully_qualified_message_name: str
) -> Union[Message, Any]:
"""
Parse a Protobuf message by proto class and returns the corresponding
partition class instance.

:param message: message received as blob
:param fully_qualified_message_name: fully qualified message name of the proto class
:return: parsed Protobuf message
"""
fully_qualified_message_name = fully_qualified_message_name.strip()
if fully_qualified_message_name in self.message_classes:
msg = self.message_classes[fully_qualified_message_name].FromString(message)
assert isinstance(msg, Message)
return msg
return None

[docs]
def get_message_class(self, fully_qualified_message_name: str) -> Any:
"""
Get Message class by fully qualified message name.

:param fully_qualified_message_name: fully qualified message name of the proto class
:return: message class
"""
fully_qualified_message_name = fully_qualified_message_name.strip()
if fully_qualified_message_name in self.message_classes:
return self.message_classes[fully_qualified_message_name]
return None

[docs]
def serialize_bytes(self, data: Union[Message, dict]) -> bytes:
"""
Parse a Protobuf message and returns the bytes.

:param data: Protobuf message
:return: bytes
"""
assert isinstance(data, Message)
message_bytes = data.SerializeToString()
assert isinstance(message_bytes, bytes)
return message_bytes

[docs]
def primary_fields(self, all_fields: bool = False) -> List[FieldDescriptor]:
"""
Return the field descriptors of the root partition class.

:param all_fields: Return all field descriptors of the root partition
class that are labeled 'repeated' if all is set to False.
Return all the field descriptors otherwise.
:return: List of field descriptors
"""
return [
f
for f in self.partition_class.DESCRIPTOR.fields
if all_fields or f.label == FieldDescriptor.LABEL_REPEATED
]

[docs]
@staticmethod
def access_field(message: Message, path: Optional[str]):
"""
Return the data at the given path.

:param message: a decoded Protobuf message
:param path: path to the field, concatenate field names with '.' to select nested fields
:return: the field referenced by the path parameter. This is either
a Message in case the field contains a nested Protobuf structure,
a simple value for integers, strings and other native types,
or a container that can be iterated for a repeated field in Protobuf.
"""
if not path:
return message
else:
current_field = message
for field in path.split("."):
current_field = getattr(current_field, field)
return current_field

def _access_message_class(self, class_name: str) -> Any:
if class_name in self.message_classes:
return self.message_classes[class_name]
else:
paths = class_name.split(".")
for i in range(len(paths), 0, -1):
current_parent_class = ".".join(paths[:i])
if current_parent_class in self.message_classes:
result_class = self.message_classes[current_parent_class]
for path in paths[i:]:
result_class = result_class.DESCRIPTOR.nested_types_by_name[
path
]._concrete_class
return result_class