From 2b1600d8c373dbe616871c014e3223bb8061f5a3 Mon Sep 17 00:00:00 2001
From: breken-ai <312387581+breken-ai@users.noreply.github.com>
Date: Fri, 25 Sep 2026 16:49:31 -0700
Subject: [PATCH 1/2] Accept edge_attr and other attr dicts in visualize()

Graph.visualize() (and Application.visualize()) pass **engine_kwargs to
graphviz.Digraph, merging dict values into Burr's defaults. Only
graph_attr and node_attr have defaults, so passing any other dict such
as edge_attr={"color": "red"} raised KeyError: 'edge_attr'.

Merge into an empty dict when Burr has no default for that key.

Generated-by: Claude Code (Claude Opus 5.5)
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
---
 burr/core/graph.py                  |  2 +-
 tests/core/test_graphviz_display.py | 16 ++++++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/burr/core/graph.py b/burr/core/graph.py
index 5889c41f6..f7de9fd47 100644
--- a/burr/core/graph.py
+++ b/burr/core/graph.py
@@ -249,7 +249,7 @@ def visualize(
         )
         for g_key, g_value in engine_kwargs.items():
             if isinstance(g_value, dict):
-                digraph_attr[g_key].update(**g_value)
+                digraph_attr.setdefault(g_key, {}).update(**g_value)
             else:
                 digraph_attr[g_key] = g_value
         digraph = graphviz.Digraph(**digraph_attr)
diff --git a/tests/core/test_graphviz_display.py b/tests/core/test_graphviz_display.py
index 7eda47cba..c10faf0d1 100644
--- a/tests/core/test_graphviz_display.py
+++ b/tests/core/test_graphviz_display.py
@@ -99,3 +99,19 @@ def test_visualize_include_state_multiline_label(reads: list, writes: list, expe
     digraph = graph.visualize(include_state=True)
 
     assert expected_label in digraph.source
+
+
+def test_visualize_engine_kwargs_attr_dicts(graph):
+    """Attribute dicts passed through ``engine_kwargs`` reach the graphviz.Digraph,
+    including ones (like ``edge_attr``) that have no Burr default to merge into."""
+    digraph = graph.visualize(
+        graph_attr={"rankdir": "LR"},
+        node_attr={"fontname": "Courier"},
+        edge_attr={"color": "red"},
+    )
+
+    assert digraph.graph_attr["rankdir"] == "LR"
+    assert digraph.graph_attr["ranksep"] == "0.4"  # Burr default is kept
+    assert digraph.node_attr["fontname"] == "Courier"
+    assert digraph.node_attr["fillcolor"] == "#b4d8e4"  # Burr default is kept
+    assert digraph.edge_attr == {"color": "red"}

From bdf29b553fa0851b999275474d7d3985d1f11913 Mon Sep 17 00:00:00 2001
From: breken-ai <312387581+breken-ai@users.noreply.github.com>
Date: Wed, 7 Oct 2026 17:27:27 -0700
Subject: [PATCH 2/2] Keep visualize() defaults when an attr dict is None

Passing graph_attr=None (or node_attr/edge_attr=None) replaced Burr's
default attributes with None, so rankdir, ranksep and the node styling
were dropped. Treat a None *_attr like an omitted one.

Also cover the user-facing path: Application.visualize(edge_attr=...)
now has its own test.

Generated-by: Claude Code (Claude Opus 5.5)
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
---
 burr/core/graph.py                  |  3 +++
 tests/core/test_graphviz_display.py | 20 ++++++++++++++++++++
 2 files changed, 23 insertions(+)

diff --git a/burr/core/graph.py b/burr/core/graph.py
index f7de9fd47..700ac4cad 100644
--- a/burr/core/graph.py
+++ b/burr/core/graph.py
@@ -248,6 +248,9 @@ def visualize(
             ),
         )
         for g_key, g_value in engine_kwargs.items():
+            if g_value is None and g_key.endswith("_attr"):
+                # Same as not passing it: keep Burr's defaults.
+                continue
             if isinstance(g_value, dict):
                 digraph_attr.setdefault(g_key, {}).update(**g_value)
             else:
diff --git a/tests/core/test_graphviz_display.py b/tests/core/test_graphviz_display.py
index c10faf0d1..fe0109221 100644
--- a/tests/core/test_graphviz_display.py
+++ b/tests/core/test_graphviz_display.py
@@ -19,6 +19,7 @@
 
 import pytest
 
+from burr.core import ApplicationBuilder
 from burr.core.graph import GraphBuilder
 
 from tests.core.test_graph import PassedInAction
@@ -115,3 +116,22 @@ def test_visualize_engine_kwargs_attr_dicts(graph):
     assert digraph.node_attr["fontname"] == "Courier"
     assert digraph.node_attr["fillcolor"] == "#b4d8e4"  # Burr default is kept
     assert digraph.edge_attr == {"color": "red"}
+
+
+def test_application_visualize_edge_attr(graph):
+    """``Application.visualize`` is the path users call; it forwards ``edge_attr``."""
+    app = ApplicationBuilder().with_graph(graph).with_entrypoint("counter").build()
+
+    digraph = app.visualize(edge_attr={"color": "red"})
+
+    assert digraph.edge_attr == {"color": "red"}
+    assert digraph.graph_attr["rankdir"] == "TB"  # Burr default is kept
+
+
+def test_visualize_none_attr_keeps_defaults(graph):
+    """Passing ``graph_attr=None`` behaves like not passing it."""
+    digraph = graph.visualize(graph_attr=None, node_attr=None, edge_attr=None)
+
+    assert digraph.graph_attr["rankdir"] == "TB"
+    assert digraph.node_attr["fillcolor"] == "#b4d8e4"
+    assert digraph.edge_attr == {}