简单的Python源码分析——StringIO

简单的py源码分析——StringIO

近几日在架构KolaWriter时，涉及到了相关的内存文本缓冲功能的实现，因此尝试着参考了一下Python中StringIO的实现方式。于是有了这篇文章www

StringIO是一个来自Python标准库io的类。它会在内存中模拟一个以w+方式打开的文件对象。你应该也听说过这样的传说，那就是对于大量的字符串拼接操作来说，使用StringIO是快于直接进行字符串相加的。这里我也进行了一个小测试:

def str_stringio(epoch: int) -> str:
    s = StringIO()
    for _ in range(epoch):
        s.write('a' * np.random.randint(1, 100))
    return s.getvalue()

def str_stringio_r(epoch: int) -> str:
    s = StringIO('a')
    for _ in range(epoch):
        s.write('a' * np.random.randint(1, 100))
    return s.getvalue()

def str_concat(epoch: int) -> str:
    s = ''
    for _ in range(epoch):
        s += 'a' * np.random.randint(1, 100)
    return s

def str_ck(epoch: int) -> str:
    for _ in range(epoch):
        'a' * np.random.randint(1, 100)
    return ''

测试环境为Python3.8.12，这里我取epoch=number=100，利用timeit计算100次运行时间。我得到了如下的结果：

E(X) S(X) min(X) max(X)
stringio 0.0223633 3.907e-07 0.0213613 0.0235901
stringio-realized 0.0239235 1.50646e-06 0.0226005 0.028129
concat 0.0243483 5.96924e-07 0.0228622 0.0267949
CK 0.0214586 3.95554e-07 0.0204597 0.0230739

取各组平均值，减去空白对照组可以得到平均净用时，再除以净用时最长的concat组，可以得到下表：

类型 stringio stringio-realized concat
净用时之比 0.313066 0.853 1

可以看到，两者的差距确实存在，利用StringIO进行字符串拼接的用时仅为使用字符串相加用时的三成。当然，我们今天的主要目的不是为了测试StringIO有多好用。相信各位也注意到了一个有趣的事情，我们的stringio-r组，在测试代码上仅比stringio组多了一个初始值，而它的净运行时间比stringio组足足多了一倍不止，这又是为什么呢？

要回答这个问题，就需要我们把目光投向cpython解释器的源码部分，来看一看stringio是如何工作的。

基础结构

StringIO的源码位于Modules/_io/stringio.c。作为一个C级对象，我们首先来看StringIO的object struct定义：

接下来的代码来自 https://github.com/python/cpython 的main分支，本文写作时的版本号为Python 3.12.0 Alpha 4。下同

typedef struct {
    PyObject_HEAD
    Py_UCS4 *buf;
    Py_ssize_t pos;
    Py_ssize_t string_size;
    size_t buf_size;

    /*  stringio 对象可以处于两种状态：正在积累或已实现。
        在累积状态下，内部缓冲区不包含任何内容，
        并且内容由嵌入式 _PyUnicodeWriter 结构给出。
        在实现状态下，内部缓冲区是有意义的，并且
        _PyUnicodeWriter 被销毁。
    */
    int state;          /* ACCUMULATING / REALIZED */
    _PyUnicodeWriter writer;

    /* 以下几个char类型的字段实际上均被视为布尔型 */
    char ok;            /* 是否已经初始化 */
    char closed;        /* 文件是否已经关闭 */
    char readuniversal; /* 是否使用decoder */
    char readtranslate; /* 是否翻译换行符 */
    PyObject *decoder;  /* type: IncrementalNewlineDecoder */
    PyObject *readnl;   /* type: str */
    PyObject *writenl;  /* type: str */

    PyObject *dict;         /* __dict__ */
    PyObject *weakreflist;  /* __weakref__ */
} stringio;

这里我给出了一个大概的字段说明。其中靠后的几个字段，如decoder等主要用来进行换行符的处理，日常使用中涉及较少，不作为本次的主要研究内容。我们主要关注的是以下两个字段：buf及writer。

CPython中的字符与字符串

