add bidi script commands

Author: navin772

py/selenium/webdriver/common/bidi/script.py

@@ -15,14 +15,207 @@
 # specific language governing permissions and limitations
 # under the License.
 
+from dataclasses import dataclass
+from typing import List, Optional
+
+from selenium.webdriver.common.bidi.common import command_builder
+
 from .log import LogEntryAdded
 from .session import Session
 
 
+class ResultOwnership:
+    """Represents the possible result ownership types."""
+
+    NONE = "none"
+    ROOT = "root"
+
+
+class RealmType:
+    """Represents the possible realm types."""
+
+    WINDOW = "window"
+    DEDICATED_WORKER = "dedicated-worker"
+    SHARED_WORKER = "shared-worker"
+    SERVICE_WORKER = "service-worker"
+    WORKER = "worker"
+    PAINT_WORKLET = "paint-worklet"
+    AUDIO_WORKLET = "audio-worklet"
+    WORKLET = "worklet"
+
+
+@dataclass
+class RealmInfo:
+    """Represents information about a realm."""
+
+    realm: str
+    origin: str
+    type: str
+    context: Optional[str] = None
+    sandbox: Optional[str] = None
+
+    @classmethod
+    def from_json(cls, json: dict) -> "RealmInfo":
+        """Creates a RealmInfo instance from a dictionary.
+
+        Parameters:
+        -----------
+            json: A dictionary containing the realm information.
+
+        Returns:
+        -------
+            RealmInfo: A new instance of RealmInfo.
+        """
+        return cls(
+            realm=json.get("realm"),
+            origin=json.get("origin"),
+            type=json.get("type"),
+            context=json.get("context"),
+            sandbox=json.get("sandbox"),
+        )
+
+
+@dataclass
+class Source:
+    """Represents the source of a script message."""
+
+    realm: str
+    context: Optional[str] = None
+
+    @classmethod
+    def from_json(cls, json: dict) -> "Source":
+        """Creates a Source instance from a dictionary.
+
+        Parameters:
+        -----------
+            json: A dictionary containing the source information.
+
+        Returns:
+        -------
+            Source: A new instance of Source.
+        """
+        return cls(
+            realm=json.get("realm"),
+            context=json.get("context"),
+        )
+
+
+@dataclass
+class EvaluateResult:
+    """Represents the result of script evaluation."""
+
+    realm: str
+    result: Optional[dict] = None
+    exception_details: Optional[dict] = None
+
+    @classmethod
+    def from_json(cls, json: dict) -> "EvaluateResult":
+        """Creates an EvaluateResult instance from a dictionary.
+
+        Parameters:
+        -----------
+            json: A dictionary containing the evaluation result.
+
+        Returns:
+        -------
+            EvaluateResult: A new instance of EvaluateResult.
+        """
+        return cls(
+            realm=json.get("realm"),
+            result=json.get("result"),
+            exception_details=json.get("exceptionDetails"),
+        )
+
+
+class ScriptMessage:
+    """Represents a script message event."""
+
+    event_class = "script.message"
+
+    def __init__(self, channel: str, data: dict, source: Source):
+        self.channel = channel
+        self.data = data
+        self.source = source
+
+    @classmethod
+    def from_json(cls, json: dict) -> "ScriptMessage":
+        """Creates a ScriptMessage instance from a dictionary.
+
+        Parameters:
+        -----------
+            json: A dictionary containing the script message.
+
+        Returns:
+        -------
+            ScriptMessage: A new instance of ScriptMessage.
+        """
+        return cls(
+            channel=json.get("channel"),
+            data=json.get("data"),
+            source=Source.from_json(json.get("source", {})),
+        )
+
+
+class RealmCreated:
+    """Represents a realm created event."""
+
+    event_class = "script.realmCreated"
+
+    def __init__(self, realm_info: RealmInfo):
+        self.realm_info = realm_info
+
+    @classmethod
+    def from_json(cls, json: dict) -> "RealmCreated":
+        """Creates a RealmCreated instance from a dictionary.
+
+        Parameters:
+        -----------
+            json: A dictionary containing the realm created event.
+
+        Returns:
+        -------
+            RealmCreated: A new instance of RealmCreated.
+        """
+        return cls(realm_info=RealmInfo.from_json(json))
+
+
+class RealmDestroyed:
+    """Represents a realm destroyed event."""
+
+    event_class = "script.realmDestroyed"
+
+    def __init__(self, realm: str):
+        self.realm = realm
+
+    @classmethod
+    def from_json(cls, json: dict) -> "RealmDestroyed":
+        """Creates a RealmDestroyed instance from a dictionary.
+
+        Parameters:
+        -----------
+            json: A dictionary containing the realm destroyed event.
+
+        Returns:
+        -------
+            RealmDestroyed: A new instance of RealmDestroyed.
+        """
+        return cls(realm=json.get("realm"))
+
+
 class Script:
+    """BiDi implementation of the script module."""
+
+    EVENTS = {
+        "message": "script.message",
+        "realm_created": "script.realmCreated",
+        "realm_destroyed": "script.realmDestroyed",
+    }
+
     def __init__(self, conn):
         self.conn = conn
         self.log_entry_subscribed = False
+        self.subscriptions = {}
+        self.callbacks = {}
 
     def add_console_message_handler(self, handler):
         self._subscribe_to_log_entries()
