sila_interoperability.communication-tests.add-features.py Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of repository Show documentation
The newest version!
#!/usr/bin/env python
"""Tool to manage Python/proto code generation"""
import subprocess
import sys
from argparse import ArgumentParser
from os.path import dirname, join

from lxml import etree

MODULE_DIR = join(dirname(__file__), "sila2_interop_communication_tester")
RESOURCE_DIR = join(MODULE_DIR, "resources")


def add_feature(fdl_path):
    # fdl validation, fdl -> proto
    feature_xml_schema = etree.XMLSchema(etree.parse(join(RESOURCE_DIR, "xsd", "FeatureDefinition.xsd")))

    # strip BOM
    with open(fdl_path, "rb") as fdl_file:
        fdl_file_content = fdl_file.read()
    if b"<" not in fdl_file_content[:10]:
        raise RuntimeError("No opening tag found in XML file")
    fdl_file_content = fdl_file_content[fdl_file_content.find(b"<") :]
    with open(fdl_path, "wb") as fdl_file:
        fdl_file.write(fdl_file_content)

    feature_xml_parser = etree.XMLParser(schema=feature_xml_schema, encoding="utf-8")
    feature_xml = etree.parse(fdl_path, parser=feature_xml_parser).getroot()
    xsl_transform = etree.XSLT(etree.parse(join(RESOURCE_DIR, "xslt", "fdl2proto.xsl")))
    proto_content = xsl_transform(feature_xml)
    feature_id = feature_xml.xpath(
        "/sila:Feature/sila:Identifier/text()", namespaces=dict(sila="http://www.sila-standard.org")
    )[0]
    proto_dir = join(RESOURCE_DIR, "proto")
    proto_file = join(proto_dir, f"{feature_id}.proto")
    with open(proto_file, "w", encoding="utf-8") as proto_fp:
        proto_fp.write(str(proto_content))
    with open(join(RESOURCE_DIR, "fdl", f"{feature_id}.sila.xml"), "wb") as fdl_out_fp:
        fdl_out_fp.write(fdl_file_content)

    # proto -> py, pyi
    stub_dir = join(MODULE_DIR, "grpc_stubs")
    process = subprocess.run(
        [
            f"{sys.executable}",
            "-m",
            "grpc_tools.protoc",
            f"--proto_path={proto_dir}",
            f"--python_out={stub_dir}",
            f"--grpc_python_out={stub_dir}",
            f"--mypy_out={stub_dir}",
            f"--mypy_grpc_out={stub_dir}",
            f"{proto_file}",
        ],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
    )
    if process.returncode != 0:
        raise RuntimeError(f"protoc failed: {process.stderr}")

    convert_to_relative_import(join(stub_dir, f"{feature_id}_pb2.py"), ["SiLAFramework"])
    convert_to_relative_import(join(stub_dir, f"{feature_id}_pb2.pyi"), ["SiLAFramework"])
    convert_to_relative_import(join(stub_dir, f"{feature_id}_pb2_grpc.py"), ["SiLAFramework", feature_id])
    convert_to_relative_import(join(stub_dir, f"{feature_id}_pb2_grpc.pyi"), ["SiLAFramework", feature_id])


def convert_to_relative_import(filepath, proto_names):
    with open(filepath, encoding="utf-8") as fp:
        content = fp.read()
    for proto_name in proto_names:
        content = content.replace(
            f"import {proto_name}_pb2",
            f"from . import {proto_name}_pb2",
        )
    with open(filepath, "w", encoding="utf-8") as fp:
        fp.write(content)


def main():
    parser = ArgumentParser("add-features", description="Add features to this tool (.sila.xml, .proto, .py, .pyi)")
    parser.add_argument("FDLs", nargs="+", help="Feature definition files (.sila.xml)")
    args = parser.parse_args()

    for fdl_path in args.FDLs:
        add_feature(fdl_path)


if __name__ == "__main__":
    main()