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 <wainersm@redhat.com>
Reviewed-by: John Snow <jsnow@redhat.com>
Message-Id: <20200204141111.3207-5-wainersm@redhat.com>
Signed-off-by: Philippe Mathieu-Daudé <philmd@redhat.com>

1 files changed, 13 insertions, 2 deletions

diff --git a/python/qemu/qmp.py b/python/qemu/qmp.py
index 0e07d80e2a..486a566da0 100644
--- 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):
         """