@@ -38,6 +231,185 @@ def remove_console_message_handler(self, id):
 
     remove_javascript_error_handler = remove_console_message_handler
 
+    def add_preload_script(
+        self,
+        function_declaration: str,
+        arguments: Optional[List[dict]] = None,
+        contexts: Optional[List[str]] = None,
+        user_contexts: Optional[List[str]] = None,
+        sandbox: Optional[str] = None,
+    ) -> str:
+        """Adds a preload script.
+
+        Parameters:
+        -----------
+            function_declaration: The function declaration to preload.
+            arguments: The arguments to pass to the function.
+            contexts: The browsing context IDs to apply the script to.
+            user_contexts: The user context IDs to apply the script to.
+            sandbox: The sandbox name to apply the script to.
+
+        Returns:
+        -------
+            str: The preload script ID.
+
+        Raises:
+        ------
+            ValueError: If both contexts and user_contexts are provided.
+        """
+        if contexts is not None and user_contexts is not None:
+            raise ValueError("Cannot specify both contexts and user_contexts")
+
+        params = {"functionDeclaration": function_declaration}
+
+        if arguments is not None:
+            params["arguments"] = arguments
+        if contexts is not None:
+            params["contexts"] = contexts
+        if user_contexts is not None:
+            params["userContexts"] = user_contexts
+        if sandbox is not None:
+            params["sandbox"] = sandbox
+
+        result = self.conn.execute(command_builder("script.addPreloadScript", params))
+        return result["script"]
+
+    def remove_preload_script(self, script_id: str) -> None:
+        """Removes a preload script.
+
+        Parameters:
+        -----------
+            script_id: The preload script ID to remove.
+        """
+        params = {"script": script_id}
+        self.conn.execute(command_builder("script.removePreloadScript", params))
+
+    def disown(self, handles: List[str], target: dict) -> None:
+        """Disowns the given handles.
+
+        Parameters:
+        -----------
+            handles: The handles to disown.
+            target: The target realm or context.
+        """
+        params = {
+            "handles": handles,
+            "target": target,
+        }
+        self.conn.execute(command_builder("script.disown", params))
+
+    def call_function(
+        self,
+        function_declaration: str,
+        await_promise: bool,
+        target: dict,
+        arguments: Optional[List[dict]] = None,
+        result_ownership: Optional[str] = None,
+        serialization_options: Optional[dict] = None,
+        this: Optional[dict] = None,
+        user_activation: bool = False,
+    ) -> EvaluateResult:
+        """Calls a provided function with given arguments in a given realm.
+
+        Parameters:
+        -----------
+            function_declaration: The function declaration to call.
+            await_promise: Whether to await promise resolution.
+            target: The target realm or context.
+            arguments: The arguments to pass to the function.
+            result_ownership: The result ownership type.
+            serialization_options: The serialization options.
+            this: The 'this' value for the function call.
+            user_activation: Whether to trigger user activation.
+
+        Returns:
+        -------
+            EvaluateResult: The result of the function call.
+        """
+        params = {
+            "functionDeclaration": function_declaration,
+            "awaitPromise": await_promise,
+            "target": target,
+            "userActivation": user_activation,
+        }
+
+        if arguments is not None:
+            params["arguments"] = arguments
+        if result_ownership is not None:
+            params["resultOwnership"] = result_ownership
+        if serialization_options is not None:
+            params["serializationOptions"] = serialization_options
+        if this is not None:
+            params["this"] = this
+
+        result = self.conn.execute(command_builder("script.callFunction", params))
+        return EvaluateResult.from_json(result)
+
+    def evaluate(
+        self,
+        expression: str,
+        target: dict,
+        await_promise: bool,
+        result_ownership: Optional[str] = None,
+        serialization_options: Optional[dict] = None,
+        user_activation: bool = False,
+    ) -> EvaluateResult:
+        """Evaluates a provided script in a given realm.
+
+        Parameters:
+        -----------
+            expression: The script expression to evaluate.
+            target: The target realm or context.
+            await_promise: Whether to await promise resolution.
+            result_ownership: The result ownership type.
+            serialization_options: The serialization options.
+            user_activation: Whether to trigger user activation.
+
+        Returns:
+        -------
+            EvaluateResult: The result of the script evaluation.
+        """
+        params = {
+            "expression": expression,
+            "target": target,
+            "awaitPromise": await_promise,
+            "userActivation": user_activation,
+        }
+
+        if result_ownership is not None:
+            params["resultOwnership"] = result_ownership
+        if serialization_options is not None:
+            params["serializationOptions"] = serialization_options
+
+        result = self.conn.execute(command_builder("script.evaluate", params))
+        return EvaluateResult.from_json(result)
+
+    def get_realms(
+        self,
+        context: Optional[str] = None,
+        type: Optional[str] = None,
+    ) -> List[RealmInfo]:
+        """Returns a list of all realms, optionally filtered.
+
+        Parameters:
+        -----------
+            context: The browsing context ID to filter by.
+            type: The realm type to filter by.
+
+        Returns:
+        -------
+            List[RealmInfo]: A list of realm information.
+        """
+        params = {}
+
+        if context is not None:
+            params["context"] = context
+        if type is not None:
+            params["type"] = type
+
+        result = self.conn.execute(command_builder("script.getRealms", params))
+        return [RealmInfo.from_json(realm) for realm in result["realms"]]
+
     def _subscribe_to_log_entries(self):
         if not self.log_entry_subscribed:
             session = Session(self.conn)