msgspec · jcrist · Sep 12, 2022 · Sep 12, 2022
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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