From e7a907e315e2c551d7fc30fb13b0406081952edf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20S=C3=A1nchez-Gallego?= <gallegoj@uw.edu>
Date: Sun, 25 Feb 2024 23:07:22 +0000
Subject: [PATCH] sdss/yao#17: buffer size does not match expected size

---
 CHANGELOG.md                    |  5 +++++
 archon/controller/controller.py | 29 +++++++++++++++++++++++------
 2 files changed, 28 insertions(+), 6 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 19b22a1..0514768 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,10 @@
 # Changelog
 
+## Next version
+
+* [yao #17](sdss/yao#17): Deal with case when the returned buffer does not match the expected size.
+
+
 ## 0.14.3 - January 12, 2024
 
 ### 🔧 Fixed
diff --git a/archon/controller/controller.py b/archon/controller/controller.py
index eba50a5..4fa7089 100644
--- a/archon/controller/controller.py
+++ b/archon/controller/controller.py
@@ -1157,8 +1157,7 @@ async def fetch(
         notifier: Optional[Callable[[str], None]] = None,
         *,
         return_buffer: Literal[False],
-    ) -> numpy.ndarray:
-        ...
+    ) -> numpy.ndarray: ...
 
     @overload
     async def fetch(
@@ -1167,8 +1166,7 @@ async def fetch(
         notifier: Optional[Callable[[str], None]] = None,
         *,
         return_buffer: Literal[True],
-    ) -> tuple[numpy.ndarray, int]:
-        ...
+    ) -> tuple[numpy.ndarray, int]: ...
 
     @overload
     async def fetch(
@@ -1176,8 +1174,7 @@ async def fetch(
         buffer_no: int = -1,
         notifier: Optional[Callable[[str], None]] = None,
         return_buffer: bool = False,
-    ) -> numpy.ndarray:
-        ...
+    ) -> numpy.ndarray: ...
 
     async def fetch(
         self,
@@ -1267,6 +1264,26 @@ async def fetch(
         # Convert to uint16 array and reshape.
         dtype = f"<u{bytes_per_pixel}"  # Buffer is little-endian
         arr = numpy.frombuffer(frame, dtype=dtype)
+
+        # See yao issue #17. In some cases the buffer size is 2 pixels short. For
+        # now if the buffer size does not match what we expect, just pad with zeros.
+        expected_size = height * width
+        if expected_size != arr.size:
+            message = (
+                "Buffer data size does not match expected size. "
+                f"Buffer size is {arr.size}; expected size is {expected_size}. "
+                "Padding with zeros."
+            )
+            notifier(message)
+            warnings.warn(message, ArchonUserWarning)
+
+            arr0 = arr.copy()
+            arr = numpy.zeros(expected_size, dtype=arr.dtype)
+            arr[: arr0.size] = arr0
+
+        else:
+            notifier(f"Buffer size is {arr.size}; expected size is {expected_size}.")
+
         arr = arr.reshape(height, width)
 
         # Turn off FETCHING bit