Update flowchart.py

Signed-off-by: Sajid Alam <[email protected]>
kedro-org · Sep 2, 2024 · 1dbfef1 · 1dbfef1
1 parent b85425a
commit 1dbfef1
Showing 1 changed file with 14 additions and 2 deletions.
diff --git a/package/kedro_viz/models/flowchart.py b/package/kedro_viz/models/flowchart.py
@@ -684,8 +684,20 @@ def set_preview(cls, _):
             return None
 
         try:
-            preview_data = cls.dataset.preview()
-            preview_type = cls.data_node.viz_metadata.get("preview_type")
+            preview_args = (
+                cls.data_node.get_preview_args() if cls.data_node.viz_metadata else None
+            )
+
+            # Generate preview data with or without preview_args
+            if preview_args:
+                preview_data = cls.dataset.preview(**preview_args)
+            else:
+                preview_data = cls.dataset.preview()
+
+            # Determine the preview type
+            preview_type = (
+                cls.data_node.viz_metadata.get("preview_type") if cls.data_node.viz_metadata else None
+            )
 
             # Validate the format based on the preview type
             if preview_type == "TablePreview":