jinja2/_speedups.c

   1 /**
   2  * jinja2._speedups
   3  * ~~~~~~~~~~~~~~~~
   4  *
   5  * This module implements a few functions in C for better performance.
   6  *
   7  * :copyright: 2008 by Armin Ronacher.
   8  * :license: BSD.
   9  */
  10
  11 #include <Python.h>
  12
  13
  14 static const char *samp = "&amp;", *slt = "&lt;", *sgt = "&gt;", *sqt = "&quot;";
  15 static Py_UNICODE *amp, *lt, *gt, *qt;
  16 static PyObject* markup;
  17
  18
  19 static int
  20 init_constants(void)
  21 {
  22         amp = ((PyUnicodeObject*)PyUnicode_DecodeASCII(samp, 5, NULL))->str;
  23         lt = ((PyUnicodeObject*)PyUnicode_DecodeASCII(slt, 4, NULL))->str;
  24         gt = ((PyUnicodeObject*)PyUnicode_DecodeASCII(sgt, 4, NULL))->str;
  25         qt = ((PyUnicodeObject*)PyUnicode_DecodeASCII(sqt, 6, NULL))->str;
  26
  27         PyObject *module = PyImport_ImportModule("jinja2.utils");
  28         if (!module)
  29                 return 0;
  30         markup = PyObject_GetAttrString(module, "Markup");
  31         Py_DECREF(module);
  32
  33         return 1;
  34 }
  35
  36 static PyObject*
  37 escape_unicode(PyUnicodeObject *in)
  38 {
  39         PyUnicodeObject *out;
  40         Py_UNICODE *outp;
  41
  42         /* First we need to figure out how long the escaped string will be */
  43         int len = 0, erepl = 0, repl = 0;
  44         Py_UNICODE *inp = in->str;
  45         while (*(inp) || in->length > inp - in->str)
  46                 switch (*inp++) {
  47                 case '&':
  48                         len += 5;
  49                         ++erepl;
  50                         break;
  51                 case '"':
  52                         len += 6;
  53                         ++erepl;
  54                         break;
  55                 case '<':
  56                 case '>':
  57                         len += 4;
  58                         ++erepl;
  59                         break;
  60                 default:
  61                         ++len;
  62                 }
  63
  64         /* Do we need to escape anything at all? */
  65         if (!erepl) {
  66                 Py_INCREF(in);
  67                 return (PyObject*)in;
  68         }
  69
  70         out = (PyUnicodeObject*)PyUnicode_FromUnicode(NULL, len);
  71         if (!out)
  72                 return NULL;
  73
  74         outp = out->str;
  75         inp = in->str;
  76         while (*(inp) || in->length > inp - in->str) {
  77                 /* copy rest of string if we have replaced everything */
  78                 if (repl == erepl) {
  79                         Py_UNICODE_COPY(outp, inp, in->length - (inp - in->str));
  80                         break;
  81                 }
  82                 /* regular replacements */
  83                 switch (*inp) {
  84                 case '&':
  85                         Py_UNICODE_COPY(outp, amp, 5);
  86                         outp += 5;
  87                         ++repl;
  88                         break;
  89                 case '"':
  90                         Py_UNICODE_COPY(outp, qt, 6);
  91                         outp += 6;
  92                         ++repl;
  93                         break;
  94                 case '<':
  95                         Py_UNICODE_COPY(outp, lt, 4);
  96                         outp += 4;
  97                         ++repl;
  98                         break;
  99                 case '>':
 100                         Py_UNICODE_COPY(outp, gt, 4);
 101                         outp += 4;
 102                         ++repl;
 103                         break;
 104                 default:
 105                         *outp++ = *inp;
 106                 };
 107                 ++inp;
 108         }
 109
 110         return (PyObject*)out;
 111 }
 112
 113
 114 static PyObject*
 115 soft_unicode(PyObject *self, PyObject *s)
 116 {
 117         if (!PyUnicode_Check(s))
 118                 return PyObject_Unicode(s);
 119         Py_INCREF(s);
 120         return s;
 121 }
 122
 123
 124 static PyObject*
 125 escape(PyObject *self, PyObject *text)
 126 {
 127         PyObject *s = NULL, *rv = NULL;
 128
 129         /* we don't have to escape integers, bools or floats */
 130         if (PyInt_CheckExact(text) || PyLong_CheckExact(text) ||
 131             PyFloat_CheckExact(text) || PyBool_Check(text) ||
 132             text == Py_None) {
 133                 PyObject *args = PyTuple_New(1);
 134                 if (!args) {
 135                         Py_DECREF(s);
 136                         return NULL;
 137                 }
 138                 PyTuple_SET_ITEM(args, 0, text);
 139                 return PyObject_CallObject(markup, args);
 140         }
 141
 142         /* if the object has an __html__ method that performs the escaping */
 143         PyObject *html = PyObject_GetAttrString(text, "__html__");
 144         if (html) {
 145                 rv = PyObject_CallObject(html, NULL);
 146                 Py_DECREF(html);
 147                 return rv;
 148         }
 149
 150         /* otherwise make the object unicode if it isn't, then escape */
 151         PyErr_Clear();
 152         if (!PyUnicode_Check(text)) {
 153                 PyObject *unicode = PyObject_Unicode(text);
 154                 if (!unicode)
 155                         return NULL;
 156                 s = escape_unicode((PyUnicodeObject*)unicode);
 157                 Py_DECREF(unicode);
 158         }
 159         else
 160                 s = escape_unicode((PyUnicodeObject*)text);
 161
 162         /* convert the unicode string into a markup object. */
 163         return PyObject_CallFunctionObjArgs(markup, (PyObject*)s, NULL);
 164 }
 165
 166
 167 static PyObject *
 168 tb_set_next(PyObject *self, PyObject *args)
 169 {
 170         PyTracebackObject *tb, *old;
 171         PyObject *next;
 172
 173         if (!PyArg_ParseTuple(args, "O!O:tb_set_next", &PyTraceBack_Type, &tb, &next))
 174                 return NULL;
 175         if (next == Py_None)
 176                 next = NULL;
 177         else if (!PyTraceBack_Check(next)) {
 178                 PyErr_SetString(PyExc_TypeError,
 179                                 "tb_set_next arg 2 must be traceback or None");
 180                 return NULL;
 181         }
 182         else
 183                 Py_INCREF(next);
 184
 185         old = tb->tb_next;
 186         tb->tb_next = (PyTracebackObject*)next;
 187         Py_XDECREF(old);
 188
 189         Py_INCREF(Py_None);
 190         return Py_None;
 191 }
 192
 193
 194 static PyMethodDef module_methods[] = {
 195         {"escape", (PyCFunction)escape, METH_O,
 196          "escape(s) -> string\n\n"
 197          "Convert the characters &, <, >, and \" in string s to HTML-safe\n"
 198          "sequences. Use this if you need to display text that might contain\n"
 199          "such characters in HTML."},
 200         {"soft_unicode", (PyCFunction)soft_unicode, METH_O,
 201          "soft_unicode(object) -> string\n\n"
 202          "Make a string unicode if it isn't already.  That way a markup\n"
 203          "string is not converted back to unicode."},
 204         {"tb_set_next", (PyCFunction)tb_set_next, METH_VARARGS,
 205          "Set the tb_next member of a traceback object."},
 206         {NULL, NULL, 0, NULL}           /* Sentinel */
 207 };
 208
 209
 210 #ifndef PyMODINIT_FUNC  /* declarations for DLL import/export */
 211 #define PyMODINIT_FUNC void
 212 #endif
 213 PyMODINIT_FUNC
 214 init_speedups(void)
 215 {
 216         if (!init_constants())
 217                 return;
 218
 219         Py_InitModule3("jinja2._speedups", module_methods, "");
 220 }