From 7dbaee0cfbdfa393101c2bae1cbee25eafb2408e Mon Sep 17 00:00:00 2001
From: "Daniel G. Taylor" <danielgtaylor@gmail.com>
Date: Thu, 10 Oct 2019 23:09:09 -0700
Subject: [PATCH] Add support for map value message types

---
 README.md                          | 23 +++++++++++++++++++++--
 betterproto/__init__.py            | 28 +++++++++++++++++++++-------
 betterproto/tests/mapmessage.json  | 10 ++++++++++
 betterproto/tests/mapmessage.proto |  9 +++++++++
 4 files changed, 61 insertions(+), 9 deletions(-)
 create mode 100644 betterproto/tests/mapmessage.json
 create mode 100644 betterproto/tests/mapmessage.proto

diff --git a/README.md b/README.md
index fb06616..d813b82 100644
--- a/README.md
+++ b/README.md
@@ -1,4 +1,22 @@
-# TODO
+# Better Protobuf / gRPC Support for Python
+
+This project aims to provide an improved experience when using Protobuf / gRPC in a modern Python environment by making use of modern language features and generating readable, understandable code. It will not support legacy features or environments. The following are supported:
+
+- Protobuf 3 & gRPC code generation
+  - Both binary & JSON serialization is built-in
+- Python 3.7+
+  - Enums
+  - Dataclasses
+  - `async`/`await`
+  - Relative imports
+  - Mypy type checking
+
+This project is heavily inspired by, and borrows functionality from:
+
+- https://github.com/eigenein/protobuf/
+- https://github.com/vmagamedov/grpclib
+
+## TODO
 
 - [x] Fixed length fields
   - [x] Packed fixed-length
@@ -7,11 +25,12 @@
 - [x] Enums
 - [x] Repeated message fields
 - [x] Maps
-  - [ ] Maps of message fields
+  - [x] Maps of message fields
 - [ ] Support passthrough of unknown fields
 - [ ] Refs to nested types
 - [ ] Imports in proto files
 - [ ] Well-known Google types
 - [ ] JSON that isn't completely naive.
 - [ ] Async service stubs
+- [ ] Python package
 - [ ] Cleanup!
diff --git a/betterproto/__init__.py b/betterproto/__init__.py
index 7d9b217..ed44e87 100644
--- a/betterproto/__init__.py
+++ b/betterproto/__init__.py
@@ -385,14 +385,13 @@ class Message(ABC):
 
         return output
 
-    def _cls_for(self, field: dataclasses.Field) -> Type:
+    def _cls_for(self, field: dataclasses.Field, index: int = 0) -> Type:
         """Get the message class for a field from the type hints."""
         module = inspect.getmodule(self)
         type_hints = get_type_hints(self, vars(module))
         cls = type_hints[field.name]
         if hasattr(cls, "__args__"):
-            print(type_hints[field.name].__args__[0])
-            cls = type_hints[field.name].__args__[0]
+            cls = type_hints[field.name].__args__[index]
         return cls
 
     def _postprocess_single(
@@ -420,11 +419,13 @@ class Message(ABC):
             elif meta.proto_type in [TYPE_MAP]:
                 # TODO: This is slow, use a cache to make it faster since each
                 #       key/value pair will recreate the class.
+                kt = self._cls_for(field, index=0)
+                vt = self._cls_for(field, index=1)
                 Entry = dataclasses.make_dataclass(
                     "Entry",
                     [
-                        ("key", Any, dataclass_field(1, meta.map_types[0], None)),
-                        ("value", Any, dataclass_field(2, meta.map_types[1], None)),
+                        ("key", kt, dataclass_field(1, meta.map_types[0], None)),
+                        ("value", vt, dataclass_field(2, meta.map_types[1], None)),
                     ],
                     bases=(Message,),
                 )
@@ -500,10 +501,18 @@ class Message(ABC):
                     v = [i for i in v if i]
                 else:
                     v = v.to_dict()
+
+                if v:
+                    output[field.name] = v
+            elif meta.proto_type == "map":
+                for k in v:
+                    if hasattr(v[k], "to_dict"):
+                        v[k] = v[k].to_dict()
+
                 if v:
                     output[field.name] = v
             elif v != field.default:
-                output[field.name] = getattr(self, field.name)
+                output[field.name] = v
         return output
 
     def from_dict(self, value: dict) -> T:
@@ -516,13 +525,18 @@ class Message(ABC):
             if field.name in value:
                 if meta.proto_type == "message":
                     v = getattr(self, field.name)
-                    print(v, value[field.name])
+                    # print(v, value[field.name])
                     if isinstance(v, list):
                         cls = self._cls_for(field)
                         for i in range(len(value[field.name])):
                             v.append(cls().from_dict(value[field.name][i]))
                     else:
                         v.from_dict(value[field.name])
+                elif meta.proto_type == "map" and meta.map_types[1] == TYPE_MESSAGE:
+                    v = getattr(self, field.name)
+                    cls = self._cls_for(field, index=1)
+                    for k in value[field.name]:
+                        v[k] = cls().from_dict(value[field.name][k])
                 else:
                     setattr(self, field.name, value[field.name])
         return self
diff --git a/betterproto/tests/mapmessage.json b/betterproto/tests/mapmessage.json
new file mode 100644
index 0000000..a944ddd
--- /dev/null
+++ b/betterproto/tests/mapmessage.json
@@ -0,0 +1,10 @@
+{
+  "items": {
+    "foo": {
+      "count": 1
+    },
+    "bar": {
+      "count": 2
+    }
+  }
+}
diff --git a/betterproto/tests/mapmessage.proto b/betterproto/tests/mapmessage.proto
new file mode 100644
index 0000000..07dcce5
--- /dev/null
+++ b/betterproto/tests/mapmessage.proto
@@ -0,0 +1,9 @@
+syntax = "proto3";
+
+message Test {
+  map<string, Nested> items = 1;
+}
+
+message Nested {
+  int32 count = 1;
+}
\ No newline at end of file