你可能会对这两个字段的类型有些陌生。这里我简单介绍一下CPython中的字符串表示形式。对于Python3来说，Python中的str类型在C中的名称不是str，也不是PyStr，而是PyUnicode。其object struct如下：

这部分的声明代码位于Include/cpython/unicodeobject.h，不需要特地看源码，在cpython解释器所在的文件夹中即可找到

/* Object format for Unicode subclasses. */
typedef struct {
    PyCompactUnicodeObject _base; /* 这里是继承了CompactUnicode这个结构体 */
    union {
        void *any;
        Py_UCS1 *latin1;
        Py_UCS2 *ucs2;
        Py_UCS4 *ucs4;
    } data;                     /* Canonical, smallest-form Unicode buffer */
} PyUnicodeObject;

在这个我们似乎看到了一个有点熟悉的类型，和buf字段相同，就是这个Py_UCS4*。Py_UCS4其实就是CPython所使用的字符类型。在虽然Python中没有单独的字符类型，但在C级还是存在这个概念的。CPython定义了3种字符类型，Py_UCS1、Py_UCS2和Py_UCS4，这里的UCS是Universal Multiple-Octet Coded Character Set的简称，后面的数字为字符编码的字节数，有1、2、4三种。Py_UCS1也就是我们一般在C语言中所使用的char类型。

_PyUnicodeWriter结构体

到这里，我们就可以理解了，buf字段的类型类似于char*，实际上就是一个C级的字符串，只不过每个字符都是是wide char（宽字符）类型。那么writer的类型_PyUnicodeWriter又是什么呢？

从字面意思简单理解的话，_PyUnicodeWriter是一个可以写入PyUnicode的接口类型。它是一个非常高效的字符串生成器，在CPython内部广为使用，是各种内部对象的repr方法及包括字符解码、格式化等字符串操作的常客。

初始化过程

在初步了解了buf和writer两个字段的作用后，我们可以发现一个奇怪的事情。一个是字符数组，另一个则是C级的字符串生成器，二者在功能上有一些重叠之处。那么，二者在StringIO中是如何相互配合工作的呢？我们把目光放到StringIO的__init__上，让我们看一看初始化的过程。

__init__在C中的具体实现函数为_io_StringIO___init___impl，签名如下：

/*[clinic input]
_io.StringIO.__init__
    initial_value as value: object(c_default="NULL") = ''
    newline as newline_obj: object(c_default="NULL") = '\n'

Text I/O implementation using an in-memory buffer.

The initial_value argument sets the value of object.  The newline
argument is like the one of TextIOWrapper's constructor.
[clinic start generated code]*/

static int
_io_StringIO___init___impl(stringio *self, PyObject *value,
                           PyObject *newline_obj);

忽略掉不关心的换行符处理部分，我们可以看以下的内容：

/* Now everything is set up, resize buffer to size of initial value,
    and copy it */
self->string_size = 0;
if (value && value != Py_None)
    value_len = PyUnicode_GetLength(value);
else
    value_len = 0;
if (value_len > 0) {
    /* This is a heuristic, for newline translation might change
        the string length. */
    if (resize_buffer(self, 0) < 0)
        return -1;
    self->state = STATE_REALIZED;
    self->pos = 0;
    if (write_str(self, value) < 0)
        return -1;
}
else {
    /* Empty stringio object, we can start by accumulating */
    if (resize_buffer(self, 0) < 0)
        return -1;
    _PyUnicodeWriter_Init(&self->writer);
    self->writer.overallocate = 1;
    self->state = STATE_ACCUMULATING;
}
self->pos = 0;

根据value的情况，初始化被分为了两种情况。第一种情况是在value为非空字符串的情况下，此时StringIO的状态字段state被设置为STATE_REALIZED。之后的函数调用write_str，实际上是StringIO.write的C级实现函数。这里是将字符串value写入了当前实例中。

而另一种情况则是调用了_PyUnicodeWriter_Init对writer字段进行了一次初始化，并将状态设置为了STATE_ACCUMULATING。

