puppeteer-ruby 0.45.6 → 0.50.0.alpha5

data/lib/puppeteer/browser.rb

@@ -1,5 +1,6 @@
+# rbs_inline: enabled
+
 require 'thread'
-require 'timeout'
 
 class Puppeteer::Browser
   include Puppeteer::DebugPrint
@@ -7,18 +8,23 @@ class Puppeteer::Browser
   include Puppeteer::IfPresent
   using Puppeteer::DefineAsyncMethod
 
-  # @param product [String|nil] 'chrome' or 'firefox'
-  # @param {!Puppeteer.Connection} connection
-  # @param {!Array<string>} contextIds
-  # @param {boolean} ignoreHTTPSErrors
-  # @param {?Puppeteer.Viewport} defaultViewport
-  # @param process [Puppeteer::BrowserRunner::BrowserProcess|NilClass]
-  # @param {function()=} closeCallback
+  # @rbs product: String? -- Browser product (chrome only)
+  # @rbs connection: Puppeteer::Connection -- CDP connection
+  # @rbs context_ids: Array[String] -- Browser context IDs
+  # @rbs ignore_https_errors: bool -- Ignore HTTPS errors
+  # @rbs default_viewport: Puppeteer::Viewport? -- Default viewport
+  # @rbs network_enabled: bool -- Whether network events are enabled
+  # @rbs process: Puppeteer::BrowserRunner::BrowserProcess? -- Browser process handle
+  # @rbs close_callback: Proc -- Close callback
+  # @rbs target_filter_callback: Proc? -- Target filter callback
+  # @rbs is_page_target_callback: Proc? -- Page target predicate
+  # @rbs return: Puppeteer::Browser -- Browser instance
   def self.create(product:,
                   connection:,
                   context_ids:,
                   ignore_https_errors:,
                   default_viewport:,
+                  network_enabled: true,
                   process:,
                   close_callback:,
                   target_filter_callback:,
@@ -29,6 +35,7 @@ class Puppeteer::Browser
       context_ids: context_ids,
       ignore_https_errors: ignore_https_errors,
       default_viewport: default_viewport,
+      network_enabled: network_enabled,
       process: process,
       close_callback: close_callback,
       target_filter_callback: target_filter_callback,
@@ -38,25 +45,34 @@ class Puppeteer::Browser
     browser
   end
 
-  # @param product [String|nil] 'chrome' or 'firefox'
-  # @param {!Puppeteer.Connection} connection
-  # @param {!Array<string>} contextIds
-  # @param {boolean} ignoreHTTPSErrors
-  # @param {?Puppeteer.Viewport} defaultViewport
-  # @param {?Puppeteer.ChildProcess} process
-  # @param {(function():Promise)=} closeCallback
+  # @rbs product: String? -- Browser product (chrome only)
+  # @rbs connection: Puppeteer::Connection -- CDP connection
+  # @rbs context_ids: Array[String] -- Browser context IDs
+  # @rbs ignore_https_errors: bool -- Ignore HTTPS errors
+  # @rbs default_viewport: Puppeteer::Viewport? -- Default viewport
+  # @rbs network_enabled: bool -- Whether network events are enabled
+  # @rbs process: Puppeteer::BrowserRunner::BrowserProcess? -- Browser process handle
+  # @rbs close_callback: Proc -- Close callback
+  # @rbs target_filter_callback: Proc? -- Target filter callback
+  # @rbs is_page_target_callback: Proc? -- Page target predicate
+  # @rbs return: void -- No return value
   def initialize(product:,
                  connection:,
                  context_ids:,
                  ignore_https_errors:,
                  default_viewport:,
+                 network_enabled: true,
                  process:,
                  close_callback:,
                  target_filter_callback:,
                  is_page_target_callback:)
-    @product = product || 'chrome'
+    @product = product ? product.to_s : 'chrome'
+    if @product != 'chrome'
+      raise ArgumentError.new("Unsupported product: #{@product}. Only 'chrome' is supported.")
+    end
     @ignore_https_errors = ignore_https_errors
     @default_viewport = default_viewport
+    @network_enabled = network_enabled
     @process = process
     @connection = connection
     @close_callback = close_callback
@@ -69,19 +85,11 @@ class Puppeteer::Browser
       @contexts[context_id] = Puppeteer::BrowserContext.new(@connection, self, context_id)
     end
 
-    if @product == 'firefox'
-      @target_manager = Puppeteer::FirefoxTargetManager.new(
-        connection: connection,
-        target_factory: method(:create_target),
-        target_filter_callback: @target_filter_callback,
-      )
-    else
-      @target_manager = Puppeteer::ChromeTargetManager.new(
-        connection: connection,
-        target_factory: method(:create_target),
-        target_filter_callback: @target_filter_callback,
-      )
-    end
+    @target_manager = Puppeteer::ChromeTargetManager.new(
+      connection: connection,
+      target_factory: method(:create_target),
+      target_filter_callback: @target_filter_callback,
+    )
   end
 
   private def default_target_filter_callback(target_info)
@@ -94,7 +102,9 @@ class Puppeteer::Browser
 
   attr_reader :is_page_target_callback
 
-  # @param event_name [Symbol] either of :disconnected, :targetcreated, :targetchanged, :targetdestroyed
+  # @rbs event_name: (String | Symbol) -- Browser event name
+  # @rbs &block: ^(untyped) -> void -- Event handler
+  # @rbs return: String -- Listener ID
   def on(event_name, &block)
     unless BrowserEmittedEvents.values.include?(event_name.to_s)
       raise ArgumentError.new("Unknown event name: #{event_name}. Known events are #{BrowserEmittedEvents.values.to_a.join(", ")}")
@@ -103,7 +113,9 @@ class Puppeteer::Browser
     super(event_name.to_s, &block)
   end
 
-  # @param event_name [Symbol]
+  # @rbs event_name: (String | Symbol) -- Browser event name
+  # @rbs &block: ^(untyped) -> void -- Event handler
+  # @rbs return: String -- Listener ID
   def once(event_name, &block)
     unless BrowserEmittedEvents.values.include?(event_name.to_s)
       raise ArgumentError.new("Unknown event name: #{event_name}. Known events are #{BrowserEmittedEvents.values.to_a.join(", ")}")
@@ -142,7 +154,7 @@ class Puppeteer::Browser
     @target_manager.remove_event_listener(*@target_manager_event_listeners)
   end
 
-  # @return [Puppeteer::BrowserRunner::BrowserProcess]
+  # @rbs return: Puppeteer::BrowserRunner::BrowserProcess? -- Browser process handle
   def process
     @process
   end
@@ -151,33 +163,52 @@ class Puppeteer::Browser
     @target_manager
   end
 
-  # @return [Puppeteer::BrowserContext]
+  # @rbs return: Puppeteer::BrowserContext -- New incognito browser context
   def create_incognito_browser_context
     result = @connection.send_message('Target.createBrowserContext')
     browser_context_id = result['browserContextId']
     @contexts[browser_context_id] = Puppeteer::BrowserContext.new(@connection, self, browser_context_id)
   end
 
+  # @rbs proxy_server: String? -- Proxy server for new context
+  # @rbs proxy_bypass_list: Array[String]? -- Proxy bypass list
+  # @rbs download_behavior: Hash[Symbol | String, untyped]? -- Download behavior options
+  # @rbs return: Puppeteer::BrowserContext -- New browser context
+  def create_browser_context(proxy_server: nil, proxy_bypass_list: nil, download_behavior: nil)
+    params = {
+      proxyServer: proxy_server,
+      proxyBypassList: proxy_bypass_list&.join(','),
+    }.compact
+    result = @connection.send_message('Target.createBrowserContext', params)
+    browser_context_id = result['browserContextId']
+    context = Puppeteer::BrowserContext.new(@connection, self, browser_context_id)
+    context.set_download_behavior(download_behavior) if download_behavior
+    @contexts[browser_context_id] = context
+  end
+
+  # @rbs return: Array[Puppeteer::BrowserContext] -- All browser contexts
   def browser_contexts
     [@default_context].concat(@contexts.values)
   end
 
-  # @return [Puppeteer::BrowserContext]
+  # @rbs return: Puppeteer::BrowserContext -- Default browser context
   def default_browser_context
     @default_context
   end
 
-  # @param context_id [String]
+  # @rbs context_id: String? -- Browser context ID
+  # @rbs return: void -- No return value
   def dispose_context(context_id)
     return unless context_id
     @connection.send_message('Target.disposeBrowserContext', browserContextId: context_id)
     @contexts.delete(context_id)
   end
 
-  class MissingBrowserContextError < StandardError ; end
+  class MissingBrowserContextError < Puppeteer::Error ; end
 
-  # @param target_info [Puppeteer::Target::TargetInfo]
-  # @param session [CDPSession|nil]
+  # @rbs target_info: Puppeteer::Target::TargetInfo -- Target info
+  # @rbs session: Puppeteer::CDPSession? -- Attached session
+  # @rbs return: Puppeteer::Target -- Created target
   def create_target(target_info, session)
     browser_context_id = target_info.browser_context_id
     context =
@@ -199,12 +230,13 @@ class Puppeteer::Browser
       session_factory: -> (auto_attach_emulated) { @connection.create_session(target_info, auto_attach_emulated: auto_attach_emulated) },
       ignore_https_errors: @ignore_https_errors,
       default_viewport: @default_viewport,
+      network_enabled: @network_enabled,
       is_page_target_callback: @is_page_target_callback,
     )
   end
 
   private def handle_attached_to_target(target)
