fix(python): validate inputs in write_png helper

PythonClient/airsim/tests/test_write_png.py

@@ -0,0 +1,38 @@
+import os
+import tempfile
+import unittest
+
+import numpy as np
+
+from airsim.utils import write_png
+
+
+class WritePngTests(unittest.TestCase):
+    def test_rejects_empty_filename(self):
+        with self.assertRaises(ValueError):
+            write_png("", np.zeros((2, 2), dtype=np.uint8))
+
+    def test_rejects_none_image(self):
+        with self.assertRaises(ValueError):
+            write_png("x.png", None)
+
+    def test_rejects_bad_ndim(self):
+        with self.assertRaises(ValueError):
+            write_png("x.png", np.zeros((2,), dtype=np.uint8))
+
+    def test_writes_when_cv2_available(self):
+        try:
+            import cv2  # noqa: F401
+        except ImportError:
+            self.skipTest("opencv not installed")
+        img = np.zeros((4, 4, 3), dtype=np.uint8)
+        img[0, 0] = (0, 0, 255)
+        with tempfile.TemporaryDirectory() as td:
+            path = os.path.join(td, "t.png")
+            write_png(path, img)
+            self.assertTrue(os.path.isfile(path))
+            self.assertGreater(os.path.getsize(path), 0)
+
+
+if __name__ == "__main__":
+    unittest.main()

PythonClient/airsim/utils.py

@@ -199,10 +199,27 @@ def write_pfm(file, image, scale=1):
 
 
 def write_png(filename, image):
-    """ image must be numpy array H X W X channels
-    """
-    import cv2      # pip install opencv-python
+    """Write image to PNG via OpenCV.
 
-    ret = cv2.imwrite(filename, image)
+    image must be a numpy array shaped H x W or H x W x channels.
+    Raises ValueError for invalid inputs and OSError if OpenCV fails to write.
+    """
+    import cv2  # pip install opencv-python
+
+    if filename is None or (isinstance(filename, str) and filename.strip() == ""):
+        raise ValueError("write_png requires a non-empty filename")
+    if image is None:
+        raise ValueError("write_png requires an image array")
+
+    arr = image if isinstance(image, np.ndarray) else np.asarray(image)
+    if not isinstance(arr, np.ndarray):
+        raise ValueError("write_png image must be convertible to a numpy.ndarray")
+    if arr.ndim not in (2, 3):
+        raise ValueError(
+            "write_png image must have shape H x W or H x W x C, got ndim=%s" % (arr.ndim,)
+        )
+
+    ret = cv2.imwrite(filename, arr)
     if not ret:
-        logging.error(f"Writing PNG file {filename} failed")
+        logging.error("Writing PNG file %s failed", filename)
+        raise OSError("OpenCV failed to write PNG file: %s" % (filename,))