到这里，我们之前测试中的关于stringio-realized组的问题有了一些初步的解释，那就是非空的初始值会让s处于一个与stringio组的s不同的状态。我简单的将这两个状态称为积累态（ACCUMULATING）与实现态（REALIZED）。要了解两个状态在实现上的差异，我们还需要往后看。

积累态与实现态

StringIO.write方法

作为StringIO中使用频率相当高，且参与初始化过程的方法，write_str这个函数理所当然的成为了我们下一步的目标。源码如下：

/* Internal routine for writing a whole PyUnicode object to the buffer of a
   StringIO object. Returns 0 on success, or -1 on error. */
static Py_ssize_t
write_str(stringio *self, PyObject *obj)
{
    Py_ssize_t len;
    PyObject *decoded = NULL;

    /* 一些前置准备工作 */

    if (self->state == STATE_ACCUMULATING) {
        if (self->string_size == self->pos) {
            /* 这里是ACCUMULATING状态下的实现，decoded就是被处理过的obj */
            if (_PyUnicodeWriter_WriteStr(&self->writer, decoded))
                goto fail;
            goto success;
        }
        if (realize(self))
            goto fail;
    }
    
    /* 这里是REALIZED状态下的实现 */
    if (self->pos + len > self->string_size) {
        if (resize_buffer(self, self->pos + len) < 0)
            goto fail;
    }

    if (self->pos > self->string_size) {
        /* In case of overseek, pad with null bytes the buffer region between
           the end of stream and the current position.

          0   lo      string_size                           hi
          |   |<---used--->|<----------available----------->|
          |   |            <--to pad-->|<---to write--->    |
          0   buf                   position

        */
        memset(self->buf + self->string_size, '\0',
               (self->pos - self->string_size) * sizeof(Py_UCS4));
    }

    /* Copy the data to the internal buffer, overwriting some of the
       existing data if self->pos < self->string_size. */
    if (!PyUnicode_AsUCS4(decoded,
                          self->buf + self->pos,
                          self->buf_size - self->pos,
                          0))
        goto fail;

success:
    /* Set the new length of the internal string if it has changed. */
    self->pos += len;
    if (self->string_size < self->pos)
        self->string_size = self->pos;

    Py_DECREF(decoded);
    return 0;

fail:
    Py_XDECREF(decoded);
    return -1;
}

这里的逻辑还是比较明确的。对于积累态（ACCUMULATING）的实现非常简单，就是调用_PyUnicodeWriter_WriteStr向writer中写入了一个Python字符串decoded，这个decoded也就是经过了前置步骤处理了换行符的obj。

实现态（REALIZED）部分，也就是在有非空字符串作为初始值的__init__函数中实际使用的部分，其实现分为了3步，第一步是在需要写入新区域是进行一个resize_buffer。这个函数的作用是调整buf缓冲区的大小使其满足需要。第二步则是对overseek的部分填充为0，第三部才是真正的通过PyUnicode API PyUnicode_AsUCS4将decoded写入缓冲区buf中。

结合之前的初始化过程，我们可以得到一个初步的结论，那就是在积累态下，StringIO的内部操作对象为writer，而在实现态下则使用buf为操作对象。

StringIO.getvalue方法

为了验证以上结论，我们可以再看一个之前在测试中使用到的方法，就是这个getvalue。C级实现的函数名为_io_StringIO_getvalue_impl。源代码如下：

/*[clinic input]
_io.StringIO.getvalue

Retrieve the entire contents of the object.
[clinic start generated code]*/

static PyObject *
_io_StringIO_getvalue_impl(stringio *self)
/*[clinic end generated code: output=27b6a7bfeaebce01 input=d23cb81d6791cf88]*/
{
    CHECK_INITIALIZED(self);
    CHECK_CLOSED(self);
    if (self->state == STATE_ACCUMULATING)
        return make_intermediate(self);
    return PyUnicode_FromKindAndData(PyUnicode_4BYTE_KIND, self->buf,
                                     self->string_size);
}

先看实现态的实现，就是调用PyUnicode_FromKindAndData将字符数组buf转为一个Python字符串。而积累态则转到了另一个函数make_intermediate：

