Use context managers in gzip_encode / gzip_decode

1 files changed, 7 insertions, 12 deletions

diff --git a/kafka/codec.py b/kafka/codec.py
index 2279200..7883158 100644
--- a/kafka/codec.py
+++ b/kafka/codec.py
@@ -24,22 +24,17 @@ def has_snappy():
 
 
 def gzip_encode(payload):
-    buffer = BytesIO()
-    handle = gzip.GzipFile(fileobj=buffer, mode="w")
-    handle.write(payload)
-    handle.close()
-    buffer.seek(0)
-    result = buffer.read()
-    buffer.close()
+    with BytesIO() as buf:
+        with gzip.GzipFile(fileobj=buf, mode="w") as gzipper:
+            gzipper.write(payload)
+        result = buf.getvalue()
     return result
 
 
 def gzip_decode(payload):
-    buffer = BytesIO(payload)
-    handle = gzip.GzipFile(fileobj=buffer, mode='r')
-    result = handle.read()
-    handle.close()
-    buffer.close()
+    with BytesIO(payload) as buf:
+        with gzip.GzipFile(fileobj=buf, mode='r') as gzipper:
+            result = gzipper.read()
     return result