# Copyright 2011 Google Inc. All Rights Reserved.
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
#
# Licensed to the PSF under a Contributor Agreement.
#
# Author: Gregory P. Smith

"""Implements the select.select() API using select.poll()."""

import itertools
import select


def Select(rlist, wlist, xlist, timeout=None):
  """A select.select() workalike implemented internally using select.poll()."""
  poll = select.poll()
  fd_object_map = {}
  for file_obj in itertools.chain(rlist, wlist, xlist):
    if not isinstance(file_obj, (int, long)):
      fd_object_map[file_obj] = file_obj.fileno()
  poll_read_flag = select.POLLIN
  poll_write_flag = select.POLLOUT
  poll_oob_flag = select.POLLPRI
  fd_flag_map = {}
  for file_obj in rlist:
    _SetFlagsInMap(file_obj, fd_flag_map, poll_read_flag)
  for file_obj in wlist:
    _SetFlagsInMap(file_obj, fd_flag_map, poll_write_flag)
  for file_obj in xlist:
    _SetFlagsInMap(file_obj, fd_flag_map, poll_oob_flag)

  # We register everything via the underlying fd for simplicity instead of
  # trying to maintain the original file object vs fd here.
  for fd, eventmask in fd_flag_map.iteritems():
    poll.register(fd, eventmask)

  events = poll.poll(timeout)

  r_ready = []
  w_ready = []
  x_ready = []
  for fd, event in events:
    if event & poll_read_flag:
      r_ready.append(fd)
    if event & poll_write_flag:
      w_ready.append(fd)
    if event & poll_oob_flag:
      x_ready.append(fd)

  # Return the original file objects for those that were specified using
  # a file object instead of an fd to start with.
  r_ready = _ConvertFdListToObjects(r_ready, fd_flag_map)
  w_ready = _ConvertFdListToObjects(w_ready, fd_flag_map)
  x_ready = _ConvertFdListToObjects(x_ready, fd_flag_map)

  return (r_ready, w_ready, x_ready)


def _SetFlagsInMap(file_obj, fd_flag_map, flags):
  """Make sure file_obj is in fd_flag_map with flags set.

  If file_obj is an object instead of a file descriptor, its fd will be
  inserted into the map using its .fileno() method.  flags are bitwise
  ORed into the existing flags for the fd in the map if any.

  Args:
    file_obj: A file object bearing a fileno() method or a file descriptor.
    fd_flag_map: A map of file descriptors to current flags for that fd.
    flags: Flags to bitwise OR into any existing flags for file_obj.
  """
  if isinstance(file_obj, (int, long)):
    fd = file_obj
  else:
    fd = file_obj.fileno()
  fd_flag_map[fd] = fd_flag_map.get(fd, 0) | flags


def _ConvertFdListToObjects(fd_list, fd_flag_map):
  """Return a list with fds in fd_list replaced with objects as appropriate.

  If a fd in fd_list appears in fd_flag_map the resulting list will contain
  the value from fd_flag_map instead of the fd; otherwise the fd.

  Args:
    fd_list: The original list to convert.
    fd_flag_map: A map of fd to file objects.

  Returns:
    A new list with fds replaced with map values when available.
  """
  result = []
  for fd in fd_list:
    if fd in fd_flag_map:
      result.append(fd_flag_map[fd])
    else:
      result.append(fd)
  return result