add methods

CHANGELOG.md

@@ -11,6 +11,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
+- Add `tab.set_content` methods to Set content to page @3mora2
+- Add `tab.expose_function` The method adds a function called name on the window object. When called, the function executes callback and returns a Promise which resolves to the return value of callback. @3mora2
+- Add `tab.expose_bindings` The method adds a function called name on the window object. When called, the function executes callback and returns a Promise which resolves to the return value of callback. If the callback returns a Promise, it will be awaited. @3mora2
+- Add `tab.add_script_tag` The method adds a script tag to page. @3mora2
+- Add `tab.add_style_tag` The method adds a style tag to page. @3mora2
+
 ### Changed
 
 ### Removed

examples/expose_function.py

@@ -0,0 +1,33 @@
+import asyncio
+import hashlib
+
+import zendriver as zd
+
+
+async def main():
+    async def sha256(text):
+        m = hashlib.sha256()
+        m.update(bytes(text, "utf8"))
+        return m.hexdigest()
+
+    async with await zd.start(headless=False) as browser:
+        page = browser.main_tab
+        await page.expose_function("sha256", sha256)
+
+        await page.set_content("""
+                <script>
+                  async function onClick() {
+                    document.querySelector('div#secret').textContent = await window.sha256('zendriver');
+                  }
+                </script>
+                <button onclick="onClick()">Click me</button>
+                <div id="secret"></div>
+            """)
+
+        await (await page.find("button")).click()
+        assert (await page.find("div#secret")).text == await sha256("zendriver")
+        print("done")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

tests/core/test_tab.py

@@ -191,7 +191,7 @@ async def test_expect_request(browser: zd.Browser) -> None:
 
     async with tab.expect_request(sample_file("groceries.html")) as request_info:
         await tab.get(sample_file("groceries.html"))
-        req = await asyncio.wait_for(request_info.value, timeout=3)
+        req = await request_info.value
         assert type(req) is zd.cdp.network.RequestWillBeSent
         assert type(req.request) is zd.cdp.network.Request
         assert req.request.url == sample_file("groceries.html")
@@ -208,7 +208,7 @@ async def test_expect_response(browser: zd.Browser) -> None:
 
     async with tab.expect_response(sample_file("groceries.html")) as response_info:
         await tab.get(sample_file("groceries.html"))
-        resp = await asyncio.wait_for(response_info.value, timeout=3)
+        resp = await response_info.value
         assert type(resp) is zd.cdp.network.ResponseReceived
         assert type(resp.response) is zd.cdp.network.Response
         assert resp.request_id is not None
@@ -225,8 +225,8 @@ async def test_expect_download(browser: zd.Browser) -> None:
     async with tab.expect_download() as download_ex:
         await tab.get(sample_file("groceries.html"))
         await (await tab.select("#download_file")).click()
-        download = await asyncio.wait_for(download_ex.value, timeout=3)
-        assert type(download) is zd.cdp.browser.DownloadWillBegin
+        download = await download_ex.value
+        assert type(download) is zd.cdp.page.DownloadWillBegin
         assert download.url is not None
 
 
@@ -246,3 +246,33 @@ async def test_intercept(browser: zd.Browser) -> None:
         assert body is not None
         # original_response = loads(body)
         # assert original_response["name"] == "Zendriver"
+
+
+async def test_expose_function(browser: zd.Browser):
+    async def sha256(text):
+        return str(text) + "sha256"
+
+    tab = browser.main_tab
+
+    await tab.expose_function("sha256", sha256)
+    await tab.set_content("""
+            <script>
+              async function onClick() {
+                document.querySelector('div#secret').textContent = await window.sha256('zendriver');
+              }
+            </script>
+            <button onclick="onClick()">Click me</button>
+            <div id="secret"></div>
+        """)
+
+    await (await tab.find("button")).click()
+    assert (await tab.find("div#secret")).text == await sha256("zendriver")
+
+
+async def test_add_script_tag(browser: zd.Browser):
+    tab = browser.main_tab
+
+    await tab.add_script_tag(content="""window.Store = {"name":"zendriver"} """)
+    assert await tab.evaluate("window.Store", return_by_value=True) == {
+        "name": "zendriver"
+    }

zendriver/core/browser.py

@@ -15,10 +15,11 @@
 import urllib.request
 import warnings
 from collections import defaultdict
-from typing import List, Tuple, Union, Any
+from typing import List, Tuple, Union, Dict, Any
 
 import asyncio_atexit
 
+from .helper import PageBinding
 from .. import cdp
 from . import tab, util
 from ._contradict import ContraDict
@@ -139,6 +140,7 @@ def __init__(self, config: Config):
         self._process_pid = None
         self._is_updating = asyncio.Event()
         self.connection = None
