From c12f12499eca6910fe8874f2caa40e93949e7f73 Mon Sep 17 00:00:00 2001
From: David Bau <david.bau@gmail.com>
Date: Mon, 28 Mar 2022 20:46:28 -0400
Subject: [PATCH] Support numpy.ndarray and torch.tensor properties.

---
 baukit/labwidget.py | 24 +++++++++++++++++++++---
 1 file changed, 21 insertions(+), 3 deletions(-)

diff --git a/baukit/labwidget.py b/baukit/labwidget.py
index e8e14f0..58beba2 100644
--- a/baukit/labwidget.py
+++ b/baukit/labwidget.py
@@ -48,7 +48,7 @@ import io
 import json
 import html
 import re
-from inspect import signature
+from inspect import signature, getmro
 from . import show
 
 
@@ -237,7 +237,7 @@ class Widget(Model):
         Returns the HTML code for the widget.
         '''
         self._viewcount += 1
-        json_data = json.dumps({
+        json_data = jsondump({
             k: v.value for k, v in vars(self).items()
             if isinstance(v, Property)})
         json_data = re.sub('</', '<\\/', json_data)
@@ -285,7 +285,7 @@ class Widget(Model):
                 colab_output.eval_js(minify(f"""
                 (window.send_{id(self)} = window.send_{id(self)} ||
                 new BroadcastChannel("channel_{id(self)}")
-                ).postMessage({json.dumps(args)});
+                ).postMessage({jsondump(args)});
                 """), ignore_result=True)
             elif WIDGET_ENV == 'jupyter':
                 if not self._comms:
@@ -1020,6 +1020,24 @@ class Image(Widget):
 # Utils
 ##########################################################################
 
+def baseclass_named(obj, *class_names):
+    '''
+    Detects if obj is a subclass of a class named clsname, without requiring import
+    of the class.
+    '''
+    for x in getmro(type(obj)):
+        if (x.__module__ + '.' + x.__name__) in class_names:
+             return True
+    return False
+
+class PermissiveEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if baseclass_named(obj, 'numpy.ndarray', 'torch.Tensor'):
+            return obj.tolist()
+        return json.JSONEncoder.default(self, obj)
+
+def jsondump(d):
+    return json.dumps(d, cls=PermissiveEncoder)
 
 def minify(t):
     # TODO: plug in some more real minification.