Issue #15022: Add pickle and comparison support to types.SimpleNamespace.

ericsnowcurrently · ericsnowcurrently · commit b5c8f927829a · 2013-02-16T16:32:39.000-07:00
diff --git a/Doc/library/types.rst b/Doc/library/types.rst
@@ -212,6 +212,8 @@ Standard names are defined for the following types:
                keys = sorted(self.__dict__)
                items = ("{}={!r}".format(k, self.__dict__[k]) for k in keys)
                return "{}({})".format(type(self).__name__, ", ".join(items))
+           def __eq__(self, other):
+               return self.__dict__ == other.__dict__
 
    ``SimpleNamespace`` may be useful as a replacement for ``class NS: pass``.
    However, for a structured record type use :func:`~collections.namedtuple`
diff --git a/Lib/test/test_types.py b/Lib/test/test_types.py
@@ -2,6 +2,7 @@
 
 from test.support import run_unittest, run_with_locale
 import collections
+import pickle
 import locale
 import sys
 import types
@@ -1077,9 +1078,19 @@ def test_repr(self):
         ns2 = types.SimpleNamespace()
         ns2.x = "spam"
         ns2._y = 5
+        name = "namespace"
 
-        self.assertEqual(repr(ns1), "namespace(w=3, x=1, y=2)")
-        self.assertEqual(repr(ns2), "namespace(_y=5, x='spam')")
+        self.assertEqual(repr(ns1), "{name}(w=3, x=1, y=2)".format(name=name))
+        self.assertEqual(repr(ns2), "{name}(_y=5, x='spam')".format(name=name))
+
+    def test_equal(self):
+        ns1 = types.SimpleNamespace(x=1)
+        ns2 = types.SimpleNamespace()
+        ns2.x = 1
+
+        self.assertEqual(types.SimpleNamespace(), types.SimpleNamespace())
+        self.assertEqual(ns1, ns2)
+        self.assertNotEqual(ns2, types.SimpleNamespace())
 
     def test_nested(self):
         ns1 = types.SimpleNamespace(a=1, b=2)
@@ -1117,11 +1128,12 @@ def test_recursive_repr(self):
         ns1.spam = ns1
         ns2.spam = ns3
         ns3.spam = ns2
+        name = "namespace"
+        repr1 = "{name}(c='cookie', spam={name}(...))".format(name=name)
+        repr2 = "{name}(spam={name}(spam={name}(...), x=1))".format(name=name)
 
-        self.assertEqual(repr(ns1),
-                         "namespace(c='cookie', spam=namespace(...))")
-        self.assertEqual(repr(ns2),
-                         "namespace(spam=namespace(spam=namespace(...), x=1))")
+        self.assertEqual(repr(ns1), repr1)
+        self.assertEqual(repr(ns2), repr2)
 
     def test_as_dict(self):
         ns = types.SimpleNamespace(spam='spamspamspam')
@@ -1144,6 +1156,14 @@ class Spam(types.SimpleNamespace):
         self.assertIs(type(spam), Spam)
         self.assertEqual(vars(spam), {'ham': 8, 'eggs': 9})
 
+    def test_pickle(self):
+        ns = types.SimpleNamespace(breakfast="spam", lunch="spam")
+
+        ns_pickled = pickle.dumps(ns)
+        ns_roundtrip = pickle.loads(ns_pickled)
+
+        self.assertEqual(ns, ns_roundtrip)
+
 
 def test_main():
     run_unittest(TypesTests, MappingProxyTests, ClassCreationTests,
diff --git a/Misc/NEWS b/Misc/NEWS
@@ -252,6 +252,8 @@ Core and Builtins
 - Issue #15111: __import__ should propagate ImportError when raised as a
   side-effect of a module triggered from using fromlist.
 
+- Issue #15022: Add pickle and comparison support to types.SimpleNamespace.
+
 Library
 -------
 
diff --git a/Objects/namespaceobject.c b/Objects/namespaceobject.c
@@ -66,16 +66,20 @@ namespace_dealloc(_PyNamespaceObject *ns)
 
 
 static PyObject *
-namespace_repr(_PyNamespaceObject *ns)
+namespace_repr(PyObject *ns)
 {
     int i, loop_error = 0;
     PyObject *pairs = NULL, *d = NULL, *keys = NULL, *keys_iter = NULL;
     PyObject *key;
     PyObject *separator, *pairsrepr, *repr = NULL;
+    const char * name;
 
-    i = Py_ReprEnter((PyObject *)ns);
+    name = (Py_TYPE(ns) == &_PyNamespace_Type) ? "namespace"
+                                               : ns->ob_type->tp_name;
+
+    i = Py_ReprEnter(ns);
     if (i != 0) {
-        return i > 0 ? PyUnicode_FromString("namespace(...)") : NULL;
+        return i > 0 ? PyUnicode_FromFormat("%s(...)", name) : NULL;
     }
 
     pairs = PyList_New(0);
@@ -127,16 +131,15 @@ namespace_repr(_PyNamespaceObject *ns)
     if (pairsrepr == NULL)
         goto error;
 
-    repr = PyUnicode_FromFormat("%s(%S)",
-                                ((PyObject *)ns)->ob_type->tp_name, pairsrepr);
+    repr = PyUnicode_FromFormat("%s(%S)", name, pairsrepr);
     Py_DECREF(pairsrepr);
 
 error:
     Py_XDECREF(pairs);
     Py_XDECREF(d);
     Py_XDECREF(keys);
     Py_XDECREF(keys_iter);
-    Py_ReprLeave((PyObject *)ns);
+    Py_ReprLeave(ns);
 
     return repr;
 }
@@ -158,14 +161,26 @@ namespace_clear(_PyNamespaceObject *ns)
 }
 
 
+static PyObject *
+namespace_richcompare(PyObject *self, PyObject *other, int op)
+{
+    if (PyObject_IsInstance(self, (PyObject *)&_PyNamespace_Type) &&
+            PyObject_IsInstance(other, (PyObject *)&_PyNamespace_Type))
+        return PyObject_RichCompare(((_PyNamespaceObject *)self)->ns_dict,
+                                   ((_PyNamespaceObject *)other)->ns_dict, op);
+    Py_INCREF(Py_NotImplemented);
+    return Py_NotImplemented;
+}
+
+
 PyDoc_STRVAR(namespace_doc,
 "A simple attribute-based namespace.\n\
 \n\
-namespace(**kwargs)");
+SimpleNamespace(**kwargs)");
 
 PyTypeObject _PyNamespace_Type = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)
-    "namespace",                                /* tp_name */
+    "types.SimpleNamespace",                    /* tp_name */
     sizeof(_PyNamespaceObject),                 /* tp_size */
     0,                                          /* tp_itemsize */
     (destructor)namespace_dealloc,              /* tp_dealloc */
@@ -188,7 +203,7 @@ PyTypeObject _PyNamespace_Type = {
     namespace_doc,                              /* tp_doc */
     (traverseproc)namespace_traverse,           /* tp_traverse */
     (inquiry)namespace_clear,                   /* tp_clear */
-    0,                                          /* tp_richcompare */
+    namespace_richcompare,                      /* tp_richcompare */
     0,                                          /* tp_weaklistoffset */
     0,                                          /* tp_iter */
     0,                                          /* tp_iternext */
