Merge pull request #11 from henribru/bugfix/timestamp-field

Fix Timestamp fields not returning DatetimeWithNanoseconds
henribru · Jan 25, 2024 · 5a1513b · 5a1513b
2 parents 0245caa + 95e5326
commit 5a1513b
Showing 1 changed file with 79 additions and 0 deletions.
diff --git a/proto-stubs/fields.pyi b/proto-stubs/fields.pyi
@@ -9,6 +9,9 @@ from typing import (
     overload,
 )
 
+from google.protobuf.timestamp_pb2 import Timestamp
+
+from .datetime_helpers import DatetimeWithNanoseconds
 from .message import Message
 from .primitives import ProtoType
 
@@ -89,6 +92,17 @@ class Field(Generic[_T]):
         optional: bool = False,
     ) -> None: ...
     @overload
+    def __init__(
+        self: Field[DatetimeWithNanoseconds],
+        proto_type: Literal[ProtoType.MESSAGE],
+        *,
+        number: int,
+        message: type[Timestamp],
+        oneof: str | None = None,
+        json_name: str | None = None,
+        optional: bool = False,
+    ) -> None: ...
+    @overload
     def __init__(
         self: Field[_T],
         proto_type: Literal[ProtoType.MESSAGE],
@@ -111,6 +125,16 @@ class Field(Generic[_T]):
         optional: bool = False,
     ) -> None: ...
     @overload
+    def __init__(
+        self: Field[DatetimeWithNanoseconds],
+        proto_type: type[Timestamp],
+        *,
+        number: int,
+        oneof: str | None = None,
+        json_name: str | None = None,
+        optional: bool = False,
+    ) -> None: ...
+    @overload
     def __init__(
         self: Field[_T],
         proto_type: type[_T],
@@ -205,6 +229,17 @@ class RepeatedField(Field[_T]):
         optional: bool = False,
     ) -> None: ...
     @overload
+    def __init__(
+        self: RepeatedField[DatetimeWithNanoseconds],
+        proto_type: Literal[ProtoType.MESSAGE],
+        *,
+        number: int,
+        message: type[Timestamp],
+        oneof: str | None = None,
+        json_name: str | None = None,
+        optional: bool = False,
+    ) -> None: ...
+    @overload
     def __init__(
         self: RepeatedField[_T],
         proto_type: Literal[ProtoType.MESSAGE],
@@ -227,6 +262,16 @@ class RepeatedField(Field[_T]):
         optional: bool = False,
     ) -> None: ...
     @overload
+    def __init__(
+        self: RepeatedField[DatetimeWithNanoseconds],
+        proto_type: type[Timestamp],
+        *,
+        number: int,
+        oneof: str | None = None,
+        json_name: str | None = None,
+        optional: bool = False,
+    ) -> None: ...
+    @overload
     def __init__(
         self: RepeatedField[_T],
         proto_type: type[_T],
@@ -306,6 +351,15 @@ class MapField(Field[_V], Generic[_K, _V]):
         number: int,
     ) -> None: ...
     @overload
+    def __init__(
+        self: MapField[int, DatetimeWithNanoseconds],
+        key_type: _IntegerProtoType,
+        value_type: Literal[ProtoType.MESSAGE],
+        *,
+        number: int,
+        message: type[Timestamp],
+    ) -> None: ...
+    @overload
     def __init__(
         self: MapField[int, _V],
         key_type: _IntegerProtoType,
@@ -324,6 +378,14 @@ class MapField(Field[_V], Generic[_K, _V]):
         enum: type[_V],
     ) -> None: ...
     @overload
+    def __init__(
+        self: MapField[int, DatetimeWithNanoseconds],
+        key_type: _IntegerProtoType,
+        value_type: type[Timestamp],
+        *,
+        number: int,
+    ) -> None: ...
+    @overload
     def __init__(
         self: MapField[int, _V],
         key_type: _IntegerProtoType,
@@ -390,6 +452,15 @@ class MapField(Field[_V], Generic[_K, _V]):
         number: int,
     ) -> None: ...
     @overload
+    def __init__(
+        self: MapField[str, DatetimeWithNanoseconds],
+        key_type: Literal[ProtoType.STRING],
+        value_type: Literal[ProtoType.MESSAGE],
+        *,
+        number: int,
+        message: type[Timestamp],
+    ) -> None: ...
+    @overload
     def __init__(
         self: MapField[str, _V],
         key_type: Literal[ProtoType.STRING],
@@ -408,6 +479,14 @@ class MapField(Field[_V], Generic[_K, _V]):
         enum: type[_V],
     ) -> None: ...
     @overload
+    def __init__(
+        self: MapField[str, DatetimeWithNanoseconds],
+        key_type: Literal[ProtoType.STRING],
+        value_type: type[Timestamp],
+        *,
+        number: int,
+    ) -> None: ...
+    @overload
     def __init__(
         self: MapField[str, _V],
         key_type: Literal[ProtoType.STRING],