Feature Request: Always return Plotly figure objects from plot_* methods to enable programmatic customization and saving

[khushthecoder]

Description

While working with the malariagen-data-python API for exploratory data analysis , I noticed a consistent API design pattern across all plotting methods that makes programmatic figure customization and exporting unnecessarily difficult.

---

The Problem

Across the codebase, there are 45+ plotting methods (e.g., plot_heterozygosity_track, plot_samples_interactive_map, plot_pca_coords, plot_frequencies_heatmap) that all share the following parameter and return logic:

def plot_something(
    ...,
    show: bool = True,
    ...
):
    fig = px.something(...)
    ...
    if show:
        fig.show(renderer=renderer)
        return None
    else:
        return fig

Because show=True by default, the fig object is consumed and None is returned unless the user explicitly passes show=False.

---

Why This Is an Issue

1. Breaks API Expectations

In most modern Python data science libraries (like matplotlib.pyplot, seaborn, or standard plotly.express), plotting functions generally return the Figure or Axes object by default, even if they also display it. Returning None by default is surprising and breaks standard visualization workflows.

2. Difficult Customization

If a user wants to quickly test a plot and then decides to update the title, tweak an axis label, or change the legend position, they discover their variable is None:

# User intuitively does this:
fig = ag3.plot_pca_coords(sample_sets="v3_wild")
# fig is now None!
 
fig.update_layout(title="My Custom Title")
# AttributeError: 'NoneType' object has no attribute 'update_layout'

They must realize the API design choice, rerun the potentially expensive data-fetching/computation step with show=False, customize the figure, and then manually call fig.show().

3. Export Friction

Saving plots for publications or reports requires the figure object. Users shouldn't have to suppress inline display (show=False) just to be able to save the plot to a file via fig.write_image().

---

Proposed Solution

Change the logic to always return the figure object, regardless of the show parameter, while maintaining full backward compatibility.

Before:

def plot_something(
    ...,
    show: bool = True,
    ...
):
    fig = px.something(...)
    ...
    if show:
        fig.show(renderer=renderer)
        return None
    else:
        return fig

After:

def plot_something(
    ...,
    show: bool = True,
    ...
):
    fig = px.something(...)
    ...
    if show:
        fig.show(renderer=renderer)
    # Always return the figure!
    return fig

---

Benefits

Backward Compatibility

Code that currently works continues to work — the inline display still happens exactly as before. No existing user workflows are broken.

Intuitive Programmatic Access

Users can now customize and export plots without re-running expensive computations:

# Displays inline AND returns figure — no show=False needed
fig = ag3.plot_pca_coords(sample_sets="v3_wild")
 
# Customize immediately
fig.update_layout(title="My Custom PCA")
 
# Export immediately
fig.write_image("pca.png")

Alignment with Ecosystem Conventions

This brings the API in line with standard Plotly and broader Python data science conventions, reducing surprise for new users coming from matplotlib, seaborn, or plotly.express.

---

Scope of Change

This change would need to be applied systematically across the ~45 affected methods in the following files:

File	Example Methods
frq_base.py	plot_frequencies_heatmap, plot_frequencies_time_series
heterozygosity.py	plot_heterozygosity_track
sample_metadata.py	plot_samples_interactive_map
dipclust.py	plot_diplotype_clustering
cnv_data.py	plot_cnv_hmm_coverage
h12.py	plot_h12_gwss_track
g123.py	plot_g123_gwss_track
fst.py	plot_fst_gwss_track

The fix is mechanical and low-risk: a one-line change per method (removing return None from the if show: branch and moving return fig outside the conditional). A search for the pattern return None following fig.show( would identify all locations reliably.

---

[khushthecoder]

Hi @jonbrenas,

Currently, every plot method defaults to show=True which calls fig.show() and return None, meaning the figure is lost. Users who want to customize, compose, or export plots must re-run the entire computation with show=False.

The fix is a one-line mechanical change per method: remove return None from the if show: branch and move return fig outside the conditional. This is fully backward compatible — existing code that ignores the return value works identically, and the inline display still happens.

The change brings the API in line with standard Plotly/matplotlib conventions where plot functions always return the figure object.

I have submitted a PR addressing this issue. Could you please review it? #1163