From 53e65adc618382562fefbf1e3dfb6534485ee5a9 Mon Sep 17 00:00:00 2001
From: Wouter Bolsterlee <uws@xs4all.nl>
Date: Fri, 18 Oct 2013 16:46:02 +0200
Subject: [PATCH] Implement explicit iterator closing (issue #19)

Added Iterator.close(), and a context manager to automatically closes an
iterator. Fixes issue #19.

Also updated the user guide, the API docs, the tests, and NEWS.
---
 NEWS.rst            |  4 ++++
 doc/api.rst         | 22 ++++++++++++++++++++--
 doc/user.rst        | 24 ++++++++++++++++++++++--
 plyvel/_plyvel.pyx  | 20 +++++++++++---------
 test/test_plyvel.py | 15 +++++++++++++++
 5 files changed, 72 insertions(+), 13 deletions(-)

diff --git a/NEWS.rst b/NEWS.rst
index 35d5e7c..6765659 100644
--- a/NEWS.rst
+++ b/NEWS.rst
@@ -6,6 +6,10 @@ Version history
 Plyvel 0.6 (not yet released)
 =============================
 
+* Allow iterators to be closed explicitly using either
+  :py:meth:`Iterator.close()` or a ``with`` block (`issue #19
+  <https://github.com/wbolster/plyvel/issues/19>`_)
+
 * Add useful ``__repr__()`` for :py:class:`DB` and :py:class:`PrefixedDB`
   instances (`issue #16 <https://github.com/wbolster/plyvel/issues/16>`_)
 
diff --git a/doc/api.rst b/doc/api.rst
index b750253..a1d1df7 100644
--- a/doc/api.rst
+++ b/doc/api.rst
@@ -65,8 +65,10 @@ LevelDB database.
       Close the database.
 
       This closes the database and releases associated resources such as open
-      file pointers and caches. Any further operations on the closed database
-      will raise a :py:exc:`RuntimeError`.
+      file pointers and caches.
+
+      Any further operations on the closed database will raise
+      :py:exc:`RuntimeError`.
 
       .. warning::
 
@@ -478,6 +480,22 @@ Directly invoking methods on the :py:class:`Iterator` returned by
       This moves the iterator to the the first key that sorts equal or before
       the specified `target` within the iterator range (`start` and `stop`).
 
+   .. py:method:: close()
+
+      Close the iterator.
+
+      This closes the iterator and releases the associated resources. Any
+      further operations on the closed iterator will raise
+      :py:exc:`RuntimeError`.
+
+      To automatically close an iterator, a context manager can be used::
+
+          with db.iterator() as it:
+              for k, v in it:
+                  pass  # do something
+
+          it.seek_to_start()  # raises RuntimeError
+
 
 Errors
 ======
diff --git a/doc/user.rst b/doc/user.rst
index df87ea0..da85df9 100644
--- a/doc/user.rst
+++ b/doc/user.rst
@@ -325,8 +325,28 @@ iterating over snapshots using the :py:meth:`Snapshot.iterator` method. This
 method works exactly the same as :py:meth:`DB.iterator`, except that it operates
 on the snapshot instead of the complete database.
 
-Advanced iterator usage
------------------------
+Closing iterators
+-----------------
+
+It is generally not required to close an iterator explicitly, since it will be
+closed when it goes out of scope (or is garbage collected). However, due to the
+way LevelDB is designed, each iterator operates on an implicit database
+snapshot, which can be an expensive resource depending on how the database is
+used during the iterator's lifetime. The :py:meth:`Iterator.close` method gives
+explicit control over when those resources are released::
+
+    >>> it = db.iterator()
+    >>> it.close()
+
+Alternatively, to ensure that an iterator is immediately closed after used, you
+can also use an iterator as a context manager using the ``with`` statement::
+
+    >>> with db.iterator() as it:
+    ...    for k, v in it:
+    ...        pass
+
+Non-linear iteration
+--------------------
 
 In the examples above, we've only used Python's standard iteration methods using
 a ``for`` loop and the :py:func:`list` constructor. This suffices for most
diff --git a/plyvel/_plyvel.pyx b/plyvel/_plyvel.pyx
index 091ae6f..63da9ee 100644
--- a/plyvel/_plyvel.pyx
+++ b/plyvel/_plyvel.pyx
@@ -716,19 +716,21 @@ cdef class Iterator:
         # Store a weak reference on the db (needed when closing db)
         db.iterators[id(self)] = self
 
-    cdef close(self):
-        # Note: this method is only for internal cleanups and hence not
-        # accessible from Python.
-        if self._iter is NULL:
-            # Already closed
-            return
-
-        del self._iter
-        self._iter = NULL
+    cpdef close(self):
+        if self._iter is not NULL:
+            del self._iter
+            self._iter = NULL
 
     def __dealloc__(self):
         self.close()
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.close()
+        return False  # propagate exceptions
+
     def __iter__(self):
         return self
 
diff --git a/test/test_plyvel.py b/test/test_plyvel.py
index 8116328..c466b53 100644
--- a/test/test_plyvel.py
+++ b/test/test_plyvel.py
@@ -332,6 +332,21 @@ def test_iteration():
             assert_equal(entry, expected)
 
 
+def test_iterator_closing():
+    with tmp_db('iteration_closing') as db:
+        db.put(b'k', b'v')
+        it = db.iterator()
+        next(it)
+        it.close()
+        assert_raises(RuntimeError, next, it)
+        assert_raises(RuntimeError, it.seek_to_stop)
+
+        with db.iterator() as it:
+            next(it)
+
+        assert_raises(RuntimeError, next, it)
+
+
 def test_iterator_return():
     with tmp_db('iteration') as db:
         db.put(b'key', b'value')