Fixing plot mode when chunk prompt triggers

src/h5forest/bindings/hist_bindings.py

@@ -9,7 +9,7 @@
 from prompt_toolkit.widgets import Label
 
 from h5forest.config import translate_key_label
-from h5forest.dataset_prompts import prompt_for_dataset_operation
+from h5forest.dataset_prompts import prompt_for_chunking_preference
 from h5forest.errors import error_handler
 from h5forest.utils import WaitIndicator
 
@@ -28,13 +28,8 @@ def select_data(event):
             app.print(f"{node.path} is not a Dataset")
             return
 
-        def run_operation(use_chunks):
-            """Set histogram data after user confirmation."""
-            # Set the text in the histogram area
-            app.hist_content.text = app.histogram_plotter.set_data_key(node)
-
-        # Prompt user if needed, then run operation
-        prompt_for_dataset_operation(app, node, run_operation)
+        # Set the text in the histogram area directly (no prompt here)
+        app.hist_content.text = app.histogram_plotter.set_data_key(node)
 
     @error_handler
     def edit_bins(event):
@@ -247,37 +242,32 @@ def plot_hist(event):
                 app.print(f"{node.path} is not a Dataset")
                 return
 
-            def run_operation(use_chunks):
-                """Set histogram data and plot after user confirmation."""
-                # Set the text in the plotting area
-                app.hist_content.text = app.histogram_plotter.set_data_key(
-                    node
-                )
-
-                # Compute and plot the histogram with wait indicator
-                with WaitIndicator(app, "Generating histogram..."):
-                    # Compute the histogram
-                    app.hist_content.text = app.histogram_plotter.compute_hist(
-                        app.hist_content.text
-                    )
-
-                    # Get the plot
-                    app.histogram_plotter.plot_and_show(app.hist_content.text)
+            # Set the text in the plotting area
+            app.hist_content.text = app.histogram_plotter.set_data_key(node)
 
-            # Prompt user if needed, then run operation
-            prompt_for_dataset_operation(app, node, run_operation)
-        else:
-            # Already have data, just plot
+        def do_plot(use_chunks):
+            """Actually perform the plot after chunking preference is set."""
             # Compute and plot the histogram with wait indicator
             with WaitIndicator(app, "Generating histogram..."):
                 # Compute the histogram
                 app.hist_content.text = app.histogram_plotter.compute_hist(
-                    app.hist_content.text
+                    app.hist_content.text, use_chunks=use_chunks
                 )
 
                 # Get the plot
                 app.histogram_plotter.plot_and_show(app.hist_content.text)
 
+        # Check if we have data selected
+        if "data" not in app.histogram_plotter.plot_params:
+            app.print("Please select a dataset first (Enter)")
+            return
+
+        # Get the node to check for chunking
+        nodes = [app.histogram_plotter.plot_params["data"]]
+
+        # Prompt for chunking preference if needed, then plot
+        prompt_for_chunking_preference(app, nodes, do_plot)
+
     @error_handler
     def save_hist(event):
         """Plot the histogram."""
@@ -291,33 +281,30 @@ def save_hist(event):
                 app.print(f"{node.path} is not a Dataset")
                 return
 
-            def run_operation(use_chunks):
-                """Set histogram data and save after user confirmation."""
-                # Set the text in the plotting area
-                app.hist_content.text = app.histogram_plotter.set_data_key(
-                    node
-                )
-
-                # Compute the histogram
-                app.hist_content.text = app.histogram_plotter.compute_hist(
-                    app.hist_content.text
-                )
-
-                # Get the plot
-                app.histogram_plotter.plot_and_save(app.hist_content.text)
+            # Set the text in the plotting area
+            app.hist_content.text = app.histogram_plotter.set_data_key(node)
 
-            # Prompt user if needed, then run operation
-            prompt_for_dataset_operation(app, node, run_operation)
-        else:
-            # Already have data, just save
+        def do_save(use_chunks):
+            """Actually save the plot after chunking preference is set."""
             # Compute the histogram
             app.hist_content.text = app.histogram_plotter.compute_hist(
-                app.hist_content.text
+                app.hist_content.text, use_chunks=use_chunks
             )
 
             # Get the plot
             app.histogram_plotter.plot_and_save(app.hist_content.text)
 
+        # Check if we have data selected
+        if "data" not in app.histogram_plotter.plot_params:
+            app.print("Please select a dataset first (Enter)")
+            return
+
+        # Get the node to check for chunking
+        nodes = [app.histogram_plotter.plot_params["data"]]
+
+        # Prompt for chunking preference if needed, then save
+        prompt_for_chunking_preference(app, nodes, do_save)
+
     @error_handler
     def reset_hist(event):
         """Reset the histogram content."""