static PyObject *
make_intermediate(stringio *self)
{
    PyObject *intermediate = _PyUnicodeWriter_Finish(&self->writer);
    self->state = STATE_REALIZED;
    if (intermediate == NULL)
        return NULL;

    _PyUnicodeWriter_Init(&self->writer);
    self->writer.overallocate = 1;
    if (_PyUnicodeWriter_WriteStr(&self->writer, intermediate)) {
        Py_DECREF(intermediate);
        return NULL;
    }
    self->state = STATE_ACCUMULATING;
    return intermediate;
}

这个函数还是非常有意思的，首先是_PyUnicodeWriter_Finish将writer中的内容生成字符串。同时，这个函数会结束writer的生命周期，使其会处于一个被销毁的状态，所以接下来又是一个_PyUnicodeWriter_Init重新初始化writer，然后将之前读出来的intermediate重新写回了writer中。通过这个过程，我们在不改变writer内容的情况下完成了一次读取。

这里我们就可以基本肯定之前的结论，两个状态下的读取与写入的确是依赖于不同的字段进行的。

状态转换

那么什么时候StringIO的状态才会发生更改呢？由累积态变为实现态可以在如下情况发生，涉及到的代码部分比较多，我做一个简单的总结：

  1. 写入时self->string_size != self->pos
  2. 读取时self->pos != 0
  3. 调用truncate、readline或__next__
  4. 函数make_intermediate出错

这里的pos字段是StringIO的写入位置。换言之，1、2的意思就是，积累态的StringIO只允许在末尾写入并从头读取，否则就会进入实现状态。

具体的转换实现函数，在之前的write_str函数源码中已经出现了一次，就是这个realize函数：

static int
realize(stringio *self)
{
    Py_ssize_t len;
    PyObject *intermediate;

    if (self->state == STATE_REALIZED)
        return 0;
    assert(self->state == STATE_ACCUMULATING);
    self->state = STATE_REALIZED;

    intermediate = _PyUnicodeWriter_Finish(&self->writer);
    if (intermediate == NULL)
        return -1;

    /* Append the intermediate string to the internal buffer.
       The length should be equal to the current cursor position.
     */
    len = PyUnicode_GET_LENGTH(intermediate);
    if (resize_buffer(self, len) < 0) {
        Py_DECREF(intermediate);
        return -1;
    }
    if (!PyUnicode_AsUCS4(intermediate, self->buf, len, 0)) {
        Py_DECREF(intermediate);
        return -1;
    }

    Py_DECREF(intermediate);
    return 0;
}

这个函数的过程就是从writer中读取出完整的字符串，并通过PyUnicode_AsUCS4写入字符数组buf中，完成了内容从嵌入式 _PyUnicodeWriter结构到内部缓冲区的转移

至于从实现态转为积累态的办法，我并没有在源码中找到。换言之，从积累态到实现态就是一张单程票，一旦转变就无法再回来了。

为什么会有这两种状态

对于这个问题，我想这也是由StringIO本身的功能决定的。首先，作为字符串拼接的工具类，其实际上是C级对象_PyUnicodeWriter在Python层面的封装，为字符串的拼接提供了高效接口，也就对应了StringIO中的积累态。另一方面，作为一个IO类，其实现了文件io接口，这也意味其需要实现文件操作位置移动的功能，而这是_PyUnicodeWriter无法做到的，因此又引入了buf字符数组来处理指针移动的问题，也就是StringIO中的实现态。

总结

在最后，在源码中其实有一段我之前没有放出的注释，我就将其简单翻译一下作为本篇的总结：stringio 对象可以处于两种状态：正在积累或已实现。在累积状态下，内部缓冲区不包含任何内容，并且内容由嵌入式 _PyUnicodeWriter 结构给出。在实现状态下，内部缓冲区是有意义的，并且 _PyUnicodeWriter 被销毁。

到此，希望本文能够带给你一点新的对于StringIO的认识。

posted @ 2023-10-12 00:00  Ovizro  阅读(351)  评论(0)  编辑  收藏  举报