Fix heatmap intensity (#1282)

* Update link

* update link, discard max_val

* Update test_heat_map.py

* Update url after rebase

* Rename temp html file utility function

* temp_html_filepath accept str and bytes

* Refactor selenium tests

* Small additional refactor of selenium fixture

* Add HeatMap with weights selenium test

* Set window size at driver load

* Print screenshot if test fails

* export canvas, not full screenshot

* Use JS resource from folium master

folium/folium.py

@@ -14,7 +14,7 @@
 from folium.raster_layers import TileLayer
 from folium.utilities import (
     _parse_size,
-    _tmp_html,
+    temp_html_filepath,
     validate_location,
     parse_options,
 )
@@ -314,7 +314,7 @@ def _to_png(self, delay=3):
             driver = webdriver.Firefox(options=options)
 
             html = self.get_root().render()
-            with _tmp_html(html) as fname:
+            with temp_html_filepath(html) as fname:
                 # We need the tempfile to avoid JS security issues.
                 driver.get('file:///{path}'.format(path=fname))
                 driver.maximize_window()

folium/plugins/heat_map.py

@@ -1,5 +1,7 @@
 # -*- coding: utf-8 -*-
 
+import warnings
+
 from branca.element import Figure, JavascriptLink
 
 from folium.map import Layer
@@ -17,7 +19,7 @@
 
 _default_js = [
     ('leaflet-heat.js',
-     'https://leaflet.github.io/Leaflet.heat/dist/leaflet-heat.js'),
+     'https://cdn.jsdelivr.net/gh/python-visualization/folium@master/folium/templates/leaflet_heat.min.js'),  # noqa
     ]
 
 
@@ -37,8 +39,6 @@ class HeatMap(Layer):
     max_zoom : default 18
         Zoom level where the points reach maximum intensity (as intensity
         scales with zoom), equals maxZoom of the map by default
-    max_val : float, default 1.
-        Maximum point intensity
     radius : int, default 25
         Radius of each "point" of the heatmap
     blur : int, default 15
