Fix keep_dot propagation in Driver display functions

DAGWorks-Inc · Oct 23, 2024 · 083dcce · 083dcce
1 parent bc7cbbf
commit 083dcce
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/hamilton/driver.py b/hamilton/driver.py
@@ -1200,6 +1200,7 @@ def display_downstream_of(
                 display_fields=show_schema,
                 custom_style_function=custom_style_function,
                 config=self.graph._config,
+                keep_dot=keep_dot,
             )
         except ImportError as e:
             logger.warning(f"Unable to import {e}", exc_info=True)
@@ -1263,6 +1264,7 @@ def display_upstream_of(
                 display_fields=show_schema,
                 custom_style_function=custom_style_function,
                 config=self.graph._config,
+                keep_dot=keep_dot,
             )
         except ImportError as e:
             logger.warning(f"Unable to import {e}", exc_info=True)