sobolevn 27 ноя в 16:40

Как работает bytearray в Python? Смотрим реализацию на C

Сложный

7 мин

4.9K

Привет! Меня зовут Никита Соболев, я core-разработчик языка программирования CPython, а так же автор серии видео про его устройство.

Сегодня я хочу рассказать, как bytearray устроен внутри.

Под катом будет про: интересные оптимизации, разные аллокаторы в CPython, C pointer math, детали устройства данной структуры данных.

Если вам такое интересно или целиком незнакомо — добро пожаловать!

Начнем с видео, а далее в текстовом формате опишем основные моменты.

Как выглядит bytearray в исходниках CPython?

Сама структура, которая используется питоном выглядит вот так:

typedef struct {
    PyObject_VAR_HEAD
    Py_ssize_t ob_alloc;   /* How many bytes allocated in ob_bytes */
    char *ob_bytes;        /* Physical backing buffer */
    char *ob_start;        /* Logical start inside ob_bytes */
    Py_ssize_t ob_exports; /* How many buffer exports */
} PyByteArrayObject;

Что здесь есть интересного?

PyObject_VAR_HEAD – макрос, который разворачивается в PyVarObject ob_base; Данный код указывает: во-первых, базовый класс (object в нашем случае), и что размеры каждого инстанса будут отличаться. Сравним bytearray() с размером 0 и bytearray(b'abc') с размером 4 байта
Py_ssize_t ob_alloc отвечает за хранение размера структуры, дальше рассмотрим волшебный метод __alloc__
char *ob_bytes – указатель на буфер с реальными данными
char *ob_start – оптимизация, указатель на логическое начало данных среди ob_bytes, далее мы рассмотрим, зачем оно нужно
Py_ssize_t ob_exports – количество раз, когда мы делали "export" объекта через Buffer протокол и метод __buffer__

Создание bytearray

Создать bytearray можно разными способами, документацию мы тут пересказывать не будем. А посмотрим мы на C-API функцию PyByteArray_FromStringAndSize (имеется в виду, конечно не питоновский str, а сишное представление), которая широко используется для создания bytearray объектов в коде. Например, при использовании bytearray.copy. Для нас она – отличный пример, потому что очень наглядная:

PyObject *
PyByteArray_FromStringAndSize(const char *bytes, Py_ssize_t size)
{
    Py_ssize_t alloc;
    /* Код упрощен для целей демонстрации. */
    PyByteArrayObject *new = PyObject_New(PyByteArrayObject, &PyByteArray_Type);
    if (size == 0) {
        new->ob_bytes = NULL;
        alloc = 0;
    }
    else {
        alloc = size + 1;
        new->ob_bytes = PyMem_Malloc(alloc);
        if (new->ob_bytes == NULL) {
            Py_DECREF(new);
            return PyErr_NoMemory();
        }
        if (bytes != NULL && size > 0)
            memcpy(new->ob_bytes, bytes, size);
        new->ob_bytes[size] = '\0';  /* Trailing null byte */
    }
    Py_SET_SIZE(new, size);
    new->ob_alloc = alloc;
    new->ob_start = new->ob_bytes;
    new->ob_exports = 0;
    return (PyObject *)new;
}

Что нас здесь может заинтересовать?

Создание объекта с длинной 0, имеет быстрый путь, ничего дополнительного мы не делаем
Если у нас есть не нулевой размер «строки», которую нам передали, то мы выделяем участок памяти на 1 байт больше, чем нам нужно, используем PyMem_Malloc. По сути мы получим просто указатель на неинициализированный кусок памяти нужного нам размера. Простой и понятный аллокатор
Далее, мы используем memcpy для копирования bytes в new->ob_bytes с размером size (на 1 меньше выделенного участка памяти)
Последним шагом мы добавляем \0 (null-byte) в последнюю незанятую позицию в ->ob_bytes
Инициализируем все нужные поля структуры

Посмотрим, что нам выдается правильный размер через магический метод __alloc__:

PyDoc_STRVAR(alloc_doc,
"B.__alloc__() -> int\n\
\n\
Return the number of bytes actually allocated.");

static PyObject *
bytearray_alloc(PyByteArrayObject *self, PyObject *Py_UNUSED(ignored))
{
    return PyLong_FromSsize_t(self->ob_alloc);
}

И проверим:

>>> b = bytearray(b'abc')
>>> b.__alloc__()
4

Все верно, ответ 4— потому что у нас есть еще один дополнительный байт – \0!

Удаляем байты из начала

Следующая интересная тема – оптимизация с ob_start. Когда мы работаем с набором байт, нам часто нужно удалить какие-то байты в начале последовательности. Например, заголовок какой-нибудь в известное количество байт. То есть, мы будем делать del b[:4] для удаления первых четырех байт.

