From a6caae23984aedb4c3855685f4e46066e4706074 Mon Sep 17 00:00:00 2001
From: Jon Mease <jon.mease@gmail.com>
Date: Sat, 1 Sep 2018 19:31:53 -0400
Subject: [PATCH] Annotated heatmap figure factory fixes and improvements

 - Fix `map object not subscriptable` error (This happened on Python 3 when specifying a custom colormap using `rgb` colors)
 - Use colorscale validator to validate specified colorscale before trying to do text color inference.
 - Handle case where some colors are specified as rgb and some as hex.
 - Add 'Cividis" to list of know lighter-to-darker colorscales.
---
 plotly/figure_factory/_annotated_heatmap.py   | 33 ++++++++++++-------
 .../test_tools/test_figure_factory.py         | 18 +++++-----
 2 files changed, 29 insertions(+), 22 deletions(-)

diff --git a/plotly/figure_factory/_annotated_heatmap.py b/plotly/figure_factory/_annotated_heatmap.py
index 5b2a8eaec9..b8ba35aa0d 100644
--- a/plotly/figure_factory/_annotated_heatmap.py
+++ b/plotly/figure_factory/_annotated_heatmap.py
@@ -3,6 +3,7 @@
 from plotly import exceptions, optional_imports
 from plotly.figure_factory import utils
 from plotly.graph_objs import graph_objs
+from plotly.validators.heatmap import ColorscaleValidator
 
 # Optional imports, may be None for users that only use our core functionality.
 np = optional_imports.get_module('numpy')
@@ -86,6 +87,11 @@ def create_annotated_heatmap(z, x=None, y=None, annotation_text=None,
     # Avoiding mutables in the call signature
     font_colors = font_colors if font_colors is not None else []
     validate_annotated_heatmap(z, x, y, annotation_text)
+
+    # validate colorscale
+    colorscale_validator = ColorscaleValidator()
+    colorscale = colorscale_validator.validate_coerce(colorscale)
+
     annotations = _AnnotatedHeatmap(z, x, y, annotation_text,
                                     colorscale, font_colors, reversescale,
                                     **kwargs).make_annotations()
@@ -112,6 +118,15 @@ def create_annotated_heatmap(z, x=None, y=None, annotation_text=None,
     return graph_objs.Figure(data=data, layout=layout)
 
 
+def to_rgb_color_list(color_str, default):
+    if 'rgb' in color_str:
+        return [int(v) for v in color_str.strip('rgb()').split(',')]
+    elif '#' in color_str:
+        return utils.hex_to_rgb(color_str)
+    else:
+        return default
+
+
 class _AnnotatedHeatmap(object):
     """
     Refer to TraceFactory.create_annotated_heatmap() for docstring
@@ -155,7 +170,7 @@ def get_text_color(self):
         colorscales = ['Greys', 'Greens', 'Blues',
                        'YIGnBu', 'YIOrRd', 'RdBu',
                        'Picnic', 'Jet', 'Hot', 'Blackbody',
-                       'Earth', 'Electric', 'Viridis']
+                       'Earth', 'Electric', 'Viridis', 'Cividis']
         # Plotly colorscales ranging from a darker shade to a lighter shade
         colorscales_reverse = ['Reds']
         if self.font_colors:
@@ -174,17 +189,11 @@ def get_text_color(self):
             min_text_color = '#000000'
             max_text_color = '#FFFFFF'
         elif isinstance(self.colorscale, list):
-            if 'rgb' in self.colorscale[0][1]:
-                min_col = map(int,
-                              self.colorscale[0][1].strip('rgb()').split(','))
-                max_col = map(int,
-                              self.colorscale[-1][1].strip('rgb()').split(','))
-            elif '#' in self.colorscale[0][1]:
-                    min_col = utils.hex_to_rgb(self.colorscale[0][1])
-                    max_col = utils.hex_to_rgb(self.colorscale[-1][1])
-            else:
-                min_col = [255, 255, 255]
-                max_col = [255, 255, 255]
+
+            min_col = to_rgb_color_list(self.colorscale[0][1],
+                                        [255, 255, 255])
+            max_col = to_rgb_color_list(self.colorscale[-1][1],
+                                        [255, 255, 255])
 
             if (min_col[0]*0.299 + min_col[1]*0.587 + min_col[2]*0.114) > 186:
                 min_text_color = '#000000'
diff --git a/plotly/tests/test_optional/test_tools/test_figure_factory.py b/plotly/tests/test_optional/test_tools/test_figure_factory.py
index 2a3d227cfa..e9a11ff0ae 100644
--- a/plotly/tests/test_optional/test_tools/test_figure_factory.py
+++ b/plotly/tests/test_optional/test_tools/test_figure_factory.py
@@ -822,16 +822,14 @@ def test_annotated_heatmap_kwargs(self):
 
         z = [[1, 0], [.25, .75], [.45, .5]]
         text = [['first', 'second'], ['third', 'fourth'], ['fifth', 'sixth']]
-        a = ff.create_annotated_heatmap(z, x=['A', 'B'],
-                                                       y=['One', 'Two',
-                                                          'Three'],
-                                                       annotation_text=text,
-                                                       colorscale=[[0,
-                                                                    '#ffffff'],
-                                                                   [1,
-                                                                    '#e6005a']]
-                                                       )
-        expected_a = {'data': [{'colorscale': [[0, '#ffffff'], [1, '#e6005a']],
+        a = ff.create_annotated_heatmap(z,
+                                        x=['A', 'B'],
+                                        y=['One', 'Two', 'Three'],
+                                        annotation_text=text,
+                                        colorscale=[[0, 'rgb(255,255,255)'],
+                                                    [1, '#e6005a']])
+        expected_a = {'data': [{'colorscale':
+                                    [[0, 'rgb(255,255,255)'], [1, '#e6005a']],
                                 'showscale': False,
                                 'type': 'heatmap',
                                 'x': ['A', 'B'],