Fix a bunch of errors in color_brewer and add unit tests

branca/scheme_base_codes.json

@@ -1 +1 @@
-{"codes": ["Spectral", "RdYlGn", "PuBu", "Accent", "OrRd", "Set1", "Set2", "Set3", "BuPu", "Dark2", "RdBu", "Oranges", "BuGn", "PiYG", "YlOrBr", "YlGn", "Pastel2", "RdPu", "Greens", "PRGn", "YlGnBu", "RdYlBu", "Paired", "BrBG", "Purples", "Reds", "Pastel1", "GnBu", "Greys", "RdGy", "YlOrRd", "PuOr", "PuRd", "Blues", "PuBuGn"]}
+{"codes": ["viridis","Spectral", "RdYlGn", "PuBu", "Accent", "OrRd", "Set1", "Set2", "Set3", "BuPu", "Dark2", "RdBu", "Oranges", "BuGn", "PiYG", "YlOrBr", "YlGn", "Pastel2", "RdPu", "Greens", "PRGn", "YlGnBu", "RdYlBu", "Paired", "BrBG", "Purples", "Reds", "Pastel1", "GnBu", "Greys", "RdGy", "YlOrRd", "PuOr", "PuRd", "Blues", "PuBuGn"]}

branca/scheme_info.json

@@ -1 +1 @@
-{"Spectral": "Diverging", "RdYlGn": "Diverging", "Set2": "Qualitative", "Accent": "Qualitative", "OrRd": "Sequential", "Set1": "Qualitative", "PuBu": "Sequential", "Set3": "Qualitative", "BuPu": "Sequential", "Dark2": "Qualitative", "RdBu": "Diverging", "BuGn": "Sequential", "PiYG": "Diverging", "YlOrBr": "Sequential", "YlGn": "Sequential", "RdPu": "Sequential", "PRGn": "Diverging", "YlGnBu": "Sequential", "RdYlBu": "Diverging", "Paired": "Qualitative", "Pastel2": "Qualitative", "Pastel1": "Qualitative", "GnBu": "Sequential", "RdGy": "Diverging", "YlOrRd": "Sequential", "PuOr": "Diverging", "PuRd": "Sequential", "BrBg": "Diverging", "PuBuGn": "Sequential"}
+{"Spectral": "Diverging", "RdYlGn": "Diverging", "Set2": "Qualitative", "Accent": "Qualitative", "OrRd": "Sequential", "Set1": "Qualitative", "PuBu": "Sequential", "Set3": "Qualitative", "BuPu": "Sequential", "Dark2": "Qualitative", "RdBu": "Diverging", "BuGn": "Sequential", "PiYG": "Diverging", "YlOrBr": "Sequential", "YlGn": "Sequential", "RdPu": "Sequential", "PRGn": "Diverging", "YlGnBu": "Sequential", "RdYlBu": "Diverging", "Paired": "Qualitative", "Pastel2": "Qualitative", "Pastel1": "Qualitative", "GnBu": "Sequential", "RdGy": "Diverging", "YlOrRd": "Sequential", "PuOr": "Diverging", "PuRd": "Sequential", "BrBG": "Diverging", "PuBuGn": "Sequential", "Greens": "Sequential", "viridis": "Sequential", "Oranges": "Sequential", "Blues": "Sequential", "Greys": "Sequential", "Reds": "Sequential", "Purples": "Sequential"}

branca/utilities.py

@@ -108,6 +108,9 @@ def color_brewer(color_code, n=6):
     maximum_n = 253
     minimum_n = 3
 
+    if not isinstance(n, int):
+        raise TypeError('n has to be an int, not a %s' % type(n))
+
     # Raise an error if the n requested is greater than the maximum.
     if n > maximum_n:
         raise ValueError('The maximum number of colors in a'
@@ -131,7 +134,7 @@ def color_brewer(color_code, n=6):
     with open(os.path.join(rootpath, '_schemes.json')) as f:
         schemes = json.loads(f.read())
 
-    with open(os.path.join(rootpath, '_cnames.json')) as f:
+    with open(os.path.join(rootpath, 'scheme_info.json')) as f:
         scheme_info = json.loads(f.read())
 
     with open(os.path.join(rootpath, 'scheme_base_codes.json')) as f:
@@ -140,10 +143,8 @@ def color_brewer(color_code, n=6):
     if base_code not in core_schemes:
         raise ValueError(base_code + ' is not a valid ColorBrewer code')
 
-    try:
-        schemes[core_color_code]
-        explicit_scheme = True
-    except KeyError:
+    explicit_scheme = True
+    if schemes.get(core_color_code) is None:
         explicit_scheme = False
 
     # Only if n is greater than the scheme length do we interpolate values.
@@ -162,10 +163,21 @@ def color_brewer(color_code, n=6):
                              ' and ' + str(max(matching_quals))
                              )
         else:
