Skip to content

Support overriding __setattr__ in Struct types #376

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
Apr 17, 2023
Merged

Support overriding __setattr__ in Struct types #376

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
47 changes: 28 additions & 19 deletions msgspec/_core.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -210,7 +210,7 @@ typedef struct {
* - the value supports GC
* - the value isn't a tuple or the object is tracked (skip tracked checks for non-tuples)
*/
#define MS_OBJ_IS_GC(x) \
#define MS_MAYBE_TRACKED(x) \
(MS_TYPE_IS_GC(Py_TYPE(x)) && \
(!PyTuple_CheckExact(x) || MS_IS_TRACKED(x)))

Expand Down Expand Up @@ -4808,12 +4808,16 @@ Struct_setattro_frozen(PyObject *self, PyObject *key, PyObject *value) {
return -1;
}

/* setattr for frozen=False, gc=True types */
/* setattr for frozen=False types */
static int
Struct_setattro_default(PyObject *self, PyObject *key, PyObject *value) {
if (PyObject_GenericSetAttr(self, key, value) < 0)
return -1;
if (value != NULL && MS_OBJ_IS_GC(value) && !MS_IS_TRACKED(self))
if (PyObject_GenericSetAttr(self, key, value) < 0) return -1;
if (
value != NULL &&
MS_OBJECT_IS_GC(self) &&
!MS_IS_TRACKED(self) &&
MS_MAYBE_TRACKED(value)
)
PyObject_GC_Track(self);
return 0;
}
Expand Down Expand Up @@ -5755,12 +5759,17 @@ StructMeta_new_inner(
((PyTypeObject *)cls)->tp_free = &PyObject_GC_Del;
}
if (info.frozen == OPT_TRUE) {
/* Frozen structs always override __setattr__ */
((PyTypeObject *)cls)->tp_setattro = &Struct_setattro_frozen;
}
else if (info.gc == OPT_FALSE) {
((PyTypeObject *)cls)->tp_setattro = &PyObject_GenericSetAttr;
}
else {
else if (
((PyTypeObject *)cls)->tp_setattro == Struct_setattro_frozen ||
((PyTypeObject *)cls)->tp_setattro == PyObject_GenericSetAttr
) {
/* Only set non-frozen __setattr__ if it hasn't been defined by the
* user. All structs will inherit `Struct_setattro_default` from the
* base Struct class, so proper use of `super()` should still result in
* this being called. */
((PyTypeObject *)cls)->tp_setattro = &Struct_setattro_default;
}

Expand Down Expand Up @@ -6549,7 +6558,7 @@ Struct_fill_in_defaults(StructMetaObject *st_type, PyObject *obj, PathNode *path
Struct_set_index(obj, i, val);
}
if (should_untrack) {
should_untrack = !MS_OBJ_IS_GC(val);
should_untrack = !MS_MAYBE_TRACKED(val);
}
}

Expand Down Expand Up @@ -6597,7 +6606,7 @@ Struct_vectorcall(PyTypeObject *cls, PyObject *const *args, size_t nargsf, PyObj
Py_INCREF(val);
*(PyObject **)addr = val;
if (should_untrack) {
should_untrack = !MS_OBJ_IS_GC(val);
should_untrack = !MS_MAYBE_TRACKED(val);
}
}

Expand Down Expand Up @@ -6643,7 +6652,7 @@ Struct_vectorcall(PyTypeObject *cls, PyObject *const *args, size_t nargsf, PyObj
Py_INCREF(val);
*(PyObject **)addr = val;
if (should_untrack) {
should_untrack = !MS_OBJ_IS_GC(val);
should_untrack = !MS_MAYBE_TRACKED(val);
}
}

Expand All @@ -6659,7 +6668,7 @@ Struct_vectorcall(PyTypeObject *cls, PyObject *const *args, size_t nargsf, PyObj
if (MS_UNLIKELY(val == NULL)) goto error;
*(PyObject **)addr = val;
if (should_untrack) {
should_untrack = !MS_OBJ_IS_GC(val);
should_untrack = !MS_MAYBE_TRACKED(val);
}
continue;
}
Expand Down Expand Up @@ -6949,7 +6958,7 @@ struct_replace(PyObject *self, PyObject *const *args, Py_ssize_t nargs, PyObject
Py_INCREF(val);
Struct_set_index(out, field_index, val);
if (should_untrack) {
should_untrack = !MS_OBJ_IS_GC(val);
should_untrack = !MS_MAYBE_TRACKED(val);
}
}

Expand All @@ -6958,7 +6967,7 @@ struct_replace(PyObject *self, PyObject *const *args, Py_ssize_t nargs, PyObject
PyObject *val = Struct_get_index(obj, i);
if (val == NULL) goto error;
if (should_untrack) {
should_untrack = !MS_OBJ_IS_GC(val);
should_untrack = !MS_MAYBE_TRACKED(val);
}
Py_INCREF(val);
Struct_set_index(out, i, val);
Expand Down Expand Up @@ -12672,7 +12681,7 @@ mpack_decode_struct_array_inner(
}
Struct_set_index(res, i, val);
if (should_untrack) {
should_untrack = !MS_OBJ_IS_GC(val);
should_untrack = !MS_MAYBE_TRACKED(val);
}
}
if (MS_UNLIKELY(size > 0)) {
Expand Down Expand Up @@ -14974,7 +14983,7 @@ json_decode_struct_array_inner(
if (MS_UNLIKELY(item == NULL)) goto error;
Struct_set_index(out, i, item);
if (should_untrack) {
should_untrack = !MS_OBJ_IS_GC(item);
should_untrack = !MS_MAYBE_TRACKED(item);
}
i++;
item_path.index++;
Expand Down Expand Up @@ -15013,7 +15022,7 @@ json_decode_struct_array_inner(
if (item == NULL) goto error;
Struct_set_index(out, i, item);
if (should_untrack) {
should_untrack = !MS_OBJ_IS_GC(item);
should_untrack = !MS_MAYBE_TRACKED(item);
}
}
Py_LeaveRecursiveCall();
Expand Down Expand Up @@ -18005,7 +18014,7 @@ from_builtins_struct_array_inner(
}
Struct_set_index(out, i, val);
if (should_untrack) {
should_untrack = !MS_OBJ_IS_GC(val);
should_untrack = !MS_MAYBE_TRACKED(val);
}
}
if (MS_UNLIKELY(size > 0)) {
Expand Down
55 changes: 53 additions & 2 deletions tests/test_struct.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1435,7 +1435,7 @@ class FrozenPoint(Struct, frozen=True):
y: int


class TestFrozen:
class TestSetAttr:
def test_frozen_objects_no_setattr(self):
p = FrozenPoint(1, 2)
with pytest.raises(AttributeError, match="immutable type: 'FrozenPoint'"):
Expand All @@ -1455,11 +1455,62 @@ def test_frozen_objects_hash_errors_if_field_unhashable(self):
with pytest.raises(TypeError):
hash(p)

def test_mutable_objects_hash_errors(self):
def test_frozen_hash_mutable_objects_hash_errors(self):
p = Point(1, 2)
with pytest.raises(TypeError, match="unhashable type"):
hash(p)

@pytest.mark.parametrize("base_gc", [True, None, False])
@pytest.mark.parametrize("base_frozen", [True, False])
@pytest.mark.parametrize("has_gc", [True, None, False])
def test_override_setattr(self, has_gc, base_gc, base_frozen):
called = False

class Base(Struct, gc=base_gc, frozen=base_frozen):
pass

class Test(Struct, gc=has_gc, frozen=False):
x: Any

def __setattr__(self, name, value):
nonlocal called
called = True
super().__setattr__(name, value)

t = Test(1)
assert not called
t.x = 2
assert called
if has_gc:
assert not gc.is_tracked(t)
t.x = [1]
assert gc.is_tracked(t)

@pytest.mark.parametrize("base_gc", [True, None, False])
@pytest.mark.parametrize("has_gc", [True, None, False])
def test_override_setattr_inherit(self, base_gc, has_gc):
called = False

class Base(Struct, gc=base_gc):
x: Any

def __setattr__(self, name, value):
nonlocal called
called = True
super().__setattr__(name, value)

class Test(Base, gc=has_gc):
pass

t = Test(1)
assert not called
t.x = 2
assert called
if has_gc:
assert not gc.is_tracked(t)
t.x = [1]
assert gc.is_tracked(t)


class TestOrderAndEq:
@staticmethod
Expand Down