Function write_python_string

python/pyfory/cpp/pyfory.cc:665–723

}

static int write_python_string(Buffer *buffer, PyObject *value) {
  if (FORY_PREDICT_FALSE(!PyUnicode_Check(value))) {
    PyErr_Format(PyExc_TypeError, "expected str, got %.200s",
                 Py_TYPE(value)->tp_name);
    return -1;
  }
  if (FORY_PREDICT_FALSE(PyUnicode_READY(value) < 0)) {
    return -1;
  }
  const Py_ssize_t length = PyUnicode_GET_LENGTH(value);
  const int kind = PyUnicode_KIND(value);
  const void *data = PyUnicode_DATA(value);
  uint64_t header = 0;
  uint32_t buffer_size = 0;

  if (kind == PyUnicode_1BYTE_KIND) {
    if (FORY_PREDICT_FALSE(length > std::numeric_limits<uint32_t>::max())) {
      PyErr_SetString(PyExc_OverflowError,
                      "string too large for fastpath encoding");
      return -1;
    }
    buffer_size = static_cast<uint32_t>(length);
    header = (static_cast<uint64_t>(length) << 2U) | 0ULL;
  } else if (kind == PyUnicode_2BYTE_KIND) {
    const uint64_t bytes = static_cast<uint64_t>(length) << 1U;
    if (FORY_PREDICT_FALSE(bytes > std::numeric_limits<uint32_t>::max())) {
      PyErr_SetString(PyExc_OverflowError,
                      "string too large for fastpath encoding");
      return -1;
    }
    buffer_size = static_cast<uint32_t>(bytes);
    // Keep wire format exactly aligned with Buffer.write_string in buffer.pyx.
    header = (static_cast<uint64_t>(length) << 3U) | 1ULL;
  } else {
    Py_ssize_t utf8_size = 0;
    const char *utf8 = PyUnicode_AsUTF8AndSize(value, &utf8_size);
    if (FORY_PREDICT_FALSE(utf8 == nullptr)) {
      return -1;
    }
    if (FORY_PREDICT_FALSE(utf8_size > std::numeric_limits<uint32_t>::max())) {
      PyErr_SetString(PyExc_OverflowError,
                      "string too large for fastpath encoding");
      return -1;
    }
    data = utf8;
    buffer_size = static_cast<uint32_t>(utf8_size);
    header = (static_cast<uint64_t>(buffer_size) << 2U) | 2ULL;
  }

  buffer->write_var_uint64(header);
  if (buffer_size == 0) {
    return 0;
  }
  const uint32_t writer_index = buffer->writer_index();
  buffer->grow(buffer_size);
  buffer->unsafe_put(writer_index, data, buffer_size);
  buffer->increase_writer_index(buffer_size);
  return 0;