diff --git a/src/betterproto/plugin/compiler.py b/src/betterproto/plugin/compiler.py new file mode 100644 index 0000000..4fd3b8f --- /dev/null +++ b/src/betterproto/plugin/compiler.py @@ -0,0 +1,40 @@ +import os.path + +try: + # betterproto[compiler] specific dependencies + import black + import jinja2 +except ImportError as err: + missing_import = err.args[0][17:-1] + print( + "\033[31m" + f"Unable to import `{missing_import}` from betterproto plugin! " + "Please ensure that you've installed betterproto as " + '`pip install "betterproto[compiler]"` so that compiler dependencies ' + "are included." + "\033[0m" + ) + raise SystemExit(1) + +from betterproto.plugin.models import OutputTemplate + + +def outputfile_compiler(output_file: OutputTemplate) -> str: + + templates_folder = os.path.abspath( + os.path.join(os.path.dirname(__file__), "..", "templates") + ) + + env = jinja2.Environment( + trim_blocks=True, + lstrip_blocks=True, + loader=jinja2.FileSystemLoader(templates_folder), + ) + template = env.get_template("template.py.j2") + + res = black.format_str( + template.render(output_file=output_file), + mode=black.FileMode(target_versions={black.TargetVersion.PY37}), + ) + + return res diff --git a/src/betterproto/plugin/parser.py b/src/betterproto/plugin/parser.py index 33991ec..85f370f 100644 --- a/src/betterproto/plugin/parser.py +++ b/src/betterproto/plugin/parser.py @@ -1,12 +1,10 @@ import itertools -import os.path import pathlib import sys from typing import List, Iterator try: # betterproto[compiler] specific dependencies - import black from google.protobuf.compiler import plugin_pb2 as plugin from google.protobuf.descriptor_pb2 import ( DescriptorProto, @@ -14,7 +12,6 @@ try: FieldDescriptorProto, ServiceDescriptorProto, ) - import jinja2 except ImportError as err: missing_import = err.args[0][17:-1] print( @@ -41,6 +38,8 @@ from betterproto.plugin.models import ( is_oneof, ) +from betterproto.plugin.compiler import outputfile_compiler + def traverse(proto_file: FieldDescriptorProto) -> Iterator: # Todo: Keep information about nested hierarchy @@ -70,16 +69,6 @@ def generate_code( ) -> None: plugin_options = request.parameter.split(",") if request.parameter else [] - templates_folder = os.path.abspath( - os.path.join(os.path.dirname(__file__), "..", "templates") - ) - - env = jinja2.Environment( - trim_blocks=True, - lstrip_blocks=True, - loader=jinja2.FileSystemLoader(templates_folder), - ) - template = env.get_template("template.py.j2") request_data = PluginRequestCompiler(plugin_request_obj=request) # Gather output packages for proto_file in request.proto_file: @@ -116,7 +105,7 @@ def generate_code( # Generate output files output_paths: pathlib.Path = set() - for output_package_name, template_data in request_data.output_packages.items(): + for output_package_name, output_package in request_data.output_packages.items(): # Add files to the response object output_path = pathlib.Path(*output_package_name.split("."), "__init__.py") @@ -126,10 +115,7 @@ def generate_code( f.name: str = str(output_path) # Render and then format the output file - f.content: str = black.format_str( - template.render(description=template_data), - mode=black.FileMode(target_versions={black.TargetVersion.PY37}), - ) + f.content: str = outputfile_compiler(output_file=output_package) # Make each output directory a package with __init__ file init_files = ( diff --git a/src/betterproto/templates/template.py.j2 b/src/betterproto/templates/template.py.j2 index fa260e4..7fd0463 100644 --- a/src/betterproto/templates/template.py.j2 +++ b/src/betterproto/templates/template.py.j2 @@ -1,26 +1,26 @@ # Generated by the protocol buffer compiler. DO NOT EDIT! -# sources: {{ ', '.join(description.input_filenames) }} +# sources: {{ ', '.join(output_file.input_filenames) }} # plugin: python-betterproto -{% for i in description.python_module_imports|sort %} +{% for i in output_file.python_module_imports|sort %} import {{ i }} {% endfor %} from dataclasses import dataclass -{% if description.datetime_imports %} -from datetime import {% for i in description.datetime_imports|sort %}{{ i }}{% if not loop.last %}, {% endif %}{% endfor %} +{% if output_file.datetime_imports %} +from datetime import {% for i in output_file.datetime_imports|sort %}{{ i }}{% if not loop.last %}, {% endif %}{% endfor %} {% endif%} -{% if description.typing_imports %} -from typing import {% for i in description.typing_imports|sort %}{{ i }}{% if not loop.last %}, {% endif %}{% endfor %} +{% if output_file.typing_imports %} +from typing import {% for i in output_file.typing_imports|sort %}{{ i }}{% if not loop.last %}, {% endif %}{% endfor %} {% endif %} import betterproto -{% if description.services %} +{% if output_file.services %} import grpclib {% endif %} -{% if description.enums %}{% for enum in description.enums %} +{% if output_file.enums %}{% for enum in output_file.enums %} class {{ enum.py_name }}(betterproto.Enum): {% if enum.comment %} {{ enum.comment }} @@ -36,7 +36,7 @@ class {{ enum.py_name }}(betterproto.Enum): {% endfor %} {% endif %} -{% for message in description.messages %} +{% for message in output_file.messages %} @dataclass class {{ message.py_name }}(betterproto.Message): {% if message.comment %} @@ -67,7 +67,7 @@ class {{ message.py_name }}(betterproto.Message): {% endfor %} -{% for service in description.services %} +{% for service in output_file.services %} class {{ service.py_name }}Stub(betterproto.ServiceStub): {% if service.comment %} {{ service.comment }} @@ -154,6 +154,6 @@ class {{ service.py_name }}Stub(betterproto.ServiceStub): {% endfor %} {% endfor %} -{% for i in description.imports|sort %} +{% for i in output_file.imports|sort %} {{ i }} {% endfor %}