Merge pull request #30 from benabel/master

Implement svg output in notebooks #29

Author: joowani

binarytree/__init__.py

@@ -5,10 +5,10 @@
 from dataclasses import dataclass
 from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
 
-from graphviz import Digraph, nohtml
 from pkg_resources import get_distribution
 
 from binarytree.exceptions import (
+ GraphvizImportError,
  NodeIndexError,
  NodeModifyError,
  NodeNotFoundError,
@@ -18,6 +18,15 @@
  TreeHeightError,
 )
 
+try:
+ from graphviz import Digraph, nohtml
+
+ GRAPHVIZ_INSTALLED = True
+except ImportError:
+ GRAPHVIZ_INSTALLED = False
+ Digraph = Any
+ from binarytree.layout import generate_svg
+
 __version__ = get_distribution("binarytree").version
 
 LEFT_FIELD = "left"
@@ -463,29 +472,32 @@ def _repr_svg_(self) -> str:
 
  .. _Jupyter notebooks: https://jupyter.org
  """
- # noinspection PyProtectedMember
- return str(self.graphviz()._repr_svg_())
+ if GRAPHVIZ_INSTALLED:
+ # noinspection PyProtectedMember
+ return str(self.graphviz()._repr_svg_())
+ else:
+ return generate_svg(self.values) # pragma: no cover
 
  def graphviz(self, *args: Any, **kwargs: Any) -> Digraph:
  """Return a graphviz.Digraph_ object representing the binary tree.
-
  This method's positional and keyword arguments are passed directly into the
  the Digraph's **__init__** method.
-
  :return: graphviz.Digraph_ object representing the binary tree.
-
+ :raise binarytree.exceptions.GraphvizImportError: If graphviz is not installed
  .. code-block:: python
-
  >>> from binarytree import tree
  >>>
  >>> t = tree()
  >>>
  >>> graph = t.graphviz() # Generate a graphviz object
  >>> graph.body # Get the DOT body
  >>> graph.render() # Render the graph