+            longest_scheme_name = base_code
+            longest_scheme_n = 0
+            for sn_name in schemes.keys():
+                if '_' not in sn_name:
+                    continue
+                if sn_name.split('_')[0] != base_code:
+                    continue
+                if int(sn_name.split('_')[1]) > longest_scheme_n:
+                    longest_scheme_name = sn_name
+                    longest_scheme_n = int(sn_name.split('_')[1])
+
             if not color_reverse:
-                color_scheme = linear_gradient(schemes.get(core_color_code), n)
+                color_scheme = linear_gradient(schemes.get(longest_scheme_name), n)
             else:
-                color_scheme = linear_gradient(schemes.get(core_color_code)[::-1], n)
+                color_scheme = linear_gradient(schemes.get(longest_scheme_name)[::-1], n)
     else:
         if not color_reverse:
             color_scheme = schemes.get(core_color_code, None)

tests/test_iframe.py

@@ -8,6 +8,7 @@
 
 import pytest
 from selenium.webdriver import Firefox
+from selenium.webdriver.common.by import By
 from selenium.webdriver.firefox.options import Options
 
 import branca.element as elem
@@ -55,7 +56,7 @@ def test_rendering_figure_notebook():
     try:
         driver.get('file://' + filepath)
         driver.switch_to.frame(0)
-        text_div = driver.find_element_by_css_selector('div')
+        text_div = driver.find_element(By.CSS_SELECTOR, 'div')
         assert text_div.text == text
     finally:
         os.remove(filepath)

tests/test_utilities.py

@@ -1,4 +1,22 @@
 import branca.utilities as ut
+import os
+from pathlib import Path
+import json
+import pytest
+
+
+rootpath = Path(os.path.dirname(os.path.abspath(__file__))) / '..' / 'branca'
+color_brewer_minimum_n = 3
+color_brewer_maximum_n = 253  # Why this limitation @ branca/utilities.py:108 ?
+
+
+# Loads schemes and their meta-data
+with open(rootpath / '_schemes.json') as f:
+    schemes = json.loads(f.read())
+with open(rootpath / 'scheme_info.json') as f:
+    scheme_info = json.loads(f.read())
+with open(rootpath / 'scheme_base_codes.json') as f:
+    core_schemes = json.loads(f.read())['codes']
 
 
 def test_color_brewer_base():
@@ -14,3 +32,62 @@ def test_color_brewer_reverse():
     scheme = ut.color_brewer('YlGnBu')
     scheme_r = ut.color_brewer('YlGnBu_r')
     assert scheme[::-1] == scheme_r
+
+
+def test_color_brewer_extendability():
+    """
+    The non-qualitative schemes should be extendable.
+
+    :see https://github.com/python-visualization/branca/issues/104
+    :see https://github.com/python-visualization/branca/issues/114
+
+    Moreover, the following error was not reported via issues:
+    * TypeError in the linear_gradient function when trying to extend any scheme. Indeed, in color_brewer, the key searched in the scheme database was not found, thus, it was passing `None` instead of a real scheme vector to linear_gradient.
+    """
+    for sname in core_schemes:
+        for n in range(color_brewer_minimum_n, color_brewer_maximum_n+1):
+            try:
+                scheme = ut.color_brewer(sname, n=n)
+            except Exception as e:
+                if scheme_info[sname] == 'Qualitative' and isinstance(e, ValueError):
+                    continue
+                raise
+
+            assert len(scheme) == n
+
+            # When we try to extend a scheme, the reverse is not always the exact reverse vector of the original one.
+            # Thus, we do not test this property!
+            _ = ut.color_brewer(sname + '_r', n=n)
+
+
+def test_color_avoid_unexpected_error():
+    """
+    We had unexpected errors by providing some scheme name with unexpected value of `n`.
+    This function tests them.
+
+    Identified errors which was not reported via issues:
+    * The scheme 'viridis' was not in the base_codes JSON;
+    * Multiple scheme hadn't any metadata in scheme_info JSON;
+    * When a `n` value provided to `color_scheme` was a float, it tried to select an unknown scheme without raising the right Exception type.
+    """
+
+    # Verify that every scheme has is present in base codes
+    scheme_names = set()
+    for sname in schemes.keys():
+        scheme_names.add(
+            sname.split('_')[0]
+        )
+    assert scheme_names == set(core_schemes)
+
+    # Verify that every scheme has a metadata
+    assert scheme_names == set(scheme_info.keys())
+
+    # Verify that every scheme can be generated in color_brewer using exotic value of `n`. Note that big but valid
+    # values are generated by test_color_brewer_extendability.
+    for sname in scheme_names:
+        for n in [-10] + list(range(-1, color_brewer_minimum_n)) + list(range(color_brewer_maximum_n+1, color_brewer_maximum_n+10)):
+            with pytest.raises(ValueError):
+                ut.color_brewer(sname, n)
+        for n in [str(color_brewer_minimum_n), float(color_brewer_minimum_n), 'abc']:
+            with pytest.raises(TypeError):
+                ut.color_brewer(sname, n)