libhwsec-foundation/utility/proto_print.py

#!/usr/bin/env python3
# Copyright 2015 The ChromiumOS Authors
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""A C++ code generator for printing protobufs which use the LITE_RUNTIME.

Normally printing a protobuf would be done with Message::DebugString(). However,
this is not available when using only MessageLite. This script generates code to
emulate Message::DebugString() without using reflection. The input must be a
valid .proto file.

Usage: proto_print.py [--package-dir=pac] [--subdir=foo] <bar.proto>

Files named print_bar_proto.h and print_bar_proto.cc will be created in the
current working directory.
"""

from __future__ import print_function

import argparse
import collections
import datetime
import os
import re
import subprocess
import sys


# Holds information about a protobuf message field.
#
# Attributes:
#   repeated: Whether the field is a repeated field.
#   type_: The type of the field. E.g. int32.
#   name: The name of the field.
Field = collections.namedtuple("Field", "repeated type_ name")


class Message:
    """Holds information about a protobuf message.

    Attributes:
        name: The name of the message.
        fields: A list of Field tuples.
    """

    def __init__(self, name):
        """Initializes a Message instance.

        Args:
            name: The protobuf message name.
        """
        self.name = name
        self.fields = []

    def AddField(self, attribute, field_type, field_name):
        """Adds a new field to the message.

        Args:
            attribute: This should be 'optional', 'required', or 'repeated'.
            field_type: The type of the field. E.g. int32.
            field_name: The name of the field.
        """
        self.fields.append(
            Field(
                repeated=attribute == "repeated",
                type_=field_type,
                name=field_name,
            )
        )


class Enum:
    """Holds information about a protobuf enum.

    Attributes:
        name: The name of the enum.
        values: A list of enum value names.
    """

    def __init__(self, name):
        """Initializes a Enum instance.

        Args:
            name: The protobuf enum name.
        """
        self.name = name
        self.values = []

    def AddValue(self, value_name):
        """Adds a new value to the enum.

        Args:
            value_name: The name of the value.
        """
        self.values.append(value_name)


class Oneof:
    """Holds information about a protobuf Oneof.

    Attributes:
        name: The name of the OnoOf.
    """

    def __init__(self, name):
        """Initializes a OnoOf instance.

        Args:
            name: The protobuf OnoOf name.
        """
        self.name = name


def ParseProto(input_file):
    """Parses a proto file and returns a tuple of parsed information.

    Args:
        input_file: The proto file to parse.

    Returns:
        A tuple in the form (package, imports, messages, enums) where
        package: A string holding the proto package.
        imports: A list of strings holding proto imports.
        messages: A list of Message objects; one for each message in the proto.
        enums: A list of Enum objects; one for each enum in the proto.
    """
    package = ""
    imports = []
    messages = []
    enums = []
    current_message_stack = []
    current_enum = None
    current_oneof = None
    package_re = re.compile(r"package\s+([.\w]+);")
    import_re = re.compile(r'import\s+"([\w]*/)*(\w+).proto";')
    message_re = re.compile(r"message\s+(\w+)\s*{")
    field_re = re.compile(r"(optional|required|repeated)\s+(\w+)\s+(\w+)\s*=")
    field_re3 = re.compile(r"(|repeated\s+)([.\w]+)\s+(\w+)\s*=")
    enum_re = re.compile(r"enum\s+(\w+)\s*{")
    enum_value_re = re.compile(r"(\w+)\s*=")
    oneof_re = re.compile(r"oneof\s+(\w+)\s*{")
    for line in input_file:
        line = line.strip()
        if not line or line.startswith("//") or line.startswith("reserved"):
            continue
        msg_match = message_re.search(line)
        enum_match = enum_re.search(line)
        oneof_match = oneof_re.search(line)
        field_match = field_re.search(line)
        field_match3 = field_re3.search(line)
        package_match = package_re.search(line)
        import_match = import_re.search(line)
        # Look for a message definition.
        if msg_match:
            prefix = ""
            if current_message_stack:
                prefix = (
                    "::".join([m.name for m in current_message_stack]) + "::"
                )
            current_message_stack.append(Message(prefix + msg_match.group(1)))
        elif current_message_stack and oneof_match:
            if current_oneof:
                raise Exception("Unsupported nested oneof")
            # TODO(b/220231404): Support the printing of oneof field.
            current_oneof = Oneof(prefix + oneof_match.group(1))
        # Look for a message field definition.
        elif current_message_stack and field_match:
            current_message_stack[-1].AddField(
                field_match.group(1),
                field_match.group(2),
                field_match.group(3),
            )
        elif (
            current_message_stack
            and field_match3
            and field_match3.group(2) != "option"
        ):
            current_message_stack[-1].AddField(
                field_match3.group(1).strip(),
                field_match3.group(2),
                field_match3.group(3),
            )
        elif enum_match:
            prefix = ""
            if current_message_stack:
                prefix = "_".join([m.name for m in current_message_stack]) + "_"
            current_enum = Enum(prefix + enum_match.group(1))
            continue
        # Look for an enum value.
        elif current_enum:
            match = enum_value_re.search(line)
            if match:
                prefix = ""
                if current_message_stack:
                    prefix = current_enum.name + "_"
                current_enum.AddValue(prefix + match.group(1))
        # Look for a package statement.
        elif package_match:
            package = package_match.group(1)
        # Look for an import statement.
        elif import_match:
            imports.append(import_match.group(2))

        # Close off the current scope. Enums first because they can't be nested.
        if line[-1] == "}":
            if current_enum:
                enums.append(current_enum)
                current_enum = None
            elif current_oneof:
                # TODO(b/220231404): Support the printing of oneof field.
                current_oneof = None
            elif current_message_stack:
                messages.append(current_message_stack.pop())
            else:
                raise Exception("Closing unknown scope")

    # Make sure we parse the proto file correctly.
    if current_enum:
        raise Exception(f'The current_enum "{current_enum.name}" is not empty')

    if current_oneof:
        raise Exception(
            f'The current_oneof "{current_oneof.name}" is not empty'
        )

    if current_message_stack:
        name_stack = [msg.name for msg in current_message_stack]
        raise Exception(f"The current_message_stack {name_stack} is not empty")

    return package, imports, messages, enums


def GenerateFileHeaders(
    proto_name,
    package,
    imports,
    subdir,
    proto_include_override,
    header_file_name,
    header_file,
    impl_file,
    package_dir,
):
    """Generates and prints file headers.

    Args:
        proto_name: The name of the proto file.
        package: The protobuf package.
        imports: A list of imported protos.
        subdir: The --subdir arg.
        proto_include_override: Include directory override for the #include
                        statement in generated code
        header_file_name: The header file name.
        header_file: The header file handle, open for writing.
        impl_file: The implementation file handle, open for writing.
        package_dir: The package directory.
    """
    if subdir:
        guard_name = "%s_%s_PRINT_%s_PROTO_H_" % (
            package_dir.upper(),
            subdir.upper(),
            proto_name.upper(),
        )
        package_with_subdir = "%s/%s" % (package_dir, subdir)
    else:
        guard_name = "%s_PRINT_%s_PROTO_H_" % (
            package_dir.upper(),
            proto_name.upper(),
        )
        package_with_subdir = package_dir
    guard_name = guard_name.replace("-", "_")
    proto_include_dir = package_with_subdir
    if proto_include_override is not None:
        proto_include_dir = proto_include_override
    namespace = package.replace(".", "::")
    includes = "\n".join(
        [
            '#include "%(package_with_subdir)s/print_%(import)s_proto.h"'
            % {
                "package_with_subdir": package_with_subdir,
                "import": current_import,
            }
            for current_import in imports
        ]
    )
    header = """\
