Simple JSON-Friendly Python Tree

gistfile1.txt

node.py

#!/usr/bin/env python
import simplejson as json

class Node(object):
    """
    Simple node to store data and build trees. This node,
    and all of its children can be easily converted into JSON
    and sent across the tubes.
    """
    def __init__(self, key, value=[]):
        if key is None: raise Exception("Node requires a key")
        self.key = key
        self.value = value
        return super(Node, self).__init__()

    @property
    def is_leaf(self):
        """ Whether or not this node has children. """
        return not type(self.value).__name__ == 'list'

    def __str__(self):
        if self.is_leaf:
            return "<Node: %s; Data: %s>" % (self.key, self.value)
        else:
            return "<Node: %s; %d Children>" % (self.key, len(self.value))

    def add_child(self, *args, **kwargs):
        """
        A generalized way to add a child. Can add add by either passing
        key and val, or another Node object.
        """
        if self.is_leaf: 
            raise Exception("Cannot add a child to leaf node %s" % self)
        if len(args) == 1 and type(args[0]).__name__ == 'Node':
            self.value.append(args[0])
        elif len(args) == 0 and 'key' in kwargs:
            val = None if 'value' not in kwargs else kwargs['value']
            self.value.append(Node(key=kwargs['key'], value=val))
        return

    @property
    def child_keys(self):
        """ Return all of the children's keys. """
        if self.is_leaf: return None
        else:
            return [node.key for node in self.value]

    @property
    def children(self):
        """ Return this node's children. """
        return self.value

    def sort(self):
        """ Sort children by key name."""
        if self.is_leaf:
            return self
        else:
            self.value = sorted(self.value, key=lambda n: n.key)
            return [node.sort() for node in self.value]

    def __iter__(self):
        """ Iterate over the values or return the value if it's a leaf. """
        if self.is_leaf: yield self
        for val in self.value:
            yield val

    def as_dict(self):
        """ Represent node as a nested dict. """
        def _dict(node):
            if node.is_leaf: 
                return {'key': node.key, 'value': node.value}
            else:
                return {'key': node.key, 'value': [_dict(n) for n in node.value]}
        return _dict(self)

    def to_json(self):
        """ Return a JSON representation of this node and all children. """
        return json.dumps(self.as_dict())

tests.py

#!/usr/bin/env python
import simplejson as json
from datetime import datetime
from django.test import TestCase
from nose.tools import eq_, ok_
from nose.tools import assert_not_equal as neq_

class NodeTests(TestCase):
    def setUp(self):
        self.nodes = []
        self.subNodes3 = [
            Node(key=True, value="a"),
            Node(key=False, value="b")
        ]
        self.subNodes2 = [
            Node(key="A", value=3),
            Node(key="C", value=2),
            Node(key="B", value=5),
            Node(key="1", value=1),
        ]
        self.subNodes1 = [
            Node(key="F", value=9),
            Node(key="R", value=self.subNodes3),
            Node(key="A", value=self.subNodes2)
        ]
        n = Node(key="root", value=self.subNodes1)
        self.node = n
        self.simple_node = Node(key="f", value=[Node(key="a",value=2)])

    def test_create(self):
        """ Test node creation and identity. """
        node = Node(key="hello")
        eq_(node.key, "hello")
        eq_(str(node.key), "hello")

    def test_primitive_sort(self):
        """ Sort a node's children. """
        self.node.sort()
        eq_(self.node.child_keys, ['A','F','R'])

        """ Make sure that its children are sorted. """
        print self.node.children
        eq_(self.node.children[0].child_keys,
            ['1','A','B','C'])

        eq_(self.node.children[2].child_keys,
            [False, True])

    def test_datetime_sort(self):
        """ Sort a node with keys that are datetimes. """
        d4 = datetime(year=2011,month=4,day=1)
        d2 = datetime(year=2011,month=2,day=1)
        d3 = datetime(year=2011,month=3,day=1)
        d1 = datetime(year=1988,month=6,day=19)
        n2 = [
            Node(key=d3, value="a"),
            Node(key=d2, value="d"),
            Node(key=d1, value="b"),
        ]
        n1 = [
            Node(key=d2, value=4),
            Node(key=d1, value=6),
            Node(key=d3, value=3),
            Node(key=d4, value=n2),
        ]
        n = Node(key="Key", value=n1)
        n.sort()
        eq_(n.child_keys, [d1,d2,d3,d4])
        eq_(n.children[3].child_keys, [d1,d2,d3])

    def test_add_child(self):
        """ Test two ways of adding a child to the node. """
        d1 = datetime(year=2011,month=5,day=1)
        d2 = datetime(year=2012, month=4, day=3)
        self.node.add_child(Node(key=d1, value=4))
        eq_(len(self.node.children), 4)

        self.node.add_child(key=d2, val=4)
        eq_(len(self.node.children), 5)

    def test_iterator(self):
        """ Test the node iterator. """
        for i,node in enumerate(self.node):
            print i, node
        eq_(i,2)

    def test_as_dict(self):
        """ Convert the Node into a dict. """
        n2 = self.simple_node
        eq_(n2.as_dict(), {'key':'f','value':[{'key':'a','value':2}]})

    def test_to_json(self):
        n2 = self.simple_node
        eq_(n2.to_json(),
            json.dumps({'key':'f','value':[{'key':'a','value':2}]}))