Merge pull request #78 from boukeversteegh/pr/google

Basic general support for Google Protobuf

Author: nat-n

README.md

@@ -256,6 +256,7 @@ Google provides several well-known message types like a timestamp, duration, and
 | `google.protobuf.duration`  | [`datetime.timedelta`][td]               | `0`                    |
 | `google.protobuf.timestamp` | Timezone-aware [`datetime.datetime`][dt] | `1970-01-01T00:00:00Z` |
 | `google.protobuf.*Value`    | `Optional[...]`                          | `None`                 |
+| `google.protobuf.*`         | `betterproto.lib.google.protobuf.*`      | `None`                 |
 
 [td]: https://docs.python.org/3/library/datetime.html#timedelta-objects
 [dt]: https://docs.python.org/3/library/datetime.html#datetime.datetime
@@ -354,6 +355,25 @@ $ pipenv run generate
 $ pipenv run test
 ```
 
+### (Re)compiling Google Well-known Types
+
+Betterproto includes compiled versions for Google's well-known types at [betterproto/lib/google](betterproto/lib/google).
+Be sure to regenerate these files when modifying the plugin output format, and validate by running the tests.
+
+Normally, the plugin does not compile any references to `google.protobuf`, since they are pre-compiled. To force compilation of `google.protobuf`, use the option `--custom_opt=INCLUDE_GOOGLE`. 
+
+Assuming your `google.protobuf` source files (included with all releases of `protoc`) are located in `/usr/local/include`, you can regenerate them as follows:
+
+```sh
+protoc \
+    --plugin=protoc-gen-custom=betterproto/plugin.py \
+    --custom_opt=INCLUDE_GOOGLE \
+    --custom_out=betterproto/lib \
+    -I /usr/local/include/ \
+    /usr/local/include/google/protobuf/*.proto
+```
+
+
 ### TODO
 
 - [x] Fixed length fields

betterproto/__init__.py

@@ -941,19 +941,23 @@ def which_one_of(message: Message, group_name: str) -> Tuple[str, Any]:
     return (field_name, getattr(message, field_name))
 
 
-@dataclasses.dataclass
-class _Duration(Message):
-    # Signed seconds of the span of time. Must be from -315,576,000,000 to
-    # +315,576,000,000 inclusive. Note: these bounds are computed from: 60
-    # sec/min * 60 min/hr * 24 hr/day * 365.25 days/year * 10000 years
-    seconds: int = int64_field(1)
-    # Signed fractions of a second at nanosecond resolution of the span of time.
-    # Durations less than one second are represented with a 0 `seconds` field and
-    # a positive or negative `nanos` field. For durations of one second or more,
-    # a non-zero value for the `nanos` field must be of the same sign as the
-    # `seconds` field. Must be from -999,999,999 to +999,999,999 inclusive.
-    nanos: int = int32_field(2)
+# Circular import workaround: google.protobuf depends on base classes defined above.
+from .lib.google.protobuf import (
+    Duration,
+    Timestamp,
+    BoolValue,
+    BytesValue,
+    DoubleValue,
+    FloatValue,
+    Int32Value,
+    Int64Value,
+    StringValue,
+    UInt32Value,
+    UInt64Value,
+)
+
 
+class _Duration(Duration):
     def to_timedelta(self) -> timedelta:
         return timedelta(seconds=self.seconds, microseconds=self.nanos / 1e3)
 
@@ -966,16 +970,7 @@ def delta_to_json(delta: timedelta) -> str:
         return ".".join(parts) + "s"
 
 
-@dataclasses.dataclass
-class _Timestamp(Message):
-    # Represents seconds of UTC time since Unix epoch 1970-01-01T00:00:00Z. Must
-    # be from 0001-01-01T00:00:00Z to 9999-12-31T23:59:59Z inclusive.
-    seconds: int = int64_field(1)
-    # Non-negative fractions of a second at nanosecond resolution. Negative
-    # second values with fractions must still have non-negative nanos values that
-    # count forward in time. Must be from 0 to 999,999,999 inclusive.
-    nanos: int = int32_field(2)
-
+class _Timestamp(Timestamp):
     def to_datetime(self) -> datetime:
         ts = self.seconds + (self.nanos / 1e9)
         return datetime.fromtimestamp(ts, tz=timezone.utc)
@@ -1016,63 +1011,18 @@ def from_dict(self: T, value: Any) -> T:
         return self
 
 
-@dataclasses.dataclass
-class _BoolValue(_WrappedMessage):
-    value: bool = bool_field(1)
-
-
-@dataclasses.dataclass
-class _Int32Value(_WrappedMessage):
-    value: int = int32_field(1)
-
-
-@dataclasses.dataclass
-class _UInt32Value(_WrappedMessage):
-    value: int = uint32_field(1)
-
-
-@dataclasses.dataclass
-class _Int64Value(_WrappedMessage):
-    value: int = int64_field(1)
-
-
-@dataclasses.dataclass
-class _UInt64Value(_WrappedMessage):
-    value: int = uint64_field(1)
-
-
-@dataclasses.dataclass
-class _FloatValue(_WrappedMessage):
-    value: float = float_field(1)
-
-
-@dataclasses.dataclass
-class _DoubleValue(_WrappedMessage):
-    value: float = double_field(1)
-
-
-@dataclasses.dataclass
-class _StringValue(_WrappedMessage):
-    value: str = string_field(1)
-
-
-@dataclasses.dataclass
-class _BytesValue(_WrappedMessage):
-    value: bytes = bytes_field(1)
-
-
 def _get_wrapper(proto_type: str) -> Type:
     """Get the wrapper message class for a wrapped type."""
     return {
-        TYPE_BOOL: _BoolValue,
-        TYPE_INT32: _Int32Value,
-        TYPE_UINT32: _UInt32Value,
-        TYPE_INT64: _Int64Value,
-        TYPE_UINT64: _UInt64Value,
-        TYPE_FLOAT: _FloatValue,
-        TYPE_DOUBLE: _DoubleValue,
-        TYPE_STRING: _StringValue,
-        TYPE_BYTES: _BytesValue,
+        TYPE_BOOL: BoolValue,
+        TYPE_INT32: Int32Value,
+        TYPE_UINT32: UInt32Value,
+        TYPE_INT64: Int64Value,
+        TYPE_UINT64: UInt64Value,
+        TYPE_FLOAT: FloatValue,
+        TYPE_DOUBLE: DoubleValue,
+        TYPE_STRING: StringValue,
+        TYPE_BYTES: BytesValue,
     }[proto_type]
 
 

betterproto/compile/__init__.py

@@ -0,0 +1,70 @@
+from typing import Dict, Type
+
+import stringcase
+
+from betterproto import safe_snake_case
+from betterproto.lib.google import protobuf as google_protobuf
+
+WRAPPER_TYPES: Dict[str, Type] = {
+    "google.protobuf.DoubleValue": google_protobuf.DoubleValue,
+    "google.protobuf.FloatValue": google_protobuf.FloatValue,
+    "google.protobuf.Int32Value": google_protobuf.Int32Value,
+    "google.protobuf.Int64Value": google_protobuf.Int64Value,
+    "google.protobuf.UInt32Value": google_protobuf.UInt32Value,
+    "google.protobuf.UInt64Value": google_protobuf.UInt64Value,
+    "google.protobuf.BoolValue": google_protobuf.BoolValue,
+    "google.protobuf.StringValue": google_protobuf.StringValue,
+    "google.protobuf.BytesValue": google_protobuf.BytesValue,
+}
+
+
+def get_ref_type(
+    package: str, imports: set, type_name: str, unwrap: bool = True
+) -> str:
+    """
+    Return a Python type name for a proto type reference. Adds the import if
+    necessary. Unwraps well known type if required.
+    """
+    # If the package name is a blank string, then this should still work
+    # because by convention packages are lowercase and message/enum types are
+    # pascal-cased. May require refactoring in the future.
+    type_name = type_name.lstrip(".")
+
+    is_wrapper = type_name in WRAPPER_TYPES
+
+    if unwrap:
+        if is_wrapper:
+            wrapped_type = type(WRAPPER_TYPES[type_name]().value)
+            return f"Optional[{wrapped_type.__name__}]"
+
+        if type_name == "google.protobuf.Duration":
+            return "timedelta"
+
+        if type_name == "google.protobuf.Timestamp":
+            return "datetime"
+
+    if type_name.startswith(package):
+        parts = type_name.lstrip(package).lstrip(".").split(".")
+        if len(parts) == 1 or (len(parts) > 1 and parts[0][0] == parts[0][0].upper()):
+            # This is the current package, which has nested types flattened.
+            # foo.bar_thing => FooBarThing
+            cased = [stringcase.pascalcase(part) for part in parts]
+            type_name = f'"{"".join(cased)}"'
+
+    # Use precompiled classes for google.protobuf.* objects
+    if type_name.startswith("google.protobuf.") and type_name.count(".") == 2:
+        type_name = type_name.rsplit(".", maxsplit=1)[1]
+        import_package = "betterproto.lib.google.protobuf"
+        import_alias = safe_snake_case(import_package)
+        imports.add(f"import {import_package} as {import_alias}")
+        return f"{import_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.
+        parts = type_name.split(".")
+        parts[-1] = stringcase.pascalcase(parts[-1])
+        imports.add(f"from .{'.'.join(parts[:-2])} import {parts[-2]}")
+        type_name = f"{parts[-2]}.{parts[-1]}"
+
+    return type_name