// Copyright %(year)s The ChromiumOS Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// THIS CODE IS GENERATED.
// Generated with command:
// %(cmd)s

#ifndef %(guard_name)s
#define %(guard_name)s

#include <string>

#include <brillo/brillo_export.h>

#include "%(proto_include_dir)s/%(proto)s.pb.h"

namespace %(namespace)s {
""" % {
        "year": datetime.date.today().year,
        "guard_name": guard_name,
        "namespace": namespace,
        "proto": proto_name,
        "proto_include_dir": proto_include_dir,
        "cmd": " ".join(sys.argv),
    }
    impl = """\
// Copyright %(year)s The ChromiumOS Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// THIS CODE IS GENERATED.
// Generated with command:
// %(cmd)s

#include "%(package_with_subdir)s/%(header_file_name)s"

#include <inttypes.h>

#include <string>

#include <base/strings/string_number_conversions.h>
#include <base/strings/stringprintf.h>

%(includes)s

namespace %(namespace)s {
""" % {
        "year": datetime.date.today().year,
        "namespace": namespace,
        "package_with_subdir": package_with_subdir,
        "header_file_name": header_file_name,
        "includes": includes,
        "cmd": " ".join(sys.argv),
    }

    header_file.write(header)
    impl_file.write(impl)


def GenerateFileFooters(
    proto_name, package, subdir, header_file, impl_file, package_dir
):
    """Generates and prints file footers.

