Merge pull request #28 from WestHealth/feature/27-add-layout-support

Feature/27 add layout support

pyvis/network.py

@@ -28,7 +28,8 @@ def __init__(self,
                  directed=False,
                  notebook=False,
                  bgcolor="#ffffff",
-                 font_color=False):
+                 font_color=False,
+                 layout=None):
         """
         :param height: The height of the canvas
         :param width: The width of the canvas
@@ -37,13 +38,15 @@ def __init__(self,
         :param notebook: True if using jupyter notebook.
         :param bgcolor: The background color of the canvas.
         :font_color: The color of the node labels text
+        :layout: Use hierarchical layout if this is set
 
         :type height: num or str
         :type width: num or str
         :type directed: bool
         :type notebook: bool
         :type bgcolor: str
         :type font_color: str
+        :type layout: bool
         """
         self.nodes = []
         self.edges = []
@@ -56,7 +59,7 @@ def __init__(self,
         self.bgcolor = bgcolor
         self.use_DOT = False
         self.dot_lang = ""
-        self.options = Options()
+        self.options = Options(layout)
         self.widget = False
         self.node_ids = []
         self.template = None
@@ -65,6 +68,7 @@ def __init__(self,
 
         if notebook:
             self.prep_notebook()
+
 
     def __str__(self):
         """
@@ -344,7 +348,7 @@ def add_edge(self, source, to, **options):
         if not edge_exists:
             e = Edge(source, to, self.directed, **options)
             self.edges.append(e.options)
-
+            
     def add_edges(self, edges):
         """
         This method serves to add multiple edges between existing nodes
@@ -721,6 +725,7 @@ def to_json(self):
         return json.dumps(self, default=lambda o: o.__dict__,
                           sort_keys=True, indent=4)
 
+
     def set_edge_smooth(self, smooth_type):
         """
         Sets the smooth.type attribute of the edges.

pyvis/options.py

@@ -75,14 +75,83 @@ def __getitem__(self, item):
         return self.__dict__[item]
 
 
+class Layout(object):
+    """
+    Acts as the camera that looks on the canvas.
+    Does the animation, zooming and focusing.
+    """
+
+    def __repr__(self):
+        return str(self.__dict__)
+
+    def __init__(self, randomSeed=None, improvedLayout=True):
+            if not randomSeed:
+                self.randomSeed = 0
+            else:
+                self.radnomSeed = randomSeed
+            self.improvedLayout = improvedLayout
+            self.hierarchical = self.Hierarchical(enabled=True)
+
+    def set_separation(self, distance):
+        """
+        The distance between the different levels.
+        """
+        self.hierarchical.levelSeparation = distance
+
+    def set_tree_spacing(self, distance):
+        """
+        Distance between different trees (independent networks). This is
+        only for the initial layout. If you enable physics, the repulsion
+        model will denote the distance between the trees.
+        """
+        self.hierarchical.treeSpacing = distance
+
+    def set_edge_minimization(self, status):
+        """
+        Method for reducing whitespace. Can be used alone or together with
+        block shifting. Enabling block shifting will usually speed up the
+        layout process. Each node will try to move along its free axis to
+        reduce the total length of it's edges. This is mainly for the
+        initial layout. If you enable physics, they layout will be determined
+        by the physics. This will greatly speed up the stabilization time
+        """
+        self.hierarchical.edgeMinimization = status
+
+    class Hierarchical(object):
+
+        def __getitem__(self, item):
+            return self.__dict__[item]
+
+        def __init__(self,
+                    enabled=False,
+                    levelSeparation=150,
+                    treeSpacing=200,
+                    blockShifting=True,
+                    edgeMinimization=True,
+                    parentCentralization=True,
+                    sortMethod='hubsize'):
+
+            self.enabled = enabled
+            self.levelSeparation = levelSeparation
+            self.treeSpacing = treeSpacing
+            self.blockShifting = blockShifting
+            self.edgeMinimization = edgeMinimization
+            self.parentCentralization = parentCentralization
+            self.sortMethod = sortMethod
+
+
 
 class Options(object):
 
     def __repr__(self):
         return str(self.__dict__)
 
-    def __init__(self):
-        # self.layout = Layout()
+    def __getitem__(self, item):
+        return self.__dict__[item]
+
+    def __init__(self, layout=None):
+        if layout:
+            self.layout = Layout()
         self.interaction = Interaction()
         self.configure = Configure()
         self.physics = Physics()

pyvis/tests/test_graph.py

@@ -302,3 +302,52 @@ def test_inherit_colors(self):
         self.assertTrue(self.g.options.edges.color.inherit)
         self.g.inherit_edge_colors_from(False)
         self.assertFalse(self.g.options.edges.color.inherit)
+
+
+class LayoutTestCase(unittest.TestCase):
+
+    def setUp(self):
+        self.g = Network(layout=True)
+
+    def test_can_enable_init(self):
+        self.assertTrue(self.g.options['layout'])
+
+    def test_layout_disabled(self):
+        self.g = Network()
+        self.assertRaises(KeyError, lambda: self.g.options['layout'])
+
+    def test_levelSeparation(self):
+        self.assertTrue(self.g.options.layout.hierarchical.levelSeparation)
+
+    def test_treeSpacing(self):
+        self.assertTrue(self.g.options.layout.hierarchical.treeSpacing)
+
+    def test_blockShifting(self):
+        self.assertTrue(self.g.options.layout.hierarchical.blockShifting)
+
+    def test_edgeMinimization(self):
+        self.assertTrue(self.g.options.layout.hierarchical.edgeMinimization)
+
+    def test_parentCentralization(self):
+        self.assertTrue(self.g.options.layout.hierarchical.parentCentralization)
+
+    def test_sortMethod(self):
+        self.assertTrue(self.g.options.layout.hierarchical.sortMethod)
+
+    def test_set_edge_minimization(self):
+        self.g.options.layout.set_separation(10)
+        self.assertTrue(self.g.options.layout.hierarchical.levelSeparation == 10)
+
+    def test_set_tree_spacing(self):
+        self.g.options.layout.set_tree_spacing(10)
+        self.assertTrue(self.g.options.layout.hierarchical.treeSpacing == 10)
+
+    def test_set_edge_minimization(self):
+        self.g.options.layout.set_edge_minimization(True)
+        self.assertTrue(self.g.options.layout.hierarchical.edgeMinimization == True)
+        self.g.options.layout.set_edge_minimization(False)
+        self.assertTrue(self.g.options.layout.hierarchical.edgeMinimization == False)
+
+
+
+