Fix python 2 support

tests/test_utils.py

@@ -1,6 +1,6 @@
-from contextlib import contextmanager
 import pytest
 from mock import Mock
+import six
 from thefuck.utils import wrap_settings, \
     memoize, get_closest, get_all_executables, replace_argument, \
     get_all_matched_commands, is_app, for_app, cache
@@ -127,11 +127,23 @@ class TestCache(object):
     def shelve(self, mocker):
         value = {}
 
-        @contextmanager
+        class _Shelve(object):
-        def _shelve(path):
+            def __init__(self, path):
-            yield value
+                pass
 
-        mocker.patch('thefuck.utils.shelve.open', new_callable=lambda: _shelve)
+            def __setitem__(self, k, v):
+                value[k] = v
+
+            def __getitem__(self, k):
+                return value[k]
+
+            def get(self, k, v=None):
+                return value.get(k, v)
+
+            def close(self):
+                return
+
+        mocker.patch('thefuck.utils.shelve.open', new_callable=lambda: _Shelve)
         return value
 
     @pytest.fixture(autouse=True)
@@ -146,26 +158,25 @@ class TestCache(object):
 
         return fn
 
-    def test_with_blank_cache(self, shelve, fn):
+    @pytest.fixture
+    def key(self):
+        if six.PY3:
+            return 'tests.test_utils.<function TestCache.fn.<locals>.fn '
+        else:
+            return 'tests.test_utils.<function fn '
+
+    def test_with_blank_cache(self, shelve, fn, key):
         assert shelve == {}
         assert fn() == 'test'
-        assert shelve == {
+        assert shelve == {key: {'etag': '0', 'value': 'test'}}
-            'tests.test_utils.<function TestCache.fn.<locals>.fn ': {
-                'etag': '0', 'value': 'test'}}
 
-    def test_with_filled_cache(self, shelve, fn):
+    def test_with_filled_cache(self, shelve, fn, key):
-        cache_value = {
+        cache_value = {key: {'etag': '0', 'value': 'new-value'}}
-            'tests.test_utils.<function TestCache.fn.<locals>.fn ': {
-                'etag': '0', 'value': 'new-value'}}
         shelve.update(cache_value)
         assert fn() == 'new-value'
         assert shelve == cache_value
 
-    def test_when_etag_changed(self, shelve, fn):
+    def test_when_etag_changed(self, shelve, fn, key):
-        shelve.update({
+        shelve.update({key: {'etag': '-1', 'value': 'old-value'}})
-            'tests.test_utils.<function TestCache.fn.<locals>.fn ': {
-                'etag': '-1', 'value': 'old-value'}})
         assert fn() == 'test'
-        assert shelve == {
+        assert shelve == {key: {'etag': '0', 'value': 'test'}}
-            'tests.test_utils.<function TestCache.fn.<locals>.fn ': {
-                'etag': '0', 'value': 'test'}}

thefuck/utils.py

@@ -2,6 +2,7 @@ from difflib import get_close_matches
 from functools import wraps
 import shelve
 from decorator import decorator
+from contextlib import closing
 import tempfile
 
 import os
@@ -176,11 +177,13 @@ def cache(*depends_on):
             return fn(*args, **kwargs)
 
         cache_path = os.path.join(tempfile.gettempdir(), '.thefuck-cache')
+        # A bit obscure, but simplest way to generate unique key for
+        # functions and methods in python 2 and 3:
         key = '{}.{}'.format(fn.__module__, repr(fn).split('at')[0])
 
         etag = '.'.join(_get_mtime(name) for name in depends_on)
 
-        with shelve.open(cache_path) as db:
+        with closing(shelve.open(cache_path)) as db:
             if db.get(key, {}).get('etag') == etag:
                 return db[key]['value']
             else: