From 43ca3e3b6bdb1ccf30ee36b54aa8570c4e4d1da4 Mon Sep 17 00:00:00 2001
From: Yann Forget <yannforget@mailbox.org>
Date: Mon, 4 Nov 2024 13:49:20 +0100
Subject: [PATCH] feat(create-cart-diagram): add more mermaid options

---
 create-cart-diagram/pipeline.py | 22 +++++++++++++++++++++-
 1 file changed, 21 insertions(+), 1 deletion(-)

diff --git a/create-cart-diagram/pipeline.py b/create-cart-diagram/pipeline.py
index 2a29a8b..0990c8e 100644
--- a/create-cart-diagram/pipeline.py
+++ b/create-cart-diagram/pipeline.py
@@ -23,6 +23,22 @@
     type=str,
     required=False,
 )
+@parameter(
+    "add_choice_labels",
+    name="Add choice labels to links",
+    help="Add valid choice labels to relationship arrows instead of yes/no",
+    type=bool,
+    required=False,
+    default=False,
+)
+@parameter(
+    "add_node_id",
+    name="Add node indexes to mermaid shapes",
+    help="Add node indexes to mermaid shapes in addition to split rule",
+    type=bool,
+    required=False,
+    default=True,
+)
 @parameter(
     "output_dir",
     name="Output directory",
@@ -31,7 +47,11 @@
     required=False,
 )
 def create_cart_diagram(
-    cart_outputs: Dataset, version_name: Optional[str], output_dir: Optional[str]
+    cart_outputs: Dataset,
+    version_name: Optional[str],
+    add_choice_labels: Optional[bool],
+    add_node_id: Optional[bool],
+    output_dir: Optional[str],
 ):
     """Create a CART diagram from CART outputs."""