import pygments
import pygments.lexers
import pygments.formatters


def ansi_highlight_code(python_code, ansi_colors):
    """Adds ANSI color code to make python_code pretty.

    Args:
      python_code (str): Python 3 code to prettify.
      ansi_colors (int): Number of colors to use, as returned by
          count_ansi_colors().

    Returns:
      str: The prettified Python code.
    """

    if ansi_colors >= 256:
        formatter = pygments.formatters.Terminal256Formatter(bg='dark')
    elif ansi_colors >= 8:
        formatter = pygments.formatters.TerminalFormatter(bg='dark')
    else:
        return python_code

    lexer = pygments.lexers.get_lexer_by_name('python3')
    return pygments.highlight(python_code, lexer, formatter)