-
  .. _graphviz.Digraph: https://graphviz.readthedocs.io/en/stable/api.html#digraph
  """
+ if not GRAPHVIZ_INSTALLED:
+ raise GraphvizImportError(
+ "Can't use graphviz method if graphviz module is not installed"
+ )
  if "node_attr" not in kwargs:
  kwargs["node_attr"] = {
  "shape": "record",
@@ -494,20 +506,14 @@ def graphviz(self, *args: Any, **kwargs: Any) -> Digraph:
  "fillcolor": "lightgray",
  "fontcolor": "black",
  }
-
  digraph = Digraph(*args, **kwargs)
-
  for node in self:
  node_id = str(id(node))
-
  digraph.node(node_id, nohtml(f"<l>|<v> {node.value}|<r>"))
-
  if node.left is not None:
  digraph.edge(f"{node_id}:l", f"{id(node.left)}:v")
-
  if node.right is not None:
  digraph.edge(f"{node_id}:r", f"{id(node.right)}:v")
-
  return digraph
 
  def pprint(self, index: bool = False, delimiter: str = "-") -> None:

binarytree/exceptions.py

@@ -28,3 +28,7 @@ class NodeValueError(BinaryTreeError):
 
 class TreeHeightError(BinaryTreeError):
  """Tree height was invalid."""
+
+
+class GraphvizImportError(BinaryTreeError):
+ """graphviz module is not installed"""

binarytree/layout.py

@@ -0,0 +1,118 @@
+""" Module containing layout related algorithms."""
+from typing import List, Tuple, Union
+
+
+def _get_coords(
+ values: List[Union[float, int, None]]
+) -> Tuple[
+ List[Tuple[int, int, Union[float, int, None]]], List[Tuple[int, int, int, int]]
+]:
+ """Generate the coordinates used for rendering the nodes and edges.
+
+ node and edges are stored as tuples in the form node: (x, y, label) and
+ edge: (x1, y1, x2, y2)
+
+ Each coordinate is relative y is the depth, x is the position of the node
+ on a level from left to right 0 to 2**depth -1
+
+ :param values: Values of the binary tree.
+ :type values: list of ints
+ :return: nodes and edges list
+ :rtype: two lists of tuples
+
+ """
+ x = 0
+ y = 0
+ nodes = []
+ edges = []
+
+ # root node
+ nodes.append((x, y, values[0]))
+ # append other nodes and their edges
+ y += 1
+ for value in values[1:]:
+ if value is not None:
+ nodes.append((x, y, value))
+ edges.append((x // 2, y - 1, x, y))
+ x += 1
+ # check if level is full
+ if x == 2 ** y:
+ x = 0
+ y += 1
+ return nodes, edges
+
+
+def generate_svg(values: List[Union[float, int, None]]) -> str:
+ """Generate a svg image from a binary tree
+
+ A simple layout is used based on a perfect tree of same height in which all
+ leaves would be regularly spaced.
+
+ :param values: Values of the binary tree.
+ :type values: list of ints
+ :return: the svg image of the tree.
+ :rtype: str
+ """
+ node_size = 16.0
+ stroke_width = 1.5
+ gutter = 0.5
+ x_scale = (2 + gutter) * node_size
+ y_scale = 3.0 * node_size
+
+ # retrieve relative coordinates
+ nodes, edges = _get_coords(values)
+ y_min = min([n[1] for n in nodes])
+ y_max = max([n[1] for n in nodes])
+
+ # generate the svg string
+ svg = f"""
+ <svg width="{x_scale * 2**y_max}" height="{y_scale * (2 + y_max)}"
+ xmlns="http://www.w3.org/2000/svg">
+ <style>
+ .bt-label {{
+ font: 300 {node_size}px sans-serif;;
+ text-align: center;
+ dominant-baseline: middle;
+ text-anchor: middle;
+ }}
+ .bt-node {{
+ fill: lightgray;
+ stroke-width: {stroke_width};
+ }}
+
+ </style>
+ <g stroke="#111">
+ """
+ # scales
+
+ def scalex(x: int, y: int) -> float:
+ depth = y_max - y
+ # offset
+ x = 2 ** (depth + 1) * x + 2 ** depth - 1
+ return 1 + node_size + x_scale * x / 2
+
+ def scaley(y: int) -> float:
+ return float(y_scale * (1 + y - y_min))
+
+ # edges
+ def svg_edge(x1: float, y1: float, x2: float, y2: float) -> str:
+ """Generate svg code for an edge"""
+ return f"""<line x1="{x1}" x2="{x2}" y1="{y1}" y2="{y2}"/>"""
+
+ for a in edges:
+ x1, y1, x2, y2 = a
+ svg += svg_edge(scalex(x1, y1), scaley(y1), scalex(x2, y2), scaley(y2))
+
+ # nodes
+ def svg_node(x: float, y: float, label: str = "") -> str:
+ """Generate svg code for a node and his label"""
+ return f"""
+ <circle class="bt-node" cx="{x}" cy="{y}" r="{node_size}"/>
+ <text class="bt-label" x="{x}" y="{y}">{label}</text>"""
+
+ for n in nodes:
+ x, y, label = n
+ svg += svg_node(scalex(x, y), scaley(y), str(label))
+
+ svg += "</g></svg>"
+ return svg

tests/test_layout.py

@@ -0,0 +1,17 @@
+import xml.etree.ElementTree as ET
+
+from binarytree.layout import _get_coords, generate_svg
+
+
+def test_get_coords():
+ values = [0, 6, 5, None, 1, 4, 2]
+ assert _get_coords(values) == (
+ [(0, 0, 0), (0, 1, 6), (1, 1, 5), (1, 2, 1), (2, 2, 4), (3, 2, 2)],
+ [(0, 0, 0, 1), (0, 0, 1, 1), (0, 1, 1, 2), (1, 1, 2, 2), (1, 1, 3, 2)],
+ )
+
+
+def test_svg():
+ svg = generate_svg([0, 1, 2])
+ svg_tree = ET.fromstring(svg)
+ assert svg_tree.tag == "{http://www.w3.org/2000/svg}svg"