From 610f01e96d3bfd9f632371de5bd6cf911a8e71ef Mon Sep 17 00:00:00 2001
From: Dana Powers <dana.powers@rd.io>
Date: Sun, 8 Mar 2015 20:01:57 -0700
Subject: Use context managers in gzip_encode / gzip_decode

---
 kafka/codec.py | 19 +++++++------------
 1 file changed, 7 insertions(+), 12 deletions(-)

(limited to 'kafka/codec.py')

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
 
 
-- 
cgit v1.2.1