def Set(assignments, callback, *args):
  """Set the stoplight using assignments and a condition callback.

  Args:
    assignments: A dict mapping returns of the callback to the stoplight value.
    callback: The callback to compute the conditional assignment.
    *args: Input arguments to the callback functions.

  Example:
    Set({1: STOPLIGHT_NORMAL, -1: STOPLIGHT_WARNING, 0: STOPLIGHT_UNAVAILABLE},
        callback)

  Returns:
    Value of the stoplight.

  Raises:
    KeyError: Raised if assignment is invalid.
  """

  key = callback(*args)
  if key not in assignments:
    raise KeyError("Cannot determine stoplight due to invalid assignment.")
  return assignments[key]