src/h5forest/bindings/plot_bindings.py

@@ -10,6 +10,7 @@
 from prompt_toolkit.widgets import Label
 
 from h5forest.config import translate_key_label
+from h5forest.dataset_prompts import prompt_for_chunking_preference
 from h5forest.errors import error_handler
 from h5forest.utils import WaitIndicator
 
@@ -211,16 +212,60 @@ def edit_plot_entry_callback():
     @error_handler
     def plot_scatter(event):
         """Plot and show scatter with mean in bins."""
-        # Make the plot with wait indicator
-        with WaitIndicator(app, "Generating scatter plot..."):
-            app.scatter_plotter.plot_and_show(app.plot_content.text)
+        # Check if we have both x and y datasets selected
+        if (
+            "x" not in app.scatter_plotter.plot_params
+            or "y" not in app.scatter_plotter.plot_params
+        ):
+            msg = "Please select both x-axis (x) and y-axis (y) datasets first"
+            app.print(msg)
+            return
 
-        app.default_focus()
+        def do_plot(use_chunks):
+            """Actually perform the plot after chunking preference is set."""
+            # Make the plot with wait indicator
+            with WaitIndicator(app, "Generating scatter plot..."):
+                app.scatter_plotter.plot_and_show(
+                    app.plot_content.text, use_chunks=use_chunks
+                )
+
+            app.default_focus()
+
+        # Get the nodes to check for chunking
+        nodes = [
+            app.scatter_plotter.plot_params["x"],
+            app.scatter_plotter.plot_params["y"],
+        ]
+
+        # Prompt for chunking preference if needed, then plot
+        prompt_for_chunking_preference(app, nodes, do_plot)
 
     @error_handler
     def save_scatter(event):
         """Save the plot."""
-        app.scatter_plotter.plot_and_save(app.plot_content.text)
+        # Check if we have both x and y datasets selected
+        if (
+            "x" not in app.scatter_plotter.plot_params
+            or "y" not in app.scatter_plotter.plot_params
+        ):
+            msg = "Please select both x-axis (x) and y-axis (y) datasets first"
+            app.print(msg)
+            return
+
+        def do_save(use_chunks):
+            """Actually save the plot after chunking preference is set."""
+            app.scatter_plotter.plot_and_save(
+                app.plot_content.text, use_chunks=use_chunks
+            )
+
+        # Get the nodes to check for chunking
+        nodes = [
+            app.scatter_plotter.plot_params["x"],
+            app.scatter_plotter.plot_params["y"],
+        ]
+
+        # Prompt for chunking preference if needed, then save
+        prompt_for_chunking_preference(app, nodes, do_save)
 
     @error_handler
     def reset(event):

src/h5forest/dataset_prompts.py

