Merge pull request #147 from autoscrape-labs/fix/script-execution

Improve script execution logic

Author: thalissonvs

pydoll/browser/tab.py

@@ -12,6 +12,7 @@
     TypeAlias,
     Union,
     cast,
+    overload,
 )
 
 import aiofiles
@@ -32,6 +33,7 @@
     IFrameNotFound,
     InvalidFileExtension,
     InvalidIFrame,
+    InvalidScriptWithElement,
     NetworkEventsNotEnabled,
     NoDialogPresent,
     NotAnIFrame,
@@ -43,9 +45,13 @@
 from pydoll.protocol.network.types import Cookie, CookieParam, NetworkLog
 from pydoll.protocol.page.events import PageEvent
 from pydoll.protocol.page.responses import CaptureScreenshotResponse, PrintToPDFResponse
-from pydoll.protocol.runtime.responses import EvaluateResponse
+from pydoll.protocol.runtime.responses import CallFunctionOnResponse, EvaluateResponse
 from pydoll.protocol.storage.responses import GetCookiesResponse
-from pydoll.utils import decode_base64_to_bytes
+from pydoll.utils import (
+    decode_base64_to_bytes,
+    has_return_outside_function,
+    is_script_already_function,
+)
 
 if TYPE_CHECKING:
     from pydoll.browser.chromium.base import Browser
@@ -554,7 +560,15 @@ async def handle_dialog(self, accept: bool, prompt_text: Optional[str] = None):
             PageCommands.handle_javascript_dialog(accept=accept, prompt_text=prompt_text)
         )
 
-    async def execute_script(self, script: str, element: Optional[WebElement] = None):
+    @overload
+    async def execute_script(self, script: str) -> EvaluateResponse: ...
+
+    @overload
+    async def execute_script(self, script: str, element: WebElement) -> CallFunctionOnResponse: ...
+
+    async def execute_script(
+        self, script: str, element: Optional[WebElement] = None
+    ) -> Union[EvaluateResponse, CallFunctionOnResponse]:
         """
         Execute JavaScript in page context.
 
@@ -565,17 +579,17 @@ async def execute_script(self, script: str, element: Optional[WebElement] = None
         Examples:
             await page.execute_script('argument.click()', element)
             await page.execute_script('argument.value = "Hello"', element)
+
+        Raises:
+            InvalidScriptWithElement: If script contains 'argument' but no element is provided.
         """
+        if 'argument' in script and element is None:
+            raise InvalidScriptWithElement('Script contains "argument" but no element was provided')
+
         if element:
-            script = script.replace('argument', 'this')
-            script = f'function(){{ {script} }}'
-            object_id = element._object_id
-            command = RuntimeCommands.call_function_on(
-                object_id=object_id, function_declaration=script, return_by_value=True
-            )
-        else:
-            command = RuntimeCommands.evaluate(expression=script)
-        return await self._execute_command(command)
+            return await self._execute_script_with_element(script, element)
+
+        return await self._execute_script_without_element(script)
 
     @asynccontextmanager
     async def expect_file_chooser(
@@ -693,6 +707,46 @@ async def callback_wrapper(event):
             event_name, function_to_register, temporary
         )
 
+    async def _execute_script_with_element(self, script: str, element: WebElement):
+        """
+        Execute script with element context.
+
+        Args:
+            script: JavaScript code to execute.
+            element: Element context (use 'argument' in script to reference).
+
+        Returns:
+            The result of the script execution.
+        """
+        if 'argument' not in script:
+            raise InvalidScriptWithElement('Script does not contain "argument"')
+
+        script = script.replace('argument', 'this')
+
+        if not is_script_already_function(script):
+            script = f'function(){{ {script} }}'
+
+        command = RuntimeCommands.call_function_on(
+            object_id=element._object_id, function_declaration=script, return_by_value=True
+        )
+        return await self._execute_command(command)
+
+    async def _execute_script_without_element(self, script: str):
+        """
+        Execute script without element context.
+
+        Args:
+            script: JavaScript code to execute.
+
+        Returns:
+            The result of the script execution.
+        """
+        if has_return_outside_function(script):
+            script = f'(function(){{ {script} }})()'
+
+        command = RuntimeCommands.evaluate(expression=script)
+        return await self._execute_command(command)
+
     async def _refresh_if_url_not_changed(self, url: str) -> bool:
         """Refresh page if URL hasn't changed."""
         current_url = await self.current_url

