finished caption support + docs

docs/api.rst

@@ -27,6 +27,13 @@ StreamQuery Object
    :members:
    :inherited-members:
 
+Caption Object
+-------------
+
+.. autoclass:: pytube.Caption
+   :members:
+   :inherited-members:
+
 Extract
 -------
 

pytube/__init__.py

@@ -15,6 +15,7 @@ __license__ = 'MIT License'
 __copyright__ = 'Copyright 2017 Nick Ficano'
 
 from pytube.logging import create_logger
+from pytube.query import CaptionQuery
 from pytube.query import StreamQuery
 from pytube.streams import Stream
 from pytube.captions import Caption

pytube/__main__.py

@@ -13,6 +13,7 @@ import json
 import logging
 
 from pytube import Caption
+from pytube import CaptionQuery
 from pytube import extract
 from pytube import mixins
 from pytube import request
@@ -163,6 +164,11 @@ class YouTube(object):
             self.fmt_streams.append(video)
 
     def initialize_caption_objects(self):
+        """Populate instances of :class:`Caption <Caption>`.
+
+        Take the unscrambled player response data, and use it to initialize
+        instances of :class:`Caption <Caption>`.
+        """
         if 'captions' not in self.player_config['args']['player_response']:
             return
         caption_tracks = (
@@ -177,7 +183,8 @@ class YouTube(object):
 
     @property
     def captions(self):
-        return self.caption_tracks
+        """Interface to query caption tracks."""
+        return CaptionQuery([c for c in self.caption_tracks])
 
     @property
     def streams(self):

pytube/captions.py

@@ -1,9 +1,35 @@
 # -*- coding: utf-8 -*-
+"""This module contrains a container for caption tracks."""
+from pytube import request
+from pytube.helpers import xml_caption_to_srt
+
+
 class Caption:
+    """Container for caption tracks."""
+
     def __init__(self, caption_track):
+        """Construct a :class:`Caption <Caption>`.
+
+        :param dict caption_track:
+            Caption track data extracted from ``watch_html``.
+        """
         self.url = caption_track.get('baseUrl')
         self.name = caption_track['name']['simpleText']
         self.code = caption_track['languageCode']
 
+    @property
+    def xml_captions(self):
+        """Download the xml caption tracks."""
+        return request.get(self.url)
+
+    def generate_srt_captions(self):
+        """Generate "SubRip Subtitle" captions.
+
+        Takes the xml captions from :meth:`~pytube.Caption.xml_captions` and
+        recompiles them into the "SubRip Subtitle" format.
+        """
+        return xml_caption_to_srt(self.xml_captions)
+
     def __repr__(self):
+        """Printable object representation."""
         return'<Caption lang="{s.name}" code="{s.code}">'.format(s=self)

pytube/compat.py

@@ -13,6 +13,12 @@ if python_version == 2:
     from urllib2 import unquote
     from urllib2 import urlopen
     from urlparse import parse_qsl
+    from HTMLParser import HTMLParser
+
+    def unescape(s):
+        """Strip HTML entries from a string."""
+        html_parser = HTMLParser()
+        return html_parser.unescape(s)
 
     def unicode(s):
         """Encode a string to utf-8."""
@@ -25,6 +31,7 @@ elif python_version == 3:
     from urllib.parse import unquote
     from urllib.parse import urlencode
     from urllib.request import urlopen
+    from html import unescape
 
     def unicode(s):
         """No-op."""

pytube/helpers.py

@@ -3,9 +3,13 @@
 from __future__ import absolute_import
 
 import logging
+import math
 import pprint
 import re
+import time
+import xml.etree.ElementTree as ElementTree
 
+from pytube.compat import unescape
 from pytube.compat import unicode
 from pytube.exceptions import RegexMatchError
 
@@ -88,3 +92,50 @@ def safe_filename(s, max_length=255):
     regex = re.compile(pattern, re.UNICODE)
     filename = regex.sub('', s)
     return unicode(filename[:max_length].rsplit(' ', 0)[0])
+
+
+def float_to_srt_time_format(d):
+    """Convert decimal durations into proper srt format.
+
+    :rtype: str
+    :returns:
+        SubRip Subtitle (str) formatted time duration.
+
+    >>> float_to_srt_time_format(3.89)
+    '00:00:03,890'
+    """
+    frac, whole = math.modf(d)
+    time_fmt = time.strftime('0%H:0%M:%S,', time.gmtime(whole))
+    ms = '{:.3f}'.format(frac).replace('0.', '')
+    return time_fmt + ms
+
+
+def xml_caption_to_srt(xml_captions):
+    """Convert xml caption tracks to "SubRip Subtitle (srt)".
+
+    :param str xml_captions:
+        XML formatted caption tracks.
+    """
+    segments = []
+    root = ElementTree.fromstring(xml_captions)
+    for i, child in enumerate(root.getchildren()):
+        text = child.text or ''
+        caption = unescape(
+            text
+            .replace('\n', ' ')
+            .replace('  ', ' '),
+        )
+        duration = float(child.attrib['dur'])
+        start = float(child.attrib['start'])
+        end = start + duration
+        sequence_number = i + 1  # convert from 0-indexed to 1.
+        line = (
+            '{seq}\n{start} --> {end}\n{text}\n'.format(
+                seq=sequence_number,
+                start=float_to_srt_time_format(start),
+                end=float_to_srt_time_format(end),
+                text=caption,
+            )
+        )
+        segments.append(line)
+    return '\n'.join(segments).strip()

pytube/query.py

@@ -1,13 +1,16 @@
 # -*- coding: utf-8 -*-
-"""This module provides a query interface for media streams."""
+"""This module provides a query interface for media streams and captions."""
 
 
 class StreamQuery:
     """Interface for querying the available media streams."""
 
     def __init__(self, fmt_streams):
-        """Construct a :class:`StreamQuery <StreamQuery>`."""
-        # list of :class:`Stream <Stream>` instances.
+        """Construct a :class:`StreamQuery <StreamQuery>`.
+
+        param list fmt_streams:
+            list of :class:`Stream <Stream>` instances.
+        """
         self.fmt_streams = fmt_streams
         self.itag_index = {int(s.itag): s for s in fmt_streams}
 
@@ -224,3 +227,33 @@ class StreamQuery:
     def all(self):
         """Get all the results represented by this query as a list."""
         return self.fmt_streams
+
+
+class CaptionQuery:
+    """Interface for querying the available captions."""
+
+    def __init__(self, captions):
+        """Construct a :class:`Caption <Caption>`.
+
+        param list captions:
+            list of :class:`Caption <Caption>` instances.
+
+        """
+        self.captions = captions
+        self.lang_code_index = {c.code: c for c in captions}
+
+    def get_by_language_code(self, lang_code):
+        """Get the :class:`Caption <Caption>` for a given ``lang_code``.
+
+        :param str lang_code:
+            The code that identifies the caption language.
+        :rtype: :class:`Caption <Caption>` or ``None``
+        :returns:
+            The :class:`Caption <Caption>` matching the given ``lang_code`` or
+            ``None`` if it does not exist.
+        """
+        return self.lang_code_index.get(lang_code)
+
+    def all(self):
+        """Get all the results represented by this query as a list."""
+        return self.captions