Add optional d3js parameter

tests/test_draw_parameters.py

@@ -0,0 +1,127 @@
+import io
+import inspect
+
+import pandas as pd
+
+import tskit_arg_visualizer
+from tskit_arg_visualizer import D3ARG, draw_D3
+
+
+def _minimal_arg_json():
+    return {
+        "data": {"nodes": [], "edges": [], "mutations": [], "breakpoints": []},
+        "width": 100,
+        "height": 100,
+        "y_axis": {"include_labels": True},
+        "edges": {"type": "line"},
+        "condense_mutations": False,
+        "label_mutations": False,
+        "tree_highlighting": True,
+        "title": "None",
+        "rotate_tip_labels": False,
+        "plot_type": "full",
+    }
+
+
+def _minimal_d3arg():
+    return D3ARG(
+        nodes=pd.DataFrame(columns=["id", "time"]),
+        edges=pd.DataFrame(columns=["source", "target", "bounds"]),
+        mutations=pd.DataFrame(columns=["position_01"]),
+        breakpoints=pd.DataFrame(
+            [
+                {"start": 0.0, "stop": 1.0, "x_pos_01": 0.0, "width_01": 1.0, "fill": "#053e4e"}
+            ]
+        ),
+        num_samples=0,
+        sample_order=[],
+        default_node_style={},
+        time_units="generations",
+    )
+
+
+class TestDrawParametersD3js:
+    def test_draw_D3_default_uses_default_url(self, monkeypatch):
+        call_args = None
+
+        def mock_display(*args, **kwargs):
+            nonlocal call_args
+            call_args = (args, kwargs)
+
+        monkeypatch.setattr(tskit_arg_visualizer, "display", mock_display)
+        draw_D3(_minimal_arg_json(), is_notebook=True)
+        html = call_args[0][0].data
+        assert "https://d3js.org/d3.v7.min" in html
+
+    def test_public_draw_apis_accept_d3js_parameter(self):
+        assert "d3js" in inspect.signature(D3ARG.draw).parameters
+        assert "d3js" in inspect.signature(D3ARG.draw_node).parameters
+        assert "d3js" in inspect.signature(D3ARG.draw_nodes).parameters
+        assert "d3js" in inspect.signature(D3ARG.draw_genome_bar).parameters
+
+    def test_draw_D3_uses_url_override(self, monkeypatch):
+        custom_url = "https://example.org/custom/d3.min"
+        call_args = None
+
+        def mock_display(*args, **kwargs):
+            nonlocal call_args
+            call_args = (args, kwargs)
+
+        monkeypatch.setattr(tskit_arg_visualizer, "display", mock_display)
+        draw_D3(_minimal_arg_json(), is_notebook=True, d3js=custom_url)
+        html = call_args[0][0].data
+        assert custom_url in html
+        assert "https://d3js.org/d3.v7.min" not in html
+
+    def test_draw_D3_inlines_file_like_content(self, monkeypatch):
+        inline_js = b"window.d3 = {version: 'inline'};"
+        call_args = None
+
+        def mock_display(*args, **kwargs):
+            nonlocal call_args
+            call_args = (args, kwargs)
+
+        monkeypatch.setattr(tskit_arg_visualizer, "display", mock_display)
+        draw_D3(_minimal_arg_json(), is_notebook=True, d3js=io.BytesIO(inline_js))
+        html = call_args[0][0].data
+        assert "<script>window.d3 = {version: 'inline'};</script>" in html
+
+    def test_draw_genome_bar_uses_d3js_override(self, monkeypatch):
+        d3arg = _minimal_d3arg()
+        custom_url = "https://example.org/alt/d3.min"
+        call_args = None
+
+        def mock_display(*args, **kwargs):
+            nonlocal call_args
+            call_args = (args, kwargs)
+
+        monkeypatch.setattr(tskit_arg_visualizer, "display", mock_display)
+        d3arg.draw_genome_bar(is_notebook=True, d3js=custom_url)
+        html = call_args[0][0].data
+        assert custom_url in html
+
+    def test_draw_genome_bar_inlines_file_like_content(self, monkeypatch):
+        d3arg = _minimal_d3arg()
+        call_args = None
+
+        def mock_display(*args, **kwargs):
+            nonlocal call_args
+            call_args = (args, kwargs)
+
+        monkeypatch.setattr(tskit_arg_visualizer, "display", mock_display)
+        d3arg.draw_genome_bar(is_notebook=True, d3js=io.StringIO("window.d3 = {};"))
+        html = call_args[0][0].data
+        assert "<script>window.d3 = {};</script>" in html
+
+    def test_draw_genome_bar_default_uses_default_url(self, monkeypatch):
+        d3arg = _minimal_d3arg()
+        call_args = None
+
+        def mock_display(*args, **kwargs):
+            nonlocal call_args
+            call_args = (args, kwargs)
+
+        monkeypatch.setattr(tskit_arg_visualizer, "display", mock_display)
+        d3arg.draw_genome_bar(is_notebook=True)
+        html = call_args[0][0].data
+        assert "https://d3js.org/d3.v7.min" in html

