merged lp:~jconti/python-apt/closeable-cache

2 files changed, 46 insertions, 4 deletions

diff --git a/apt/cache.py b/apt/cache.py
index 86a788bb..adc936ef 100644
--- a/apt/cache.py
+++ b/apt/cache.py
@@ -42,6 +42,9 @@ class FetchFailedException(IOError):
 class LockFailedException(IOError):
     """Exception that is thrown when locking fails."""
 
+class CacheClosedException(Exception):
+    """Exception that is thrown when the cache is used after close()."""
+
 
 class Cache(object):
     """Dictionary-like package cache.
@@ -172,6 +175,19 @@ class Cache(object):
         progress.done()
         self._run_callbacks("cache_post_open")
 
+    def close(self):
+        """ Close the package cache """
+        del self._records
+        self._records = None
+
+    def __enter__(self):
+        """ Enter the with statement """
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """ Exit the with statement """
+        self.close()
+
     def __getitem__(self, key):
         """ look like a dictionary (get key) """
         try:
@@ -238,6 +254,8 @@ class Cache(object):
     @property
     def required_download(self):
         """Get the size of the packages that are required to download."""
+        if self._records is None:
+            raise CacheClosedException("Cache object used after close() called")
         pm = apt_pkg.PackageManager(self._depcache)
         fetcher = apt_pkg.Acquire()
         pm.get_archives(fetcher, self._list, self._records)
@@ -288,6 +306,8 @@ class Cache(object):
 
     def _fetch_archives(self, fetcher, pm):
         """ fetch the needed archives """
+        if self._records is None:
+            raise CacheClosedException("Cache object used after close() called")
 
         # get lock
         lockfile = apt_pkg.config.find_dir("Dir::Cache::Archives") + "lock"
diff --git a/tests/test_apt_cache.py b/tests/test_apt_cache.py
index 448aed75..fc429cf3 100644
--- a/tests/test_apt_cache.py
+++ b/tests/test_apt_cache.py
@@ -7,19 +7,21 @@
 # are permitted in any medium without royalty provided the copyright
 # notice and this notice are preserved.
 """Unit tests for verifying the correctness of check_dep, etc in apt_pkg."""
+
+import glob
+import logging
 import os
+import shutil
+import sys
 import tempfile
 import unittest
 
 from test_all import get_library_dir
-import sys
 sys.path.insert(0, get_library_dir())
 
 import apt
 import apt_pkg
-import shutil
-import glob
-import logging
+
 
 def if_sources_list_is_readable(f):
     def wrapper(*args, **kwargs):
@@ -29,6 +31,16 @@ def if_sources_list_is_readable(f):
             logging.warn("skipping '%s' because sources.list is not readable" % f)
     return wrapper
 
+
+def get_open_file_descriptors():
+    try:
+        fds = os.listdir("/proc/self/fd")
+    except OSError:
+        logging.warn("failed to list /proc/self/fd")
+        return set([])
+    return set(map(int, fds))
+
+
 class TestAptCache(unittest.TestCase):
     """ test the apt cache """
 
@@ -70,6 +82,16 @@ class TestAptCache(unittest.TestCase):
                 self.assertTrue(len(r['Description']) > 0)
                 self.assertTrue(str(r).startswith('Package: %s\n' % pkg.shortname))
 
+    @if_sources_list_is_readable
+    def test_closeable_cache(self):
+        fd_0 = get_open_file_descriptors()
+        cache = apt.Cache()
+        opened_fd = get_open_file_descriptors().difference(fd_0)
+        cache.close()
+        fd_1 = get_open_file_descriptors()
+        unclosed_fd = opened_fd.intersection(fd_1)
+        self.assertEqual(unclosed_fd, set())
+
     def test_get_provided_packages(self):
         apt.apt_pkg.config.set("Apt::architecture", "i386")
         cache = apt.Cache(rootdir="./data/test-provides/")