Как оно будет работать? Посмотрим на bytearray_setslice_linear, который вызывается из bytearray_setslice:

static int
bytearray_setslice_linear(PyByteArrayObject *self,
                          Py_ssize_t lo, Py_ssize_t hi,
                          char *bytes, Py_ssize_t bytes_len)
{
    Py_ssize_t avail = hi - lo;
    char *buf = PyByteArray_AS_STRING(self);
    Py_ssize_t growth = bytes_len - avail;
    int res = 0;
    assert(avail >= 0);

    if (growth < 0) {
        if (!_canresize(self))
            return -1;

        if (lo == 0) {
            /* Shrink the buffer by advancing its logical start */
            self->ob_start -= growth;
            /*
              0   lo               hi             old_size
              |   |<----avail----->|<-----tail------>|
              |      |<-bytes_len->|<-----tail------>|
              0    new_lo         new_hi          new_size
            */
        }
        else {
            /*
              0   lo               hi               old_size
              |   |<----avail----->|<-----tomove------>|
              |   |<-bytes_len->|<-----tomove------>|
              0   lo         new_hi              new_size
            */
            memmove(buf + lo + bytes_len, buf + hi,
                    Py_SIZE(self) - hi);
        }
        // ...
    }
    // ...
}

Давайте посмотрим на такой питон код, который вызовет код выше:

>>> b = bytearray(b'abc')
>>> del b[0]
>>> b
bytearray(b'bc')
>>> b.__alloc__()
4

Смотрите: мы удаляем нулевой объект в bytearray, но почему-то не происходит изменения размера в self->ob_alloc. Ответ в коде выше на C. Что там происходит?

Если мы сталкиваемся с «отрицательным ростом» размера объекта, то мы смотрим, откуда начинается изменение: с нуля или нет
Если мы удаем, скажем del b[2:4], то нам приходится идти по сложному пути и двигать память при помощи memmove
Если мы удаляем объекты с нуля, то мы просто перемещаем указатель ob_start на нужное количество байт

Тут нужно сделать оговорку про self->ob_start -= growth; Тут немного нужно разобраться, что такое C pointer math. Потому что тут мы отнимаем число от указателя (точнее прибавляем, потому что growth отрицательный)! На самом деле мы делаем что-то вроде: self->ob_start -= (growth * sizeof(char)) . То есть, мы прибавляем размер данных внутри указателя, что позволяет нам, например, переходить по значениям внутри массива. Пример (отсюда):

#include <stdio.h>

int main() {
  int int_arr[] = {12, 23, 45};
  int *ptrArr = int_arr;

  printf("Value at ptrArr: %d\n", *ptrArr);  // 12

  // Добавляем 2 к указателю:
  ptrArr = ptrArr + 2;  // ptrArr + (2 * sizeof(int))
  // Мы с 12 перевели на 2 значения вперед: до 45
  printf("Value at ptrArr after adding 2: %d\n", *ptrArr);  // 45

  return 0;
}

Следовательно, когда мы делаем self->ob_start -= growth при growth = -1, бы добавляем: self->ob_start = self.ob_start - (-1 * sizeof(char)), где sizeof(char) является 1. И мы просто сдвигаем указатель в случае del b[0] на один байт. Смотрите:

static inline char* PyByteArray_AS_STRING(PyObject *op)
{
    PyByteArrayObject *self = _PyByteArray_CAST(op);
    if (Py_SIZE(self)) {
        return self->ob_start;
    }
    return _PyByteArray_empty_string;
}

Почти везде в bytearray мы далее используем self->ob_start через PyByteArray_AS_STRING. И потому такая оптимизация дает нам O(1) удаление с начала последовательности!

Добавляем новые значения через .append()

Сам по себе bytearray.append очень простой. Но, в нем есть самая интересная часть:

if (PyByteArray_Resize((PyObject *)self, n + 1) < 0)
    return NULL;
PyByteArray_AS_STRING(self)[n] = item;

Давайте разбираться с ней. Как происходит увеличение bytearray? Сначала мы выбираем стратегию увеличения размера: будем ли мы делать пере-аллокацию? Или добавим только нужное количество байт?

