[py] Add type annotations to BiDi module

[mzyndul]

User description

🔗 Related Issues

Improves type coverage for the BiDi browser module to pass mypy --strict.

💥 What does this PR do?

Adds type annotations to selenium/webdriver/common/bidi/

Also adds type narrowing assertions in webdriver.py for properties that pass WebSocketConnection to now-typed BiDi classes.

🔧 Implementation Notes

TYPE_CHECKING pattern: WebSocketConnection is imported under TYPE_CHECKING to avoid runtime circular imports while still providing type hints. This matches the pattern used in sibling modules (session.py, script.py, input.py).

cast() usage: The conn.execute() method returns Any (untyped upstream). Using cast(str, result["userContext"]) documents the expected type per BiDi spec and satisfies mypy --strict.

💡 Additional Considerations

Decision needed: Should the assertions be moved outside the if block? Example:

Current (inside if):

if not self._websocket_connection:
    self._start_bidi()
    assert self._websocket_connection is not None

Alternative:

if not self._websocket_connection:
    self._start_bidi()

assert self._websocket_connection is not None

Both pass mypy and tox checks. The "outside" pattern is slightly more defensive but adds an assertion on every property access, so I decided to go with the first one.

🔄 Types of changes

• Cleanup (formatting, renaming)

---

PR Type

Enhancement

---

Description

• Add comprehensive type annotations to BiDi module for mypy --strict compliance

• Import WebSocketConnection under TYPE_CHECKING to avoid circular imports

• Use cast() for untyped upstream results to document expected types per BiDi spec

• Add type narrowing assertions in webdriver.py properties for WebSocketConnection

• Improve type hints for dataclass methods, callbacks, and dictionary parameters

---

Diagram Walkthrough

flowchart LR
  A["BiDi Modules<br/>browser.py, input.py,<br/>script.py, session.py"] -->|"Add type hints"| B["Typed Parameters<br/>& Return Values"]
  C["WebSocketConnection"] -->|"TYPE_CHECKING import"| B
  B -->|"cast() for untyped<br/>upstream results"| D["mypy --strict<br/>Compliance"]
  E["webdriver.py"] -->|"Add assertions"| F["Type Narrowing<br/>for properties"]
  F -->|"Ensures type safety"| D

Loading

File Walkthrough

Relevant files

Type annotations

browser.py Add type annotations to Browser class                                        py/selenium/webdriver/common/bidi/browser.py Add from __future__ import annotations and TYPE_CHECKING import for WebSocketConnection Type annotate __init__ method with WebSocketConnection parameter and None return type Update from_dict classmethod to use dict[str, Any] instead of generic dict Use cast(str, result["userContext"]) for untyped upstream result Improve type hint for destination_folder parameter to str | os.PathLike[str] +10/-5    input.py Add comprehensive type annotations to Input module              py/selenium/webdriver/common/bidi/input.py Add from __future__ import annotations and TYPE_CHECKING import for WebSocketConnection Import Callable from collections.abc for callback type hints Type annotate all dataclass __init__ and to_dict methods with return types Update dictionary parameters to use dict[str, Any] throughout Type annotate Input class __init__, callbacks, and handler methods with proper Callable signatures Add type hints for subscriptions and callbacks dictionaries +34/-28  script.py Add type annotations to Script module                                        py/selenium/webdriver/common/bidi/script.py Add from __future__ import annotations and TYPE_CHECKING imports for WebDriver and WebSocketConnection Import Callable from collections.abc for handler type hints Type annotate Script class __init__ with WebDriver optional parameter Add type hints for handler methods with Callable signatures for console and error handlers Use cast(str, result["script"]) for untyped upstream result in preload script Type annotate all internal methods with proper dict and list type parameters Update classmethod return types to remove string quotes +40/-30  session.py Add type annotations to Session module                                      py/selenium/webdriver/common/bidi/session.py Add from __future__ import annotations and TYPE_CHECKING import for WebSocketConnection Import Generator from collections.abc for subscribe/unsubscribe return types Type annotate Session __init__ with WebSocketConnection parameter Add Generator return type hints for subscribe and unsubscribe methods Type annotate status method to return dict[str, Any] Add type hints for method parameters including variadic events and optional browsing_contexts +18/-6    webdriver.py Add type narrowing assertions for BiDi properties                py/selenium/webdriver/remote/webdriver.py Add type narrowing assertions after _start_bidi() calls in script, browser, _session, and input properties Assertions ensure WebSocketConnection is not None before passing to typed BiDi classes Maintains defensive programming pattern while satisfying mypy type checker +4/-0

---

[qodo-code-review]

PR Compliance Guide 🔍

