Implement split method

atpalmer · atpalmer · commit f6312b546b5a · 2021-03-31T14:45:22.000-05:00
diff --git a/src/cstring.c b/src/cstring.c
@@ -580,6 +580,33 @@ PyObject *cstring_startswith(PyObject *self, PyObject *args) {
     return PyBool_FromLong(cmp == 0);
 }
 
+PyDoc_STRVAR(strip__doc__, "");
+PyObject *cstring_strip(PyObject *self, PyObject *args) {
+    PyObject *charsobj = NULL;
+    if(!PyArg_ParseTuple(args, "|O", &charsobj))
+        return NULL;
+
+    const char *chars = " \t\n\v\f\r";
+
+    if(charsobj && charsobj != Py_None) {
+        if(!PyUnicode_Check(charsobj))
+            return _bad_argument_type(charsobj);
+        chars = PyUnicode_AsUTF8(charsobj);
+    }
+
+    const char *start = CSTRING_VALUE(self);
+    while(strchr(chars, *start))
+        ++start;
+
+    const char *end = &CSTRING_LAST_BYTE(self) - 1;
+    while(strchr(chars, *end))
+        --end;
+
+    Py_ssize_t newsize = end - start + 1;
+
+    return _cstring_new(Py_TYPE(self), start, newsize);
+}
+
 PyDoc_STRVAR(endswith__doc__, "");
 PyObject *cstring_endswith(PyObject *self, PyObject *args) {
     struct _substr_params params;
@@ -680,7 +707,7 @@ static PyMethodDef cstring_methods[] = {
     /* TODO: split */
     /* TODO: splitlines */
     {"startswith", cstring_startswith, METH_VARARGS, startswith__doc__},
-    /* TODO: strip */
+    {"strip", cstring_strip, METH_VARARGS, strip__doc__},
     {"swapcase", cstring_swapcase, METH_NOARGS, swapcase__doc__},
     /* TODO: title */
     /* TODO: translate */
diff --git a/test/test_methods.py b/test/test_methods.py
@@ -190,6 +190,21 @@ def test_rindex_missing():
         return target.rindex('lo', 0, 4)
 
 
+def test_strip_noargs():
+    target = cstring('\r\n\n\n\t  hello, world  \t\r\n\n  ')
+    assert target.strip() == cstring('hello, world')
+
+
+def test_strip_None():
+    target = cstring('\r\n\n\n\t  hello, world  \t\r\n\n  ')
+    assert target.strip(None) == cstring('hello, world')
+
+
+def test_strip_arg():
+    target = cstring('hello, world')
+    assert target.strip('held') == cstring('o, wor')
+
+
 def test_startswith():
     target = cstring('hello, world')
     assert target.startswith('hello,') is True