int
PyByteArray_Resize(PyObject *self, Py_ssize_t requested_size)
{
    void *sval;
    PyByteArrayObject *obj = ((PyByteArrayObject *)self);
    /* All computations are done unsigned to avoid integer overflows
       (see issue #22335). */
    size_t alloc = (size_t) obj->ob_alloc;
    size_t logical_offset = (size_t) (obj->ob_start - obj->ob_bytes);
    size_t size = (size_t) requested_size;

    if (size + logical_offset + 1 <= alloc) {
        /* Current buffer is large enough to host the requested size,
           decide on a strategy. */
        if (size < alloc / 2) {
            /* Major downsize; resize down to exact size */
            alloc = size + 1;
        }
        else {
            /* Minor downsize; quick exit */
            Py_SET_SIZE(self, size);
            PyByteArray_AS_STRING(self)[size] = '\0'; /* Trailing null */
            return 0;
        }
    }
    else {
        /* Need growing, decide on a strategy */
        if (size <= alloc * 1.125) {
            /* Moderate upsize; overallocate similar to list_resize() */
            alloc = size + (size >> 3) + (size < 9 ? 3 : 6);
        }
        else {
            /* Major upsize; resize up to exact size */
            alloc = size + 1;
        }
    }
    // ...
  }

Что мы видим?

Если объект уменьшается, то мы можем пойти по долгому пути (в случае если запрашиваемый размер в два раза меньше текущего), чтобы сохранить память, или просто сделать быструю замену размера и установить \0 в конец. Так быстрее, но требует больше «лишней» памяти
Если объект увеличивается, то мы можем выделить дополнительную память, которая пока не очень нужна (в случае если запрашиваемый размер в 1.125 раза больше текущего), или просто добавить требуемое количество памяти

Далее, нам нужно выполнить саму аллокацию / реаллокацию памяти и установить нужные поля в объекте. Тут мы будем смотреть на наличие ob_start в logical_offset:

    // ...
    if (logical_offset > 0) {
        sval = PyMem_Malloc(alloc);
        if (sval == NULL) {
            PyErr_NoMemory();
            return -1;
        }
        memcpy(sval, PyByteArray_AS_STRING(self),
               Py_MIN((size_t)requested_size, (size_t)Py_SIZE(self)));
        PyMem_Free(obj->ob_bytes);
    }
    else {
        sval = PyMem_Realloc(obj->ob_bytes, alloc);
        if (sval == NULL) {
            PyErr_NoMemory();
            return -1;
        }
    }

    obj->ob_bytes = obj->ob_start = sval;
    Py_SET_SIZE(self, size);
    obj->ob_alloc = alloc;
    obj->ob_bytes[size] = '\0'; /* Trailing null byte */

    return 0;
}

Если ob_start больше ob_bytes, то мы имеем дело с «оптимизированным» случаем из примера в прошлой главе. Тогда нам нужно создать полностью новый участок памяти нового размера при помощи PyMem_Malloc. И скопировать память в новое место из PyByteArray_AS_STRING(self) (self->ob_start) правильного размера. То есть, не потребуются уже отрошенные куски памяти. Такая операция, очевидно, долгая
В противном случае (если мы не оптимизировали ничего и оба указателя совпадают), то мы используем PyMem_Realloc для добавления нужного количества байт к self->ob_bytes, что значительно быстрее

Итого, данное объяснение позволяет нам ответить на вопрос: какой код быстрее?

b1 = bytearray(b'abc')
b1.append(ord(b'd'))
del b1[0]

b2 = bytearray(b'abc')
del b2[0]
b2.append(ord(b'd'))

Да, правильно. Первый код быстрее. Замеры:

» pyperf timeit -s 'b = bytearray(b"abc"); c = ord(b"d")' '                   
b.append(c); del b[0]'
.....................
Mean +- std dev: 33.1 ns +- 0.1 ns
                                                                                                                                                
» pyperf timeit -s 'b = bytearray(b"abc"); c = ord(b"d")' '
del b[0]; b.append(c)'
.....................
Mean +- std dev: 41.3 ns +- 0.4 ns

Выводы

Мы обсудили много интересного. Но много всего осталось за кадром. Почему, например, код вида

>>> b = bytearray(b'1234')
>>> del b[:4:2]
>>> b
bytearray(b'24')
>>> b.__alloc__()
5

работает так, как он работает? Ответ прячется в bytearray_ass_subscript.

Если нравится такой контект, забегайте ко мне в телеграм канал. Там мы регулярно обсуждаем устройство технологий. С любовью и глубоким погружением. В том числе и пример выше.

А что-то (как, например, место bytearray среди интерфейсов питона и ob_exports) я более подробно объяснил в приложенном видео.

Ну а вывод от текущего только один – если вы понимаете, как работает ваш код, то вы можете писать более оптимальные и быстрые программы! :)

Теги:

Хабы:

Как работает bytearray в Python? Смотрим реализацию на C

Как выглядит bytearray в исходниках CPython?

Создание bytearray

Удаляем байты из начала

Добавляем новые значения через .append()

Выводы

Публикации

Истории

Работа

Ближайшие события