Support explicit concatenation of template and str

Author: lysnikolaou

Objects/templateobject.c

@@ -50,18 +50,92 @@ template_repr(templateobject *self)
                                 self->args);
 }
 
+static PyObject *
+template_add_template_str(templateobject *template, PyUnicodeObject *str, int templateleft)
+{
+    Py_ssize_t templatesize = PyTuple_GET_SIZE(template->args);
+
+    PyObject *tuple = PyTuple_New(templatesize + 1);
+    if (!tuple) {
+        return NULL;
+    }
+
+    Py_ssize_t i = 0;
+    Py_ssize_t j = 0;
+    if (!templateleft) {
+        PyTuple_SET_ITEM(tuple, i++, Py_NewRef(str));
+    }
+    for (j = 0; j < templatesize; j++) {
+        PyTuple_SET_ITEM(tuple, i + j, Py_NewRef(PyTuple_GET_ITEM(template->args, j)));
+    }
+    if (templateleft) {
+        PyTuple_SET_ITEM(tuple, i + j, Py_NewRef(str));
+    }
+
+    PyObject *newtemplate = PyObject_CallOneArg((PyObject *) &_PyTemplate_Type, tuple);
+    Py_DECREF(tuple);
+    return newtemplate;
+}
+
+static PyObject *
+template_add_templates(templateobject *self, templateobject *other)
+{
+    Py_ssize_t selfsize = PyTuple_GET_SIZE(self->args);
+    Py_ssize_t othersize = PyTuple_GET_SIZE(other->args);
+
+    PyObject *tuple = PyTuple_New(selfsize + othersize);
+    if (!tuple) {
+        return NULL;
+    }
+
+    Py_ssize_t i;
+    for (i = 0; i < selfsize; i++) {
+        PyTuple_SET_ITEM(tuple, i, Py_NewRef(PyTuple_GET_ITEM(self->args, i)));
+    }
+    for (Py_ssize_t j = 0; j < othersize; j++) {
+        PyTuple_SET_ITEM(tuple, i + j, Py_NewRef(PyTuple_GET_ITEM(other->args, j)));
+    }
+
+    PyObject *newtemplate = PyObject_CallOneArg((PyObject *) &_PyTemplate_Type, tuple);
+    Py_DECREF(tuple);
+    return newtemplate;
+}
+
+static PyObject *
+template_add(PyObject *self, PyObject *other)
+{
+    if (PyObject_TypeCheck(self, &_PyTemplate_Type) &&
+            PyObject_TypeCheck(other, &_PyTemplate_Type)) {
+        return template_add_templates((templateobject *) self, (templateobject *) other);
+    }
+    else if (PyObject_TypeCheck(self, &_PyTemplate_Type) && PyUnicode_Check(other)) {
+        return template_add_template_str((templateobject *) self, (PyUnicodeObject *) other, 1);
+    }
+    else if (PyUnicode_Check(self) && PyObject_TypeCheck(other, &_PyTemplate_Type)) {
+        return template_add_template_str((templateobject *) other, (PyUnicodeObject *) self, 0);
+    }
+    else {
+        Py_RETURN_NOTIMPLEMENTED;
+    }
+}
+
 static PyMemberDef template_members[] = {
     {"args", Py_T_OBJECT_EX, offsetof(templateobject, args), Py_READONLY, "Args"},
     {NULL}
 };
 
+static PyNumberMethods template_as_number = {
+    .nb_add = template_add
+};
+
 PyTypeObject _PyTemplate_Type = {
     PyVarObject_HEAD_INIT(NULL, 0)
     .tp_name = "templatelib.Template",
     .tp_doc = PyDoc_STR("Template object"),
     .tp_basicsize = sizeof(templateobject),
     .tp_itemsize = 0,
     .tp_flags = Py_TPFLAGS_DEFAULT | Py_TPFLAGS_BASETYPE | _Py_TPFLAGS_MATCH_SELF,
+    .tp_as_number = &template_as_number,
     .tp_new = (newfunc) template_new,
     .tp_dealloc = (destructor) template_dealloc,
     .tp_repr = (reprfunc) template_repr,