Compile well-known types (#35)

* Compile well-known types

* Don't import well-known types from the lib

* Add duration and timestamp to well-known types

Author: AdrienVannson

.github/workflows/ci.yml

@@ -62,4 +62,4 @@ jobs:
 
       - name: Generate code from proto files
         shell: bash
-        run: poetry run python -m tests.generate -v
+        run: poetry run poe generate

docs/development.md

@@ -2,20 +2,17 @@
 
 ## Recompiling the lib proto files
 
-After some updates in the compiler, it might be useful to recompile the standard Google proto files. As the proto files
-are distributed with `protoc`, their path might depend on your installation.
+After some updates in the compiler, it might be useful to recompile the standard Google proto files used by the
+compiler. As the proto files are distributed with `protoc`, their path might depend on your installation.
 
 ```bash
 mkdir lib
 protoc \
-    --python_betterproto2_opt=INCLUDE_GOOGLE \
     --python_betterproto2_out=lib \
     -I /usr/include/ \
     /usr/include/google/protobuf/*.proto
 ```
 
-The generated files should be distributed in the `betterproto2` package.
-
 !!! warning
     These proto files are written with the `proto2` syntax, which is not supported by betterproto. For the compiler to
     work, you need to manually patch the generated file to mark the field `oneof_index` in `Field` and

pyproject.toml

@@ -68,9 +68,42 @@ cmd = "pytest"
 help = "Run tests"
 
 [tool.poe.tasks.generate]
-script = "tests.generate:main"
+sequence = ["_generate_tests", "_generate_tests_lib"]
 help = "Generate test cases"
 
+[tool.poe.tasks._generate_tests]
+script = "tests.generate:main"
+
+[tool.poe.tasks._generate_tests_lib]
+shell = """
+python -m grpc.tools.protoc \
+    --python_betterproto2_out=tests/output_betterproto \
+    google/protobuf/any.proto \
+    google/protobuf/api.proto \
+    google/protobuf/duration.proto \
+    google/protobuf/empty.proto \
+    google/protobuf/field_mask.proto \
+    google/protobuf/source_context.proto \
+    google/protobuf/struct.proto \
+    google/protobuf/timestamp.proto \
+    google/protobuf/type.proto \
+    google/protobuf/wrappers.proto
+
+python -m grpc.tools.protoc \
+    --python_betterproto2_out=tests/output_betterproto_pydantic \
+    --python_betterproto2_opt=pydantic_dataclasses \
+    google/protobuf/any.proto \
+    google/protobuf/api.proto \
+    google/protobuf/duration.proto \
+    google/protobuf/empty.proto \
+    google/protobuf/field_mask.proto \
+    google/protobuf/source_context.proto \
+    google/protobuf/struct.proto \
+    google/protobuf/timestamp.proto \
+    google/protobuf/type.proto \
+    google/protobuf/wrappers.proto
+"""
+
 [tool.poe.tasks.typecheck]
 cmd  = "pyright src"
 help = "Typecheck the code with Pyright"
@@ -103,25 +136,6 @@ help = "Check the code with the Ruff linter"
 cmd = "mkdocs serve"
 help = "Serve the documentation locally"
 
-[tool.poe.tasks.generate_lib]
-shell = """
-mkdir -p tests/output_lib
-python -m grpc.tools.protoc \
-    --python_betterproto2_opt=INCLUDE_GOOGLE \
-    --python_betterproto2_out=tests/output_lib \
-    google/protobuf/any.proto \
-    google/protobuf/api.proto \
-    google/protobuf/duration.proto \
-    google/protobuf/empty.proto \
-    google/protobuf/field_mask.proto \
-    google/protobuf/source_context.proto \
-    google/protobuf/struct.proto \
-    google/protobuf/timestamp.proto \
-    google/protobuf/type.proto \
-    google/protobuf/wrappers.proto
-"""
-help = "Compile the well-known proto types"
-
 [build-system]
 requires = ["poetry-core>=1.0.0,<2"]
 build-backend = "poetry.core.masonry.api"

src/betterproto2_compiler/compile/importing.py

@@ -98,14 +98,6 @@ def get_type_reference(
     py_package: list[str] = source_package.split(".") if source_package else []
     py_type: str = pythonize_class_name(source_type)
 
-    compiling_google_protobuf = current_package == ["google", "protobuf"]
-    importing_google_protobuf = py_package == ["google", "protobuf"]
-    if importing_google_protobuf and not compiling_google_protobuf:
-        py_package = ["betterproto2", "lib"] + (["pydantic"] if settings.pydantic_dataclasses else []) + py_package
-
-    if py_package[:1] == ["betterproto2"]:
-        return reference_absolute(imports, py_package, py_type)
-
     if py_package == current_package:
         return reference_sibling(py_type)
 

src/betterproto2_compiler/known_types/__init__.py

@@ -1,10 +1,14 @@
 from collections.abc import Callable
 
 from .any import Any
+from .duration import Duration
+from .timestamp import Timestamp
 
 # For each (package, message name), lists the methods that should be added to the message definition.
 # The source code of the method is read from the `known_types` folder. If imports are needed, they can be directly added
 # to the template file: they will automatically be removed if not necessary.
 KNOWN_METHODS: dict[tuple[str, str], list[Callable]] = {
     ("google.protobuf", "Any"): [Any.pack, Any.unpack, Any.to_dict],
+    ("google.protobuf", "Timestamp"): [Timestamp.from_datetime, Timestamp.to_datetime, Timestamp.timestamp_to_json],
+    ("google.protobuf", "Duration"): [Duration.from_timedelta, Duration.to_timedelta, Duration.delta_to_json],
 }

src/betterproto2_compiler/known_types/duration.py

@@ -0,0 +1,25 @@
+import datetime
+
+from betterproto2.lib.std.google.protobuf import Duration as VanillaDuration
+
+
+class Duration(VanillaDuration):
+    @classmethod
+    def from_timedelta(
+        cls, delta: datetime.timedelta, *, _1_microsecond: datetime.timedelta = datetime.timedelta(microseconds=1)
+    ) -> "Duration":
+        total_ms = delta // _1_microsecond
+        seconds = int(total_ms / 1e6)
+        nanos = int((total_ms % 1e6) * 1e3)
+        return cls(seconds, nanos)
+
+    def to_timedelta(self) -> datetime.timedelta:
+        return datetime.timedelta(seconds=self.seconds, microseconds=self.nanos / 1e3)
+
+    @staticmethod
+    def delta_to_json(delta: datetime.timedelta) -> str:
+        parts = str(delta.total_seconds()).split(".")
+        if len(parts) > 1:
+            while len(parts[1]) not in (3, 6, 9):
+                parts[1] = f"{parts[1]}0"
+        return f"{'.'.join(parts)}s"

src/betterproto2_compiler/known_types/timestamp.py

@@ -0,0 +1,45 @@
+import datetime
+
+from betterproto2.lib.std.google.protobuf import Timestamp as VanillaTimestamp
+
+
+class Timestamp(VanillaTimestamp):
+    @classmethod
+    def from_datetime(cls, dt: datetime.datetime) -> "Timestamp":
+        # manual epoch offset calulation to avoid rounding errors,
+        # to support negative timestamps (before 1970) and skirt
+        # around datetime bugs (apparently 0 isn't a year in [0, 9999]??)
+        offset = dt - datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc)
+        # below is the same as timedelta.total_seconds() but without dividing by 1e6
+        # so we end up with microseconds as integers instead of seconds as float
+        offset_us = (offset.days * 24 * 60 * 60 + offset.seconds) * 10**6 + offset.microseconds
+        seconds, us = divmod(offset_us, 10**6)
+        return cls(seconds, us * 1000)
+
+    def to_datetime(self) -> datetime.datetime:
+        # datetime.fromtimestamp() expects a timestamp in seconds, not microseconds
+        # if we pass it as a floating point number, we will run into rounding errors
+        # see also #407
+        offset = datetime.timedelta(seconds=self.seconds, microseconds=self.nanos // 1000)
+        return datetime.datetime(1970, 1, 1, tzinfo=datetime.timezone.utc) + offset
+
+    @staticmethod
+    def timestamp_to_json(dt: datetime.datetime) -> str:
+        nanos = dt.microsecond * 1e3
+        if dt.tzinfo is not None:
+            # change timezone aware datetime objects to utc
+            dt = dt.astimezone(datetime.timezone.utc)
+        copy = dt.replace(microsecond=0, tzinfo=None)
+        result = copy.isoformat()
+        if (nanos % 1e9) == 0:
+            # If there are 0 fractional digits, the fractional
+            # point '.' should be omitted when serializing.
+            return f"{result}Z"
+        if (nanos % 1e6) == 0:
+            # Serialize 3 fractional digits.
+            return f"{result}.{int(nanos // 1e6) :03d}Z"
+        if (nanos % 1e3) == 0:
+            # Serialize 6 fractional digits.
+            return f"{result}.{int(nanos // 1e3) :06d}Z"
+        # Serialize 9 fractional digits.
+        return f"{result}.{nanos:09d}"

src/betterproto2_compiler/plugin/models.py

@@ -189,7 +189,6 @@ class OutputTemplate:
     messages: dict[str, "MessageCompiler"] = field(default_factory=dict)
     enums: dict[str, "EnumDefinitionCompiler"] = field(default_factory=dict)
     services: dict[str, "ServiceCompiler"] = field(default_factory=dict)
-    output: bool = True
 
     settings: Settings
 

src/betterproto2_compiler/plugin/parser.py

@@ -107,11 +107,6 @@ def generate_code(request: CodeGeneratorRequest) -> CodeGeneratorResponse:
         # Add this input file to the output corresponding to this package
         request_data.output_packages[output_package_name].input_files.append(proto_file)
 
-        if proto_file.package == "google.protobuf" and "INCLUDE_GOOGLE" not in plugin_options:
-            # If not INCLUDE_GOOGLE,
-            # skip outputting Google's well-known types
-            request_data.output_packages[output_package_name].output = False
-
     # Read Messages and Enums
     # We need to read Messages before Services in so that we can
     # get the references to input/output messages for each service
@@ -147,9 +142,6 @@ def generate_code(request: CodeGeneratorRequest) -> CodeGeneratorResponse:
     # Generate output files
     output_paths: set[pathlib.Path] = set()
     for output_package_name, output_package in request_data.output_packages.items():
-        if not output_package.output:
-            continue
-
         # Add files to the response object
         output_path = pathlib.Path(*output_package_name.split("."), "__init__.py")
         output_paths.add(output_path)