Skip to content

Commit aa8af83

Browse files
skrawczelijahbenizzy
authored andcommitted
Fixes regression for visualization stripping path
Viz would strip paths like `./foo/bar/dag.png` to just `dag.png`. This fixes that.
1 parent e570014 commit aa8af83

File tree

8 files changed

+9
-10
lines changed

8 files changed

+9
-10
lines changed

examples/hello_world/a_path.dot.png

-36 KB
Binary file not shown.

examples/hello_world/a_path.png

35.9 KB
Loading

examples/hello_world/my_dag.dot.png

-67 KB
Binary file not shown.

examples/hello_world/my_dag.png

66.9 KB
Loading

examples/hello_world/my_full_dag.png

64.1 KB
Loading

examples/hello_world/my_script.py

+3-4
Original file line numberDiff line numberDiff line change
@@ -33,12 +33,11 @@
3333
print(df.to_string())
3434

3535
# To visualize do `pip install "sf-hamilton[visualization]"` if you want these to work
36-
dr.visualize_execution(output_columns, "./my_dag.dot", {"format": "png"})
36+
dr.visualize_execution(output_columns, "./my_dag.png")
3737
dr.visualize_path_between(
3838
"spend_mean",
3939
"spend_zero_mean_unit_variance",
40-
"./a_path.dot",
41-
{"format": "png"},
40+
"./a_path.png",
4241
strict_path_visualization=False,
4342
)
44-
# dr.display_all_functions("./my_full_dag.dot", {"format": "png"})
43+
dr.display_all_functions("./my_full_dag.png")

hamilton/graph.py

+2-4
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88

99
import inspect
1010
import logging
11-
import pathlib
11+
import os.path
1212
import uuid
1313
from enum import Enum
1414
from types import ModuleType
@@ -856,12 +856,10 @@ def display(
856856
)
857857
kwargs = {"view": False, "format": "png"} # default format = png
858858
if output_file_path: # infer format from path
859-
suffix = pathlib.Path(output_file_path).suffix
859+
output_file_path, suffix = os.path.splitext(output_file_path)
860860
if suffix != "":
861861
inferred_format = suffix.partition(".")[-1]
862862
kwargs.update(format=inferred_format)
863-
# remove suffix if exist because dot.render() will append the format
864-
output_file_path = str(pathlib.Path(output_file_path).stem)
865863
if render_kwargs and isinstance(render_kwargs, dict): # accept explicit format
866864
kwargs.update(render_kwargs)
867865
if output_file_path:

tests/test_graph.py

+4-2
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
import inspect
2+
import os
23
import pathlib
34
import uuid
45
from itertools import permutations
@@ -815,7 +816,7 @@ def _styling_function(*, node, node_class):
815816

816817
@pytest.mark.parametrize("show_legend", [(True), (False)])
817818
def test_function_graph_display_legend(show_legend: bool, tmp_path: pathlib.Path):
818-
dot_file_path = tmp_path / "dag"
819+
dot_file_path = tmp_path / "dag.png"
819820
fg = graph.FunctionGraph.from_modules(tests.resources.dummy_functions, config={"b": 1, "c": 2})
820821

821822
fg.display(
@@ -824,7 +825,8 @@ def test_function_graph_display_legend(show_legend: bool, tmp_path: pathlib.Path
824825
render_kwargs={"view": False},
825826
show_legend=show_legend,
826827
)
827-
dot = dot_file_path.open("r").read()
828+
dot_file = pathlib.Path(os.path.splitext(str(dot_file_path))[0])
829+
dot = dot_file.open("r").read()
828830

829831
found_legend = "cluster__legend" in dot
830832
assert found_legend is show_legend

0 commit comments

Comments
 (0)