@@ -62,7 +62,7 @@ class HeatMap(Layer):
         """)
 
     def __init__(self, data, name=None, min_opacity=0.5, max_zoom=18,
-                 max_val=1.0, radius=25, blur=15, gradient=None,
+                 radius=25, blur=15, gradient=None,
                  overlay=True, control=True, show=True, **kwargs):
         super(HeatMap, self).__init__(name=name, overlay=overlay,
                                       control=control, show=show)
@@ -72,10 +72,13 @@ def __init__(self, data, name=None, min_opacity=0.5, max_zoom=18,
                      for line in data]
         if np.any(np.isnan(self.data)):
             raise ValueError('data may not contain NaNs.')
+        if kwargs.pop('max_val', None):
+            warnings.warn('The `max_val` parameter is no longer necessary. '
+                          'The largest intensity is calculated automatically.',
+                          stacklevel=2)
         self.options = parse_options(
             min_opacity=min_opacity,
             max_zoom=max_zoom,
-            max=max_val,
             radius=radius,
             blur=blur,
             gradient=gradient,

folium/utilities.py

@@ -429,12 +429,12 @@ def normalize(rendered):
 
 
 @contextmanager
-def _tmp_html(data):
+def temp_html_filepath(data):
     """Yields the path of a temporary HTML file containing data."""
     filepath = ''
     try:
         fid, filepath = tempfile.mkstemp(suffix='.html', prefix='folium_')
-        os.write(fid, data.encode('utf8'))
+        os.write(fid, data.encode('utf8') if isinstance(data, str) else data)
         os.close(fid)
         yield filepath
     finally:

tests/plugins/test_heat_map.py

@@ -28,7 +28,7 @@ def test_heat_map():
     out = normalize(m._parent.render())
 
     # We verify that the script import is present.
-    script = '<script src="https://leaflet.github.io/Leaflet.heat/dist/leaflet-heat.js"></script>'  # noqa
+    script = '<script src="https://cdn.jsdelivr.net/gh/python-visualization/folium@master/folium/templates/leaflet_heat.min.js"></script>'  # noqa
     assert script in out
 
     # We verify that the script part is correct.
@@ -38,7 +38,6 @@ def test_heat_map():
                 {
                     minOpacity: {{this.min_opacity}},
                     maxZoom: {{this.max_zoom}},
-                    max: {{this.max_val}},
                     radius: {{this.radius}},
                     blur: {{this.blur}},
                     gradient: {{this.gradient}}

tests/selenium/conftest.py

@@ -0,0 +1,56 @@
+
+import pytest
+from selenium.webdriver import Chrome, ChromeOptions
+from selenium.webdriver.common.by import By
+from selenium.webdriver.support.ui import WebDriverWait
+from selenium.webdriver.support.expected_conditions import visibility_of_element_located
+
+
+@pytest.fixture(scope='session')
+def driver():
+    """Pytest fixture that yields a Selenium WebDriver instance"""
+    driver = DriverFolium()
+    try:
+        yield driver
+    finally:
+        driver.quit()
+
+
+class DriverFolium(Chrome):
+    """Selenium WebDriver wrapper that adds folium test specific features."""
+
+    def __init__(self):
+        options = ChromeOptions()
+        options.add_argument('--no-sandbox')
+        options.add_argument('--disable-dev-shm-usage')
+        options.add_argument('--disable-gpu')
+        options.add_argument('--headless')
+        options.add_argument("--window-size=1024,768")
+        super().__init__(options=options)
+
+    def get_file(self, filepath):
+        self.clean_window()
+        super().get('file://' + filepath)
+
+    def clean_window(self):
+        """Make sure we have a fresh window (without restarting the browser)."""
+        # open new tab
+        self.execute_script('window.open();')
+        # close old tab
+        self.close()
+        # switch to new tab
+        self.switch_to.window(self.window_handles[0])
+
+    def verify_js_logs(self):
+        """Raise an error if there are errors in the browser JS console."""
+        logs = self.get_log('browser')
+        for log in logs:
+            if log['level'] == 'SEVERE':
+                msg = ' '.join(log['message'].split()[2:])
+                raise RuntimeError('Javascript error: "{}".'.format(msg))
+
+    def wait_until(self, css_selector, timeout=10):
+        """Wait for and return the element(s) selected by css_selector."""
+        wait = WebDriverWait(self, timeout=timeout)
+        is_visible = visibility_of_element_located((By.CSS_SELECTOR, css_selector))
+        return wait.until(is_visible)

tests/selenium/test_heat_map_selenium.py

@@ -0,0 +1,43 @@
+import base64
+import os
+
+import folium
+from folium.plugins.heat_map import HeatMap
+from folium.utilities import temp_html_filepath
+
+
+def test_heat_map_with_weights(driver):
+    """Verify that HeatMap uses weights in data correctly.
+
+    This test will fail in non-headless mode because window size will be different.
+
+    """
+    m = folium.Map((0.5, 0.5), zoom_start=8, tiles=None)
+    HeatMap(
+        # make four dots with different weights: 1, 1, 1.5 and 2.
+        data=[
+            (0, 0, 1.5),
+            (0, 1, 1),
+            (1, 0, 1),
+            (1, 1, 2),
+        ],
+        radius=70,
+        blur=50,
+    ).add_to(m)
+    html = m.get_root().render()
+    with temp_html_filepath(html) as filepath:
+        driver.get_file(filepath)
+        assert driver.wait_until('.folium-map')
+        driver.verify_js_logs()
+    canvas = driver.wait_until('canvas.leaflet-heatmap-layer')
+    assert canvas
+    # get the canvas as a PNG base64 string
+    canvas_base64 = driver.execute_script(
+        "return arguments[0].toDataURL('image/png').substring(21);", canvas)
+    screenshot = base64.b64decode(canvas_base64)
+    path = os.path.dirname(__file__)
+    with open(os.path.join(path, 'test_heat_map_selenium_screenshot.png'), 'rb') as f:
+        screenshot_expected = f.read()
+    if hash(screenshot) != hash(screenshot_expected):
+        print(screenshot)
+        assert False, 'screenshot is not as expected'

tests/selenium/test_selenium.py

@@ -7,41 +7,9 @@
 
 import nbconvert
 import pytest
-from selenium.webdriver import Chrome, ChromeOptions
 from selenium.common.exceptions import UnexpectedAlertPresentException
-from selenium.webdriver.common.by import By
-from selenium.webdriver.support.ui import WebDriverWait
-from selenium.webdriver.support.expected_conditions import visibility_of_element_located
 
-
-def create_driver():
-    """Create a Selenium WebDriver instance."""
-    options = ChromeOptions()
-    options.add_argument('--no-sandbox')
-    options.add_argument('--disable-dev-shm-usage')
-    options.add_argument('--disable-gpu')
-    options.add_argument('--headless')
-    driver = Chrome(options=options)
-    return driver
-
-
-@pytest.fixture(scope='module')
-def driver():
-    """Pytest fixture that yields a Selenium WebDriver instance"""
-    driver = create_driver()
-    try:
-        yield driver
-    finally:
-        driver.quit()
-
-
-def clean_window(driver):
-    # open new tab
-    driver.execute_script('window.open();')
-    # close old tab
-    driver.close()
-    # switch to new tab
-    driver.switch_to.window(driver.window_handles[0])
+from folium.utilities import temp_html_filepath
 
 
 def find_notebooks():
@@ -59,22 +27,15 @@ def find_notebooks():
 @pytest.mark.parametrize('filepath', find_notebooks())
 def test_notebook(filepath, driver):
     for filepath_html in get_notebook_html(filepath):
-        clean_window(driver)
-        driver.get('file://' + filepath_html)
-        wait = WebDriverWait(driver, timeout=10)
-        map_is_visible = visibility_of_element_located((By.CSS_SELECTOR, '.folium-map'))
+        driver.get_file(filepath_html)
         try:
-            assert wait.until(map_is_visible)
+            assert driver.wait_until('.folium-map')
         except UnexpectedAlertPresentException:
             # in Plugins.ipynb we get an alert about geolocation permission
             # for some reason it cannot be closed or avoided, so just ignore it
             print('skipping', filepath_html, 'because of alert')
             continue
-        logs = driver.get_log('browser')
-        for log in logs:
-            if log['level'] == 'SEVERE':
-                msg = ' '.join(log['message'].split()[2:])
-                raise RuntimeError('Javascript error: "{}".'.format(msg))
+        driver.verify_js_logs()
 
 
 def get_notebook_html(filepath_notebook, execute=True):
@@ -93,15 +54,9 @@ def get_notebook_html(filepath_notebook, execute=True):
     parser.feed(body)
     iframes = parser.iframes
 
-    for i, iframe in enumerate(iframes):
-        filepath_html = filepath_notebook.replace('.ipynb', '.{}.html'.format(i))
-        filepath_html = os.path.abspath(filepath_html)
-        with open(filepath_html, 'wb') as f:
-            f.write(iframe)
-        try:
+    for iframe in iframes:
+        with temp_html_filepath(iframe) as filepath_html:
             yield filepath_html
-        finally:
-            os.remove(filepath_html)
 
 
 class IframeParser(HTMLParser):