~chipaca/ytd/trunk

"""

YouTubeDownloader

(for now! eventually, any-video-site-out-there downloader ;-p

"""

import cgi

from cStringIO import StringIO

import logging

import os

import re

import sys

import tempfile

import urllib

from htmlentitydefs import name2codepoint

import glib

import gtk

from async_downloader import AsyncDownloader

logging.basicConfig(level=logging.DEBUG)

# each extension has its own format, and so its own extension

EXTS = {'37': 'mp4',

        '22': 'mp4',

        '35': 'flv',

        '18': 'mp4',

        '34': 'flv',

        '17': '3gp',

        '0': 'flv',

        '5': 'flv',

        '6': 'flv',

        '13': '3gp',

        }

FMT_MAP = {'mobile': ['17', '13'],

           'low': ['34', '5', '0'],

           'medium': ['18', '6'],

           'high': ['35'],

           '720p': ['22'],

           '1080p': ['37'],

           }

def match2char(match):

    """

    Transform a entity (in the match's groups) to its unicode character

    """

    num, name = match.groups()

    if num is not None:

        num = int(num)

        try:

            return unichr(num)

        except (ValueError, OverflowError):

            return match.group()

    else:

        codepoint = name2codepoint.get(name)

        if codepoint is not None:

            return unichr(codepoint)

        else:

            return match.group()

def de_entify(string):

    """

    Strip entities from a string

    """

    return re.sub(r'&(?:#(\d+)|(\w+));', match2char, string)

# pylint: disable-msg=E1101