    Args:
        proto_name: The name of the proto file.
        package: The protobuf package.
        subdir: The --subdir arg.
        header_file: The header file handle, open for writing.
        impl_file: The implementation file handle, open for writing.
        package_dir: The package directory.
    """
    namespace = package.replace(".", "::")
    if subdir:
        guard_name = "%s_%s_PRINT_%s_PROTO_H_" % (
            package_dir.upper(),
            subdir.upper(),
            proto_name.upper(),
        )
    else:
        guard_name = "%s_PRINT_%s_PROTO_H_" % (
            package_dir.upper(),
            proto_name.upper(),
        )
    header = """

}  // namespace %(namespace)s

#endif  // %(guard_name)s
""" % {
        "guard_name": guard_name,
        "namespace": namespace,
    }
    impl = """
}  // namespace %(namespace)s
""" % {
        "namespace": namespace
    }

    header_file.write(header)
    impl_file.write(impl)


def GenerateEnumPrinter(enum, header_file, impl_file):
    """Generates and prints a function to print an enum value.

    Args:
        enum: An Enum instance.
        header_file: The header file handle, open for writing.
        impl_file: The implementation file handle, open for writing.
    """
    declare = """
std::string GetProtoDebugStringWithIndent(%(name)s value, int indent_size);
BRILLO_EXPORT std::string GetProtoDebugString(%(name)s value);""" % {
        "name": enum.name
    }
    define_begin = """
std::string GetProtoDebugString(%(name)s value) {
  return GetProtoDebugStringWithIndent(value, 0);
}

std::string GetProtoDebugStringWithIndent(%(name)s value, int indent_size) {
""" % {
        "name": enum.name
    }
    define_end = """
  return "<unknown>";
}
"""
    condition = """
  if (value == %(value_name)s) {
    return "%(value_name)s";
  }"""

    header_file.write(declare)
    impl_file.write(define_begin)
    for value_name in enum.values:
        impl_file.write(condition % {"value_name": value_name})
    impl_file.write(define_end)


def GenerateMessagePrinter(message, header_file, impl_file):
    """Generates and prints a function to print a message.

    Args:
        message: A Message instance.
        header_file: The header file handle, open for writing.
        impl_file: The implementation file handle, open for writing.
    """
    declare = """
std::string GetProtoDebugStringWithIndent(const %(name)s& value,
                                          int indent_size);
BRILLO_EXPORT std::string GetProtoDebugString(const %(name)s& value);""" % {
        "name": message.name
    }
    define_begin = """
std::string GetProtoDebugString(const %(name)s& value) {
  return GetProtoDebugStringWithIndent(value, 0);
}

