[pypy-svn] pypy collections-module: Write defaultdict... cheating.

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: collections-module
Changeset: r41986:2a2352811697
Date: 2011-02-15 18:31 +0100
http://bitbucket.org/pypy/pypy/changeset/2a2352811697/

Log:	Write defaultdict... cheating.

diff --git a/pypy/module/_collections/__init__.py b/pypy/module/_collections/__init__.py
--- a/pypy/module/_collections/__init__.py
+++ b/pypy/module/_collections/__init__.py
@@ -3,7 +3,9 @@
 from pypy.interpreter.mixedmodule import MixedModule
 
 class Module(MixedModule):
-    appleveldefs = {}
+    appleveldefs = {
+        'defaultdict': 'app_defaultdict.defaultdict',
+        }
 
     interpleveldefs = {
         'deque' : 'interp_deque.W_Deque',

diff --git a/pypy/module/_collections/app_defaultdict.py b/pypy/module/_collections/app_defaultdict.py
new file mode 100644
--- /dev/null
+++ b/pypy/module/_collections/app_defaultdict.py
@@ -0,0 +1,56 @@
+# NOT_RPYTHON
+
+# For now this is here, living at app-level.
+#
+# The issue is that for now we don't support writing interp-level
+# subclasses of Wrappable that inherit at app-level from a type like
+# 'dict'.
+
+
+class defaultdict(dict):
+    
+    def __init__(self, *args, **kwds):
+        self.default_factory = None
+        if 'default_factory' in kwds:
+            self.default_factory = kwds.pop('default_factory')
+        elif len(args) > 0 and (callable(args[0]) or args[0] is None):
+            self.default_factory = args[0]
+            args = args[1:]
+        super(defaultdict, self).__init__(*args, **kwds)
+ 
+    def __missing__(self, key):
+        # from defaultdict docs
+        if self.default_factory is None: 
+            raise KeyError(key)
+        self[key] = value = self.default_factory()
+        return value
+
+    def __repr__(self, recurse=set()):
+        # XXX not thread-safe, but good enough
+        if id(self) in recurse:
+            return "defaultdict(...)"
+        try:
+            recurse.add(id(self))
+            return "defaultdict(%s, %s)" % (repr(self.default_factory), super(defaultdict, self).__repr__())
+        finally:
+            recurse.remove(id(self))
+
+    def copy(self):
+        return type(self)(self, default_factory=self.default_factory)
+    
+    def __copy__(self):
+        return self.copy()
+
+    def __reduce__(self):
+        """
+        __reduce__ must return a 5-tuple as follows:
+
+           - factory function
+           - tuple of args for the factory function
+           - additional state (here None)
+           - sequence iterator (here None)
+           - dictionary iterator (yielding successive (key, value) pairs
+
+           This API is used by pickle.py and copy.py.
+        """
+        return (type(self), (self.default_factory,), None, None, self.iteritems())

diff --git a/pypy/module/_collections/test/test_defaultdict.py b/pypy/module/_collections/test/test_defaultdict.py
new file mode 100644
--- /dev/null
+++ b/pypy/module/_collections/test/test_defaultdict.py
@@ -0,0 +1,18 @@
+import py
+from pypy.conftest import gettestobjspace
+
+class AppTestBasic:
+    def setup_class(cls):
+        cls.space = gettestobjspace(usemodules=['_collections'])
+
+    def test_basics(self):
+        from _collections import defaultdict
+        d = defaultdict(list)
+        l = d[5]
+        d[5].append(42)
+        d[5].append(43)
+        assert l == [42, 43]
+        l2 = []
+        d[5] = l2
+        d[5].append(44)
+        assert l == [42, 43] and l2 == [44]