QOL fixes (#141)

- Add missing type annotations
- Various style improvements
- Use constants more consistently
- enforce black on benchmark code

src/betterproto/plugin/parser.py

@@ -1,7 +1,7 @@
 import itertools
 import pathlib
 import sys
-from typing import List, Iterator
+from typing import TYPE_CHECKING, Iterator, List, Tuple, Union, Set
 
 try:
     # betterproto[compiler] specific dependencies
@@ -13,10 +13,9 @@ try:
         ServiceDescriptorProto,
     )
 except ImportError as err:
-    missing_import = err.args[0][17:-1]
     print(
         "\033[31m"
-        f"Unable to import `{missing_import}` from betterproto plugin! "
+        f"Unable to import `{err.name}` from betterproto plugin! "
         "Please ensure that you've installed betterproto as "
         '`pip install "betterproto[compiler]"` so that compiler dependencies '
         "are included."
@@ -24,26 +23,32 @@ except ImportError as err:
     )
     raise SystemExit(1)
 
-from betterproto.plugin.models import (
-    PluginRequestCompiler,
-    OutputTemplate,
-    MessageCompiler,
-    FieldCompiler,
-    OneOfFieldCompiler,
-    MapEntryCompiler,
+from .compiler import outputfile_compiler
+from .models import (
     EnumDefinitionCompiler,
+    FieldCompiler,
+    MapEntryCompiler,
+    MessageCompiler,
+    OneOfFieldCompiler,
+    OutputTemplate,
+    PluginRequestCompiler,
     ServiceCompiler,
     ServiceMethodCompiler,
     is_map,
     is_oneof,
 )
 
-from betterproto.plugin.compiler import outputfile_compiler
+if TYPE_CHECKING:
+    from google.protobuf.descriptor import Descriptor
 
 
-def traverse(proto_file: FieldDescriptorProto) -> Iterator:
+def traverse(
+    proto_file: FieldDescriptorProto,
+) -> "itertools.chain[Tuple[Union[str, EnumDescriptorProto], List[int]]]":
     # Todo: Keep information about nested hierarchy
-    def _traverse(path, items, prefix=""):
+    def _traverse(
+        path: List[int], items: List["Descriptor"], prefix=""
+    ) -> Iterator[Tuple[Union[str, EnumDescriptorProto], List[int]]]:
         for i, item in enumerate(items):
             # Adjust the name since we flatten the hierarchy.
             # Todo: don't change the name, but include full name in returned tuple
@@ -104,7 +109,7 @@ def generate_code(
                 read_protobuf_service(service, index, output_package)
 
     # Generate output files
-    output_paths: pathlib.Path = set()
+    output_paths: Set[pathlib.Path] = set()
     for output_package_name, output_package in request_data.output_packages.items():
 
         # Add files to the response object
@@ -112,20 +117,17 @@ def generate_code(
         output_paths.add(output_path)
 
         f: response.File = response.file.add()
-        f.name: str = str(output_path)
+        f.name = str(output_path)
 
         # Render and then format the output file
-        f.content: str = outputfile_compiler(output_file=output_package)
+        f.content = outputfile_compiler(output_file=output_package)
 
     # Make each output directory a package with __init__ file
-    init_files = (
-        set(
-            directory.joinpath("__init__.py")
-            for path in output_paths
-            for directory in path.parents
-        )
-        - output_paths
-    )
+    init_files = {
+        directory.joinpath("__init__.py")
+        for path in output_paths
+        for directory in path.parents
+    } - output_paths
 
     for init_file in init_files:
         init = response.file.add()