Merge pull request #161 from nficano/feature-caption-support

docs/api.rst

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

pytube/__init__.py

@@ -15,8 +15,10 @@ __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
 from pytube.__main__ import YouTube
 
 logger = create_logger()

pytube/__main__.py

@@ -12,6 +12,8 @@ from __future__ import absolute_import
 import json
 import logging
 
+from pytube import Caption
+from pytube import CaptionQuery
 from pytube import extract
 from pytube import mixins
 from pytube import request
@@ -59,6 +61,7 @@ class YouTube(object):
         self.player_config = None  # inline js in the html containing streams
 
         self.fmt_streams = []  # list of :class:`Stream <Stream>` instances
+        self.caption_tracks = []
 
         # video_id part of /watch?v=<video_id>
         self.video_id = extract.video_id(url)
@@ -115,6 +118,8 @@ class YouTube(object):
         # build instances of :class:`Stream <Stream>`
         self.initialize_stream_objects(progressive_fmts)
         self.initialize_stream_objects(adaptive_fmts)
+
+        self.initialize_caption_objects()
         logger.info('init finished successfully')
 
     def prefetch(self):
@@ -158,6 +163,29 @@ 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 = (
+            self.player_config['args']
+            ['player_response']
+            ['captions']
+            ['playerCaptionsTracklistRenderer']
+            ['captionTracks']
+        )
+        for caption_track in caption_tracks:
+            self.caption_tracks.append(Caption(caption_track))
+
+    @property
+    def captions(self):
+        """Interface to query caption tracks."""
+        return CaptionQuery([c for c in self.caption_tracks])
+
     @property
     def streams(self):
         """Interface to query both adaptive (DASH) and progressive streams."""

pytube/captions.py

@@ -0,0 +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

setup.cfg

@@ -3,7 +3,7 @@ commit = True
 tag = True
 current_version = 7.0.9
 parse = (?P<major>\d+)\.(?P<minor>\d+)\.(?P<patch>\d+)(\-(?P<release>[a-z]+))?
-serialize = 
+serialize =
 	{major}.{minor}.{patch}
 
 [metadata]
@@ -15,6 +15,5 @@ description-file = README.md
 
 [coverage:run]
 source = pytube
-omit = 
+omit =
 	pytube/compat.py
-