@@ -5,6 +5,116 @@
 """
 
 
+def prompt_for_chunking_preference(app, nodes, operation_callback):
+    """
+    Prompt user about chunking preference for plotting/histogram operations.
+
+    This function checks if any of the provided nodes are chunked, and
+    prompts the user to decide whether to use chunked processing.
+
+    Args:
+        app (H5Forest):
+            The main application instance.
+        nodes (list):
+            List of Node objects to check for chunking.
+        operation_callback (callable):
+            Function to call with signature: operation_callback(use_chunks)
+            where use_chunks is True to use chunked processing.
+    """
+    # Check if config has always_chunk enabled
+    if app.config.always_chunk_datasets():
+        operation_callback(use_chunks=True)
+        return
+
+    # Check if any node is chunked
+    has_chunked = any(node.is_chunked for node in nodes)
+
+    # If no chunked data, proceed without chunking
+    if not has_chunked:
+        operation_callback(use_chunks=False)
+        return
+
+    # At this point, we have chunked data and need to prompt
+    # Calculate memory footprint for all nodes
+    import h5py
+    import numpy as np
+
+    total_size_gb = 0.0
+    total_chunks = 0
+
+    try:
+        for node in nodes:
+            if node.is_chunked:
+                with h5py.File(node.filepath, "r") as hdf:
+                    dataset = hdf[node.path]
+                    uncompressed_bytes = dataset.size * dataset.dtype.itemsize
+                    total_size_gb += uncompressed_bytes / (10**9)
+
+                    if dataset.chunks:
+                        num_chunks = int(
+                            np.prod(
+                                [
+                                    np.ceil(s / c)
+                                    for s, c in zip(
+                                        dataset.shape, dataset.chunks
+                                    )
+                                ]
+                            )
+                        )
+                        total_chunks += num_chunks
+    except (OSError, FileNotFoundError, KeyError):
+        # If file access fails, use node's nbytes as estimate
+        for node in nodes:
+            if node.is_chunked:
+                total_size_gb += node.nbytes / (10**9)
+        total_chunks = "unknown"
+
+    def on_yes_chunk():
+        """User wants chunk-by-chunk processing."""
+        app.default_focus()
+        operation_callback(use_chunks=True)
+
+    def on_no_chunk():
+        """User wants to load all at once."""
+        app.default_focus()
+
+        def on_yes_load_all():
+            """User confirms loading all at once."""
+            app.default_focus()
+            operation_callback(use_chunks=False)
+
+        def on_no_load_all():
+            """User wants to abort."""
+            app.default_focus()
+            app.print("Operation aborted.")
+
+        # Ask second question
+        app.prompt_yn(
+            "Should we load all at once? (If not, we will abort) [y/n]:",
+            on_yes_load_all,
+            on_no_load_all,
+        )
+
+    # Build the prompt message
+    if isinstance(total_chunks, int) and total_chunks > 0:
+        prompt_msg = (
+            f"Chunked Dataset found ({total_size_gb:.2f} GB, "
+            f"{total_chunks} chunks). "
+            f"Should we process chunk by chunk? [y/n]:"
+        )
+    else:
+        prompt_msg = (
+            f"Chunked Dataset found ({total_size_gb:.2f} GB). "
+            f"Should we process chunk by chunk? [y/n]:"
+        )
+
+    app.prompt_yn(
+        prompt_msg,
+        on_yes_chunk,
+        on_no_chunk,
+    )
+
+
 def prompt_for_chunked_dataset(
     app, node, operation_callback, return_to_normal=True
 ):

src/h5forest/h5_forest.py

@@ -166,6 +166,7 @@ def _init(self, hdf5_filepath, use_default_config=False):
         self._flag_plotting_mode = False
         self._flag_hist_mode = False
         self._flag_search_mode = False
+        self._flag_in_prompt = False  # Track if we're in a prompt_yn dialog
 
         # Timer for debouncing search input
         self.search_timer = None
@@ -640,6 +641,7 @@ def return_to_normal_mode(self):
         self._flag_plotting_mode = False
         self._flag_hist_mode = False
         self._flag_search_mode = False
+        self._flag_in_prompt = False  # Track if we're in a prompt_yn dialog
         self.mode_title.update_title("Normal Mode")
 
     def _init_text_areas(self):
@@ -1110,6 +1112,9 @@ def prompt_yn(self, prompt, on_yes, on_no):
             on_no (callable):
                 Callback to execute when user presses 'n'.
         """
+        # Set the prompt flag to True
+        self._flag_in_prompt = True
+
         # Set the prompt message
         self.input_buffer_content.text = prompt
         self.mini_buffer_content.text = ""
@@ -1120,6 +1125,9 @@ def prompt_yn(self, prompt, on_yes, on_no):
 
         def cleanup():
             """Remove temporary keybindings and clear prompt."""
+            # Clear the prompt flag first
+            self._flag_in_prompt = False
+
             for handler in handlers_to_remove:
                 try:
                     self.kb.bindings.remove(handler)
@@ -1143,11 +1151,19 @@ def handle_escape(event):
             cleanup()
             on_no()
 
-        # Add temporary keybindings for y/n/escape
-        # These bindings should always be active during the prompt
-        y_handler = self.kb.add("y")(handle_yes)
-        n_handler = self.kb.add("n")(handle_no)
-        esc_handler = self.kb.add("escape")(handle_escape)
+        # Add temporary keybindings for y/n/escape with filter
+        # These bindings are ONLY active when _flag_in_prompt is True
+        from prompt_toolkit.filters import Condition
+
+        y_handler = self.kb.add(
+            "y", filter=Condition(lambda: self._flag_in_prompt)
+        )(handle_yes)
+        n_handler = self.kb.add(
+            "n", filter=Condition(lambda: self._flag_in_prompt)
+        )(handle_no)
+        esc_handler = self.kb.add(
+            "escape", filter=Condition(lambda: self._flag_in_prompt)
+        )(handle_escape)
 
         handlers_to_remove.extend([y_handler, n_handler, esc_handler])