pydoll/exceptions.py

@@ -239,3 +239,15 @@ class NetworkEventsNotEnabled(PydollException):
     """Raised when network events are not enabled."""
 
     message = 'Network events not enabled'
+
+
+class ScriptException(PydollException):
+    """Base class for exceptions related to JavaScript execution."""
+
+    message = 'A script execution error occurred'
+
+
+class InvalidScriptWithElement(ScriptException):
+    """Raised when a script contains 'argument' but no element is provided."""
+
+    message = 'Script contains "argument" but no element was provided'

pydoll/utils.py

@@ -1,6 +1,7 @@
 import base64
 import logging
 import os
+import re
 
 import aiohttp
 
@@ -70,3 +71,96 @@ def validate_browser_paths(paths: list[str]) -> str:
         if os.path.exists(path) and os.access(path, os.X_OK):
             return path
     raise InvalidBrowserPath(f'No valid browser path found in: {paths}')
+
+
+def clean_script_for_analysis(script: str) -> str:
+    """
+    Clean JavaScript code by removing comments and string literals.
+
+    This helps avoid false positives when analyzing script structure.
+
+    Args:
+        script: JavaScript code to clean.
+
+    Returns:
+        str: Cleaned script with comments and strings removed.
+    """
+    # Remove line comments
+    cleaned = re.sub(r'//.*?$', '', script, flags=re.MULTILINE)
+    # Remove block comments
+    cleaned = re.sub(r'/\*.*?\*/', '', cleaned, flags=re.DOTALL)
+    # Remove double quoted strings
+    cleaned = re.sub(r'"[^"]*"', '""', cleaned)
+    # Remove single quoted strings
+    cleaned = re.sub(r"'[^']*'", "''", cleaned)
+    # Remove template literals
+    cleaned = re.sub(r'`[^`]*`', '``', cleaned)
+
+    return cleaned
+
+
+def is_script_already_function(script: str) -> bool:
+    """
+    Check if a JavaScript script is already wrapped in a function.
+
+    Args:
+        script: JavaScript code to analyze.
+
+    Returns:
+        bool: True if script is already a function, False otherwise.
+    """
+    cleaned_script = clean_script_for_analysis(script)
+
+    function_pattern = r'^\s*function\s*\([^)]*\)\s*\{'
+    arrow_function_pattern = r'^\s*\([^)]*\)\s*=>\s*\{'
+
+    return bool(
+        re.match(function_pattern, cleaned_script.strip())
+        or re.match(arrow_function_pattern, cleaned_script.strip())
+    )
+
+
+def has_return_outside_function(script: str) -> bool:
+    """
+    Check if a JavaScript script has return statements outside of functions.
+
+    Args:
+        script: JavaScript code to analyze.
+
+    Returns:
+        bool: True if script has return outside function, False otherwise.
+    """
+    cleaned_script = clean_script_for_analysis(script)
+
+    # If already a function, no need to check
+    if is_script_already_function(cleaned_script):
+        return False
+
+    # Look for 'return' statements
+    return_pattern = r'\breturn\b'
+    if not re.search(return_pattern, cleaned_script):
+        return False
+
+    # Check if return is inside a function by counting braces
+    lines = cleaned_script.split('\n')
+    brace_count = 0
+    in_function = False
+
+    for line in lines:
+        # Check for function declarations
+        if re.search(r'\bfunction\b', line) or re.search(r'=>', line):
+            in_function = True
+
+        # Count braces
+        brace_count += line.count('{') - line.count('}')
+
+        # Check for return statement
+        if re.search(return_pattern, line):
+            if not in_function or brace_count <= 0:
+                return True
+
+        # Reset function flag if we're back to top level
+        if brace_count <= 0:
+            in_function = False
+
+    return False