Keep map only in png screenshot

folium/folium.py

@@ -330,7 +330,8 @@ def _to_png(self, delay=3, driver=None):
                 driver.get('file:///{path}'.format(path=fname))
                 driver.maximize_window()
                 time.sleep(delay)
-                png = driver.get_screenshot_as_png()
+                div = driver.find_element('class name', 'folium-map')
+                png = div.screenshot_as_png
                 driver.quit()
             self._png_image = png
         return self._png_image

tests/test_repr.py

@@ -5,7 +5,6 @@
 """
 
 import io
-import sys
 
 import PIL.Image
 
@@ -47,16 +46,22 @@ def test__repr_png_no_image(m):
     assert png is None
 
 
-@pytest.mark.xfail
 def test__repr_png_is_bytes(m_png):
     png = m_png._repr_png_()
     assert isinstance(png, bytes)
 
 
-@pytest.mark.xfail
-@pytest.mark.skipif(sys.version_info < (3, 0),
-                    reason="Doesn't work on Python 2.7.")
 def test_valid_png(m_png):
     png = m_png._repr_png_()
     img = PIL.Image.open(io.BytesIO(png))
-    isinstance(img, PIL.PngImagePlugin.PngImageFile)
+    assert isinstance(img, PIL.PngImagePlugin.PngImageFile)
+
+
+def test_valid_png_size(m_png):
+    from folium.utilities import _parse_size
+    w = h = 500
+    m_png.width = _parse_size(w)
+    m_png.height = _parse_size(h)
+    png = m_png._repr_png_()
+    img = PIL.Image.open(io.BytesIO(png))
+    assert img.size == (w, h)