From 41164f421daeb1c4d775fefa0b76dd7b77d18690 Mon Sep 17 00:00:00 2001 From: boukeversteegh Date: Thu, 28 May 2020 09:10:26 +0200 Subject: [PATCH 1/3] #9 Support importing wellknown types Struct and Empty --- betterproto/plugin.py | 14 ++++++++++++++ betterproto/tests/inputs/config.py | 3 +-- .../googletypes_service_returns_empty.proto | 11 ----------- .../googletypes_service_returns_googletype.proto | 12 ++++++++++++ 4 files changed, 27 insertions(+), 13 deletions(-) delete mode 100644 betterproto/tests/inputs/googletypes_service_returns_empty/googletypes_service_returns_empty.proto create mode 100644 betterproto/tests/inputs/googletypes_service_returns_googletype/googletypes_service_returns_googletype.proto diff --git a/betterproto/plugin.py b/betterproto/plugin.py index 8eecef2b0..ae706544c 100755 --- a/betterproto/plugin.py +++ b/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, @@ -44,6 +46,11 @@ }, ) +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 @@ -82,6 +89,13 @@ def get_ref_type( 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. diff --git a/betterproto/tests/inputs/config.py b/betterproto/tests/inputs/config.py index cf4c99671..47d1a16d7 100644 --- a/betterproto/tests/inputs/config.py +++ b/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", } diff --git a/betterproto/tests/inputs/googletypes_service_returns_empty/googletypes_service_returns_empty.proto b/betterproto/tests/inputs/googletypes_service_returns_empty/googletypes_service_returns_empty.proto deleted file mode 100644 index c45469141..000000000 --- a/betterproto/tests/inputs/googletypes_service_returns_empty/googletypes_service_returns_empty.proto +++ /dev/null @@ -1,11 +0,0 @@ -syntax = "proto3"; - -import "google/protobuf/empty.proto"; - -service Test { - rpc Send (RequestMessage) returns (google.protobuf.Empty) { - } -} - -message RequestMessage { -} \ No newline at end of file diff --git a/betterproto/tests/inputs/googletypes_service_returns_googletype/googletypes_service_returns_googletype.proto b/betterproto/tests/inputs/googletypes_service_returns_googletype/googletypes_service_returns_googletype.proto new file mode 100644 index 000000000..05a53a815 --- /dev/null +++ b/betterproto/tests/inputs/googletypes_service_returns_googletype/googletypes_service_returns_googletype.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 { +} \ No newline at end of file From 67a628907e4090d28d97d04291c4f76b981cd28e Mon Sep 17 00:00:00 2001 From: boukeversteegh Date: Thu, 28 May 2020 09:17:05 +0200 Subject: [PATCH 2/3] #9 Support importing wellknown types Value and ListValue --- betterproto/plugin.py | 2 ++ betterproto/tests/inputs/config.py | 2 +- .../googletypes_service_returns_googletype.proto | 2 ++ 3 files changed, 5 insertions(+), 1 deletion(-) diff --git a/betterproto/plugin.py b/betterproto/plugin.py index ae706544c..81adddbf5 100755 --- a/betterproto/plugin.py +++ b/betterproto/plugin.py @@ -49,6 +49,8 @@ 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, + 'google.protobuf.Value': google.protobuf.struct_pb2.Value, + 'google.protobuf.ListValue': google.protobuf.struct_pb2.ListValue, }) diff --git a/betterproto/tests/inputs/config.py b/betterproto/tests/inputs/config.py index 47d1a16d7..5af995d6e 100644 --- a/betterproto/tests/inputs/config.py +++ b/betterproto/tests/inputs/config.py @@ -10,7 +10,7 @@ "oneof_enum", # 63 "casing_message_field_uppercase", # 11 "namespace_keywords", # 70 - "namespace_builtin_types" # 53 + "namespace_builtin_types", # 53 } services = { diff --git a/betterproto/tests/inputs/googletypes_service_returns_googletype/googletypes_service_returns_googletype.proto b/betterproto/tests/inputs/googletypes_service_returns_googletype/googletypes_service_returns_googletype.proto index 05a53a815..11348be25 100644 --- a/betterproto/tests/inputs/googletypes_service_returns_googletype/googletypes_service_returns_googletype.proto +++ b/betterproto/tests/inputs/googletypes_service_returns_googletype/googletypes_service_returns_googletype.proto @@ -6,6 +6,8 @@ import "google/protobuf/struct.proto"; service Test { rpc GetEmpty (RequestMessage) returns (google.protobuf.Empty); rpc GetStruct (RequestMessage) returns (google.protobuf.Struct); + rpc GetListValue (RequestMessage) returns (google.protobuf.ListValue); + rpc GetValue (RequestMessage) returns (google.protobuf.Value); } message RequestMessage { From ac523201ee8a6cda1ecedbbd28fd971396097e2b Mon Sep 17 00:00:00 2001 From: boukeversteegh Date: Thu, 28 May 2020 09:45:04 +0200 Subject: [PATCH 3/3] Blacken --- betterproto/plugin.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/betterproto/plugin.py b/betterproto/plugin.py index 81adddbf5..7f66b173c 100755 --- a/betterproto/plugin.py +++ b/betterproto/plugin.py @@ -46,12 +46,15 @@ }, ) -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, - 'google.protobuf.Value': google.protobuf.struct_pb2.Value, - 'google.protobuf.ListValue': google.protobuf.struct_pb2.ListValue, -}) +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, + "google.protobuf.Value": google.protobuf.struct_pb2.Value, + "google.protobuf.ListValue": google.protobuf.struct_pb2.ListValue, + }, +) def get_ref_type( @@ -96,7 +99,7 @@ def get_ref_type( 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}' + return f"{type_module_alias}.{type_name}" if "." in type_name: # This is imported from another package. No need