mkadvice-0.2.1

#!/usr/bin/env python
# -*- coding: utf-8 -*-

from argparse import ArgumentParser
from collections import Counter
from cStringIO import StringIO
from Queue import Queue
from threading import Lock, Thread

from PIL import Image, ImageDraw, ImageFont
from progressbar import Bar, ProgressBar, Percentage, SimpleProgress
from requests import get as req_get

DEFAULT_OUT = 'output.jpg'
FONT = '/usr/share/fonts/ubuntu-font-family/Ubuntu-C.ttf'
WIDTH = 214  # Poster width
HIGHT = 317  # Poster hight

class Movie:
    def __init__(self, imdb_id, title, genre, poster=None):
        self.imdb_id = imdb_id
        self.title = title
        self.genre = genre
        self.poster = poster

    def dl_poster(self):
        payload = {'id': self.imdb_id, 'plot': 'none'}
        json = req_get('http://imdbapi.org/', params=payload).json()
        link = json.get('poster')
        if link:
            self.poster = req_get(link).content


def worker(mov_queue, pbar, pbar_lock):
    """Downloads posters."""
    while True:
        movie = mov_queue.get()
        movie.dl_poster()
        mov_queue.task_done()
        with pbar_lock:
            pbar.update(mov_queue.qsize())


def make_caption(title):
    """Returns image with caption."""
    fake_poster = Image.new('RGB', (WIDTH, HIGHT), 'lightgrey')
    caption = ImageDraw.Draw(fake_poster)
    caption_font = ImageFont.truetype(FONT, 24)
    caption.text((5, 10), title, fill='black', font=caption_font)
    return fake_poster


def main(input_file, output_file):
    # Read movies from file
    movies = [Movie(*(i.strip().capitalize() for i in line.split(',', 2)))
              for line in open(input_file, 'r') if not line.isspace()]

    # Get unique genres and maximum number of posters in column
    genres = Counter(mov.genre for mov in movies)
    unique_genres = genres.keys()
    unique_genres.sort()
    max_posters = max(genres.values())

    # Create progress bar, lock, queue and thread pool
    widgets = [Percentage(), Bar(left='[', right=']'), SimpleProgress()]
    pbar = ProgressBar(maxval=len(movies), widgets=widgets).start()
    pbar_lock = Lock()
    mov_queue = Queue(len(movies))
    threads_num = (len(movies) if len(movies) < 30 else 30)
    for i in xrange(threads_num):
        download = Thread(target=worker, args=(mov_queue, pbar, pbar_lock))
        download.daemon = True
        download.start()

    for mov in movies:
        mov_queue.put(mov)

    mov_queue.join()
    pbar.finish()

    # Create canvas, add genres and posters to canvas
    canv_width = 20 + (WIDTH+20) * len(unique_genres)
    canv_hight = 100 + HIGHT*max_posters + 20
    canvas = Image.new('RGB', (canv_width, canv_hight), 'black')
    draw = ImageDraw.Draw(canvas)
    genre_font = ImageFont.truetype(FONT, 42)
    shift = 0
    for genre in unique_genres:
        draw.text((20 + shift, 20), genre, fill='white', font=genre_font)
        shift += (WIDTH+20)

    p_counter = Counter()  # Current number of posters in each column (genre)
    y_counter = Counter()  # Current y value in each column
    x, y = (0, 0)
    for mov in movies:
        if mov.poster:
            poster = Image.open(StringIO(mov.poster))
            if poster.size != (WIDTH, HIGHT):
                poster = poster.resize((WIDTH, HIGHT))
        else:
                poster = make_caption(mov.title)

        # Set the (x, y) in appliance with genre
        x = 20 + (WIDTH+20) * unique_genres.index(mov.genre)
        y = 100 + y_counter[mov.genre]
        p_counter[mov.genre] += 1
        y_counter[mov.genre] = HIGHT * p_counter[mov.genre]

        canvas.paste(poster, (x, y))
    canvas.save(output_file)


if __name__ == '__main__':
    parser = ArgumentParser(description='Downloads posters and makes "advice"')
    parser.add_argument('input_file',
                        metavar='FILE',
                        help='Text file containing movies (id, title, genre)')
    parser.add_argument('-o',
                        dest='output_file',
                        metavar='OUTPUT',
                        default=DEFAULT_OUT,
                        help='Name of out image (default: '+ DEFAULT_OUT +')')
    args = parser.parse_args()
    main(args.input_file, args.output_file)