Add pickle support for tstrings (#75)

Author: sobolevn

Lib/string/templatelib.py

@@ -1,3 +1,23 @@
 """Support for template string literals (t-strings)."""
 
 from _templatelib import Interpolation, Template
+
+__all__ = [
+    "Interpolation",
+    "Template",
+]
+
+def _template_unpickle(*args):
+    import itertools
+
+    if len(args) != 2:
+        raise ValueError('Template expects tuple of length 2 to unpickle')
+
+    strings, interpolations = args
+    parts = []
+    for string, interpolation in itertools.zip_longest(strings, interpolations):
+        if string is not None:
+            parts.append(string)
+        if interpolation is not None:
+            parts.append(interpolation)
+    return Template(*parts)

Lib/test/test_string/_support.py

@@ -2,7 +2,7 @@
 from string.templatelib import Interpolation
 
 
-class TStringTestCase(unittest.TestCase):
+class TStringBaseCase:
     def assertTStringEqual(self, t, strings, interpolations):
         """Test template string literal equality.
 
@@ -23,7 +23,7 @@ def assertTStringEqual(self, t, strings, interpolations):
                 continue
 
             if len(exp) == 3:
-                self.assertEqual((i.value, i.expression, i.conversion),  exp)
+                self.assertEqual((i.value, i.expression, i.conversion), exp)
                 self.assertEqual(i.format_spec, '')
                 continue
 

Lib/test/test_string/test_templatelib.py

@@ -1,9 +1,22 @@
+import pickle
+import unittest
 from string.templatelib import Template, Interpolation
 
-from test.test_string._support import TStringTestCase, fstring
+from test.test_string._support import TStringBaseCase, fstring
 
 
-class TestTemplate(TStringTestCase):
+class TestTemplate(unittest.TestCase, TStringBaseCase):
+
+    def test_common(self):
+        self.assertEqual(type(t'').__name__, 'Template')
+        self.assertEqual(type(t'').__qualname__, 'Template')
+        self.assertEqual(type(t'').__module__, 'string.templatelib')
+
+        a = 'a'
+        i = t'{a}'.interpolations[0]
+        self.assertEqual(type(i).__name__, 'Interpolation')
+        self.assertEqual(type(i).__qualname__, 'Interpolation')
+        self.assertEqual(type(i).__module__, 'string.templatelib')
 
     def test_basic_creation(self):
         # Simple t-string creation
@@ -53,3 +66,44 @@ def test_creation_interleaving(self):
             t, ('', '', ''), [('Maria', 'name'), ('Python', 'language')]
         )
         self.assertEqual(fstring(t), 'MariaPython')
+
+    def test_pickle_template(self):
+        user = 'test'
+        for template in (
+            t'',
+            t"No values",
+            t'With inter {user}',
+            t'With ! {user!r}',
+            t'With format {1 / 0.3:.2f}',
+            Template(),
+            Template('a'),
+            Template(Interpolation('Nikita', 'name', None, '')),
+            Template('a', Interpolation('Nikita', 'name', 'r', '')),
+        ):
+            for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+                with self.subTest(proto=proto, template=template):
+                    pickled = pickle.dumps(template, protocol=proto)
+                    unpickled = pickle.loads(pickled)
+
+                    self.assertEqual(unpickled.values, template.values)
+                    self.assertEqual(fstring(unpickled), fstring(template))
+
+    def test_pickle_interpolation(self):
+        for interpolation in (
+            Interpolation('Nikita', 'name', None, ''),
+            Interpolation('Nikita', 'name', 'r', ''),
+            Interpolation(1/3, 'x', None, '.2f'),
+        ):
+            for proto in range(pickle.HIGHEST_PROTOCOL + 1):
+                with self.subTest(proto=proto, interpolation=interpolation):
+                    pickled = pickle.dumps(interpolation, protocol=proto)
+                    unpickled = pickle.loads(pickled)
+
+                    self.assertEqual(unpickled.value, interpolation.value)
+                    self.assertEqual(unpickled.expression, interpolation.expression)
+                    self.assertEqual(unpickled.conversion, interpolation.conversion)
+                    self.assertEqual(unpickled.format_spec, interpolation.format_spec)
+
+
+if __name__ == '__main__':
+    unittest.main()

Lib/test/test_tstring.py

@@ -1,9 +1,9 @@
 import unittest
 
-from test.test_string._support import TStringTestCase, fstring
+from test.test_string._support import TStringBaseCase, fstring
 
 
-class TestTString(TStringTestCase):
+class TestTString(unittest.TestCase, TStringBaseCase):
     def test_string_representation(self):
         # Test __repr__
         t = t"Hello"

Objects/interpolationobject.c

@@ -125,6 +125,22 @@ static PyMemberDef interpolation_members[] = {
     {NULL}
 };
 
+static PyObject*
+interpolation_reduce(PyObject *op, PyObject *Py_UNUSED(dummy))
+{
+    interpolationobject *self = interpolationobject_CAST(op);
+    return Py_BuildValue("(O(OOOO))", (PyObject *)Py_TYPE(op),
+                         self->value, self->expression,
+                         self->conversion, self->format_spec);
+}
+
+static PyMethodDef interpolation_methods[] = {
+    {"__reduce__", interpolation_reduce, METH_VARARGS,
+     PyDoc_STR("__reduce__() -> (cls, state)")},
+
+    {NULL,      NULL},
+};
+
 PyTypeObject _PyInterpolation_Type = {
     PyVarObject_HEAD_INIT(NULL, 0)
     .tp_name = "string.templatelib.Interpolation",
@@ -139,6 +155,7 @@ PyTypeObject _PyInterpolation_Type = {
     .tp_free = PyObject_GC_Del,
     .tp_repr = interpolation_repr,
     .tp_members = interpolation_members,
+    .tp_methods = interpolation_methods,
     .tp_traverse = interpolation_traverse,
 };
 

Objects/templateobject.c

@@ -114,7 +114,10 @@ template_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
             last_was_str = 0;
         }
         else {
-            PyErr_SetString(PyExc_TypeError, "Template.__new__ *args need to be of type 'str' or 'Interpolation'");
+            PyErr_Format(
+                PyExc_TypeError,
+                "Template.__new__ *args need to be of type 'str' or 'Interpolation', got %T",
+                item);
             return NULL;
         }
     }
@@ -418,14 +421,43 @@ static PyMemberDef template_members[] = {
 };
 
 static PyGetSetDef template_getset[] = {
-    {"values", template_values_get, NULL, "Values of interpolations", NULL},
+    {"values", template_values_get, NULL,
+     PyDoc_STR("Values of interpolations"), NULL},
     {NULL},
 };
 
 static PySequenceMethods template_as_sequence = {
     .sq_concat = _PyTemplate_Concat,
 };
 
+static PyObject*
+template_reduce(PyObject *op, PyObject *Py_UNUSED(dummy))
+{
+    PyObject *mod = PyImport_ImportModule("string.templatelib");
+    if (mod == NULL) {
+        return NULL;
+    }
+    PyObject *func = PyObject_GetAttrString(mod, "_template_unpickle");
+    Py_DECREF(mod);
+    if (func == NULL) {
+        return NULL;
+    }
+
+    templateobject *self = templateobject_CAST(op);
+    PyObject *result = Py_BuildValue("O(OO)",
+                                     func,
+                                     self->strings,
+                                     self->interpolations);
+
+    Py_DECREF(func);
+    return result;
+}
+
+static PyMethodDef template_methods[] = {
+    {"__reduce__", template_reduce, METH_NOARGS, NULL},
+    {NULL,      NULL},
+};
+
 PyTypeObject _PyTemplate_Type = {
     PyVarObject_HEAD_INIT(NULL, 0)
     .tp_name = "string.templatelib.Template",
@@ -441,6 +473,7 @@ PyTypeObject _PyTemplate_Type = {
     .tp_free = PyObject_GC_Del,
     .tp_repr = template_repr,
     .tp_members = template_members,
+    .tp_methods = template_methods,
     .tp_getset = template_getset,
     .tp_iter = template_iter,
     .tp_traverse = template_traverse,