From f7faef799050602ea5b288b82387284f6b188a8e Mon Sep 17 00:00:00 2001
From: Rafique Anwar <rafique.anwar@gmail.com>
Date: Sat, 2 Jan 2021 19:36:52 +1100
Subject: [PATCH 1/2] fix error on retrieving result

---
 ipyvolume/headless.py | 26 +++++++++++++++++++++-----
 1 file changed, 21 insertions(+), 5 deletions(-)

diff --git a/ipyvolume/headless.py b/ipyvolume/headless.py
index 3dd2b0f8..4c0a24ca 100644
--- a/ipyvolume/headless.py
+++ b/ipyvolume/headless.py
@@ -49,7 +49,20 @@ def _screenshot_data(
     # else:
     #     assert isinstance(fig, ipv.Figure)
 
-    chrome = PyChromeDevTools.ChromeInterface(headless_kwargs)
+
+    try:
+        host = "localhost"
+        port = 9222
+
+        if headless_kwargs.get("host"):
+            host = headless_kwargs.get("host")
+
+        chrome = PyChromeDevTools.ChromeInterface(host=host, port=port)
+        #chrome = PyChromeDevTools.ChromeInterface(host="localhost")
+    except:
+        print("Falied to connect to headless chrome, please make sure it's running on host={} and port={}".format(host, port))
+        raise 
+
     chrome.Network.enable()
     chrome.Page.enable()
     chrome.Page.navigate(url=html_filename)
@@ -59,16 +72,19 @@ def _screenshot_data(
     result = chrome.Runtime.evaluate(expression='ipvss()')
     tries = 0
     while tries < 10:
+        #print(str(result))
         try:
-            url = result['result']['result']['value']
+            url = result[0]['result']['result']['value']
             return url
-        except:
-            if 'ipvss' in result['result']['result']['description']:
+        except Exception as ex:
+            #print(str(ex))
+            #raise ex
+            if 'ipvss' in result[0]['result']['result']['description']:
                 tries += 1
                 time.sleep(0.5)
             else:
                 print('error getting result, return value was:', result)
-                raise
+                raise ex
 
 
 def _main():

From 0a230ee1923aab245f4abd38c8cf9b9b6c5386cc Mon Sep 17 00:00:00 2001
From: Rafique Anwar <rafique.anwar@gmail.com>
Date: Tue, 5 Jan 2021 15:46:30 +1100
Subject: [PATCH 2/2] clean up

---
 ipyvolume/headless.py | 16 +++-------------
 1 file changed, 3 insertions(+), 13 deletions(-)

diff --git a/ipyvolume/headless.py b/ipyvolume/headless.py
index 4c0a24ca..8b10ea2e 100644
--- a/ipyvolume/headless.py
+++ b/ipyvolume/headless.py
@@ -42,7 +42,7 @@ def _get_browser():
 
 
 def _screenshot_data(
-    html_filename, timeout_seconds=10, output_widget=None, format="png", width=None, height=None, fig=None, **headless_kwargs):
+    html_filename, timeout_seconds=10, output_widget=None, format="png", width=None, height=None, fig=None, host="localhost", port=9222, **headless_kwargs):
     # browser = _get_browser()
     # if fig is None:
     #     fig = gcf()
@@ -51,14 +51,7 @@ def _screenshot_data(
 
 
     try:
-        host = "localhost"
-        port = 9222
-
-        if headless_kwargs.get("host"):
-            host = headless_kwargs.get("host")
-
         chrome = PyChromeDevTools.ChromeInterface(host=host, port=port)
-        #chrome = PyChromeDevTools.ChromeInterface(host="localhost")
     except:
         print("Falied to connect to headless chrome, please make sure it's running on host={} and port={}".format(host, port))
         raise 
@@ -72,19 +65,16 @@ def _screenshot_data(
     result = chrome.Runtime.evaluate(expression='ipvss()')
     tries = 0
     while tries < 10:
-        #print(str(result))
         try:
             url = result[0]['result']['result']['value']
             return url
-        except Exception as ex:
-            #print(str(ex))
-            #raise ex
+        except Exception:
             if 'ipvss' in result[0]['result']['result']['description']:
                 tries += 1
                 time.sleep(0.5)
             else:
                 print('error getting result, return value was:', result)
-                raise ex
+                raise
 
 
 def _main():