Merge pull request #5 from brunorosilva/hotfix/general-colorscale

fix global colorscale problems

plotly_calplot/calplot.py

@@ -2,6 +2,10 @@
 from plotly import graph_objects as go
 from plotly.subplots import make_subplots
 
+from plotly_calplot.layout_formatter import (
+    apply_general_colorscaling,
+    showscale_of_heatmaps,
+)
 from plotly_calplot.single_year_calplot import year_calplot
 from plotly_calplot.utils import fill_empty_with_zeros
 
@@ -22,6 +26,7 @@ def calplot(
     month_lines: bool = True,
     total_height: int = None,
     space_between_plots: float = 0.08,
+    showscale: bool = False,
 ) -> go.Figure:
     """
     Yearly Calendar Heatmap
@@ -78,6 +83,10 @@ def calplot(
 
     space_between_plots: float = 0.08
         controls the vertical space between the plots
+
+    showscale: bool = False
+        if True, a color legend will be created.
+        Thanks to @ghhar98!
     """
     unique_years = data[x].dt.year.unique()
     unique_years_amount = len(unique_years)
@@ -120,4 +129,8 @@ def calplot(
             total_height=total_height,
         )
 
+    fig = apply_general_colorscaling(data, y, fig)
+    if showscale:
+        fig = showscale_of_heatmaps(fig)
+
     return fig

plotly_calplot/layout_formatter.py

@@ -105,3 +105,14 @@ def update_plot_with_current_layout(
     fig.update_layout(width=width, height=total_height)
     fig.add_traces(cplt, rows=[(row + 1)] * len(cplt), cols=[1] * len(cplt))
     return fig
+
+
+def apply_general_colorscaling(data: pd.DataFrame, y: str, fig: go.Figure) -> go.Figure:
+    return fig.update_traces(selector=dict(type="heatmap"), zmax=data[y].max(), zmin=0)
+
+
+def showscale_of_heatmaps(fig: go.Figure) -> go.Figure:
+    return fig.update_traces(
+        showscale=True,
+        selector=dict(type="heatmap"),
+    )