From 7443fad90f1a749dac3aaef215fed8c8cf25aec1 Mon Sep 17 00:00:00 2001
From: Torrin Jones <torrinj@gmail.com>
Date: Sat, 7 Jan 2017 10:39:03 -0800
Subject: [PATCH] Allow use of the with statement

---
 asynchronousfilereader/__init__.py   |  7 +++++++
 tests/test_asynchronousfilereader.py | 10 ++++++++++
 2 files changed, 17 insertions(+)

diff --git a/asynchronousfilereader/__init__.py b/asynchronousfilereader/__init__.py
index e8aa35e..f6db7e2 100644
--- a/asynchronousfilereader/__init__.py
+++ b/asynchronousfilereader/__init__.py
@@ -39,6 +39,13 @@ def __init__(self, fd, queue=None, autostart=True):
         if autostart:
             self.start()
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if self.is_alive():
+            self.join()
+
     def run(self):
         """
         The body of the tread: read lines and put them on the queue.
diff --git a/tests/test_asynchronousfilereader.py b/tests/test_asynchronousfilereader.py
index 82d3e2b..49aa9b3 100644
--- a/tests/test_asynchronousfilereader.py
+++ b/tests/test_asynchronousfilereader.py
@@ -22,6 +22,16 @@ def test_simple(self):
 
         self.assertEqual(['line1\n', 'line2\n'], lines)
 
+    def test_with(self):
+        data_stream = StringIO('line1\nline2\n')
+        with AsynchronousFileReader(data_stream) as reader:
+            lines = []
+            while not reader.eof():
+                for line in reader.readlines():
+                    lines.append(line)
+
+        self.assertEqual(['line1\n', 'line2\n'], lines)
+
 
 if __name__ == '__main__':
     unittest.main()