From 14fcc5d485e255cb935b5834e90e8f6ed7a149f6 Mon Sep 17 00:00:00 2001
From: Lucien Morey <lucienp@live.com.au>
Date: Sun, 10 Nov 2024 05:08:45 +1100
Subject: [PATCH] generate packing for python messages (#1535)

Generate packet serialization in Python, too.
---
 .../generated/MultiTargetPNPResultSerde.py    |  13 ++-
 .../generated/PhotonPipelineMetadataSerde.py  |  19 ++-
 .../generated/PhotonPipelineResultSerde.py    |  18 ++-
 .../generated/PhotonTrackedTargetSerde.py     |  41 ++++++-
 .../photonlibpy/generated/PnpResultSerde.py   |  20 +++-
 .../generated/TargetCornerSerde.py            |  13 ++-
 photon-lib/py/photonlibpy/packet.py           | 109 +++++++++++++++++-
 photon-serde/generate_messages.py             |   1 +
 photon-serde/message_data_types.yaml          |   6 +
 photon-serde/messages.yaml                    |   1 +
 photon-serde/templates/ThingSerde.py.jinja    |  29 +++++
 11 files changed, 263 insertions(+), 7 deletions(-)

diff --git a/photon-lib/py/photonlibpy/generated/MultiTargetPNPResultSerde.py b/photon-lib/py/photonlibpy/generated/MultiTargetPNPResultSerde.py
index 080e1ba19..56705202d 100644
--- a/photon-lib/py/photonlibpy/generated/MultiTargetPNPResultSerde.py
+++ b/photon-lib/py/photonlibpy/generated/MultiTargetPNPResultSerde.py
@@ -21,14 +21,25 @@
 ###############################################################################
 
 from ..targeting import *
+from ..packet import Packet
 
 
 class MultiTargetPNPResultSerde:
-
     # Message definition md5sum. See photon_packet.adoc for details
     MESSAGE_VERSION = "541096947e9f3ca2d3f425ff7b04aa7b"
     MESSAGE_FORMAT = "PnpResult:ae4d655c0a3104d88df4f5db144c1e86 estimatedPose;int16 fiducialIDsUsed[?];"
 
+    @staticmethod
+    def pack(value: "MultiTargetPNPResult") -> "Packet":
+        ret = Packet()
+
+        # estimatedPose is of non-intrinsic type PnpResult
+        ret.encodeBytes(PnpResult.photonStruct.pack(value.estimatedPose).getData())
+
+        # fiducialIDsUsed is a custom VLA!
+        ret.encodeShortList(value.fiducialIDsUsed)
+        return ret
+
     @staticmethod
     def unpack(packet: "Packet") -> "MultiTargetPNPResult":
         ret = MultiTargetPNPResult()
diff --git a/photon-lib/py/photonlibpy/generated/PhotonPipelineMetadataSerde.py b/photon-lib/py/photonlibpy/generated/PhotonPipelineMetadataSerde.py
index f8e68f123..99a994186 100644
--- a/photon-lib/py/photonlibpy/generated/PhotonPipelineMetadataSerde.py
+++ b/photon-lib/py/photonlibpy/generated/PhotonPipelineMetadataSerde.py
@@ -21,14 +21,31 @@
 ###############################################################################
 
 from ..targeting import *
+from ..packet import Packet
 
 
 class PhotonPipelineMetadataSerde:
-
     # Message definition md5sum. See photon_packet.adoc for details
     MESSAGE_VERSION = "ac0a45f686457856fb30af77699ea356"
     MESSAGE_FORMAT = "int64 sequenceID;int64 captureTimestampMicros;int64 publishTimestampMicros;int64 timeSinceLastPong;"
 
+    @staticmethod
+    def pack(value: "PhotonPipelineMetadata") -> "Packet":
+        ret = Packet()
+
+        # sequenceID is of intrinsic type int64
+        ret.encodeLong(value.sequenceID)
+
+        # captureTimestampMicros is of intrinsic type int64
+        ret.encodeLong(value.captureTimestampMicros)
+
+        # publishTimestampMicros is of intrinsic type int64
+        ret.encodeLong(value.publishTimestampMicros)
+
+        # timeSinceLastPong is of intrinsic type int64
+        ret.encodeLong(value.timeSinceLastPong)
+        return ret
+
     @staticmethod
     def unpack(packet: "Packet") -> "PhotonPipelineMetadata":
         ret = PhotonPipelineMetadata()
diff --git a/photon-lib/py/photonlibpy/generated/PhotonPipelineResultSerde.py b/photon-lib/py/photonlibpy/generated/PhotonPipelineResultSerde.py
index d54917cf8..a1c186595 100644
--- a/photon-lib/py/photonlibpy/generated/PhotonPipelineResultSerde.py
+++ b/photon-lib/py/photonlibpy/generated/PhotonPipelineResultSerde.py
@@ -21,14 +21,30 @@
 ###############################################################################
 
 from ..targeting import *
+from ..packet import Packet
 
 
 class PhotonPipelineResultSerde:
-
     # Message definition md5sum. See photon_packet.adoc for details
     MESSAGE_VERSION = "4b2ff16a964b5e2bf04be0c1454d91c4"
     MESSAGE_FORMAT = "PhotonPipelineMetadata:ac0a45f686457856fb30af77699ea356 metadata;PhotonTrackedTarget:cc6dbb5c5c1e0fa808108019b20863f1 targets[?];optional MultiTargetPNPResult:541096947e9f3ca2d3f425ff7b04aa7b multitagResult;"
 
+    @staticmethod
+    def pack(value: "PhotonPipelineResult") -> "Packet":
+        ret = Packet()
+
+        # metadata is of non-intrinsic type PhotonPipelineMetadata
+        ret.encodeBytes(
+            PhotonPipelineMetadata.photonStruct.pack(value.metadata).getData()
+        )
+
+        # targets is a custom VLA!
+        ret.encodeList(value.targets, PhotonTrackedTarget.photonStruct)
+
+        # multitagResult is optional! it better not be a VLA too
+        ret.encodeOptional(value.multitagResult, MultiTargetPNPResult.photonStruct)
+        return ret
+
     @staticmethod
     def unpack(packet: "Packet") -> "PhotonPipelineResult":
         ret = PhotonPipelineResult()
diff --git a/photon-lib/py/photonlibpy/generated/PhotonTrackedTargetSerde.py b/photon-lib/py/photonlibpy/generated/PhotonTrackedTargetSerde.py
index 1a54e7b53..2eea5a6ad 100644
--- a/photon-lib/py/photonlibpy/generated/PhotonTrackedTargetSerde.py
+++ b/photon-lib/py/photonlibpy/generated/PhotonTrackedTargetSerde.py
@@ -21,14 +21,53 @@
 ###############################################################################
 
 from ..targeting import *
+from ..packet import Packet
 
 
 class PhotonTrackedTargetSerde:
-
     # Message definition md5sum. See photon_packet.adoc for details
     MESSAGE_VERSION = "cc6dbb5c5c1e0fa808108019b20863f1"
     MESSAGE_FORMAT = "float64 yaw;float64 pitch;float64 area;float64 skew;int32 fiducialId;int32 objDetectId;float32 objDetectConf;Transform3d bestCameraToTarget;Transform3d altCameraToTarget;float64 poseAmbiguity;TargetCorner:16f6ac0dedc8eaccb951f4895d9e18b6 minAreaRectCorners[?];TargetCorner:16f6ac0dedc8eaccb951f4895d9e18b6 detectedCorners[?];"
 
+    @staticmethod
+    def pack(value: "PhotonTrackedTarget") -> "Packet":
+        ret = Packet()
+
+        # yaw is of intrinsic type float64
+        ret.encodeDouble(value.yaw)
+
+        # pitch is of intrinsic type float64
+        ret.encodeDouble(value.pitch)
+
+        # area is of intrinsic type float64
+        ret.encodeDouble(value.area)
+
+        # skew is of intrinsic type float64
+        ret.encodeDouble(value.skew)
+
+        # fiducialId is of intrinsic type int32
+        ret.encodeInt(value.fiducialId)
+
+        # objDetectId is of intrinsic type int32
+        ret.encodeInt(value.objDetectId)
+
+        # objDetectConf is of intrinsic type float32
+        ret.encodeFloat(value.objDetectConf)
+
+        ret.encodeTransform(value.bestCameraToTarget)
+
+        ret.encodeTransform(value.altCameraToTarget)
+
+        # poseAmbiguity is of intrinsic type float64
+        ret.encodeDouble(value.poseAmbiguity)
+
+        # minAreaRectCorners is a custom VLA!
+        ret.encodeList(value.minAreaRectCorners, TargetCorner.photonStruct)
+
+        # detectedCorners is a custom VLA!
+        ret.encodeList(value.detectedCorners, TargetCorner.photonStruct)
+        return ret
+
     @staticmethod
     def unpack(packet: "Packet") -> "PhotonTrackedTarget":
         ret = PhotonTrackedTarget()
diff --git a/photon-lib/py/photonlibpy/generated/PnpResultSerde.py b/photon-lib/py/photonlibpy/generated/PnpResultSerde.py
index a9d4cb81e..44cffc368 100644
--- a/photon-lib/py/photonlibpy/generated/PnpResultSerde.py
+++ b/photon-lib/py/photonlibpy/generated/PnpResultSerde.py
@@ -21,14 +21,32 @@
 ###############################################################################
 
 from ..targeting import *
+from ..packet import Packet
 
 
 class PnpResultSerde:
-
     # Message definition md5sum. See photon_packet.adoc for details
     MESSAGE_VERSION = "ae4d655c0a3104d88df4f5db144c1e86"
     MESSAGE_FORMAT = "Transform3d best;Transform3d alt;float64 bestReprojErr;float64 altReprojErr;float64 ambiguity;"
 
+    @staticmethod
+    def pack(value: "PnpResult") -> "Packet":
+        ret = Packet()
+
+        ret.encodeTransform(value.best)
+
+        ret.encodeTransform(value.alt)
+
+        # bestReprojErr is of intrinsic type float64
+        ret.encodeDouble(value.bestReprojErr)
+
+        # altReprojErr is of intrinsic type float64
+        ret.encodeDouble(value.altReprojErr)
+
+        # ambiguity is of intrinsic type float64
+        ret.encodeDouble(value.ambiguity)
+        return ret
+
     @staticmethod
     def unpack(packet: "Packet") -> "PnpResult":
         ret = PnpResult()
diff --git a/photon-lib/py/photonlibpy/generated/TargetCornerSerde.py b/photon-lib/py/photonlibpy/generated/TargetCornerSerde.py
index ef1814a35..0c57b42ee 100644
--- a/photon-lib/py/photonlibpy/generated/TargetCornerSerde.py
+++ b/photon-lib/py/photonlibpy/generated/TargetCornerSerde.py
@@ -21,14 +21,25 @@
 ###############################################################################
 
 from ..targeting import *
+from ..packet import Packet
 
 
 class TargetCornerSerde:
-
     # Message definition md5sum. See photon_packet.adoc for details
     MESSAGE_VERSION = "16f6ac0dedc8eaccb951f4895d9e18b6"
     MESSAGE_FORMAT = "float64 x;float64 y;"
 
+    @staticmethod
+    def pack(value: "TargetCorner") -> "Packet":
+        ret = Packet()
+
+        # x is of intrinsic type float64
+        ret.encodeDouble(value.x)
+
+        # y is of intrinsic type float64
+        ret.encodeDouble(value.y)
+        return ret
+
     @staticmethod
     def unpack(packet: "Packet") -> "TargetCorner":
         ret = TargetCorner()
diff --git a/photon-lib/py/photonlibpy/packet.py b/photon-lib/py/photonlibpy/packet.py
index a4e375d43..9a4e1fadd 100644
--- a/photon-lib/py/photonlibpy/packet.py
+++ b/photon-lib/py/photonlibpy/packet.py
@@ -22,7 +22,7 @@ import wpilib
 
 
 class Packet:
-    def __init__(self, data: bytes):
+    def __init__(self, data: bytes = b""):
         """
         * Constructs an empty packet.
         *
@@ -198,3 +198,110 @@ class Packet:
             return serde.unpack(self)
         else:
             return None
+
+    def _encodeGeneric(self, packFormat, value):
+        """
+        Append bytes to the packet data buffer.
+        """
+        self.packetData = self.packetData + struct.pack(packFormat, value)
+        self.size = len(self.packetData)
+
+    def encode8(self, value: int):
+        """
+        Encodes a single byte and appends it to the packet.
+        """
+        self._encodeGeneric("<b", value)
+
+    def encode16(self, value: int):
+        """
+        Encodes a short (2 bytes) and appends it to the packet.
+        """
+        self._encodeGeneric("<h", value)
+
+    def encodeInt(self, value: int):
+        """
+        Encodes an int (4 bytes) and appends it to the packet.
+        """
+        self._encodeGeneric("<l", value)
+
+    def encodeFloat(self, value: float):
+        """
+        Encodes a float (4 bytes) and appends it to the packet.
+        """
+        self._encodeGeneric("<f", value)
+
+    def encodeLong(self, value: int):
+        """
+        Encodes a long (8 bytes) and appends it to the packet.
+        """
+        self._encodeGeneric("<q", value)
+
+    def encodeDouble(self, value: float):
+        """
+        Encodes a double (8 bytes) and appends it to the packet.
+        """
+        self._encodeGeneric("<d", value)
+
+    def encodeBoolean(self, value: bool):
+        """
+        Encodes a boolean as a single byte and appends it to the packet.
+        """
+        self.encode8(1 if value else 0)
+
+    def encodeDoubleArray(self, values: list[float]):
+        """
+        Encodes an array of doubles and appends it to the packet.
+        """
+        self.encode8(len(values))
+        for value in values:
+            self.encodeDouble(value)
+
+    def encodeShortList(self, values: list[int]):
+        """
+        Encodes a list of shorts, with length prefixed as a single byte.
+        """
+        self.encode8(len(values))
+        for value in values:
+            self.encode16(value)
+
+    def encodeTransform(self, transform: Transform3d):
+        """
+        Encodes a Transform3d (translation and rotation) and appends it to the packet.
+        """
+        # Encode Translation3d part (x, y, z)
+        self.encodeDouble(transform.translation().x)
+        self.encodeDouble(transform.translation().y)
+        self.encodeDouble(transform.translation().z)
+
+        # Encode Rotation3d as Quaternion (w, x, y, z)
+        quaternion = transform.rotation().getQuaternion()
+        self.encodeDouble(quaternion.W())
+        self.encodeDouble(quaternion.X())
+        self.encodeDouble(quaternion.Y())
+        self.encodeDouble(quaternion.Z())
+
+    def encodeList(self, values: list[Any], serde: Type):
+        """
+        Encodes a list of items using a specific serializer and appends it to the packet.
+        """
+        self.encode8(len(values))
+        for item in values:
+            packed = serde.pack(item)
+            self.packetData = self.packetData + packed.getData()
+            self.size = len(self.packetData)
+
+    def encodeOptional(self, value: Optional[Any], serde: Type):
+        """
+        Encodes an optional value using a specific serializer.
+        """
+        if value is None:
+            self.encodeBoolean(False)
+        else:
+            self.encodeBoolean(True)
+            packed = serde.pack(value)
+            self.packetData = self.packetData + packed.getData()
+            self.size = len(self.packetData)
+
+    def encodeBytes(self, value: bytes):
+        self.packetData = self.packetData + value
+        self.size = len(self.packetData)
diff --git a/photon-serde/generate_messages.py b/photon-serde/generate_messages.py
index 2b4922dc0..2efca5281 100644
--- a/photon-serde/generate_messages.py
+++ b/photon-serde/generate_messages.py
@@ -46,6 +46,7 @@ class MessageType(TypedDict):
     # C++ helpers
     cpp_include: str
     # python shim types
+    python_encode_shim: str
     python_decode_shim: str
     # Java import name
     java_import: str
diff --git a/photon-serde/message_data_types.yaml b/photon-serde/message_data_types.yaml
index b6a846e4c..85fdc608b 100644
--- a/photon-serde/message_data_types.yaml
+++ b/photon-serde/message_data_types.yaml
@@ -5,29 +5,35 @@ bool:
   java_type: bool
   cpp_type: bool
   java_decode_method: decodeBoolean
+  java_encode_shim: encodeBoolean
 int16:
   len: 2
   java_type: short
   cpp_type: int16_t
   java_decode_method: decodeShort
   java_list_decode_method: decodeShortList
+  java_encode_shim: encodeShort
 int32:
   len: 4
   java_type: int
   cpp_type: int32_t
   java_decode_method: decodeInt
+  java_encode_shim: encodeBoolean
 int64:
   len: 8
   java_type: long
   cpp_type: int64_t
   java_decode_method: decodeLong
+  java_encode_shim: encodeLong
 float32:
   len: 4
   java_type: float
   cpp_type: float
   java_decode_method: decodeFloat
+  java_encode_shim: encodeFloat
 float64:
   len: 8
   java_type: double
   cpp_type: double
   java_decode_method: decodeDouble
+  java_encode_shim: encodeDouble
diff --git a/photon-serde/messages.yaml b/photon-serde/messages.yaml
index cbdf4e90e..6c4e461a3 100644
--- a/photon-serde/messages.yaml
+++ b/photon-serde/messages.yaml
@@ -17,6 +17,7 @@
   cpp_type: frc::Transform3d
   cpp_include: "<frc/geometry/Transform3d.h>"
   python_decode_shim: packet.decodeTransform
+  python_encode_shim: encodeTransform
   java_import: edu.wpi.first.math.geometry.Transform3d
   # shim since we expect fields to at least exist
   fields: []
diff --git a/photon-serde/templates/ThingSerde.py.jinja b/photon-serde/templates/ThingSerde.py.jinja
index 335adc038..f6e441b47 100644
--- a/photon-serde/templates/ThingSerde.py.jinja
+++ b/photon-serde/templates/ThingSerde.py.jinja
@@ -21,6 +21,7 @@
 ###############################################################################
 
 from ..targeting import *
+from ..packet import Packet
 
 class {{ name }}Serde:
 
@@ -28,6 +29,34 @@ class {{ name }}Serde:
     MESSAGE_VERSION = "{{ message_hash }}"
     MESSAGE_FORMAT = "{{ message_fmt }}"
 
+    @staticmethod
+    def pack(value: '{{ name }}' ) -> 'Packet':
+        ret = Packet()
+{% for field in fields -%}
+{%- if field.type | is_shimmed %}
+        ret.{{ get_message_by_name(field.type).python_encode_shim}}(value.{{ field.name }})
+{%- elif field.optional == True %}
+        # {{ field.name }} is optional! it better not be a VLA too
+        ret.encodeOptional(value.{{ field.name }}, {{ field.type }}.photonStruct)
+{%- elif field.vla == True and not field.type | is_intrinsic %}
+        # {{ field.name }} is a custom VLA!
+        ret.encodeList(value.{{ field.name }}, {{ field.type }}.photonStruct)
+{%- elif field.vla == True and field.type | is_intrinsic %}
+        # {{ field.name }} is a custom VLA!
+        ret.encode{{ type_map[field.type].java_type.title() }}List(value.{{ field.name }})
+{%- elif field.type | is_intrinsic %}
+        # {{ field.name }} is of intrinsic type {{ field.type }}
+        ret.{{ type_map[field.type].java_encode_shim }}(value.{{field.name}})
+{%- else %}
+        # {{ field.name }} is of non-intrinsic type {{ field.type }}
+        ret.encodeBytes({{ field.type }}.photonStruct.pack(value.{{field.name}}).getData())
+{%- endif %}
+{%- if not loop.last %}
+{% endif -%}
+{% endfor%}
+        return ret
+
+
     @staticmethod
     def unpack(packet: 'Packet') -> '{{ name }}':
         ret = {{ name }}()