From 99b1e692a2e99a4cfc19d97a698cb614e545b7ee Mon Sep 17 00:00:00 2001 From: Jim Crist-Harif Date: Mon, 12 Sep 2022 00:44:56 -0500 Subject: [PATCH] Add rich support This adds support for `rich.print` to `msgspec.Struct` and `msgspec.Meta` objects. This allows rich to pretty-print nested objects containing these types. This was pretty straightforward to support, hooray for well defined interfaces. --- docs/source/structs.rst | 2 ++ msgspec/__init__.pyi | 2 ++ msgspec/_core.c | 58 ++++++++++++++++++++++++++++++++++ tests/basic_typing_examples.py | 9 ++++++ tests/test_constraints.py | 12 +++++++ tests/test_struct.py | 16 ++++++++++ 6 files changed, 99 insertions(+) diff --git a/docs/source/structs.rst b/docs/source/structs.rst index f7a27ca9..a1d27a5d 100644 --- a/docs/source/structs.rst +++ b/docs/source/structs.rst @@ -41,6 +41,7 @@ annotations: - ``__copy__`` - ``__eq__`` & ``__ne__`` - ``__match_args__`` (for Python 3.10+'s `pattern matching`_) +- ``__rich_repr__`` (for pretty printing support with rich_) .. code-block:: python @@ -686,3 +687,4 @@ collected (leading to a memory leak). .. _reference counting: https://en.wikipedia.org/wiki/Reference_counting .. _cyclic garbage collector: https://devguide.python.org/garbage_collector/ .. _tagged unions: https://en.wikipedia.org/wiki/Tagged_union +.. _rich: https://rich.readthedocs.io/en/stable/pretty.html diff --git a/msgspec/__init__.pyi b/msgspec/__init__.pyi index eb88e21c..c18374df 100644 --- a/msgspec/__init__.pyi +++ b/msgspec/__init__.pyi @@ -54,6 +54,7 @@ class Struct(metaclass=__StructMeta): gc: bool = True, weakref: bool = False, ) -> None: ... + def __rich_repr__(self) -> Iterable[Tuple[str, Any]]: ... def defstruct( name: str, @@ -116,6 +117,7 @@ class Meta: description: Final[Union[str, None]] examples: Final[Union[list, None]] extra_json_schema: Final[Union[dict, None]] + def __rich_repr__(self) -> Iterable[Tuple[str, Any]]: ... class MsgspecError(Exception): ... class EncodeError(MsgspecError): ... diff --git a/msgspec/_core.c b/msgspec/_core.c index 4dad5cad..6e376f40 100644 --- a/msgspec/_core.c +++ b/msgspec/_core.c @@ -1724,6 +1724,35 @@ Meta_repr(Meta *self) { return strbuilder_build(&builder); } +static PyObject * +Meta_rich_repr(PyObject *py_self, PyObject *args) { + Meta *self = (Meta *)py_self; + PyObject *out = PyList_New(0); + if (out == NULL) goto error; +#define DO_REPR(field) do { \ + if (self->field != NULL) { \ + PyObject *part = Py_BuildValue("(UO)", #field, self->field); \ + if (part == NULL || (PyList_Append(out, part) < 0)) goto error;\ + } } while(0) + DO_REPR(gt); + DO_REPR(ge); + DO_REPR(lt); + DO_REPR(le); + DO_REPR(multiple_of); + DO_REPR(pattern); + DO_REPR(min_length); + DO_REPR(max_length); + DO_REPR(title); + DO_REPR(description); + DO_REPR(examples); + DO_REPR(extra_json_schema); +#undef DO_REPR + return out; +error: + Py_XDECREF(out); + return NULL; +} + static int _meta_richcompare_part(PyObject *left, PyObject *right) { if ((left == NULL) != (right == NULL)) { @@ -1812,6 +1841,11 @@ Meta_hash(Meta *self) { return (acc == (Py_uhash_t)-1) ? 1546275796 : acc; } +static PyMethodDef Meta_methods[] = { + {"__rich_repr__", Meta_rich_repr, METH_NOARGS, "rich repr"}, + {NULL, NULL}, +}; + static PyMemberDef Meta_members[] = { {"gt", T_OBJECT, offsetof(Meta, gt), READONLY, NULL}, {"ge", T_OBJECT, offsetof(Meta, ge), READONLY, NULL}, @@ -1838,6 +1872,7 @@ static PyTypeObject Meta_Type = { .tp_traverse = (traverseproc) Meta_traverse, .tp_clear = (inquiry) Meta_clear, .tp_dealloc = (destructor) Meta_dealloc, + .tp_methods = Meta_methods, .tp_members = Meta_members, .tp_repr = (reprfunc) Meta_repr, .tp_richcompare = (richcmpfunc) Meta_richcompare, @@ -5869,6 +5904,28 @@ Struct_reduce(PyObject *self, PyObject *args) return out; } +static PyObject * +Struct_rich_repr(PyObject *self, PyObject *args) { + PyObject *fields = StructMeta_GET_FIELDS(Py_TYPE(self)); + Py_ssize_t nfields = PyTuple_GET_SIZE(fields); + + PyObject *out = PyTuple_New(nfields); + if (out == NULL) goto error; + + for (Py_ssize_t i = 0; i < nfields; i++) { + PyObject *field = PyTuple_GET_ITEM(fields, i); + PyObject *val = Struct_get_index(self, i); + if (val == NULL) goto error; + PyObject *part = PyTuple_Pack(2, field, val); + if (part == NULL) goto error; + PyTuple_SET_ITEM(out, i, part); + } + return out; +error: + Py_XDECREF(out); + return NULL; +} + static PyObject * StructMixin_fields(PyObject *self, void *closure) { PyObject *out = ((StructMetaObject *)Py_TYPE(self))->struct_fields; @@ -5893,6 +5950,7 @@ StructMixin_defaults(PyObject *self, void *closure) { static PyMethodDef Struct_methods[] = { {"__copy__", Struct_copy, METH_NOARGS, "copy a struct"}, {"__reduce__", Struct_reduce, METH_NOARGS, "reduce a struct"}, + {"__rich_repr__", Struct_rich_repr, METH_NOARGS, "rich repr"}, {NULL, NULL}, }; diff --git a/tests/basic_typing_examples.py b/tests/basic_typing_examples.py index c939e92c..e0544e0e 100644 --- a/tests/basic_typing_examples.py +++ b/tests/basic_typing_examples.py @@ -187,6 +187,9 @@ class Point(msgspec.Struct): a.x = a.x + b.y repr(a) + for name, val in a.__rich_repr__(): + print(f"{name} = {val}") + def check_struct_attributes() -> None: class Point(msgspec.Struct): @@ -299,6 +302,12 @@ def check_meta_equal() -> None: print("ok") +def check_meta_methods() -> None: + c = msgspec.Meta() + for name, val in c.__rich_repr__(): + print(f"{name} = {val}") + + ########################################################## # Raw # ########################################################## diff --git a/tests/test_constraints.py b/tests/test_constraints.py index e04228e5..c319041e 100644 --- a/tests/test_constraints.py +++ b/tests/test_constraints.py @@ -104,6 +104,18 @@ def test_repr_multiple_fields(self): c = Meta(gt=0, lt=1) assert repr(c) == "msgspec.Meta(gt=0, lt=1)" + def test_rich_repr_empty(self): + assert Meta().__rich_repr__() == [] + + @pytest.mark.parametrize("field", FIELDS) + def test_rich_repr_one_field(self, field): + m = Meta(**{field: FIELDS[field]}) + assert m.__rich_repr__() == [(field, FIELDS[field])] + + def test_rich_repr_multiple_fields(self): + m = Meta(gt=0, lt=1) + assert m.__rich_repr__() == [("gt", 0), ("lt", 1)] + def test_equality(self): assert_eq(Meta(), Meta()) assert_ne(Meta(), None) diff --git a/tests/test_struct.py b/tests/test_struct.py index 45377846..ad3f59ed 100644 --- a/tests/test_struct.py +++ b/tests/test_struct.py @@ -480,6 +480,22 @@ class Test(Struct): repr(t) +def test_struct_rich_repr(): + assert Struct().__rich_repr__() == () + + class Test(Struct): + a: int + b: str + + t = Test(1, "hello") + + assert t.__rich_repr__() == (("a", 1), ("b", "hello")) + + del t.b + with pytest.raises(AttributeError): + t.__rich_repr__() + + def test_struct_copy(): x = copy.copy(Struct()) assert type(x) is Struct