[pypy-commit] pypy identity-dict-strategy: first version of the strategy for instances which compares by identity. Broken in case we mutate the class after they are already in the dict

[antocuni]

Author: Antonio Cuni <anto.cuni at gmail.com>
Branch: identity-dict-strategy
Changeset: r45758:beb63ba8a93a
Date: 2011-07-20 10:07 +0200
http://bitbucket.org/pypy/pypy/changeset/beb63ba8a93a/

Log:	first version of the strategy for instances which compares by
	identity. Broken in case we mutate the class after they are already
	in the dict

diff --git a/pypy/objspace/std/dictmultiobject.py b/pypy/objspace/std/dictmultiobject.py
--- a/pypy/objspace/std/dictmultiobject.py
+++ b/pypy/objspace/std/dictmultiobject.py
@@ -157,11 +157,15 @@
        return self.erase(None)
 
     def switch_to_correct_strategy(self, w_dict, w_key):
-        #XXX implement other strategies later
+        trackcomparebyidentity = self.space.config.objspace.std.trackcomparebyidentity
         if type(w_key) is self.space.StringObjectCls:
             self.switch_to_string_strategy(w_dict)
-        elif self.space.is_w(self.space.type(w_key), self.space.w_int):
+            return
+        w_type = self.space.type(w_key)
+        if self.space.is_w(w_type, self.space.w_int):
             self.switch_to_int_strategy(w_dict)
+        elif trackcomparebyidentity and w_type.compares_by_identity():
+            self.switch_to_identity_strategy(w_dict)
         else:
             self.switch_to_object_strategy(w_dict)
 
@@ -177,6 +181,12 @@
         w_dict.strategy = strategy
         w_dict.dstorage = storage
 
+    def switch_to_identity_strategy(self, w_dict):
+        strategy = self.space.fromcache(IdentityDictStrategy)
+        storage = strategy.get_empty_storage()
+        w_dict.strategy = strategy
+        w_dict.dstorage = storage
+
     def switch_to_object_strategy(self, w_dict):
         strategy = self.space.fromcache(ObjectDictStrategy)
         storage = strategy.get_empty_storage()
@@ -338,7 +348,6 @@
 
     def getitem(self, w_dict, w_key):
         space = self.space
-
         if self.is_correct_type(w_key):
             return self.unerase(w_dict.dstorage).get(self.unwrap(w_key), None)
         elif self._never_equal_to(space.type(w_key)):
@@ -404,6 +413,23 @@
     def keys(self, w_dict):
         return self.unerase(w_dict.dstorage).keys()
 
+
+class IdentityDictStrategy(ObjectDictStrategy):
+    """
+    Strategy for custom instances which compares by identity (i.e., the
+    default unless you override __hash__, __eq__ or __cmp__).  The storage is
+    just a normal RPython dict, which has already the correct by-identity
+    semantics.
+    """
+
+    def is_correct_type(self, w_obj):
+        w_type = self.space.type(w_obj)
+        return w_type.compares_by_identity()
+
+    def get_empty_storage(self):
+        return self.erase({})
+
+
 class StringDictStrategy(AbstractTypedStrategy, DictStrategy):
 
     erase, unerase = rerased.new_erasing_pair("string")
diff --git a/pypy/objspace/std/test/test_dictmultiobject.py b/pypy/objspace/std/test/test_dictmultiobject.py
--- a/pypy/objspace/std/test/test_dictmultiobject.py
+++ b/pypy/objspace/std/test/test_dictmultiobject.py
@@ -1105,3 +1105,60 @@
     fakespace = FakeSpace()
     d = fakespace.newdict(module=True)
     assert type(d.strategy) is StringDictStrategy
+
+
+class AppTestIdentityDict(object):
+    def setup_class(cls):
+        cls.space = gettestobjspace(**{"objspace.std.trackcomparebyidentity": True})
+        if option.runappdirect:
+            py.test.skip("__repr__ doesn't work on appdirect")
+
+    def w_uses_identity_strategy(self, obj):
+        import __pypy__
+        return "IdentityDictStrategy" in __pypy__.internal_repr(obj)
+
+    def test_use_strategy(self):
+        class X(object):
+            pass
+        d = {}
+        x = X()
+        d[x] = 1
+        assert self.uses_identity_strategy(d)
+        assert d[x] == 1
+
+    def test_bad_item(self):
+        class X(object):
+            pass
+        class Y(object):
+            def __hash__(self):
+                return 32
+
+        d = {}
+        x = X()
+        y = Y()
+        d[x] = 1
+        assert self.uses_identity_strategy(d)
+        d[y] = 2
+        assert not self.uses_identity_strategy(d)
+        assert d[x] == 1
+        assert d[y] == 2
+
+    def test_bad_key(self):
+        class X(object):
+            pass
+        d = {}
+        x = X()
+
+        class Y(object):
+            def __hash__(self):
+                return hash(x) # to make sure we do x == y
+
+            def __eq__(self, other):
+                return True
+
+        y = Y()
+        d[x] = 1
+        assert self.uses_identity_strategy(d)
+        assert d[y] == 1
+        assert not self.uses_identity_strategy(d)
+