-    if target.initialized_promise.value!
+    if target.initialized_promise.wait
       emit_event(BrowserEmittedEvents::TargetCreated, target)
       target.browser_context.emit_event(BrowserContextEmittedEvents::TargetCreated, target)
     end
@@ -213,7 +245,7 @@ class Puppeteer::Browser
   private def handle_detached_from_target(target)
     target.ignore_initialize_callback_promise
     target.closed_callback
-    if target.initialized_promise.value!
+    if target.initialized_promise.wait
       emit_event(BrowserEmittedEvents::TargetDestroyed, target)
       target.browser_context.emit_event(BrowserContextEmittedEvents::TargetDestroyed, target)
     end
@@ -233,20 +265,21 @@ class Puppeteer::Browser
     emit_event('targetdiscovered', target_info)
   end
 
-  # @return [String]
+  # @rbs return: String -- WebSocket endpoint URL
   def ws_endpoint
     @connection.url
   end
 
+  # @rbs return: Puppeteer::Page -- New page in default context
   def new_page
     @default_context.new_page
   end
 
-  class MissingTargetError < StandardError ; end
-  class CreatePageError < StandardError ; end
+  class MissingTargetError < Puppeteer::Error ; end
+  class CreatePageError < Puppeteer::Error ; end
 
-  # @param {?string} contextId
-  # @return {!Promise<!Puppeteer.Page>}
+  # @rbs context_id: String? -- Browser context ID
+  # @rbs return: Puppeteer::Page -- Created page
   def create_page_in_context(context_id)
     create_target_params = {
       url: 'about:blank',
@@ -258,7 +291,7 @@ class Puppeteer::Browser
     unless target
       raise MissingTargetError.new("Missing target for page (id = #{target_id})")
     end
-    unless target.initialized_promise.value!
+    unless target.initialized_promise.wait
       raise CreatePageError.new("Failed to create target for page (id = #{target_id})")
     end
     page = target.page
@@ -270,12 +303,14 @@ class Puppeteer::Browser
 
   # All active targets inside the Browser. In case of multiple browser contexts, returns
   # an array with all the targets in all browser contexts.
+  # @rbs return: Array[Puppeteer::Target] -- Active targets
   def targets
     @target_manager.available_targets.values.select { |target| target.initialized? }
   end
 
 
   # The target associated with the browser.
+  # @rbs return: Puppeteer::Target -- Browser target
   def target
     targets.find { |target| target.type == 'browser' } or raise 'Browser target is not found'
   end
@@ -285,29 +320,30 @@ class Puppeteer::Browser
     @target_manager.available_targets[target_id]
   end
 
-  # @param predicate [Proc(Puppeteer::Target -> Boolean)]
-  # @return [Puppeteer::Target]
+  # @rbs predicate: Proc -- Predicate for target matching
+  # @rbs timeout: Numeric? -- Timeout in milliseconds
+  # @rbs return: Puppeteer::Target -- Matching target
   def wait_for_target(predicate:, timeout: nil)
     timeout_helper = Puppeteer::TimeoutHelper.new('target', timeout_ms: timeout, default_timeout_ms: 30000)
     existing_target = targets.find { |target| predicate.call(target) }
     return existing_target if existing_target
 
     event_listening_ids = []
-    target_promise = resolvable_future
+    target_promise = Async::Promise.new
     event_listening_ids << add_event_listener(BrowserEmittedEvents::TargetCreated) do |target|
       if predicate.call(target)
-        target_promise.fulfill(target)
+        target_promise.resolve(target)
       end
     end
     event_listening_ids << add_event_listener(BrowserEmittedEvents::TargetChanged) do |target|
       if predicate.call(target)
-        target_promise.fulfill(target)
+        target_promise.resolve(target)
       end
     end
 
     begin
       timeout_helper.with_timeout do
-        target_promise.value!
+        target_promise.wait
       end
     ensure
       remove_event_listener(*event_listening_ids)
@@ -316,43 +352,49 @@ class Puppeteer::Browser
 
   # @!method async_wait_for_target(predicate:, timeout: nil)
   #
-  # @param predicate [Proc(Puppeteer::Target -> Boolean)]
   define_async_method :async_wait_for_target
 
-  # @return {!Promise<!Array<!Puppeteer.Page>>}
+  # @rbs return: Array[Puppeteer::Page] -- All pages across contexts
   def pages
     browser_contexts.flat_map(&:pages)
   end
 
-  # @return [String]
+  # @rbs return: String -- Browser version string
   def version
     Version.fetch(@connection).product
   end
 
-  # @return [String]
+  # @rbs return: String -- Browser user agent string
   def user_agent
     Version.fetch(@connection).user_agent
   end
 
+  # @rbs return: void -- No return value
   def close
     @close_callback.call
     disconnect
   end
 
+  # @rbs return: void -- No return value
   def disconnect
     @target_manager.dispose
     @connection.dispose
   end
 
+  # @rbs return: bool -- Whether the browser is connected
   def connected?
     !@connection.closed?
   end
 
   class Version
+    # @rbs connection: Puppeteer::Connection -- CDP connection
+    # @rbs return: Puppeteer::Browser::Version -- Browser version info
     def self.fetch(connection)
       new(connection.send_message('Browser.getVersion'))
     end
 
+    # @rbs hash: Hash[String, String] -- Version payload
+    # @rbs return: void -- No return value
     def initialize(hash)
       @protocol_version = hash['protocolVersion']
       @product = hash['product']

data/lib/puppeteer/browser_connector.rb

@@ -12,7 +12,11 @@ class Puppeteer::BrowserConnector
   # @return [Puppeteer::Browser]
   def connect_to_browser
     version = Puppeteer::Browser::Version.fetch(connection)
-    product = version.product.downcase.include?('firefox') ? 'firefox' : 'chrome'
+    product_name = version.product.to_s.downcase
+    if product_name.include?('firefox')
+      raise Puppeteer::Error.new('Firefox CDP support has been removed. Use puppeteer-bidi for Firefox automation.')
+    end
+    product = 'chrome'
 
     result = connection.send_message('Target.getBrowserContexts')
     browser_context_ids = result['browserContextIds']
@@ -23,6 +27,7 @@ class Puppeteer::BrowserConnector
       context_ids: browser_context_ids,
       ignore_https_errors: @browser_options.ignore_https_errors?,
       default_viewport: @browser_options.default_viewport,
+      network_enabled: @browser_options.network_enabled,
       process: nil,
       close_callback: -> { connection.send_message('Browser.close') },
       target_filter_callback: @browser_options.target_filter,
@@ -46,7 +51,12 @@ class Puppeteer::BrowserConnector
   # @return [Puppeteer::Connection]
   private def connect_with_browser_ws_endpoint(browser_ws_endpoint)
     transport = Puppeteer::WebSocketTransport.create(browser_ws_endpoint)
-    Puppeteer::Connection.new(browser_ws_endpoint, transport, @browser_options.slow_mo)
+    Puppeteer::Connection.new(
+      browser_ws_endpoint,
+      transport,
+      @browser_options.slow_mo,
+      protocol_timeout: @browser_options.protocol_timeout,
+    )
   end
 
   # @return [Puppeteer::Connection]
@@ -62,6 +72,11 @@ class Puppeteer::BrowserConnector
 
   # @return [Puppeteer::Connection]
   private def connect_with_transport(transport)
-    Puppeteer::Connection.new('', transport, @browser_options.slow_mo)
+    Puppeteer::Connection.new(
+      '',
+      transport,
+      @browser_options.slow_mo,
+      protocol_timeout: @browser_options.protocol_timeout,
+    )
   end
 end

data/lib/puppeteer/browser_context.rb

@@ -1,3 +1,6 @@
+require 'async/semaphore'
+require 'uri'
+
 class Puppeteer::BrowserContext
   include Puppeteer::EventCallbackable
   using Puppeteer::DefineAsyncMethod
@@ -9,10 +12,39 @@ class Puppeteer::BrowserContext
     @connection = connection
     @browser = browser
     @id = context_id
+    @closed = false
+    @screenshot_semaphore = nil
+    @screenshot_operations_count = 0
   end
 
   attr_reader :id
 
+  class ScreenshotGuard
+    def initialize(semaphore, on_release: nil)
+      @semaphore = semaphore
+      @on_release = on_release
+      @released = false
+    end
+
+    def release
+      return if @released
+
+      @released = true
+      @semaphore.release
+      @on_release&.call
+    end
+    alias_method :close, :release
+  end
+
+  def ==(other)
+    other = other.__getobj__ if other.is_a?(Puppeteer::ReactorRunner::Proxy)
+    return true if equal?(other)
+    return false unless other.is_a?(Puppeteer::BrowserContext)
+    return false if @id.nil? || other.id.nil?
+
+    @id == other.id
+  end
+
   # @param event_name [Symbol] either of :disconnected, :targetcreated, :targetchanged, :targetdestroyed
   def on(event_name, &block)
     unless BrowserContextEmittedEvents.values.include?(event_name.to_s)
@@ -36,6 +68,25 @@ class Puppeteer::BrowserContext
     @browser.targets.select { |target| target.browser_context == self }
   end
 
+  def start_screenshot
+    semaphore = @screenshot_semaphore || Async::Semaphore.new(1)
+    @screenshot_semaphore = semaphore
+    @screenshot_operations_count += 1
+    semaphore.acquire
+    ScreenshotGuard.new(semaphore, on_release: lambda {
+      @screenshot_operations_count -= 1
+      @screenshot_semaphore = nil if @screenshot_operations_count.zero?
+    })
+  end
+
+  def wait_for_screenshot_operations
+    semaphore = @screenshot_semaphore
+    return nil unless semaphore
+
+    semaphore.acquire
+    ScreenshotGuard.new(semaphore)
+  end
+
   # @param predicate [Proc(Puppeteer::Target -> Boolean)]
   # @return [Puppeteer::Target]
   def wait_for_target(predicate:, timeout: nil)
@@ -51,16 +102,21 @@ class Puppeteer::BrowserContext
   define_async_method :async_wait_for_target
 
   # @return {!Promise<!Array<!Puppeteer.Page>>}
-  def pages
+  def pages(include_all: false)
     targets.select { |target|
-      target.type == 'page' || (target.type == 'other' && @browser.is_page_target_callback&.call(target.target_info))
-    }.map(&:page).reject { |page| !page }
+      target.type == 'page' ||
+        ((target.type == 'other' || include_all) && @browser.is_page_target_callback&.call(target.target_info))
+    }.map(&:page).compact
   end
 
   def incognito?
     !!@id
   end
 
