Add tests and alter changelog

hyanwong · hyanwong · commit 02778c44adf2 · 2020-04-28T18:42:51.000+01:00
diff --git a/python/CHANGELOG.rst b/python/CHANGELOG.rst
@@ -6,6 +6,10 @@ In development
 
 **New features**
 
+- Add classes to SVG drawings to allow easy adjustment and styling, and document the new
+  draw_svg() methods. This also fixes a bug when using SVG drawings in Jupyter notebooks
+  (:user:`hyanwong`, :pr:`555`)
+
 - Add an optional node traversal order in ``tskit.Tree`` that uses the minimum
   lexicographic order of leaf nodes visited. This ordering (``"minlex_postorder"``)
   adds more determinism because it constraints the order in which children of
diff --git a/python/requirements/CI/requirements.txt b/python/requirements/CI/requirements.txt
@@ -22,4 +22,5 @@ sphinx==2.4.4
 sphinx-argparse==0.2.5
 sphinx-issues==1.2.0
 sphinx_rtd_theme==0.4.3
-svgwrite==1.4
+svgwrite==1.4
+xmlunittest==0.5.0
diff --git a/python/requirements/development.txt b/python/requirements/development.txt
@@ -22,4 +22,5 @@ sphinx==2.4.4 #Pinned as breathe v3 compatibility is rough for now.
 sphinx-argparse
 sphinx-issues
 sphinx_rtd_theme
-svgwrite
+svgwrite
+xmlunittest
diff --git a/python/tests/test_drawing.py b/python/tests/test_drawing.py
@@ -30,6 +30,7 @@
 import xml.etree
 
 import msprime
+from xmlunittest import XmlTestCase
 
 import tests.tsutil as tsutil
 import tskit
@@ -44,18 +45,20 @@ def get_binary_tree(self):
         ts = msprime.simulate(10, random_seed=1, mutation_rate=1)
         return next(ts.trees())
 
-    def get_nonbinary_tree(self):
+    def get_nonbinary_ts(self, sample_size=10):
         demographic_events = [
             msprime.SimpleBottleneck(time=0.1, population=0, proportion=0.5)
         ]
