Degrades to copy on mutable types - Cheap is the mutable types are already COW.

1 files changed, 90 insertions, 9 deletions

diff --git a/lib/bb/COW.py b/lib/bb/COW.py
index ac0e21c74..d34dd5f24 100644
--- a/lib/bb/COW.py
+++ b/lib/bb/COW.py
@@ -8,9 +8,24 @@ Please Note: Be careful when using mutable types (ie Dict and Lists). The copy o
 
 from inspect import getmro
 
-class COWMeta(type):
+import copy
+import types, sets
+types.ImmutableTypes = tuple([ \
+    types.BooleanType, \
+    types.ComplexType, \
+    types.FloatType, \
+    types.IntType, \
+    types.LongType, \
+    types.NoneType, \
+    types.TupleType, \
+    sets.ImmutableSet] + \
+    list(types.StringTypes))
+
+MUTABLE = "_mutable__"
+ 
+class COWDictMeta(type):
     def __str__(cls):
-        return "<COW Level: %i Current Keys: %i>" % (cls.__count__, len(cls.__dict__))
+        return "<COWDict Level: %i Current Keys: %i>" % (cls.__count__, len(cls.__dict__))
     __repr__ = __str__
 
     def cow(cls):
@@ -19,16 +34,50 @@ class COWMeta(type):
         return C
 
     def __setitem__(cls, key, value):
+        if not isinstance(value, types.ImmutableTypes):
+            key += MUTABLE
         setattr(cls, key, value)
+    
+    def __getmutable__(cls, key):
+        """
+        This gets called when you do a "o.b" and b doesn't exist on o.
+
+        IE When the type is mutable.
+        """
+        nkey = key + MUTABLE
+        # Do we already have a copy of this on us?
+        if nkey in cls.__dict__:
+            return cls.__dict__[nkey]
+
+        r = getattr(cls, nkey)
+        try:
+            r = r.copy()
+        except NameError, e:
+            r = copy.copy(r)
+        setattr(cls, nkey, r)
+        return r
+
     def __getitem__(cls, key):
-        return getattr(cls, key)
-    def haskey(cls, key):
-        return hasattr(cls, key)
+        try:
+            try:
+                return getattr(cls, key)
+            except AttributeError:
+                # Degrade to slow mode if type is not immutable, 
+                # If the type is also a COW this will be cheap anyway
+                return cls.__getmutable__(key)
+        except AttributeError, e:
+            raise KeyError(e)
+
+    def has_key(cls, key):
+        return (hasattr(cls, key) or hasattr(cls, key + MUTABLE))
 
     def iter(cls, type):
         for key in dir(cls):
             if key.startswith("__"):
-                continue
+                    continue
+
+            if key.endswith(MUTABLE):
+                key = key[:-len(MUTABLE)]
 
             if type == "keys":
                 yield key
@@ -46,26 +95,58 @@ class COWMeta(type):
         return cls.iter("items")
     copy = cow
 
-class COWBase(object):
-    __metaclass__ = COWMeta
+class COWSetMeta(COWDictMeta):
+    def __str__(cls):
+        return "<COWSet Level: %i Current Keys: %i>" % (cls.__count__, len(cls.__dict__))
+    __repr__ = __str__
+
+    def cow(cls):
+        class C(cls):
+            __count__ = cls.__count__ + 1
+        return C
+
+    def add(cls, key, value):
+        cls.__setitem__(hash(key), value)
+
+class COWDictBase(object):
+    __metaclass__ = COWDictMeta
     __count__ = 0
 
+
 if __name__ == "__main__":
-    a = COWBase
+    a = COWDictBase.copy()
     print a
     a['a'] = 'a'
+    a['dict'] = {}
+
+    print "ha"
+    hasattr(a, "a")
+    print "ha"
     
     b = a.copy()
     print b
     b['b'] = 'b'
 
+    for x in a.iteritems():
+        print x
+    print "--"
     for x in b.iteritems():
         print x
     print
 
+    b['dict']['a'] = 'b'
     b['a'] = 'c'
 
+    for x in a.iteritems():
+        print x
+    print "--"
     for x in b.iteritems():
         print x
     print
 
+    try:
+        b['dict2']
+    except KeyError, e:
+        print "Okay!"
+
+