msgspec · jcrist · Jun 12, 2021 · Jun 9, 2021 · Jun 10, 2021 · Jun 10, 2021
diff --git a/docs/source/api.rst b/docs/source/api.rst
@@ -24,6 +24,13 @@ Struct
     :members:
 
 
+Ext
+---
+
+.. autoclass:: Ext
+    :members:
+
+
 Functions
 ---------
 

diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -41,7 +41,7 @@ Highlights
   excellent editor integration.
 - ``msgspec`` is **flexible**. Unlike other libraries like ``msgpack`` or
   ``json``, ``msgspec`` natively supports a wider range of Python builtin
-  types.
+  types. Support for additional types can also be added through :ref:`extensions`.
 - ``msgspec`` supports :ref:`"schema evolution" <schema-evolution>`. Messages can
   be sent between clients with different schemas without error.
 
@@ -104,6 +104,10 @@ Msgspec currently supports serializing/deserializing the following types:
 - `set`
 - `enum.Enum`
 - `msgspec.Struct`
+- `msgspec.Ext`
+
+Support for serializing additional types can be added through the use of the
+``default`` callback on the `Encoder`, or by defining custom :ref:`extensions`.
 
 .. _typed-deserialization:
 
@@ -123,6 +127,7 @@ mapped to Python types as follows:
 - ``bin``: `bytes`
 - ``array``: `list` or `tuple` [#tuple]_
 - ``map``: `dict`
+- ``ext``: `msgspec.Ext`
 
 .. [#tuple] Tuples are only used when the array type must be hashable (e.g.
    keys in a ``dict`` or ``set``). All other array types are deserialized as
@@ -190,6 +195,7 @@ acceptible:
 - `set` / `typing.Set`
 - `typing.Any`
 - `typing.Optional`
+- `msgspec.Ext`
 - `enum.Enum` derived types
 - `enum.IntEnum` derived types
 - `msgspec.Struct` derived types
@@ -280,6 +286,129 @@ deserialization, it also can improve performance. Depending on the schema,
 deserializing a message into a `Struct` can be *roughly twice as fast* as
 deserializing it into a `dict`.
 
+.. _extensions:
+
+Extensions
+~~~~~~~~~~
+
+The MessagePack specification provides support for defining custom
+`Extensions <https://github.com/msgpack/msgpack/blob/master/spec.md#extension-types>`__.
+Extensions consist of:
+
+- An integer code (between 0 and 127, inclusive) representing the "type" of the
+  extension.
+- An arbitrary byte buffer of data (up to ``(2^32) - 1`` in length).
+
+By default extensions are serialized to/from `Ext` objects.
+
+.. code-block:: python
+
+    >>> ext = msgspec.Ext(1, b"some data")  # an extension object, with type code 1
+    >>> msg = msgspec.encode(ext)
+    >>> ext2 = msgspec.decode(msg)
+    >>> ext == ext2  # deserializes as an Ext object
+    True
+
+While manually creating `Ext` objects from buffers can be useful, usually the
+user wants to map extension types to/from their own custom objects. This can be
+accomplished by defining two callback functions:
+
+- ``default`` in `Encoder`, for transforming custom types into values
+  that ``msgspec`` already knows how to serialize.
+- ``ext_hook`` in `Decoder`, for converting extensions back into those
+  custom types.
+
+These should have the following signatures:
+
+.. code-block:: python
+
+    def default(obj: Any) -> Any:
+        """Given an object that msgspec doesn't know how to serialize by
+        default, convert it into an object that it does know how to
+        serialize"""
+        pass
+
+    def ext_hook(code: int, data: memoryview) -> Any:
+        """Given an extension type code and data buffer, deserialize whatever
+        custom object the extension type represents"""
+        pass
+
+
+For example, perhaps you wanted to serialize `complex` number objects as an
+extension type.  These objects can be represented as tuples of two floats (one
+"real" and one "imaginary"). If we represent each float as 8 bytes (a
+"double"), then any complex number can be fully represented by a 16 byte
+buffer.
+
+.. code-block::
+
+    +---------+---------+
+    |  real   |  imag   |
+    +---------+---------+
+      8 bytes   8 bytes 
+
+
+Here we define ``default`` and ``ext_hook`` callbacks to convert `complex`
+objects to/from this binary representation as a MessagePack extension.
+
+.. code-block:: python
+
+    import msgspec
+    import struct
+    from typing import Any
+
+    # All extension types need a unique integer designator so the decoder knows
+    # which type they're decoding. Here we arbitrarily choose 1, but any integer
+    # between 0 and 127 (inclusive) would work.
+    COMPLEX_TYPE_CODE = 1
+
+    def default(obj: Any) -> Any:
+        if isinstance(obj, complex):
+            # encode the complex number into a 16 byte buffer
+            data = struct.pack('dd', obj.real, obj.imag)
+
+            # Return an `Ext` object so msgspec serializes it as an extension type.
+            return msgspec.Ext(COMPLEX_TYPE_CODE, data)
+        else:
+            # Raise a TypeError for other types
+            raise TypeError(f"Objects of type {type(obj)} are not supported")
+
+
+    def ext_hook(code: int, data: memoryview) -> Any:
+        if code == COMPLEX_TYPE_CODE:
+            # This extension type represents a complex number, decode the data
+            # buffer accordingly.
+            real, imag = struct.unpack('dd', data)
+            return complex(real, imag)
+        else:
+            # Raise a TypeError for other extension type codes
+            raise TypeError(f"Extension type code {code} is not supported")
+
+
+    # Create an encoder and a decoder using the custom callbacks
+    enc = msgspec.Encoder(default=default)
+    dec = msgspec.Decoder(ext_hook=ext_hook)
+
+    # Define a message that contains complex numbers
+    msg = {"roots": [0, 0.75, 1 + 0.5j, 1 - 0.5j]}
+
+    # Encode and decode the message to show that things work
+    buf = enc.encode(msg)
+    msg2 = dec.decode(buf)
+    assert msg == msg2  # True
+
+
+.. note::
+
+    Note that the ``data`` argument to ``ext_hook`` is a `memoryview`. This
+    view is attached to the larger buffer containing the complete message being
+    decoded. As such, you'll want to ensure that you don't keep a reference to
+    the underlying buffer, otherwise you may accidentally persist the larger
+    message buffer around for longer than necessary, resulting in increased
+    memory usage.
+
+
+
 .. _schema-evolution:
 
 Schema Evolution
@@ -299,6 +428,8 @@ For schema evolution to work smoothly, you need to follow a few guidelines:
 
 1. Any new fields on a `Struct` must specify default values.
 2. Don't change the type annotations for existing messages or fields
+3. Don't change the type codes or implementations for any defined
+   :ref:`Extensions`
 
 For example, suppose we wanted to add a new ``email`` field to our ``Person``
 struct. To do so, we add it at the end of the definition, with a default value.
@@ -334,7 +465,6 @@ efficiently skipped without decoding.
     >>> new_dec.decode(old_msg) # deserializing an old msg with a new decoder
     Person2(first='Harry', last='Potter', address='4 Privet Drive', phone=None, email=None)
 
-
 .. _type annotations: https://docs.python.org/3/library/typing.html
 .. _quickle: https://jcristharif.com/quickle/
 .. _pickle: https://docs.python.org/3/library/pickle.html

diff --git a/msgspec/__init__.py b/msgspec/__init__.py
@@ -2,6 +2,7 @@
     Struct,
     Encoder,
     Decoder,
+    Ext,
     MsgspecError,
     EncodingError,
     DecodingError,