From 96626cae1f906d83cc621ae5cddfb22c6ed5b6c9 Mon Sep 17 00:00:00 2001
From: Mitchell Goff <mitchellgoffpc@gmail.com>
Date: Tue, 20 Jun 2023 16:55:51 -0700
Subject: [PATCH] Added VisionIpcBufExtra properties to python VisionIpcClient
 (#479)

* Added VisionIpcBufExtra properties to python VisionIpcClient

* Check frame_id in test_send_single_buffer
---
 visionipc/tests/test_visionipc.py | 36 +++++++++++++++++++++++++++++++
 visionipc/visionipc_pyx.pyx       | 15 ++++++++++++-
 2 files changed, 50 insertions(+), 1 deletion(-)

diff --git a/visionipc/tests/test_visionipc.py b/visionipc/tests/test_visionipc.py
index 4c693cc..17f83ac 100644
--- a/visionipc/tests/test_visionipc.py
+++ b/visionipc/tests/test_visionipc.py
@@ -60,3 +60,39 @@ class TestVisionIpc(unittest.TestCase):
     recv_buf = client.recv()
     self.assertIsNot(recv_buf, None)
     self.assertEqual(recv_buf.view('<i4')[0], 1234)
+    self.assertEqual(client.frame_id, 1337)
+
+  def test_no_conflate(self):
+    self.create_vipc_server("camerad", VisionStreamType.VISION_STREAM_ROAD)
+    client = VisionIpcClient("camerad", VisionStreamType.VISION_STREAM_ROAD, False)
+    self.assertTrue(client.connect(True))
+    zmq_sleep()
+
+    buf = np.zeros(100 * 150, dtype=np.uint8)
+    self.server.send(VisionStreamType.VISION_STREAM_ROAD, buf, frame_id=1)
+    self.server.send(VisionStreamType.VISION_STREAM_ROAD, buf, frame_id=2)
+
+    recv_buf = client.recv()
+    self.assertIsNot(recv_buf, None)
+    self.assertEqual(client.frame_id, 1)
+
+    recv_buf = client.recv()
+    self.assertIsNot(recv_buf, None)
+    self.assertEqual(client.frame_id, 2)
+
+  def test_conflate(self):
+    self.create_vipc_server("camerad", VisionStreamType.VISION_STREAM_ROAD)
+    client = VisionIpcClient("camerad", VisionStreamType.VISION_STREAM_ROAD, True)
+    self.assertTrue(client.connect(True))
+    zmq_sleep()
+
+    buf = np.zeros(100 * 150, dtype=np.uint8)
+    self.server.send(VisionStreamType.VISION_STREAM_ROAD, buf, frame_id=1)
+    self.server.send(VisionStreamType.VISION_STREAM_ROAD, buf, frame_id=2)
+
+    recv_buf = client.recv()
+    self.assertIsNot(recv_buf, None)
+    self.assertEqual(client.frame_id, 2)
+
+    recv_buf = client.recv()
+    self.assertIs(recv_buf, None)
diff --git a/visionipc/visionipc_pyx.pyx b/visionipc/visionipc_pyx.pyx
index 2ec7888..8f99eee 100644
--- a/visionipc/visionipc_pyx.pyx
+++ b/visionipc/visionipc_pyx.pyx
@@ -65,6 +65,7 @@ cdef class VisionIpcServer:
 cdef class VisionIpcClient:
   cdef cppVisionBuf * buf
   cdef cppVisionIpcClient * client
+  cdef VisionIpcBufExtra extra
 
   def __cinit__(self, string name, VisionStreamType stream, bool conflate):
     self.client = new cppVisionIpcClient(name, stream, conflate, NULL, NULL)
@@ -89,8 +90,20 @@ cdef class VisionIpcClient:
   def uv_offset(self):
     return None if not self.buf else self.buf.uv_offset
 
+  @property
+  def frame_id(self):
+    return self.extra.frame_id
+
+  @property
+  def timestamp_sof(self):
+    return self.extra.timestamp_sof
+
+  @property
+  def timestamp_eof(self):
+    return self.extra.timestamp_eof
+
   def recv(self, int timeout_ms=100):
-    self.buf = self.client.recv(NULL, timeout_ms)
+    self.buf = self.client.recv(&self.extra, timeout_ms)
     if not self.buf:
       return None
     cdef cnp.ndarray dat = np.empty(self.buf.len, dtype=np.uint8)