std::string GetProtoDebugStringWithIndent(const %(name)s& value,
                                          int indent_size) {
  std::string indent(indent_size, ' ');
  std::string output = base::StringPrintf("[%%s] {\\n",
                                          value.GetTypeName().c_str());
""" % {
        "name": message.name
    }
    define_end = """
  output += indent + "}";
  return output;
}
"""
    # As long as the field has a has_name method, call it to check whether the
    # field exists. We don't need to print it if it doesn't.
    singular_field = """
  []<typename T>(const T& value, int indent_size,
                 const std::string& indent, std::string& output) {
    if constexpr (requires(T t) { t.has_%(name)s(); }) {
      if (!value.has_%(name)s()) {
        return;
      }
    }
    output += indent + "  %(name)s: ";
    base::StringAppendF(&output, %(format)s);
    output += "\\n";
  }(value, indent_size, indent, output);"""
    repeated_field = """
  output += indent + "  %(name)s: {";
  for (int i = 0; i < value.%(name)s_size(); ++i) {
    if (i > 0) {
      output += ",";
    }
    output += "\\n    " + indent;
    base::StringAppendF(&output, %(format)s);
    if (i == value.%(name)s_size() - 1) {
      output += "\\n  " + indent;
    }
  }
  output += "}\\n";"""
    singular_field_get = "value.%(name)s()"
    repeated_field_get = "value.%(name)s(i)"
    formats = {
        "bool": '"%%s", %(value)s ? "true" : "false"',
        "int32": '"%%" PRId32, %(value)s',
        "int64": '"%%" PRId64, %(value)s',
        "uint32": '"%%" PRIu32 " (0x%%08" PRIX32 ")", %(value)s, %(value)s',
        "uint64": '"%%" PRIu64 " (0x%%016" PRIX64 ")", %(value)s, %(value)s',
        "string": '"%%s", %(value)s.c_str()',
        "bytes": """"%%s", base::HexEncode(%(value)s.data(),
                                         %(value)s.size()).c_str()""",
    }
    subtype_format = (
        '"%%s", GetProtoDebugStringWithIndent(%(value)s, '
        "indent_size + %(indent_incr)s).c_str()"
    )

    header_file.write(declare)
    impl_file.write(define_begin)
    for field in message.fields:
        if field.repeated:
            value_get = repeated_field_get % {"name": field.name}
            field_code = repeated_field
        else:
            value_get = singular_field_get % {"name": field.name}
            field_code = singular_field
        if field.type_ in formats:
            value_format = formats[field.type_] % {"value": value_get}
        else:
            value_format = subtype_format % {
                "value": value_get,
                "indent_incr": 4 if field.repeated else 2,
            }
        impl_file.write(
            field_code % {"name": field.name, "format": value_format}
        )
    impl_file.write(define_end)


def FormatFile(filename):
    subprocess.call(["clang-format", "-i", "-style=Chromium", filename])


def main():
    parser = argparse.ArgumentParser(description="print proto code generator")
    parser.add_argument(
        "--package-dir",
        default="",
        help=("Override for the package directory name"),
    )
    parser.add_argument(
        "--subdir",
        default="",
        help=(
            "The subdirectory under which the generated "
            + "file will reside in"
        ),
    )
    parser.add_argument(
        "--proto-include-override",
        default=None,
        help=(
            "Override for the include path of *.pb.h in" + "generated header"
        ),
    )
    parser.add_argument(
        "--output-dir", default=".", help=("The output directory")
    )
    parser.add_argument("input_files", nargs="*")
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)

    for input_path in args.input_files:
        with open(input_path, encoding="utf-8") as input_file:
            package, imports, messages, enums = ParseProto(input_file)
        package_dir = package
        if args.package_dir != "":
            package_dir = args.package_dir
        proto_name = os.path.basename(input_path).rsplit(".", 1)[0]
        header_file_name = "print_%s_proto.h" % proto_name
        impl_file_name = "print_%s_proto.cc" % proto_name
        header_file_path = os.path.join(args.output_dir, header_file_name)
        impl_file_path = os.path.join(args.output_dir, impl_file_name)
        with open(header_file_path, "w", encoding="utf-8") as header_file:
            with open(impl_file_path, "w", encoding="utf-8") as impl_file:
                GenerateFileHeaders(
                    proto_name,
                    package,
                    imports,
                    args.subdir,
                    args.proto_include_override,
                    header_file_name,
                    header_file,
                    impl_file,
                    package_dir,
                )
                for enum in enums:
                    GenerateEnumPrinter(enum, header_file, impl_file)
                for message in messages:
                    GenerateMessagePrinter(message, header_file, impl_file)
                GenerateFileFooters(
                    proto_name,
                    package,
                    args.subdir,
                    header_file,
                    impl_file,
                    package_dir,
                )
        FormatFile(header_file_path)
        FormatFile(impl_file_path)


if __name__ == "__main__":
    main()