danielgtaylor#9 Support importing wellknown types Struct and Empty

Author: boukeversteegh

betterproto/plugin.py

@@ -28,6 +28,8 @@
 from betterproto.casing import safe_snake_case
 
 import google.protobuf.wrappers_pb2 as google_wrappers
+import google.protobuf.empty_pb2
+import google.protobuf.struct_pb2
 
 WRAPPER_TYPES: Dict[str, Optional[Type]] = defaultdict(lambda: None, {
     'google.protobuf.DoubleValue': google_wrappers.DoubleValue,
@@ -41,6 +43,11 @@
     'google.protobuf.BytesValue': google_wrappers.BytesValue,
 })
 
+WELLKNOWN_TYPES: Dict[str, Optional[Type]] = defaultdict(lambda: None, {
+    'google.protobuf.Empty': google.protobuf.empty_pb2.Empty,
+    'google.protobuf.Struct': google.protobuf.struct_pb2.Struct,
+})
+
 
 def get_ref_type(package: str, imports: set, type_name: str, unwrap: bool = True) -> str:
     """
@@ -77,6 +84,13 @@ def get_ref_type(package: str, imports: set, type_name: str, unwrap: bool = True
             cased = [stringcase.pascalcase(part) for part in parts]
             type_name = f'"{"".join(cased)}"'
 
+    if type_name in WELLKNOWN_TYPES:
+        type_module = WELLKNOWN_TYPES[type_name].__module__
+        type_name = WELLKNOWN_TYPES[type_name].__name__
+        type_module_alias = safe_snake_case(type_module)
+        imports.add(f"import {type_module} as {type_module_alias}")
+        return f'{type_module_alias}.{type_name}'
+
     if "." in type_name:
         # This is imported from another package. No need
         # to use a forward ref and we need to add the import.

betterproto/tests/inputs/config.py

@@ -8,7 +8,6 @@
     "import_circular_dependency",  # failing because of other bugs now
     "import_packages_same_name",  # 25
     "oneof_enum",  # 63
-    "googletypes_service_returns_empty",  # 9
     "casing_message_field_uppercase",  # 11
     "namespace_keywords",  # 70
     "namespace_builtin_types"  # 53
@@ -19,5 +18,5 @@
     "googletypes_response_embedded",
     "service",
     "import_service_input_message",
-    "googletypes_service_returns_empty",
+    "googletypes_service_returns_googletype",
 }

betterproto/tests/inputs/googletypes_service_returns_empty/googletypes_service_returns_empty.proto

@@ -0,0 +1,12 @@
+syntax = "proto3";
+
+import "google/protobuf/empty.proto";
+import "google/protobuf/struct.proto";
+
+service Test {
+    rpc GetEmpty (RequestMessage) returns (google.protobuf.Empty);
+    rpc GetStruct (RequestMessage) returns (google.protobuf.Struct);
+}
+
+message RequestMessage {
+}