class YouTubeDownloader(object):

    """

    An app that lets you download a video from youtube

    """

    def __init__(self, url=None):

        builder = gtk.Builder()

        builder.set_translation_domain('YouTubeDownloader')

        builder.add_from_file('ytd.glade')

        builder.connect_signals(self)

        for obj in builder.get_objects():

            name = getattr(obj, 'name', None)

            if name is None and isinstance(obj, gtk.Buildable):

                # work around bug lp:507739

                name = gtk.Buildable.get_name(obj)

            if name is None:

                logging.warn("%s has no name (??)", obj)

            else:

                setattr(self, name.replace('-', '_'), obj)

        # work around a bug somewhere

        self.entry_filename.set_property('secondary-icon-sensitive', False)

        self.entry_url.set_property('secondary-icon-sensitive', False)

        # random stuff we need to cart around

        self.entry_url_timer = None

        self.entry_filename_timer = None

        self._thumb_tempfile = None

        self.urlmap = None

        self.current_fmt = None

        self.file_size = None

    def run(self, url=None):

        """

        Start the app

        """

        self.entry_url.set_text(url)

        glib.timeout_add_seconds(0, self.get_info)

        try:

            gtk.main()

        except KeyboardInterrupt:

            pass

    def download(self):

        """

        Download the full video (at the quality selected)

        """

        self.vbox1.set_sensitive(False)

        self.button_ok.set_sensitive(False)

        self.entry_filename_timer = glib.timeout_add(250,

                                                     self.entry_filename_pulse)

        AsyncDownloader(self.urlmap[self.current_fmt],

                        self.partial_factory,

                        self._download_download_done_cb,

                        self.error,

                        query_info_cb=self._download_query_info_cb,

                        tick_cb=self._download_tick_cb).start()

    def partial_factory(self):

        """

        Create a .partial file in a reasonably safe way

        """

        filename = os.path.join(self.filechooserbutton.get_current_folder(),

                                self.entry_filename.get_text() + '.part')

        filedes = os.open(filename, os.O_CREAT | os.O_EXCL | os.O_WRONLY, 0600)

        # XXX this can still fail in older versions of NFS

        fileobj = os.fdopen(filedes, 'w')

        return fileobj

    def _download_query_info_cb(self, qinfo):

        """

        Callback for when the query_info arrives

        """

        self.file_size = qinfo.get_size()

    def _download_tick_cb(self, fileobj):

        """

        Callback for when we get a chunk of data

        """

        if self.entry_filename_timer is not None:

            glib.source_remove(self.entry_filename_timer)

            self.entry_filename_timer = None

        fraction = 1.0 * fileobj.tell() / self.file_size

        self.entry_filename.set_progress_fraction(fraction)

    def _download_download_done_cb(self, fileobj):

        """

        Callback for when the download finishes

        """

        fileobj.flush()

        os.fsync(fileobj.fileno())

        filename = os.path.join(self.filechooserbutton.get_current_folder(),

                                self.entry_filename.get_text())

        os.rename(filename + '.part', filename)

        logging.debug('downloaded movie to %s', filename)

        self.clear_progress()

        self.vbox1.set_sensitive(True)

        self.button_ok.set_sensitive(True)

    def get_video_id(self):

        """

        Extract the video id from the url

        """

        url = self.entry_url.get_text()

        url = re.sub(r'^http://(?:www\.|m\.)youtube\.com/.*next_url=([^&]+)',

                     lambda m: urllib.unquote_plus(m.group(1)), url)

        match = re.match(r'http://(?:www\.|m\.)?youtube\.com/'

                         r'(?:watch\?v=|v/)([^<>?&,]+)',

                         url)

        if match is None:

            return None

        else:

            return match.group(1)

    def get_thumb(self, uri):

        """

        Download the thumbnail of the video

        """

        AsyncDownloader(uri, tempfile.NamedTemporaryFile,

                        self._thumb_download_done_cb, self.error).start()

    def _thumb_download_done_cb(self, fileobj):

        """

        Callback for when the thumbnail download finished

        """

        fileobj.flush()

        logging.debug("downloaded thumb into %s", fileobj.name)

        self.img_preview.set_from_file(fileobj.name)

        self.clear_progress()

    def clear_progress(self):

        """

        Turn off any progress indicators

        """

        if self.entry_url_timer is not None:

            glib.source_remove(self.entry_url_timer)

            self.entry_url_timer = None

        self.entry_url.set_progress_fraction(0.)

        if self.entry_filename_timer is not None:

            glib.source_remove(self.entry_filename_timer)

            self.entry_filename_timer = None

        self.entry_filename.set_progress_fraction(0.)

    def error(self, exc):

        """

        Show and log an error

        """

        msg = str(exc)

        logging.exception(msg)

        self.main_window.set_sensitive(False)

        self.error_dialog.format_secondary_text(msg)

        self.error_dialog.run()

        self.error_dialog.hide()

        self.main_window.set_sensitive(True)

        self.vbox1.set_sensitive(True)

        self.button_ok.set_sensitive(True)

        self.clear_progress()

    def get_info_result(self, urls, title, thumbnail=None):

        """

        Called with the parsed result of the call to get_info

        Enables/disables the appropriate radiobuttons (depening on the

        qualities available), and gets the thumbnail

        """

        logging.debug("get_info answered with: title: %r  thumb: %r  urls: %r",

                      title, thumbnail, urls)

        title = title.replace('/', u'\u2044')

        urls = re.findall(r'(\d+)\|([^,]+)', urls)

        self.urlmap = dict(urls)

        logging.debug("movie has formats %s", sorted(self.urlmap.keys()))

        logging.info("unknown formats: %s",

                     sorted(set(self.urlmap) - set(EXTS)))

        self.hbox_url_info.set_sensitive(True)

        best = None

        for name in '1080p', '720p', 'high', 'medium', 'low', 'mobile':

            radiobtn = getattr(self, 'radiobutton_' + name)

            radiobtn.set_active(False)

            radiobtn.fmt = None

            for fmt in FMT_MAP[name]:

                has_fmt = fmt in self.urlmap

                radiobtn.set_sensitive(has_fmt)

                if has_fmt:

                    radiobtn.fmt = fmt

                    if best is None:

                        best = fmt

                        radiobtn.set_active(True)

                    break

        if best is None:

            raise ValueError("No known formats available")

        self.current_fmt = best

        self.entry_filename.set_text(title + '.' + EXTS[best])

        if thumbnail:

            self.get_thumb(thumbnail)

        else:

            self.img_preview.set_from_icon_name('b0rken', gtk.ICON_SIZE_DND)

            self.clear_progress()

    def scrape(self):

        """

        Download the page, so it gets scraped (get_info failed)

        """

        AsyncDownloader(self.entry_url.get_text(),

                        StringIO,

                        self._scrape_download_done).start()

    def _scrape_download_done(self, buf):

        """

        extract the necessary info out of the page

        """

        try:

            body = buf.getvalue()

            match = re.search(r"'SWF_ARGS' *: *({.*?})", body, re.S)

            if not match:

                # XXX check for error-box (cf. youtubedown)

                raise ValueError("Unable to scrape the HTML")

            match = re.search(r'"fmt_url_map": "(.*?)"', match.group(1), re.S)

            if not match:

                raise ValueError("no fmt_url_map found")

            urls = urllib.unquote(de_entify(match.group(1)))

            match = re.search(r"'VIDEO_TITLE': '(.*?)'", body, re.S)

            if not match:

                raise ValueError("No title found")

            title = de_entify(match.group(1)).strip()

            self.get_info_result(urls, title)

        except StandardError, exc:

            self.error(exc)

        self.vbox1.set_sensitive(True)

    def get_info(self):

        """

        Download the "get_info" page for the video, which is easyer to

        parse than the "real" page

        """

        self.vbox1.set_sensitive(False)

        video_id = self.get_video_id()

        if video_id is None:

            self.error('Unable to find an ID in the URL %r'

                       % (self.entry_url.get_text(),))

        else:

            self.entry_url_timer = glib.timeout_add(250,

                                                    self.entry_url_pulse)

            info_url = "http://www.youtube.com/get_video_info?video_id=" \

                + video_id

            AsyncDownloader(info_url, StringIO,

                            self._get_info_download_done_cb, self.error).start()

        return False

    def _get_info_download_done_cb(self, fileobj):

        """

        The "get_info" page finished downloading

        """

        try:

            result = de_entify(fileobj.getvalue())

            result = dict(cgi.parse_qsl(result))

            urls = result.get('fmt_url_map', result.get('fmt_stream_map'))

            if urls is None:

                self.scrape()

                return

            if 'title' not in result:

                raise RuntimeError("No title found")

            title = de_entify(result['title']).strip()

            self.get_info_result(urls, title, result.get('thumbnail_url'))

        except StandardError, exc:

            self.error(exc)

        self.vbox1.set_sensitive(True)

    def entry_filename_pulse(self):

        """

        pulse the filename entry progressbar

        """

        self.entry_filename.progress_pulse()

        return True

    def entry_url_pulse(self):

        """

        pulse the url entry progressbar

        """

        self.entry_url.progress_pulse()

        return True

    def on_radiobutton_fmt_toggled(self, widget):

        """

        Callback for when one of the format radiobuttons is toggled

        This sets the appropriate extension on the filename

        """

        if widget.fmt is None:

            return

        this_ext = EXTS[widget.fmt]

        if widget.get_active():

            self.current_fmt = widget.fmt

            text = self.entry_filename.get_text() + this_ext

        else:

            text = self.entry_filename.get_text().replace(this_ext, '')

        self.entry_filename.set_text(text)

    def on_entry_filename_activate(self, *ignored):

        """

        Callback for when the filename entry is activated: Start the download

        """

        return self.download()

    on_entry_filename_icon_press = on_entry_filename_activate

    def on_entry_url_activate(self, *ignored):

        """

        Callback for when the url entry is activated: get the metadata

        """

        return self.get_info()

    on_entry_url_icon_press = on_entry_url_activate

    def on_main_window_response(self, ignored, response):

        """

        Callback for handling a dialog response (either cancel close or ok)

        cancel should abort and close

        close should just close

        ok should start the download

        """

        if response == gtk.RESPONSE_OK:

            self.download()

        else:

            if response not in (gtk.RESPONSE_DELETE_EVENT, gtk.RESPONSE_CANCEL):

                desc = "???"

                for key, val in vars(gtk).iteritems():

                    if key.startswith('RESPONSE_') and val == response:

                        desc = key

                logging.error("Unknown response %s (%s)", response, desc)

            gtk.main_quit()

    def on_entry_url_changed(self, widget):

        """

        Callback for when the url entry changed

        Used to give feedback about the validity / utility of the url

        """

        widget.set_property('secondary-icon-sensitive',

                            self.get_video_id() is not None)

    def on_entry_filename_changed(self, widget):

        """

        Callback for when the filename entry changed

        Used to give feedback about the validity / utility of the filename

        """

        widget.set_property('secondary-icon-sensitive', bool(widget.get_text()))

if __name__ == '__main__':

    YouTubeDownloader().run(sys.argv[1] if len(sys.argv) > 1 else None)