Update jit_graph and tests

2 files changed, 44 insertions, 29 deletions

diff --git a/networkx/readwrite/json_graph/jit.py b/networkx/readwrite/json_graph/jit.py
index b9472301..cd390c78 100644
--- a/networkx/readwrite/json_graph/jit.py
+++ b/networkx/readwrite/json_graph/jit.py
@@ -1,7 +1,14 @@
+#    Copyright (C) 2011-2016 by
+#    Aric Hagberg <hagberg@lanl.gov>
+#    Dan Schult <dschult@colgate.edu>
+#    Pieter Swart <swart@lanl.gov>
+#    All rights reserved.
+#    BSD license.
+
 """
 Read and write NetworkX graphs as JavaScript InfoVis Toolkit (JIT) format JSON.
 
-See the JIT documentation and examples at http://thejit.org
+See the `JIT documentation`_ for more examples.
 
 Format
 ------
@@ -23,27 +30,21 @@ var json = [
 
   'other nodes go here...'
 ];
+.. _JIT documentation: http://thejit.org
 """
 
-#    Copyright (C) 2011-2013 by
-#    Aric Hagberg <hagberg@lanl.gov>
-#    Dan Schult <dschult@colgate.edu>
-#    Pieter Swart <swart@lanl.gov>
-#    All rights reserved.
-#    BSD license.
-
-__all__ = ['jit_graph', 'jit_data']
-
-import networkx as nx
 import json
+import networkx as nx
+from networkx.utils.decorators import not_implemented_for
 
+__all__ = ['jit_graph', 'jit_data']
 
 def jit_graph(data):
     """Read a graph from JIT JSON.
 
     Parameters
     ----------
-    data : JIT JSON string
+    data : JSON Graph Object
 
     Returns
     -------
@@ -51,12 +52,14 @@ def jit_graph(data):
     """
     G = nx.Graph()
     for node in data:
-        for adj in node['adjacencies']:
-            G.add_edge(node['id'], adj['nodeTo'], **adj['data'])
-            G.add_node(node['id'], **node['data'])
+        G.add_node(node['id'], **node['data'])
+        if node.get('adjacencies') is not None:
+            for adj in node['adjacencies']:
+                G.add_edge(node['id'], adj['nodeTo'], **adj['data'])
     return G
 
 
+@not_implemented_for('multigraph')
 def jit_data(G, indent=None):
     """Return data in JIT JSON format.
 
@@ -64,6 +67,12 @@ def jit_data(G, indent=None):
     ----------
     G : NetworkX Graph
 
+    indent: optional, default=None
+        If indent is a non-negative integer, then JSON array elements and object
+        members will be pretty-printed with that indent level. An indent level
+        of 0, or negative, will only insert newlines. None (the default) selects
+        the most compact representation.
+
     Returns
     -------
     data: JIT JSON string
@@ -75,10 +84,7 @@ def jit_data(G, indent=None):
             "name": node
         }
         # node data
-        if G.node[node]:
-            json_node["data"] = G.node[node]
-        else:
-            json_node["data"] = {}
+        json_node["data"] = G.node[node]
         # adjacencies
         if G[node]:
             json_node["adjacencies"] = []
@@ -87,10 +93,7 @@ def jit_data(G, indent=None):
                     "nodeTo": neighbour,
                 }
                 # adjacency data
-                if G.edge[node][neighbour]:
-                    adjacency["data"] = G.edge[node][neighbour]
-                else:
-                    adjacency["data"] = {}
+                adjacency["data"] = G.edge[node][neighbour]
                 json_node["adjacencies"].append(adjacency)
         json_graph.append(json_node)
     return json.dumps(json_graph, indent=indent)
diff --git a/networkx/readwrite/json_graph/tests/test_jit.py b/networkx/readwrite/json_graph/tests/test_jit.py
index 3e4482b8..00397fcc 100644
--- a/networkx/readwrite/json_graph/tests/test_jit.py
+++ b/networkx/readwrite/json_graph/tests/test_jit.py
@@ -3,15 +3,27 @@ from nose.tools import assert_true
 import networkx as nx
 from networkx.readwrite.json_graph import jit_data, jit_graph
 
-class TestJIT:
 
-    def test_jit_graph(self):
+class TestJIT(object):
+    def test_jit(self):
         G = nx.Graph()
-        G.add_node('Node1', node_data = 'foobar')
-        G.add_node('Node3', node_data = 'bar')
-        G.add_edge('Node1', 'Node2', weight = 9, something='isSomething')
-        G.add_edge('Node2', 'Node3', weight = 4, something='isNotSomething')
+        G.add_node('Node1', node_data='foobar')
+        G.add_node('Node3', node_data='bar')
+        G.add_node('Node4')
+        G.add_edge('Node1', 'Node2', weight=9, something='isSomething')
+        G.add_edge('Node2', 'Node3', weight=4, something='isNotSomething')
         G.add_edge('Node1', 'Node2')
         d = jit_data(G)
         K = jit_graph(json.loads(d))
         assert_true(nx.is_isomorphic(G, K))
+
+    def test_jit_2(self):
+        G = nx.Graph()
+        G.add_node(1, node_data=3)
+        G.add_node(3, node_data=0)
+        G.add_edge(1, 2, weight=9, something=0)
+        G.add_edge(2, 3, weight=4, something=3)
+        G.add_edge(1, 2)
+        d = jit_data(G)
+        K = jit_graph(json.loads(d))
+        assert_true(nx.is_isomorphic(G, K))