#!/usr/bin/env python
# coding: utf-8

# Custom JSON serialization in IPython
# 
# This is an illustration of what would be needed for objects to define their own custom JSON serialization (to strings) in IPython.
# 
# IPython currently supports objects defining their own serialization to *JSONable* types, but not JSON strings without potentially expensive double-encoding.
# 
# First, we define a simple JSON serialization function that allows this functionality via `_raw_json_()` methods:

# In[1]:


import json
from jupyter_client.jsonutil import date_default


def json_strings(obj, default=date_default):
    """Yield string fragments that should result in JSON after joining.
    
    Yielding fragments allows objects that define `_raw_json_()`
    to return their own JSON representation.
    """
    
    if hasattr(obj, '_raw_json_'):
        # If an object has a `_raw_json_` method, call it instead
        # _raw_json_ should return a JSON-serialized string
        yield obj._raw_json_()
    
    elif isinstance(obj, (list, tuple)):
        yield '['
        first = True
        for item in obj:
            if not first:
                yield ','
            else:
                first = False

            for s in json_strings(item):
                yield s
        yield ']'

    elif isinstance(obj, dict):
        yield '{'
        first = True
        for key, value in obj.items():
            if not first:
                yield ','
            else:
                first = False
            
            yield json.dumps(key)
            yield ': '
            for s in json_strings(value):
                yield s
        yield '}'
    else:
        yield json.dumps(obj, default=default)

def my_dumps(obj):
    """Serialize (to bytes) with JSON
    
    allowing objects to define their own JSON serialization.
    """
    return ''.join(s for s in json_strings(obj)).encode('utf8')


# In[2]:


my_dumps({
    'a': [
        5,
        'b',
        (1,'x')
    ]
})


# Next, we create a toy object that defines `_raw_json_` returning a JSON string
# and an `_ipython_display_` method for displaying itself using this.

# In[3]:


from IPython.display import display


class MyObject:
    
    def __init__(self, value):
        self.value = value
    
    def _raw_json_(self):
        """I know how to JSON-serialize myself"""
        
        return json.dumps({
            'classname': self.__class__.__name__,
            'value': self.value,
        })
    
    def _ipython_display_(self):
        display({'application/myobject+json': obj}, raw=True)


# In[4]:


obj = MyObject(10)
my_dumps(obj)


# But this doesn't get through IPython,
# because our JSON serialization isn't in use:

# In[5]:


obj


# To get something like this to work,
# we will need a public API for allowing objects to register their own JSON serializers,
# and ensure that custom objects arrive at the custom serializer.
# 
# The first step is to register `my_dumps` as the serializer for all messages.
# This will allow `_raw_json_`-having objects
# to be present in the message object and get serialized by their own definition:

# In[6]:


session = get_ipython().kernel.session
session.pack = my_dumps


# But we run into a problem where IPython is checking types before passing things to JSON.
# For now, we can disable type checking in json_clean,
# which will need to be modified to make this work as an official API:

# In[7]:


# disable json_clean
from ipykernel import jsonutil
jsonutil.json_clean = lambda obj: obj


# Now we can display the object,
# which will show up in the document:

# In[8]:


obj


# In[9]:


import nbformat
nb = nbformat.read('custom-json.ipynb', as_version=4)
nb.cells[-2].outputs