From b5230d00eb13babb34c96b702c9e197c8b0a040b Mon Sep 17 00:00:00 2001
From: Thorsten Groh <groh.thorsten@gmail.com>
Date: Wed, 2 Mar 2022 16:52:19 +0100
Subject: [PATCH] Added xarray object encoding and decoding.

---
 sipyco/pyon.py           | 32 ++++++++++++++++++++++++++++++--
 sipyco/test/test_pyon.py | 19 +++++++++++++++++++
 2 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/sipyco/pyon.py b/sipyco/pyon.py
index d75af79..53e4fb2 100644
--- a/sipyco/pyon.py
+++ b/sipyco/pyon.py
@@ -24,6 +24,7 @@
 import tempfile
 
 import numpy
+import xarray
 import pybase64 as base64
 
 
@@ -42,7 +43,9 @@
     slice: "slice",
     Fraction: "fraction",
     OrderedDict: "ordereddict",
-    numpy.ndarray: "nparray"
+    numpy.ndarray: "nparray",
+    xarray.DataArray: "xrarray",
+    xarray.Dataset: "xrdataset"
 }
 
 _numpy_scalar = {
@@ -164,6 +167,21 @@ def encode_npscalar(self, x):
         r += base64.b64encode(x.data).decode()
         r += "\")"
         return r
+    
+    def encode_xrarray(self, x):
+        d = x.to_dict(data=False)
+        d["data"] = x.values
+        for k in d["coords"]:
+            d["coords"][k]["data"] = x.coords[k].values
+        return "xrarray(" + self.encode_dict(d) + ")"
+        
+    def encode_xrdataset(self, x):
+        d = x.to_dict(data=False)
+        for k in d["data_vars"]:
+            d["data_vars"][k]["data"] = x.data_vars[k].values
+        for k in d["coords"]:
+            d["coords"][k]["data"] = x.coords[k].values
+        return "xrdataset(" + self.encode_dict(d) + ")"
 
     def encode(self, x):
         ty = _encode_map.get(type(x), None)
@@ -189,6 +207,14 @@ def _npscalar(ty, data):
     return numpy.frombuffer(base64.b64decode(data), dtype=ty)[0]
 
 
+def _xrarray(d):
+    return xarray.DataArray.from_dict(d)
+
+
+def _xrdataset(d):
+    return xarray.Dataset.from_dict(d)
+
+
 _eval_dict = {
     "__builtins__": {},
 
@@ -202,7 +228,9 @@ def _npscalar(ty, data):
     "Fraction": Fraction,
     "OrderedDict": OrderedDict,
     "nparray": _nparray,
-    "npscalar": _npscalar
+    "npscalar": _npscalar,
+    "xrarray": _xrarray,
+    "xrdataset": _xrdataset,
 }
 
 
diff --git a/sipyco/test/test_pyon.py b/sipyco/test/test_pyon.py
index 25a2217..1819e55 100644
--- a/sipyco/test/test_pyon.py
+++ b/sipyco/test/test_pyon.py
@@ -4,6 +4,7 @@
 from collections import OrderedDict
 
 import numpy as np
+import xarray as xr
 
 from sipyco import pyon
 
@@ -21,6 +22,11 @@
     "od": OrderedDict(zip("abc", range(3))),
 }
 
+_pyon_test_xarray_object = xr.DataArray(
+    data=np.random.randn(5),
+    coords=[("x", np.arange(0, 5))],
+    attrs={"foo": "bar"},
+)
 
 class PYON(unittest.TestCase):
     def test_encdec(self):
@@ -46,6 +52,18 @@ def test_encdec_array_order(self):
         array = np.reshape(np.arange(6), (2, 3), order='F')
         np.testing.assert_array_equal(
             array, pyon.decode(pyon.encode(array)))
+        
+    def test_encdec_xarray(self):
+        """Test encoding & decoding of xarray objects:"""
+        for enc in pyon.encode, lambda x: pyon.encode(x, True):
+            xr.testing.assert_equal(
+                pyon.decode(enc(_pyon_test_xarray_object)),
+                _pyon_test_xarray_object
+            )
+            xr.testing.assert_equal(
+                pyon.decode(enc(_pyon_test_xarray_object.to_dataset(name="a"))),
+                _pyon_test_xarray_object.to_dataset(name="a")
+            )
 
 
 _json_test_object = {
@@ -62,3 +80,4 @@ def test_encdec(self):
             for dec in pyon.decode, json.loads:
                 self.assertEqual(dec(enc(_json_test_object)),
                                  _json_test_object)
+