class Selenium::WebDriver::TargetLocator

def window(id)
  if block_given?
    original = begin
      @bridge.window_handle
    rescue Error::NoSuchWindowError
      nil
    end
    unless @bridge.window_handles.include? id
      raise Error::NoSuchWindowError, "The specified identifier '#{id}' is not found in the window handle list"
    end
    @bridge.switch_to_window id
    begin
      returned = yield
    ensure
      current_handles = @bridge.window_handles
      original = current_handles.first unless current_handles.include? original
      @bridge.switch_to_window original
      returned
    end
  else
    @bridge.switch_to_window id
  end
end