+        self._pageBindings: Dict[str, PageBinding] = dict()
         logger.debug("Session object initialized: %s" % vars(self))
 
     @property

zendriver/core/errors.py

@@ -0,0 +1,31 @@
+"""Exceptions for zendriver package."""
+
+
+class ZendriverError(Exception):
+    """Base exception for zendriver."""
+
+    pass
+
+
+class BrowserError(ZendriverError):
+    """Exception raised from browser."""
+
+    pass
+
+
+class ElementHandleError(ZendriverError):
+    """ElementHandle related exception."""
+
+    pass
+
+
+class NetworkError(ZendriverError):
+    """Network/Protocol related exception."""
+
+    pass
+
+
+class PageError(ZendriverError):
+    """Page/Frame related exception."""
+
+    pass

zendriver/core/helper.py

@@ -0,0 +1,128 @@
+"""Helper functions."""
+
+import asyncio
+import json
+import logging
+from typing import Any, Callable
+
+
+from .parse_evaluation_result import parse_evaluation_result_value
+from .raw_bindings_controller import raw_bindings_controller_source
+
+logger = logging.getLogger(__name__)
+
+
+def evaluation_string(fun: str, *args: Any, **kwargs) -> str:
+    """Convert function and arguments to str."""
+    _args = ", ".join([json.dumps("undefined" if arg is None else arg) for arg in args])
+    expr = f"({fun})({_args})"
+    return expr
+
+
+def init_script(source: str):
+    return f"(() => {{\n{source}\n}})();"
+
+
+class BindingSource:
+    def __init__(self, page, browser, execution_context_id):
+        self.page = page
+        self.browser = browser
+        self.execution_context_id = execution_context_id
+
+
+class PageBinding:
+    kController = "__zendriver__binding__controller__"
+    kBindingName = "__zendriver__binding__"
+
+    @staticmethod
+    def create_init_script():
+        source = f"""
+        (() => {{
+            const module = {{}};
+            {raw_bindings_controller_source}
+            const property = '{PageBinding.kController}';
+            if (!globalThis[property])
+                globalThis[property] = new (module.exports.BindingsController())(globalThis, '{PageBinding.kBindingName}');
+        }})();
+        """
+        return init_script(source)
+
+    def __init__(self, name: str, function: Callable, needs_handle: bool):
+        self.name = name
+        self.function = function
+        self.needs_handle = needs_handle
+        self.init_script = init_script(
+            f"globalThis['{PageBinding.kController}'].addBinding({json.dumps(name)}, {str(needs_handle).lower()})"
+        )
+
+        self.cleanup_script = (
+            f"globalThis['{PageBinding.kController}'].removeBinding({json.dumps(name)})"
+        )
+
+    @staticmethod
+    async def dispatch(page, event, browser):
+        execution_context_id = event.execution_context_id
+        data = json.loads(event.payload)
+        name = data["name"]
+        seq = data["seq"]
+        serialized_args = data["serializedArgs"]
+        binding = page.get_binding(name)
+        if not binding:
+            raise Exception(f'Function "{name}" is not exposed')
+        if not isinstance(serialized_args, list):
+            raise Exception(
+                "serializedArgs is not an array. This can happen when Array.prototype.toJSON is defined incorrectly"
+            )
+        try:
+            if binding.needs_handle:
+                try:
+                    handle = await page.evaluate(
+                        evaluation_string(
+                            "arg => globalThis['{}'].takeBindingHandle(arg)".format(
+                                PageBinding.kController
+                            ),
+                            {"name": name, "seq": seq},
+                        )
+                    )
+                except Exception:
+                    handle = None
+                result = binding.function(
+                    BindingSource(
+                        page=page,
+                        browser=browser,
+                        execution_context_id=execution_context_id,
+                    ),
+                    handle=handle,
+                )
+            else:
+                args = [parse_evaluation_result_value(a) for a in serialized_args]
+                result = binding.function(
+                    BindingSource(
+                        page=page,
+                        browser=browser,
+                        execution_context_id=execution_context_id,
+                    ),
+                    *args,
+                )
+
+            if asyncio.iscoroutine(result):
+                result = await result
+
+            await page.evaluate(
+                evaluation_string(
+                    "arg => globalThis['{}'].deliverBindingResult(arg)".format(
+                        PageBinding.kController
+                    ),
+                    {"name": name, "seq": seq, "result": result},
+                )
+            )
+        except Exception as error:
+            logger.error(error)
+            await page.evaluate(
+                evaluation_string(
+                    "arg => globalThis['{}'].deliverBindingResult(arg)".format(
+                        PageBinding.kController
+                    ),
+                    {"name": name, "seq": seq, "error": str(error)},
+                )
+            )