python/qemu: qmp: Make QEMUMonitorProtocol a context manager

This implement the __enter__ and __exit__ functions on
QEMUMonitorProtocol class so that it can be used on 'with'
statement and the resources will be free up on block end:

with QEMUMonitorProtocol(socket_path) as qmp:
    qmp.connect()
    qmp.command('query-status')

Signed-off-by: Wainer dos Santos Moschetta <[email protected]>
Reviewed-by: John Snow <[email protected]>
Message-Id: <20200204141111[email protected]>
Signed-off-by: Philippe Mathieu-Daudé <[email protected]>

diff --git a/python/qemu/qmp.py b/python/qemu/qmp.py
index 0e07d80e2ad332167ee32b23df4238898bbee233..486a566da0758285a50f39aa3ad5b3f072d44c2b 100644 (file)
--- a/python/qemu/qmp.py
+++ b/python/qemu/qmp.py
@@ -139,6 +139,15 @@ class QEMUMonitorProtocol:
                 raise QMPConnectError("Error while reading from socket")
             self.__sock.settimeout(None)
 
+    def __enter__(self):
+        # Implement context manager enter function.
+        return self
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        # Implement context manager exit function.
+        self.close()
+        return False
+
     def connect(self, negotiate=True):
         """
         Connect to the QMP Monitor and perform capabilities negotiation.
@@ -265,8 +274,10 @@ class QEMUMonitorProtocol:
         """
         Close the socket and socket file.
         """
-        self.__sock.close()
-        self.__sockfile.close()
+        if self.__sock:
+            self.__sock.close()
+        if self.__sockfile:
+            self.__sockfile.close()
 
     def settimeout(self, timeout):
         """