Below is a summary of compliance checks for this PR:

Security Compliance
🟢	No security concerns identified No security vulnerabilities detected by AI analysis. Human verification advised for critical code.
Ticket Compliance
⚪	🎫 No ticket provided Create ticket/issue
Codebase Duplication Compliance
⚪	Codebase context is not defined Follow the guide to enable codebase context checks.
Custom Compliance
🟢	Generic: Comprehensive Audit Trails Objective: To create a detailed and reliable record of critical system actions for security analysis and compliance. Status: Passed Learn more about managing compliance generic rules or creating your own custom rules
	Generic: Meaningful Naming and Self-Documenting Code Objective: Ensure all identifiers clearly express their purpose and intent, making code self-documenting Status: Passed Learn more about managing compliance generic rules or creating your own custom rules
	Generic: Secure Error Handling Objective: To prevent the leakage of sensitive system information through error messages while providing sufficient detail for internal debugging. Status: Passed Learn more about managing compliance generic rules or creating your own custom rules
	Generic: Secure Logging Practices Objective: To ensure logs are useful for debugging and auditing without exposing sensitive information like PII, PHI, or cardholder data. Status: Passed Learn more about managing compliance generic rules or creating your own custom rules
	Generic: Security-First Input Validation and Data Handling Objective: Ensure all data inputs are validated, sanitized, and handled securely to prevent vulnerabilities Status: Passed Learn more about managing compliance generic rules or creating your own custom rules
⚪	Generic: Robust Error Handling and Edge Case Management Objective: Ensure comprehensive error handling that provides meaningful context and graceful degradation Status: Assertion-based checks: The new assert self._websocket_connection is not None statements can fail without actionable context and are skipped under Python -O, potentially passing None into BiDi constructors without a controlled error path. Referred Code def script(self) -> Script: if not self._websocket_connection: self._start_bidi() assert self._websocket_connection is not None if not self._script: self._script = Script(self._websocket_connection, self) return self._script def _start_bidi(self) -> None: if self.caps.get("webSocketUrl"): ws_url = self.caps.get("webSocketUrl") else: raise WebDriverException("Unable to find url to connect to from capabilities") if not isinstance(self.command_executor, RemoteConnection): raise WebDriverException("command_executor must be a RemoteConnection instance for BiDi support") self._websocket_connection = WebSocketConnection( ws_url, ... (clipped 183 lines) Learn more about managing compliance generic rules or creating your own custom rules
Update

Compliance status legend

🟢 - Fully Compliant
🟡 - Partial Compliant
🔴 - Not Compliant
⚪ - Requires Further Human Verification
🏷️ - Compliance label

[qodo-code-review]

PR Code Suggestions ✨

Explore these optional code suggestions:

Category	Suggestion	Impact
Learned best practice	Validate remote response fields Validate that the response contains userContext and that it is a str before returning; avoid using cast as a substitute for runtime checks on remote data. py/selenium/webdriver/common/bidi/browser.py [208-209] result = self.conn.execute(command_builder("browser.createUserContext", params)) -return cast(str, result["userContext"]) +user_context = result.get("userContext") +if not isinstance(user_context, str): + raise ValueError("Invalid response: missing/invalid 'userContext'") +return user_context Apply / Chat Suggestion importance[1-10]: 6 __ Why: Relevant best practice - Add explicit validation/guards at integration boundaries (e.g., remote BiDi responses) instead of blindly trusting untyped fields.	Low
Update

[cgoldberg]

I don't understand the assertions.

For example:

if not self._websocket_connection:
    self._start_bidi()
    assert self._websocket_connection is not None

Since _start_bidi() sets the WebSocket connection, there is no way this assertion could ever fail. Is mypy complaining about that?

In general, we should almost never use assert in code that is not in the test suite. It just raises an unhelpful exception. If something has an invalid type or there is a situation where something can't run, we should be raising TypeError or some other exception along with a useful message. I don't like the idea of just sprinkling around assertions to keep the type checker happy.

[mzyndul]

Yes, mypy complains because it doesn't understand that _start_bidi() guarantees the connection is set:

  selenium/webdriver/remote/webdriver.py:1161: error: Argument 1 to "Session" has incompatible type "WebSocketConnection | None"; expected "WebSocketConnection"  [arg-type]

I agree that assert isn't ideal here. I could use cast() instead:

  self._bidi_session = Session(cast(WebSocketConnection, self._websocket_connection))

Alternatively, _start_bidi() could return the connection, allowing the type to flow naturally.

[mzyndul]

@cgoldberg I replaced asserts with "cast" in py/selenium/webdriver/remote/webdriver.py