-        ts = msprime.simulate(
-            10,
+        return msprime.simulate(
+            sample_size,
             recombination_rate=5,
             mutation_rate=10,
             demographic_events=demographic_events,
             random_seed=1,
         )
-        for t in ts.trees():
+
+    def get_nonbinary_tree(self, sample_size=10):
+        for t in self.get_nonbinary_ts(sample_size).trees():
             for u in t.nodes():
                 if len(t.children(u)) > 2:
                     return t
@@ -1134,16 +1137,44 @@ def test_max_tree_height(self):
                 t.draw_text(max_tree_height=bad_max_tree_height)
 
 
-class TestDrawSvg(TestTreeDraw):
+class TestDrawSvg(TestTreeDraw, XmlTestCase):
     """
     Tests the SVG tree drawing.
     """
 
-    def verify_basic_svg(self, svg, width=200, height=200):
+    def verify_basic_svg(self, svg, width=200, height=200, new_routine=True):
+        prefix = "{http://www.w3.org/2000/svg}"
         root = xml.etree.ElementTree.fromstring(svg)
-        self.assertEqual(root.tag, "{http://www.w3.org/2000/svg}svg")
+        self.assertEqual(root.tag, prefix + "svg")
         self.assertEqual(width, int(root.attrib["width"]))
         self.assertEqual(height, int(root.attrib["height"]))
+        if new_routine:
+            # Verify the class structure of the svg
+            root_group = root.find(prefix + "g")
+            self.assertIn("class", root_group.attrib)
+            self.assertRegexpMatches(
+                root_group.attrib["class"], r"\b(tree|tree-sequence)\b"
+            )
+            if "tree-sequence" in root_group.attrib["class"]:
+                trees = root_group.find(prefix + "g")
+                self.assertIn("class", trees.attrib)
+                self.assertRegexpMatches(trees.attrib["class"], r"\btrees\b")
+                first_tree = trees.find(prefix + "g")
+                self.assertIn("class", first_tree.attrib)
+                self.assertRegexpMatches(first_tree.attrib["class"], r"\btree\b")
+            else:
+                first_tree = root_group
+            # Check that we have edges, symbols, and labels groups
+            for group in first_tree.findall(prefix + "g"):
+                self.assertIn("class", group.attrib)
+                cls = group.attrib["class"]
+                self.assertRegexpMatches(cls, r"\b(edges|symbols|labels)\b")
+                if "symbols" in cls or "labels" in cls:
+                    # Check that we have nodes & mutations subgroups
+                    for subgroup in group.findall(prefix + "g"):
+                        self.assertIn("class", subgroup.attrib)
+                        subcls = subgroup.attrib["class"]
+                        self.assertRegexpMatches(subcls, r"\b(nodes|mutations)\b")
 
     def test_draw_file(self):
         t = self.get_binary_tree()
@@ -1177,7 +1208,7 @@ def test_draw_file(self):
     def test_draw_defaults(self):
         t = self.get_binary_tree()
         svg = t.draw()
-        self.verify_basic_svg(svg)
+        self.verify_basic_svg(svg, new_routine=False)
         svg = t.draw_svg()
         self.verify_basic_svg(svg)
 
@@ -1483,3 +1514,156 @@ def test_tree_height_scale(self):
         for bad_scale in [0, "", "NOT A SCALE"]:
             with self.assertRaises(ValueError):
                 ts.draw_svg(tree_height_scale=bad_scale)
+
+    def test_known_svg_tree(self):
+        tree = self.get_nonbinary_tree(sample_size=3)
+        svg = tree.draw_svg(
+            root_svg_attributes={"id": "AB"}, style=".edges {stroke: red}"
+        )
+        expected_svg = b"""<?xml version="1.0" encoding="utf-8" ?>
+<svg baseProfile="full" height="200" id="AB" version="1.1" width="200"
+ xmlns="http://www.w3.org/2000/svg" xmlns:ev="http://www.w3.org/2001/xml-events"
+ xmlns:xlink="http://www.w3.org/1999/xlink">
+  <defs><style type="text/css"><![CDATA[.edges {stroke: red}]]></style></defs>
+  <g class="tree t0">
+    <g class="edges" fill="none" stroke="black">
+      <path class="p4 c0" d="M 55.0 170.0 V 30.0 H 100.0"/>
+      <path class="p4 c1" d="M 100.0 170.0 V 30.0 H 100.0"/>
+      <path class="p4 c2" d="M 145.0 170.0 V 30.0 H 100.0"/>
+    </g>
+    <g class="symbols">
+      <g class="nodes">
+        <circle class="n4" cx="100.0" cy="30.0" r="3"/>
+        <circle class="n0 sample" cx="55.0" cy="170.0" r="3"/>
+        <circle class="n1 sample" cx="100.0" cy="170.0" r="3"/>
+        <circle class="n2 sample" cx="145.0" cy="170.0" r="3"/>
+      </g>
+      <g class="mutations" fill="red"/>
+    </g>
+    <g class="labels" dominant-baseline="middle" font-size="14">
+      <g class="nodes">
+        <g text-anchor="start"/>
+        <g text-anchor="middle">
+          <g transform="translate(100.0, 25.0)"><text class="n4">4</text></g>
+          <g transform="translate(55.0, 190.0)"><text class="n0 sample">0</text></g>
+          <g transform="translate(100.0, 190.0)"><text class="n1 sample">1</text></g>
+          <g transform="translate(145.0, 190.0)"><text class="n2 sample">2</text></g>
+        </g>
+        <g text-anchor="end"/>
+      </g>
+      <g class="mutations" font-style="italic">
+        <g text-anchor="start"/>
+        <g text-anchor="end"/>
+      </g>
+    </g>
+  </g>
+</svg>"""
+        self.assertXmlEquivalentOutputs(svg, expected_svg)
+
+    def test_known_svg_ts(self):
+        ts = self.get_nonbinary_ts(sample_size=3)
+        svg = ts.draw_svg(
+            root_svg_attributes={"id": "AB"}, style=".edges {stroke: red}"
+        )
+        self.verify_basic_svg(svg, width=200 * ts.num_trees)
+        expected_svg = b"""<?xml version="1.0" encoding="utf-8" ?>
+<svg baseProfile="full" id="AB" height="200" version="1.1" width="400"
+ xmlns="http://www.w3.org/2000/svg" xmlns:ev="http://www.w3.org/2001/xml-events"
+ xmlns:xlink="http://www.w3.org/1999/xlink">
+  <defs><style type="text/css"><![CDATA[.edges {stroke: red}]]></style></defs>
+  <g class="tree-sequence">
+    <g class="trees">
+      <g class="tree t0" transform="translate(20 15)">
+        <g class="edges" fill="none" stroke="black">
+          <path class="p4 c0" d="M 50.0 140.0 V 30.0 H 90.0"/>
+          <path class="p4 c1" d="M 90.0 140.0 V 30.0 H 90.0"/>
+          <path class="p4 c2" d="M 130.0 140.0 V 30.0 H 90.0"/>
+        </g>
+        <g class="symbols">
+          <g class="nodes">
+            <circle class="n4" cx="90.0" cy="30.0" r="3"/>
+            <circle class="n0 sample" cx="50.0" cy="140.0" r="3"/>
+            <circle class="n1 sample" cx="90.0" cy="140.0" r="3"/>
+            <circle class="n2 sample" cx="130.0" cy="140.0" r="3"/>
+          </g>
+          <g class="mutations" fill="red"/>
+        </g>
+        <g class="labels" dominant-baseline="middle" font-size="14">
+          <g class="nodes">
+            <g text-anchor="start"/>
+            <g text-anchor="middle">
+              <g transform="translate(90.0, 25.0)"><text class="n4">4</text></g>
+              <g transform="translate(50.0, 160.0)"><text class="n0 sample">0</text></g>
+              <g transform="translate(90.0, 160.0)"><text class="n1 sample">1</text></g>
+              <g transform="translate(130.0, 160.0)"><text class="n2 sample">2</text></g>
+            </g>
+            <g text-anchor="end"/>
+          </g>
+          <g class="mutations" font-style="italic">
+            <g text-anchor="start"/>
+            <g text-anchor="end"/>
+          </g>
+        </g>
+      </g>
+      <g class="tree t1" transform="translate(200.0 15)">
+        <g class="edges" fill="none" stroke="black">
+          <path class="p4 c1" d="M 50.0 140.0 V 30.0 H 80.0"/>
+          <path class="p4 c3" d="M 110.0 100.38696392754377 V 30.0 H 80.0"/>
+          <path class="p3 c0" d="M 90.0 140.0 V 100.38696392754377 H 110.0"/>
+          <path class="p3 c2" d="M 130.0 140.0 V 100.38696392754377 H 110.0"/>
+        </g>
+        <g class="symbols">
+          <g class="nodes">
+            <circle class="n4" cx="80.0" cy="30.0" r="3"/>
+            <circle class="n1 sample" cx="50.0" cy="140.0" r="3"/>
+            <circle class="n3" cx="110.0" cy="100.38696392754377" r="3"/>
+            <circle class="n0 sample" cx="90.0" cy="140.0" r="3"/>
+            <circle class="n2 sample" cx="130.0" cy="140.0" r="3"/>
+          </g>
+          <g class="mutations" fill="red">
+            <rect class="m0 s0 n1" height="6" transform="translate(-3 -3)" width="6"
+             x="50.0" y="85.0"/>
+          </g>
+        </g>
+        <g class="labels" dominant-baseline="middle" font-size="14">
+          <g class="nodes">
+            <g text-anchor="start">
+              <g transform="translate(115.0, 95.38696392754377)">
+                <text class="n3">3</text>
+              </g>
+            </g>
+            <g text-anchor="middle">
+              <g transform="translate(80.0, 25.0)"><text class="n4">4</text></g>
+              <g transform="translate(50.0, 160.0)"><text class="n1 sample">1</text></g>
+              <g transform="translate(90.0, 160.0)"><text class="n0 sample">0</text></g>
+              <g transform="translate(130.0, 160.0)"><text class="n2 sample">2</text></g>
+            </g>
+            <g text-anchor="end"/>
+          </g>
+          <g class="mutations" font-style="italic">
+            <g text-anchor="start"/>
+            <g text-anchor="end">
+              <g transform="translate(45.0, 89.0)"><text class="m0 s0 n1">0</text></g>
+            </g>
+          </g>
+        </g>
+      </g>
+    </g>
+    <g class="axis">
+      <line stroke="black" x1="20" x2="380" y1="180" y2="180"/>
+      <line stroke="black" x1="20" x2="20" y1="175" y2="185"/>
+      <g transform="translate(20, 200)">
+        <text font-size="14" font-weight="bold" text-anchor="middle">0.00</text>
+      </g>
+      <line stroke="black" x1="200.0" x2="200.0" y1="175" y2="185"/>
+      <g transform="translate(200.0, 200)">
+        <text font-size="14" font-weight="bold" text-anchor="middle">0.16</text>
+      </g>
+      <line stroke="black" x1="380.0" x2="380.0" y1="175" y2="185"/>
+      <g transform="translate(380.0, 200)">
+        <text font-size="14" font-weight="bold" text-anchor="middle">1.00</text>
+      </g>
+    </g>
+  </g>
+</svg>"""
+        self.assertXmlEquivalentOutputs(svg, expected_svg)
