Add a class for svg strings

Author: hyanwong

docs/python-api.md

@@ -1509,3 +1509,9 @@ Also see the {ref}`sec_python_api_metadata` summary.
     :members:
 ```
 
+#### The {class}`SVGString` class
+```{eval-rst}
+.. autoclass:: SVGString
+    :members:
+    :private-members: _repr_svg_
+```

python/CHANGELOG.rst

@@ -2,6 +2,12 @@
 [0.5.1] - 2022-0X-XX
 --------------------
 
+**Changes**
+
+- SVG drawing routines now return a special string object that is automatically
+  rendered in a Jupyter notebook (:user:`hyanwong`, :pr:`2377`)
+
+
 --------------------
 [0.5.0] - 2022-06-22
 --------------------

python/tests/test_drawing.py

@@ -1446,6 +1446,12 @@ def test_max_time(self):
             with pytest.raises(ValueError):
                 t.draw_text(max_time=bad_max_time)
 
+    def test_no_repr_svg(self):
+        tree = self.get_simple_ts().first()
+        output = tree.draw(format="unicode")
+        with pytest.raises(AttributeError, match="no attribute"):
+            output._repr_svg_()
+
 
 class TestDrawSvg(TestTreeDraw, xmlunittest.XmlTestMixin):
     """
@@ -1490,6 +1496,15 @@ def verify_basic_svg(self, svg, width=200, height=200, num_trees=1):
             cls = group.attrib["class"]
             assert re.search(r"\broot\b", cls)
 
+    def test_repr_svg(self):
+        ts = self.get_simple_ts()
+        svg = ts.draw_svg()
+        assert str(svg) == svg._repr_svg_()
+        svg = ts.first().draw_svg()
+        assert str(svg) == svg._repr_svg_()
+        svg = ts.first().draw(format="svg")
+        assert str(svg) == svg._repr_svg_()
+
     def test_draw_to_file(self, tmp_path):
         # NB: to view output files for testing changes to drawing code, it is possible
         # to save to a fixed directory using e.g. `pytest --basetemp=/tmp/svgtest ...`

python/tskit/__init__.py

@@ -84,6 +84,7 @@
     all_tree_labellings,
     TopologyCounter,
 )
+from tskit.drawing import SVGString  # NOQA
 from tskit.exceptions import *  # NOQA
 from tskit.util import *  # NOQA
 from tskit.metadata import *  # NOQA

python/tskit/drawing.py

@@ -94,6 +94,16 @@ def linear_transform(self, y):
         return self.plot_min - (y - self.min_time) * y_scale
 
 
+class SVGString(str):
+    "A string containing an SVG representation"
+
+    def _repr_svg_(self):
+        """
+        Simply return the SVG string: called by jupyter notebooks to render trees.
+        """
+        return self
+
+
 def check_orientation(orientation):
     if orientation is None:
         orientation = TOP
@@ -415,7 +425,7 @@ def remap_style(original_map, new_key, none_value):
             mutation_attrs=mutation_attrs,
             order=order,
         )
-        return tree.drawing.tostring()
+        return SVGString(tree.drawing.tostring())
 
     else:
         if width is not None:

python/tskit/trees.py

@@ -1794,7 +1794,7 @@ def draw_svg(
             from the same tree sequence may share some plotted mutations.
 
         :return: An SVG representation of a tree.
-        :rtype: str
+        :rtype: SVGString
         """
         draw = drawing.SvgTree(
             self,
@@ -1824,7 +1824,7 @@ def draw_svg(
         if path is not None:
             # TODO: removed the pretty here when this is stable.
             draw.drawing.saveas(path, pretty=True)
-        return output
+        return drawing.SVGString(output)
 
     def draw(
         self,
@@ -6502,7 +6502,7 @@ def draw_svg(
             at each y tickmark.
 
         :return: An SVG representation of a tree sequence.
-        :rtype: str
+        :rtype: SVGString
 
         .. note::
             Technically, x_lim[0] specifies a *minimum* value for the start of the X
@@ -6541,7 +6541,7 @@ def draw_svg(
         if path is not None:
             # TODO remove the 'pretty' when we are done debugging this.
             draw.drawing.saveas(path, pretty=True)
-        return output
+        return drawing.SVGString(output)
 
     def draw_text(
         self,