Skip to content

Support creating a Raw buffer from a str #252

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jcrist merged 1 commit into from
Jan 2, 2023
Merged

Support creating a Raw buffer from a str #252

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions msgspec/__init__.pyi
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,9 +91,9 @@ def defstruct(
# places where an object that implements the buffer protocol is valid
class Raw(bytes):
@overload
def __init__(self) -> None: ...
def __new__(self) -> "Raw": ...
@overload
def __init__(self, msg: bytes) -> None: ...
def __new__(self, msg: Union[bytes, str]) -> "Raw": ...
def copy(self) -> "Raw": ...

class Meta:
Expand Down
31 changes: 19 additions & 12 deletions msgspec/_core.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1012,7 +1012,7 @@ typedef struct Raw {
PyObject *base;
char *buf;
Py_ssize_t len;
bool wraps_bytes;
bool is_view;
} Raw;

static PyObject *
Expand All @@ -1024,7 +1024,14 @@ Raw_New(PyObject *msg) {
out->base = msg;
out->buf = PyBytes_AS_STRING(msg);
out->len = PyBytes_GET_SIZE(msg);
out->wraps_bytes = true;
out->is_view = false;
}
else if (PyUnicode_CheckExact(msg)) {
out->base = msg;
out->buf = (char *)unicode_str_and_size(msg, &out->len);
if (out->buf == NULL) return NULL;
Py_INCREF(msg);
out->is_view = false;
}
else {
Py_buffer buffer;
Expand All @@ -1035,13 +1042,13 @@ Raw_New(PyObject *msg) {
out->base = buffer.obj;
out->buf = buffer.buf;
out->len = buffer.len;
out->wraps_bytes = false;
out->is_view = true;
}
return (PyObject *)out;
}

PyDoc_STRVAR(Raw__doc__,
"Raw(msg=None, /)\n"
"Raw(msg=b"", /)\n"
"--\n"
"\n"
"A buffer containing an encoded message.\n"
Expand All @@ -1058,10 +1065,10 @@ PyDoc_STRVAR(Raw__doc__,
"\n"
"Parameters\n"
"----------\n"
"msg : bytes, bytearray, or memoryview, optional\n"
" A byte buffer containing an encoded message. One of bytes, bytearray,\n"
" memoryview, or any object that implements the buffer protocol. If not\n"
" present, defaults to a zero-length bytes object (``b""``)."
"msg : bytes, bytearray, memoryview, or str, optional\n"
" A buffer containing an encoded message. One of bytes, bytearray, memoryview,\n"
" str, or any object that implements the buffer protocol. If not present,\n"
" defaults to a zero-length bytes object (``b""``)."
);
static PyObject *
Raw_new(PyTypeObject *type, PyObject *args, PyObject *kwargs) {
Expand Down Expand Up @@ -1103,7 +1110,7 @@ static void
Raw_dealloc(Raw *self)
{
if (self->base != NULL) {
if (self->wraps_bytes) {
if (!self->is_view) {
Py_DECREF(self->base);
}
else {
Expand All @@ -1130,7 +1137,7 @@ Raw_FromView(PyObject *buffer_obj, char *data, Py_ssize_t len) {
out->base = buffer.obj;
out->buf = data;
out->len = len;
out->wraps_bytes = false;
out->is_view = true;
return (PyObject *)out;
}

Expand Down Expand Up @@ -1179,7 +1186,7 @@ static PySequenceMethods Raw_as_sequence = {
static PyObject *
Raw_reduce(Raw *self, PyObject *unused)
{
if (self->wraps_bytes) {
if (!self->is_view) {
return Py_BuildValue("O(O)", &Raw_Type, self->base);
}
return Py_BuildValue("O(y#)", &Raw_Type, self->buf, self->len);
Expand All @@ -1200,7 +1207,7 @@ PyDoc_STRVAR(Raw_copy__doc__,
static PyObject *
Raw_copy(Raw *self, PyObject *unused)
{
if (self->wraps_bytes) {
if (!self->is_view) {
Py_INCREF(self);
return (PyObject *)self;
}
Expand Down
1 change: 1 addition & 0 deletions tests/basic_typing_examples.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -378,6 +378,7 @@ def check_raw_constructor() -> None:
r2 = msgspec.Raw(b"test")
r3 = msgspec.Raw(bytearray(b"test"))
r4 = msgspec.Raw(memoryview(b"test"))
r2 = msgspec.Raw("test")


def check_raw_copy() -> None:
Expand Down
19 changes: 15 additions & 4 deletions tests/test_raw.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,17 +14,18 @@ def test_raw_noargs():
assert not r


@pytest.mark.parametrize("type", [bytes, bytearray, memoryview])
@pytest.mark.parametrize("type", [bytes, bytearray, memoryview, str])
def test_raw_constructor(type):
r = msgspec.Raw(type(b"test"))
msg = "test" if type is str else type(b"test")
r = msgspec.Raw(msg)
assert bytes(r) == b"test"
assert len(r) == 4
assert r


def test_raw_constructor_errors():
with pytest.raises(TypeError):
msgspec.Raw("test")
msgspec.Raw(1)

with pytest.raises(TypeError):
msgspec.Raw(msg=b"test")
Expand Down Expand Up @@ -68,13 +69,23 @@ def test_raw_copy():
assert ref() is None


def test_raw_pickle():
def test_raw_pickle_bytes():
orig_buffer = b"test"
r = msgspec.Raw(orig_buffer)
o = r.__reduce__()
assert o == (msgspec.Raw, (b"test",))
assert o[1][0] is orig_buffer


def test_raw_pickle_str():
orig_buffer = "test"
r = msgspec.Raw(orig_buffer)
o = r.__reduce__()
assert o == (msgspec.Raw, ("test",))
assert o[1][0] is orig_buffer


def test_raw_pickle_view():
r = msgspec.Raw(memoryview(b"test")[:3])
o = r.__reduce__()
assert o == (msgspec.Raw, (b"tes",))
Expand Down