+  def closed?
+    @closed || !@browser.browser_contexts.include?(self)
+  end
+
   WEB_PERMISSION_TO_PROTOCOL = {
     'geolocation' => 'geolocation',
     'midi' => 'midi',
@@ -105,9 +161,23 @@ class Puppeteer::BrowserContext
     end
   end
 
+  # @param download_behavior [Hash]
+  def set_download_behavior(download_behavior)
+    behavior = hash_value(download_behavior, 'policy')
+    download_path = hash_value(download_behavior, 'downloadPath', 'download_path')
+    @connection.send_message('Browser.setDownloadBehavior', {
+      behavior: behavior,
+      downloadPath: download_path,
+      browserContextId: @id,
+    }.compact)
+  end
+
   # @return [Future<Puppeteer::Page>]
   def new_page
+    guard = wait_for_screenshot_operations
     @browser.create_page_in_context(@id)
+  ensure
+    guard&.release
   end
 
   # @return [Browser]
@@ -120,5 +190,128 @@ class Puppeteer::BrowserContext
       raise 'Non-incognito profiles cannot be closed!'
     end
     @browser.dispose_context(@id)
+    @closed = true
+  end
+
+  # @return [Array<Hash>]
+  def cookies
+    params = { browserContextId: @id }.compact
+    response = @connection.send_message('Storage.getCookies', params)
+    response.fetch('cookies', []).map do |cookie|
+      normalized = cookie.dup
+      partition_key = cookie['partitionKey']
+      if partition_key
+        normalized['partitionKey'] = convert_partition_key_from_cdp(partition_key)
+      end
+      normalized['sameParty'] = cookie['sameParty'] || false
+      normalized
+    end
+  end
+
+  # @param cookies [Array<Hash>]
+  def set_cookie(*cookies)
+    items = cookies.map do |cookie|
+      normalized = normalize_cookie_hash(cookie)
+      partition_key = normalized.delete('partitionKey') || normalized.delete('partition_key')
+      normalized['partitionKey'] = convert_partition_key_for_cdp(partition_key) if partition_key
+      normalized
+    end
+    @connection.send_message('Storage.setCookies', {
+      browserContextId: @id,
+      cookies: items,
+    }.compact)
+  end
+
+  # @param cookies [Array<Hash>]
+  def delete_cookie(*cookies)
+    items = cookies.map do |cookie|
+      normalized = normalize_cookie_hash(cookie)
+      normalized['expires'] = 1
+      normalized
+    end
+    set_cookie(*items)
+  end
+
+  # @param filters [Array<Hash>]
+  def delete_matching_cookies(*filters)
+    cookies_to_delete = cookies.select do |cookie|
+      filters.any? do |filter|
+        filter_name = hash_value(filter, 'name')
+        next false unless filter_name == cookie['name']
+
+        filter_domain = hash_value(filter, 'domain')
+        next true if filter_domain && filter_domain == cookie['domain']
+
+        filter_path = hash_value(filter, 'path')
+        next true if filter_path && filter_path == cookie['path']
+
+        filter_partition_key = hash_value(filter, 'partitionKey', 'partition_key')
+        if filter_partition_key && cookie['partitionKey']
+          if cookie['partitionKey'].is_a?(String)
+            raise Puppeteer::Error.new('Unexpected string partition key')
+          end
+
+          cookie_partition_source_origin = hash_value(cookie['partitionKey'], 'sourceOrigin', 'source_origin')
+          filter_partition_source_origin =
+            if filter_partition_key.is_a?(String)
+              filter_partition_key
+            else
+              hash_value(filter_partition_key, 'sourceOrigin', 'source_origin')
+            end
+
+          next true if filter_partition_source_origin == cookie_partition_source_origin
+        end
+
+        filter_url = hash_value(filter, 'url')
+        if filter_url
+          url = URI.parse(filter_url)
+          url_path = url.path.to_s.empty? ? '/' : url.path
+          next true if url.hostname == cookie['domain'] && url_path == cookie['path']
+        end
+
+        true
+      end
+    end
+
+    delete_cookie(*cookies_to_delete)
+  end
+
+  private def normalize_cookie_hash(cookie)
+    cookie.each_with_object({}) do |(key, value), normalized|
+      normalized[key.to_s] = value
+    end
+  end
+
+  private def hash_value(hash, *keys)
+    return nil unless hash
+
+    keys.each do |key|
+      return hash[key] if hash.key?(key)
+      return hash[key.to_sym] if key.is_a?(String) && hash.key?(key.to_sym)
+      return hash[key.to_s] if key.is_a?(Symbol) && hash.key?(key.to_s)
+    end
+    nil
+  end
+
+  private def convert_partition_key_for_cdp(partition_key)
+    return nil if partition_key.nil?
+    return { topLevelSite: partition_key, hasCrossSiteAncestor: false } if partition_key.is_a?(String)
+
+    source_origin = hash_value(partition_key, 'sourceOrigin', 'source_origin')
+    has_cross_site_ancestor = hash_value(partition_key, 'hasCrossSiteAncestor', 'has_cross_site_ancestor')
+    {
+      topLevelSite: source_origin,
+      hasCrossSiteAncestor: has_cross_site_ancestor.nil? ? false : has_cross_site_ancestor,
+    }
+  end
+
+  private def convert_partition_key_from_cdp(partition_key)
+    return nil if partition_key.nil?
+    return partition_key if partition_key.is_a?(String)
+
+    {
+      'sourceOrigin' => hash_value(partition_key, 'topLevelSite', 'top_level_site'),
+      'hasCrossSiteAncestor' => hash_value(partition_key, 'hasCrossSiteAncestor', 'has_cross_site_ancestor'),
+    }.compact
   end
 end