From e34c21c2a006e4ac950a519e1e662750f3e09e79 Mon Sep 17 00:00:00 2001
From: Georg Brandl <georg@python.org>
Date: Sun, 30 Mar 2008 20:20:39 +0000
Subject: Make AST nodes pickleable.

---
 Lib/test/test_ast.py | 14 ++++++++++++++
 Parser/asdl_c.py     | 29 +++++++++++++++++++++++++++--
 Python/Python-ast.c  | 29 +++++++++++++++++++++++++++--
 3 files changed, 68 insertions(+), 4 deletions(-)

diff --git a/Lib/test/test_ast.py b/Lib/test/test_ast.py
index e068b0a..9d2bd66 100644
--- a/Lib/test/test_ast.py
+++ b/Lib/test/test_ast.py
@@ -166,6 +166,20 @@ class AST_Tests(unittest.TestCase):
         # this used to fail because Sub._fields was None
         x = _ast.Sub()
 
+    def test_pickling(self):
+        import pickle
+        mods = [pickle]
+        try:
+            import cPickle
+            mods.append(cPickle)
+        except ImportError:
+            pass
+        protocols = [0, 1, 2]
+        for mod in mods:
+            for protocol in protocols:
+                for ast in (compile(i, "?", "exec", 0x400) for i in exec_tests):
+                    ast2 = mod.loads(mod.dumps(ast, protocol))
+                    self.assertEquals(to_tuple(ast2), to_tuple(ast))
 
 def test_main():
     test_support.run_unittest(AST_Tests)
diff --git a/Parser/asdl_c.py b/Parser/asdl_c.py
index 29e2547..7e12ea6 100755
--- a/Parser/asdl_c.py
+++ b/Parser/asdl_c.py
@@ -629,9 +629,34 @@ ast_type_init(PyObject *self, PyObject *args, PyObject *kw)
     return res;
 }
 
+/* Pickling support */
+static PyObject *
+ast_type_reduce(PyObject *self, PyObject *unused)
+{
+    PyObject *res;
+    PyObject *dict = PyObject_GetAttrString(self, "__dict__");
+    if (dict == NULL) {
+        if (PyErr_ExceptionMatches(PyExc_AttributeError))
+            PyErr_Clear();
+        else
+            return NULL;
+    }
+    if (dict) {
+        res = Py_BuildValue("O()O", Py_TYPE(self), dict);
+        Py_DECREF(dict);
+        return res;
+    }
+    return Py_BuildValue("O()", Py_TYPE(self));
+}
+
+static PyMethodDef ast_type_methods[] = {
+    {"__reduce__", ast_type_reduce, METH_NOARGS, NULL},
+    {NULL}
+};
+
 static PyTypeObject AST_type = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)
-    "AST",
+    "_ast.AST",
     sizeof(PyObject),
     0,
     0,                       /* tp_dealloc */
@@ -657,7 +682,7 @@ static PyTypeObject AST_type = {
     0,                       /* tp_weaklistoffset */
     0,                       /* tp_iter */
     0,                       /* tp_iternext */
-    0,                       /* tp_methods */
+    ast_type_methods,        /* tp_methods */
     0,                       /* tp_members */
     0,                       /* tp_getset */
     0,                       /* tp_base */
diff --git a/Python/Python-ast.c b/Python/Python-ast.c
index ba310e6..f958145 100644
--- a/Python/Python-ast.c
+++ b/Python/Python-ast.c
@@ -420,9 +420,34 @@ ast_type_init(PyObject *self, PyObject *args, PyObject *kw)
     return res;
 }
 
+/* Pickling support */
+static PyObject *
+ast_type_reduce(PyObject *self, PyObject *unused)
+{
+    PyObject *res;
+    PyObject *dict = PyObject_GetAttrString(self, "__dict__");
+    if (dict == NULL) {
+        if (PyErr_ExceptionMatches(PyExc_AttributeError))
+            PyErr_Clear();
+        else
+            return NULL;
+    }
+    if (dict) {
+        res = Py_BuildValue("O()O", Py_TYPE(self), dict);
+        Py_DECREF(dict);
+        return res;
+    }
+    return Py_BuildValue("O()", Py_TYPE(self));
+}
+
+static PyMethodDef ast_type_methods[] = {
+    {"__reduce__", ast_type_reduce, METH_NOARGS, NULL},
+    {NULL}
+};
+
 static PyTypeObject AST_type = {
     PyVarObject_HEAD_INIT(&PyType_Type, 0)
-    "AST",
+    "_ast.AST",
     sizeof(PyObject),
     0,
     0,                       /* tp_dealloc */
@@ -448,7 +473,7 @@ static PyTypeObject AST_type = {
     0,                       /* tp_weaklistoffset */
     0,                       /* tp_iter */
     0,                       /* tp_iternext */
-    0,                       /* tp_methods */
+    ast_type_methods,        /* tp_methods */
     0,                       /* tp_members */
     0,                       /* tp_getset */
     0,                       /* tp_base */
-- 
cgit v0.12