Fix draw_svg warning for mutations above unplotted nodes

Author: duncanMR

python/tests/test_drawing.py

@@ -1,6 +1,6 @@
 # MIT License
 #
-# Copyright (c) 2018-2022 Tskit Developers
+# Copyright (c) 2018-2023 Tskit Developers
 # Copyright (C) 2017 University of Oxford
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
@@ -1615,10 +1615,12 @@ def test_draw_multiroot(self):
 
     def test_draw_mutations_over_roots(self):
         t = self.get_mutations_over_roots_tree()
-        svg = t.draw()
-        self.verify_basic_svg(svg)
-        svg = t.draw_svg()
-        self.verify_basic_svg(svg)
+        with pytest.warns(UserWarning, match="nodes which are not present"):
+            svg = t.draw()
+            self.verify_basic_svg(svg)
+        with pytest.warns(UserWarning, match="nodes which are not present"):
+            svg = t.draw_svg()
+            self.verify_basic_svg(svg)
 
     def test_draw_unary(self):
         t = self.get_unary_node_tree()
@@ -2742,20 +2744,23 @@ def test_known_svg_ts_mutation_times_logscale(self, overwrite_viz, draw_plotbox)
             svg, "ts_mut_times_logscale.svg", overwrite_viz, width=200 * ts.num_trees
         )
 
-    def test_known_svg_ts_mut_no_edges(self, overwrite_viz, draw_plotbox, caplog):
+    def test_known_svg_ts_mut_no_edges(self, overwrite_viz, draw_plotbox):
         # An example with some muts on axis but not on a visible node
         ts = msprime.simulate(10, random_seed=2, mutation_rate=1)
         tables = ts.dump_tables()
         tables.edges.clear()
         tables.mutations.time = np.full_like(tables.mutations.time, tskit.UNKNOWN_TIME)
         ts_no_edges = tables.tree_sequence()
-        svg = ts_no_edges.draw_svg(debug_box=draw_plotbox)
-        assert "not present in the displayed tree" in caplog.text
-        self.verify_known_svg(
-            svg, "ts_mutations_no_edges.svg", overwrite_viz, width=200 * ts.num_trees
-        )
+        with pytest.warns(UserWarning, match="nodes which are not present"):
+            svg = ts_no_edges.draw_svg(debug_box=draw_plotbox)
+            self.verify_known_svg(
+                svg,
+                "ts_mutations_no_edges.svg",
+                overwrite_viz,
+                width=200 * ts.num_trees,
+            )
 
-    def test_known_svg_ts_timed_mut_no_edges(self, overwrite_viz, draw_plotbox, caplog):
+    def test_known_svg_ts_timed_mut_no_edges(self, overwrite_viz, draw_plotbox):
         # An example with some muts on axis but not on a visible node
         ts = msprime.simulate(10, random_seed=2, mutation_rate=1)
         tables = ts.dump_tables()
@@ -2764,14 +2769,15 @@ def test_known_svg_ts_timed_mut_no_edges(self, overwrite_viz, draw_plotbox, capl
             ts.num_mutations, dtype=tables.mutations.time.dtype
         )
         ts_no_edges = tables.tree_sequence()
-        svg = ts_no_edges.draw_svg(debug_box=draw_plotbox)
-        assert "not present in the displayed tree" in caplog.text
-        self.verify_known_svg(
-            svg,
-            "ts_mutations_timed_no_edges.svg",
-            overwrite_viz,
-            width=200 * ts.num_trees,
-        )
+
+        with pytest.warns(UserWarning, match="nodes which are not present"):
+            svg = ts_no_edges.draw_svg(debug_box=draw_plotbox)
+            self.verify_known_svg(
+                svg,
+                "ts_mutations_timed_no_edges.svg",
+                overwrite_viz,
+                width=200 * ts.num_trees,
+            )
 
     def test_known_svg_ts_multiroot(self, overwrite_viz, draw_plotbox, caplog):
         tables = wf.wf_sim(

python/tskit/drawing.py

@@ -1,6 +1,6 @@
 # MIT License
 #
-# Copyright (c) 2018-2022 Tskit Developers
+# Copyright (c) 2018-2023 Tskit Developers
 # Copyright (c) 2015-2017 University of Oxford
 #
 # Permission is hereby granted, free of charge, to any person obtaining a copy
@@ -25,7 +25,6 @@
 """
 import collections
 import itertools
-import logging
 import math
 import numbers
 import operator
@@ -1365,9 +1364,11 @@ def __init__(
                     else:
                         unplotted.append(mutation.id + self.offsets.mutation)
         if len(unplotted) > 0:
-            logging.warning(
+            warnings.warn(
                 f"Mutations {unplotted} are above nodes which are not present in the "
-                "displayed tree, so are not plotted on the topology."
+                "displayed tree, so are not plotted on the topology.",
+                UserWarning,
+                stacklevel=2,
             )
         self.left_extent = tree.interval.left
         self.right_extent = tree.interval.right