tskit_arg_visualizer/__init__.py

@@ -1,4 +1,4 @@
-import collections
+from collections import namedtuple
 import itertools
 import json
 import math
@@ -23,7 +23,7 @@
 )
 
 
-default_mutation_styles = {
+DEFAULT_MUTATION_STYLES = {
     "size": 5,
     "unknown_time": {
         "fill": "gold",
@@ -41,6 +41,11 @@
     },
 }
 
+DEFAULT_D3JS_URL = "https://d3js.org/d3.v7.min"
+
+D3jsSource = namedtuple("D3jsSource", ["url", "inline_script"])
+IncludedInfo = namedtuple("IncludedInfo", ["nodes", "edges", "mutations", "breakpoints"])
+
 @dataclass
 class IncludedObjects:
     """Stores the IDs of the various objects included in a drawing."""
@@ -97,7 +102,25 @@ def running_in_notebook():
             return False  # Other type (?)
     except NameError:
         return False      # Probably standard Python interpreter
-
+
+
+def resolve_d3js_source(d3js) -> D3jsSource:
+    """
+    Returns a tuple in which exactly one field is meaningful:
+    - url is a string URL (and inline_script is ""), or
+    - url is None and inline_script is a "<script>…</script>" HTML fragment.
+    """
+    if d3js is None:
+        return D3jsSource(url=DEFAULT_D3JS_URL, inline_script="")
+    try:
+        d3_content = d3js.read()
+    except AttributeError:
+        return D3jsSource(url=str(d3js), inline_script="")
+    if isinstance(d3_content, bytes):
+        d3_content = d3_content.decode("utf-8")
+    return D3jsSource(url=None, inline_script=f"<script>{d3_content}</script>")
+
+
 def calculate_evenly_distributed_positions(num_elements, start=0, end=1, round_to=0):
     """Returns a list of `num_elements` evenly distributed positions on a given `length`
 
@@ -179,21 +202,23 @@ def convert_time_to_position(t, min_time, max_time, scale, unique_times, h_spaci
     return (1-(t-min_time)/time_range) * (height-100) + y_shift
 
 
-def draw_D3(arg_json, styles=None, is_notebook=None):
+def draw_D3(arg_json, styles=None, is_notebook=None, d3js=None):
+    d3js_source = resolve_d3js_source(d3js)
     if is_notebook is None:
         is_notebook = running_in_notebook()
     arg_json["source"] = json.dumps(arg_json.copy())  # first escape the plain json data
+    arg_json["d3_url"] = d3js_source.url  # Could be None if an inline script is used
     arg_json = {k: json.dumps(v) for k, v in arg_json.items()}  # now escape all
     arg_json["divnum"] = str(random.randint(0,9999999999))
     arg_id = "arg_" + arg_json['divnum']
     JS_text = Template((
-        '<div id="{}" class="d3arg" style="min-width:{}px; min-height:{}px;"></div>'
+        '<div id="{}" class="d3arg" style="min-width:{}px; min-height:{}px;"></div>$d3_inline_script'
         '<script>$main_text</script>'
     ).format(arg_id, float(arg_json["width"]) + 40, float(arg_json["height"]) + 80))
     with open(os.path.dirname(__file__) + "/visualizer.js", "r") as visualizerjs:
         main_text_template = Template(visualizerjs.read())
     main_text = main_text_template.safe_substitute(arg_json)
-    html = JS_text.safe_substitute({'main_text': main_text})
+    html = JS_text.safe_substitute({'main_text': main_text, 'd3_inline_script': d3js_source.inline_script})
     with open(os.path.dirname(__file__) + "/visualizer.css", "r") as css:
         general_styles = css.read()
     specific_styles = ""
@@ -666,12 +691,12 @@ def _convert_edges_table(ts, recombination_nodes_to_merge, progressbar=None):
                         # an edge. Essentially, giving them a false time just for plotting.
                         middle = (new_edge[3] + new_edge[4]) / 2
                         plot_time = middle + random.uniform(-(new_edge[3]-middle),(new_edge[3]-middle))
-                        fill = default_mutation_styles["unknown_time"]["fill"]
-                        stroke = default_mutation_styles["unknown_time"]["stroke"]
+                        fill = DEFAULT_MUTATION_STYLES["unknown_time"]["fill"]
+                        stroke = DEFAULT_MUTATION_STYLES["unknown_time"]["stroke"]
                     else:
                         plot_time = mut.time
-                        fill = default_mutation_styles["known_time"]["fill"]
-                        stroke = default_mutation_styles["known_time"]["stroke"]
+                        fill = DEFAULT_MUTATION_STYLES["known_time"]["fill"]
+                        stroke = DEFAULT_MUTATION_STYLES["known_time"]["stroke"]
                     inherited_state = site.ancestral_state if mut.parent == tskit.NULL else ts.mutation(mut.parent).derived_state
                     mutations.append({
                         "edge": new_edge[0],
@@ -687,7 +712,7 @@ def _convert_edges_table(ts, recombination_nodes_to_merge, progressbar=None):
                         "derived": mut.derived_state,
                         "fill": fill,
                         "stroke": stroke,
-                        "size": default_mutation_styles["size"],
+                        "size": DEFAULT_MUTATION_STYLES["size"],
                     })
         mutations_output = pd.DataFrame(mutations, columns=["edge","source","target","time","plot_time","site_id","position","position_01","ancestral","inherited","derived","fill","stroke","size"])
         return edges_output, mutations_output
@@ -840,11 +865,11 @@ def reset_all_mutation_styles(self):
         """Resets mutation styles to default
         """
         is_unknown = tskit.is_unknown_time(self.mutations.time)
-        self.mutations["size"] = default_mutation_styles["size"]
-        self.mutations["fill"] = default_mutation_styles["known_time"]["fill"]
-        self.mutations["stroke"] = default_mutation_styles["known_time"]["stroke"]
-        self.mutations.loc[is_unknown, "fill"] = default_mutation_styles["unknown_time"]["fill"]
-        self.mutations.loc[is_unknown, "stroke"] = default_mutation_styles["unknown_time"]["stroke"]
+        self.mutations["size"] = DEFAULT_MUTATION_STYLES["size"]
+        self.mutations["fill"] = DEFAULT_MUTATION_STYLES["known_time"]["fill"]
+        self.mutations["stroke"] = DEFAULT_MUTATION_STYLES["known_time"]["stroke"]
+        self.mutations.loc[is_unknown, "fill"] = DEFAULT_MUTATION_STYLES["unknown_time"]["fill"]
+        self.mutations.loc[is_unknown, "stroke"] = DEFAULT_MUTATION_STYLES["unknown_time"]["stroke"]
 
     def set_mutation_styles(self, styles):
         """Individually control the styling of each mutation.
@@ -1210,8 +1235,8 @@ def _prepare_json(
                                 "site_id": list(muts["site_id"]),
                                 "mutation_id": list(muts.index),
                                 "x_pos": list(x_pos),
-                                "fill": default_mutation_styles["condensed"]["fill"],
-                                "stroke": default_mutation_styles["condensed"]["stroke"],
+                                "fill": DEFAULT_MUTATION_STYLES["condensed"]["fill"],
+                                "stroke": DEFAULT_MUTATION_STYLES["condensed"]["stroke"],
                                 "active": False,
                                 "label": "⨉"+str(muts.shape[0]),
                                 "content": "<br>".join(muts.content),
@@ -1459,6 +1484,7 @@ def draw(
             styles=None,
             preamble=None,
             save_filename=None,
+            d3js=None,
         ):
         """Draws the D3ARG using D3.js by sending a custom JSON object to visualizer.js 
 
@@ -1533,6 +1559,12 @@ def draw(
         save_filename : str
             Filename to use when selecting "Download as" in the visualization
             (default=None, treated as "tskit_arg_visualizer")
+        d3js : optional
+            Source for loading D3.js. If None, uses the default URL
+            https://d3js.org/d3.v7.min. Otherwise, this is first treated as a file-like
+            object by trying `.read()`: if successful, the returned JavaScript source is
+            embedded directly in the output HTML (bytes are decoded as UTF-8). If `.read()`
+            is unavailable, the value is converted to `str(...)` and used as a URL.
 
         Returns
         -------
@@ -1574,7 +1606,7 @@ def draw(
             preamble=preamble,
             save_filename=save_filename,
         )
-        info = draw_D3(arg_json=arg, styles=styles, is_notebook=is_notebook)
+        info = draw_D3(arg_json=arg, styles=styles, is_notebook=is_notebook, d3js=d3js)
         info.included.nodes = included_nodes["id"].tolist()
         return info
 
@@ -1699,7 +1731,7 @@ def subset_graph(self, seed_nodes, depth):
         included_breakpoints.append(current_region) # make sure to append the last region
         included_breakpoints = pd.DataFrame(included_breakpoints)
 
-        return collections.namedtuple('IncludedInfo', ['nodes', 'edges', 'mutations', 'breakpoints'])(
+        return IncludedInfo(
             included_nodes, included_edges, included_mutations, included_breakpoints
         )
 
@@ -1723,6 +1755,7 @@ def draw_node(
             styles=None,
             preamble=None,
             save_filename=None,
+            d3js=None,
         ):
         """Draws a subgraph of the D3ARG using D3.js by sending a custom JSON object to visualizer.js.
 
@@ -1789,6 +1822,12 @@ def draw_node(
         save_filename : str
             Filename to use when selecting "Download as" in the visualization
             (default=None, treated as "tskit_arg_visualizer")
+        d3js : optional
+            Source for loading D3.js. If None, uses the default URL
+            https://d3js.org/d3.v7.min. Otherwise, this is first treated as a file-like
+            object by trying `.read()`: if successful, the returned JavaScript source is
+            embedded directly in the output HTML (bytes are decoded as UTF-8). If `.read()`
+            is unavailable, the value is converted to `str(...)` and used as a URL.
 
         Returns
         -------
@@ -1825,7 +1864,7 @@ def draw_node(
             preamble=preamble,
             save_filename=save_filename,
         )
-        info = draw_D3(arg_json=arg, styles=styles, is_notebook=is_notebook)
+        info = draw_D3(arg_json=arg, styles=styles, is_notebook=is_notebook, d3js=d3js)
         info.included.nodes = included.nodes["id"].tolist()
         return info
 
@@ -1839,6 +1878,7 @@ def draw_genome_bar(
             windows=None,
             show_mutations=False,
             is_notebook=None,
+            d3js=None,
         ):
         """Draws a genome bar for the D3ARG using D3.js
 
@@ -1857,10 +1897,16 @@ def draw_genome_bar(
             whether it is being called in a notebook environment. This may not work in
             some untested environments, in which case you may wish to set this explicitly
             to True (to force a notebook display) or False (to force a standalone HTML page).
+        d3js : optional
+            Source for loading D3.js. If None, uses the default URL
+            https://d3js.org/d3.v7.min. Otherwise, this is first treated as a file-like
+            object by trying `.read()`: if successful, the returned JavaScript source is
+            embedded directly in the output HTML (bytes are decoded as UTF-8). If `.read()`
+            is unavailable, the value is converted to `str(...)` and used as a URL.
         """
         if is_notebook is None:
             is_notebook = running_in_notebook()
-
+        d3js_source = resolve_d3js_source(d3js)
         transformed_bps = self.breakpoints.loc[:,:]
         transformed_bps["x_pos"] = transformed_bps["x_pos_01"] * width
         transformed_bps["width"] = transformed_bps["width_01"] * width
@@ -1897,13 +1943,15 @@ def draw_genome_bar(
         }
 
         genome_bar_json["source"] = genome_bar_json.copy()
+        genome_bar_json["d3_url"] = d3js_source.url  # Could be None if an inline script is used. Will be escaped below
+        genome_bar_json = {k: json.dumps(v) for k, v in genome_bar_json.items()}
         genome_bar_json["divnum"] = str(random.randint(0,9999999999))
-        JS_text = Template("<div id='genome_bar_" + genome_bar_json['divnum'] + "'class='d3arg' style='min-width:" + str(genome_bar_json["width"]+40) + "px; min-height:180px;'></div><script>$main_text</script>")
+        JS_text = Template("<div id='genome_bar_" + genome_bar_json['divnum'] + "'class='d3arg' style='min-width:" + str(json.loads(genome_bar_json["width"])+40) + "px; min-height:180px;'></div>$d3_inline_script<script>$main_text</script>")
         breakpointsjs = open(os.path.dirname(__file__) + "/alternative_plots/genome_bar.js", "r")
         main_text_template = Template(breakpointsjs.read())
         breakpointsjs.close()
         main_text = main_text_template.safe_substitute(genome_bar_json)
-        html = JS_text.safe_substitute({'main_text': main_text})
+        html = JS_text.safe_substitute({'main_text': main_text, 'd3_inline_script': d3js_source.inline_script})
         css = open(os.path.dirname(__file__) + "/visualizer.css", "r")
         styles = css.read()
         css.close()
@@ -1912,7 +1960,7 @@ def draw_genome_bar(
         else:
             with tempfile.NamedTemporaryFile("w", delete=False, suffix=".html") as f:
                 url = "file://" + f.name
-                f.write("<!DOCTYPE html><html><head><style>"+styles+"</style><script src='https://cdn.rawgit.com/eligrey/canvas-toBlob.js/f1a01896135ab378aa5c0118eadd81da55e698d8/canvas-toBlob.js'></script><script src='https://cdn.rawgit.com/eligrey/FileSaver.js/e9d941381475b5df8b7d7691013401e171014e89/FileSaver.min.js'></script><script src='https://d3js.org/d3.v7.min.js'></script></head><body>" + html + "</body></html>")
+                f.write("<!DOCTYPE html><html><head><style>"+styles+"</style><script src='https://cdn.rawgit.com/eligrey/canvas-toBlob.js/f1a01896135ab378aa5c0118eadd81da55e698d8/canvas-toBlob.js'></script><script src='https://cdn.rawgit.com/eligrey/FileSaver.js/e9d941381475b5df8b7d7691013401e171014e89/FileSaver.min.js'></script></head><body>" + html + "</body></html>")
             webbrowser.open(url, new=2)