From 193a96dbeb6d68b3a5df2b0e6c78b1b996ef802f Mon Sep 17 00:00:00 2001
From: David Lutton <David.Lutton@gmail.com>
Date: Sat, 13 Apr 2024 15:20:42 +0100
Subject: [PATCH] Add screenshot to AgilentE8357A

---
 src/labtoolkit/NetworkAnalyser/AgilentE8357A.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/src/labtoolkit/NetworkAnalyser/AgilentE8357A.py b/src/labtoolkit/NetworkAnalyser/AgilentE8357A.py
index ff145c4..d4ac31a 100644
--- a/src/labtoolkit/NetworkAnalyser/AgilentE8357A.py
+++ b/src/labtoolkit/NetworkAnalyser/AgilentE8357A.py
@@ -2,6 +2,8 @@
 from ..SCPI import SCPI
 import numpy as np
 import pandas as pd
+import PIL.Image as Image
+import io
 
 
 class AgilentE8357A(IEEE488, SCPI):
@@ -142,7 +144,20 @@ def reference(self):
         """
         return self.query('SENSe:ROSCillator:SOURce?')
 
-
+    def screenshot(self):
+        self.write('HCOPY:FILE \'screen.png\'')
+
+        image = Image.open(io.BytesIO(
+            self.query_binary_values(
+                ':MMEMory:TRANsfer? \'screen.png\'',
+                datatype='B', 
+                is_big_endian=False, 
+                container=bytearray
+            )
+        ))
+        self.write('MMEMory:DELete \'screen.png\'')
+        return image
+    
     '''
     @property
     def form(self):