PYON: add plugin support

setup.py

@@ -7,7 +7,7 @@
     url="https://m-labs.hk/artiq",
     description="Simple Python communications",
     license="LGPLv3+",
-    install_requires=["setuptools", "numpy", "pybase64"],
+    install_requires=["setuptools", "numpy", "pluggy", "pybase64"],
     packages=find_packages(),
     entry_points={
         "console_scripts": [

sipyco/__init__.py

@@ -0,0 +1,4 @@
+import pluggy
+
+hookimpl = pluggy.HookimplMarker("sipyco")
+"""Marker to be imported and used in plugins (and for own implementations)."""

sipyco/hookspecs.py

@@ -0,0 +1,21 @@
+import typing
+
+import pluggy
+
+
+hookspec = pluggy.HookspecMarker("sipyco")
+
+
+@hookspec(firstresult=True)
+def sipyco_pyon_encode(value: typing.Any, pretty: bool, indent_level: int) -> str:
+    """Encode a python object to a PYON string."""
+
+
+@hookspec
+def sipyco_pyon_decoders() -> typing.Sequence[typing.Tuple[str, typing.Any]]:
+    """Return elements for the decoding dictionary (passed to eval).
+
+    The return value should be a sequence of tuples, to allow one plugin function
+    to return multiple types that it can decode. Each tuple is the 'name' of the
+    value (in PYON) and the Python value (e.g. ``('pi', np.pi)``).
+    """

sipyco/plugins.py

@@ -0,0 +1,15 @@
+import pluggy
+
+import sipyco.hookspecs as hookspecs
+import sipyco.pyon as pyon
+
+
+def get_plugin_manager() -> pluggy.PluginManager:
+    """Get the PluginManager for sipyco plugins.
+
+    You can call the plugin hooks via this manager."""
+    pm = pluggy.PluginManager("sipyco")
+    pm.add_hookspecs(hookspecs)
+    pm.register(pyon)
+    pm.load_setuptools_entrypoints("sipyco")
+    return pm

sipyco/pyon.py

@@ -14,18 +14,24 @@
 that JSON does not support Numpy and more generally cannot be extended with
 other data types while keeping a concise syntax. Here we can use the Python
 function call syntax to express special data types.
-"""
 
+PYON can be extended via encoding & decoding plugins to whatever types you would like to serialize.
+An example can be found in :mod:`sipyco.test.test_pyon_plugin`.
+"""
 
+import itertools
+import os
+import tempfile
 from operator import itemgetter
 from fractions import Fraction
 from collections import OrderedDict
-import os
-import tempfile
 
 import numpy
 import pybase64 as base64
 
+import sipyco
+import sipyco.plugins as plugin
+
 
 _encode_map = {
     type(None): "none",
@@ -42,7 +48,7 @@
     slice: "slice",
     Fraction: "fraction",
     OrderedDict: "ordereddict",
-    numpy.ndarray: "nparray"
+    numpy.ndarray: "nparray",
 }
 
 _numpy_scalar = {
@@ -66,12 +72,12 @@
 
 
 class _Encoder:
-    def __init__(self, pretty):
+    def __init__(self, pretty, indent_level=0):
         self.pretty = pretty
-        self.indent_level = 0
+        self.indent_level = indent_level
 
     def indent(self):
-        return "    "*self.indent_level
+        return "    " * self.indent_level
 
     def encode_none(self, x):
         return "null"
@@ -94,22 +100,22 @@ def encode_bytes(self, x):
 
     def encode_tuple(self, x):
         if len(x) == 1:
-            return "(" + self.encode(x[0]) + ", )"
+            return "(" + encode(x[0], self.pretty, self.indent_level) + ", )"
         else:
             r = "("
-            r += ", ".join([self.encode(item) for item in x])
+            r += ", ".join([encode(item, self.pretty, self.indent_level) for item in x])
             r += ")"
             return r
 
     def encode_list(self, x):
         r = "["
-        r += ", ".join([self.encode(item) for item in x])
+        r += ", ".join([encode(item, self.pretty, self.indent_level) for item in x])
         r += "]"
         return r
 
     def encode_set(self, x):
         r = "{"
-        r += ", ".join([self.encode(item) for item in x])
+        r += ", ".join([encode(item, self.pretty, self.indent_level) for item in x])
         r += "}"
         return r
 
@@ -121,8 +127,14 @@ def encode_dict(self, x):
 
         r = "{"
         if not self.pretty or len(x) < 2:
-            r += ", ".join([self.encode(k) + ": " + self.encode(v)
-                           for k, v in items()])
+            r += ", ".join(
+                [
+                    encode(k, self.pretty, self.indent_level)
+                    + ": "
+                    + encode(v, self.pretty, self.indent_level)
+                    for k, v in items()
+                ]
+            )
         else:
             self.indent_level += 1
             r += "\n"
@@ -131,7 +143,12 @@ def encode_dict(self, x):
                 if not first:
                     r += ",\n"
                 first = False
-                r += self.indent() + self.encode(k) + ": " + self.encode(v)
+                r += (
+                    self.indent()
+                    + encode(k, self.pretty, self.indent_level)
+                    + ": "
+                    + encode(v, self.pretty, self.indent_level)
+                )
             r += "\n"  # no ','
             self.indent_level -= 1
             r += self.indent()
@@ -142,40 +159,63 @@ def encode_slice(self, x):
         return repr(x)
 
     def encode_fraction(self, x):
-        return "Fraction({}, {})".format(self.encode(x.numerator),
-                                         self.encode(x.denominator))
+        return "Fraction({}, {})".format(
+            encode(x.numerator, self.pretty, self.indent_level),
+            encode(x.denominator, self.pretty, self.indent_level),
+        )
 
     def encode_ordereddict(self, x):
-        return "OrderedDict(" + self.encode(list(x.items())) + ")"
+        return (
+            "OrderedDict("
+            + encode(list(x.items()), self.pretty, self.indent_level)
+            + ")"
+        )
 
     def encode_nparray(self, x):
         x = numpy.ascontiguousarray(x)
         r = "nparray("
-        r += self.encode(x.shape) + ", "
-        r += self.encode(x.dtype.str) + ", b\""
+        r += encode(x.shape, self.pretty, self.indent_level) + ", "
+        r += encode(x.dtype.str, self.pretty, self.indent_level) + ", b\""
         r += base64.b64encode(x.data).decode()
         r += "\")"
         return r
 
     def encode_npscalar(self, x):
         r = "npscalar("
-        r += self.encode(x.dtype.str) + ", b\""
+        r += encode(x.dtype.str, self.pretty, self.indent_level) + ", b\""
         r += base64.b64encode(x.data).decode()
         r += "\")"
         return r
 
     def encode(self, x):
         ty = _encode_map.get(type(x), None)
         if ty is None:
-            raise TypeError("`{!r}` ({}) is not PYON serializable"
-                            .format(x, type(x)))
+            raise TypeError("`{!r}` ({}) is not PYON serializable".format(x, type(x)))
         return getattr(self, "encode_" + ty)(x)
 
 
-def encode(x, pretty=False):
+@sipyco.hookimpl
+def sipyco_pyon_encode(value, pretty=False, indent_level=0):
+    """Default PYON encoder implementation."""
+    try:
+        return _Encoder(pretty=pretty, indent_level=indent_level).encode(value)
+    except TypeError:
+        return None
+
+
+def encode(x, pretty=False, indent_level=0):
     """Serializes a Python object and returns the corresponding string in
     Python syntax."""
-    return _Encoder(pretty).encode(x)
+    pm = plugin.get_plugin_manager()
+
+    func_val = pm.hook.sipyco_pyon_encode(
+        value=x, pretty=pretty, indent_level=indent_level
+    )
+
+    if func_val is None:
+        raise TypeError("`{!r}` ({}) is not PYON serializable".format(x, type(x)))
+    else:
+        return func_val
 
 
 def _nparray(shape, dtype, data):
@@ -201,21 +241,35 @@ def _npscalar(ty, data):
     "Fraction": Fraction,
     "OrderedDict": OrderedDict,
     "nparray": _nparray,
-    "npscalar": _npscalar
+    "npscalar": _npscalar,
 }
 
 
+@sipyco.hookimpl
+def sipyco_pyon_decoders():
+    """Default Sipyco PYON decoding valid types."""
+    return _eval_dict.items()
+
+
 def decode(s):
     """Parses a string in the Python syntax, reconstructs the corresponding
     object, and returns it."""
-    return eval(s, _eval_dict, {})
+    pm = plugin.get_plugin_manager()
+    _decode_eval_dict = dict(
+        itertools.chain.from_iterable(
+            filter(lambda r: r is not None, pm.hook.sipyco_pyon_decoders())
+        )
+    )
+    return eval(s, _decode_eval_dict, {})
 
 
 def store_file(filename, x):
     """Encodes a Python object and writes it to the specified file."""
     contents = encode(x, True)
     directory = os.path.abspath(os.path.dirname(filename))
-    with tempfile.NamedTemporaryFile("w", dir=directory, delete=False, encoding="utf-8") as f:
+    with tempfile.NamedTemporaryFile(
+        "w", dir=directory, delete=False, encoding="utf-8"
+    ) as f:
         f.write(contents)
         f.write("\n")
         tmpname = f.name

sipyco/test/test_pyon_plugin.py

@@ -0,0 +1,58 @@
+import dataclasses
+
+import pluggy
+import pytest
+
+import sipyco
+import sipyco.hookspecs as hookspecs
+import sipyco.plugins as plugin
+import sipyco.pyon as pyon
+
+
+@dataclasses.dataclass
+class Point:
+    x: float
+    y: float
+
+
+class TestPyonPlugin:
+    @sipyco.hookimpl
+    def sipyco_pyon_encode(value, pretty=False):
+        if isinstance(value, Point):
+            return repr(value)
+
+    @sipyco.hookimpl
+    def sipyco_pyon_decoders():
+        return [("Point", Point)]
+
+
+def test_pyon_plugin_fail_without_plugin():
+    with pytest.raises(TypeError):
+        pyon.encode(Point(3, 4))
+
+
+def pyon_extra_plugin():
+    pm = pluggy.PluginManager("sipyco")
+    pm.add_hookspecs(sipyco.hookspecs)
+    pm.register(pyon)
+    pm.load_setuptools_entrypoints("sipyco")
+    pm.register(TestPyonPlugin)
+    return pm
+
+
+def test_pyon_plugin_encode(monkeypatch):
+    monkeypatch.setattr(plugin, "get_plugin_manager", pyon_extra_plugin)
+    assert pyon.encode(Point(2, 3)) == "Point(x=2, y=3)"
+
+
+def test_pyon_plugin_encode_decode(monkeypatch):
+    monkeypatch.setattr(plugin, "get_plugin_manager", pyon_extra_plugin)
+    test_value = Point(2.5, 3.4)
+    assert pyon.decode(pyon.encode(test_value)) == test_value
+
+
+def test_pyon_nested_encode(monkeypatch):
+    """Tests that nested items will be properly encoded."""
+    monkeypatch.setattr(plugin, "get_plugin_manager", pyon_extra_plugin)
+    test_value = {"first": Point(2.5, {"nothing": 0})}
+    assert pyon.decode(pyon.encode(test_value)) == test_value