youtube_dl/utils.py

   1 #!/usr/bin/env python
   2 # -*- coding: utf-8 -*-
   3
   4 import gzip
   5 import htmlentitydefs
   6 import HTMLParser
   7 import locale
   8 import os
   9 import re
  10 import sys
  11 import zlib
  12 import urllib2
  13 import email.utils
  14 import json
  15
  16 try:
  17         import cStringIO as StringIO
  18 except ImportError:
  19         import StringIO
  20
  21 std_headers = {
  22         'User-Agent': 'Mozilla/5.0 (X11; Linux x86_64; rv:10.0) Gecko/20100101 Firefox/10.0',
  23         'Accept-Charset': 'ISO-8859-1,utf-8;q=0.7,*;q=0.7',
  24         'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,*/*;q=0.8',
  25         'Accept-Encoding': 'gzip, deflate',
  26         'Accept-Language': 'en-us,en;q=0.5',
  27 }
  28
  29 try:
  30     compat_str = unicode # Python 2
  31 except NameError:
  32     compat_str = str
  33
  34 def preferredencoding():
  35         """Get preferred encoding.
  36
  37         Returns the best encoding scheme for the system, based on
  38         locale.getpreferredencoding() and some further tweaks.
  39         """
  40         def yield_preferredencoding():
  41                 try:
  42                         pref = locale.getpreferredencoding()
  43                         u'TEST'.encode(pref)
  44                 except:
  45                         pref = 'UTF-8'
  46                 while True:
  47                         yield pref
  48         return yield_preferredencoding().next()
  49
  50
  51 def htmlentity_transform(matchobj):
  52         """Transforms an HTML entity to a Unicode character.
  53
  54         This function receives a match object and is intended to be used with
  55         the re.sub() function.
  56         """
  57         entity = matchobj.group(1)
  58
  59         # Known non-numeric HTML entity
  60         if entity in htmlentitydefs.name2codepoint:
  61                 return unichr(htmlentitydefs.name2codepoint[entity])
  62
  63         # Unicode character
  64         mobj = re.match(ur'(?u)#(x?\d+)', entity)
  65         if mobj is not None:
  66                 numstr = mobj.group(1)
  67                 if numstr.startswith(u'x'):
  68                         base = 16
  69                         numstr = u'0%s' % numstr
  70                 else:
  71                         base = 10
  72                 return unichr(long(numstr, base))
  73
  74         # Unknown entity in name, return its literal representation
  75         return (u'&%s;' % entity)
  76
  77 HTMLParser.locatestarttagend = re.compile(r"""<[a-zA-Z][-.a-zA-Z0-9:_]*(?:\s+(?:(?<=['"\s])[^\s/>][^\s/=>]*(?:\s*=+\s*(?:'[^']*'|"[^"]*"|(?!['"])[^>\s]*))?\s*)*)?\s*""", re.VERBOSE) # backport bugfix
  78 class IDParser(HTMLParser.HTMLParser):
  79         """Modified HTMLParser that isolates a tag with the specified id"""
  80         def __init__(self, id):
  81                 self.id = id
  82                 self.result = None
  83                 self.started = False
  84                 self.depth = {}
  85                 self.html = None
  86                 self.watch_startpos = False
  87                 self.error_count = 0
  88                 HTMLParser.HTMLParser.__init__(self)
  89
  90         def error(self, message):
  91                 if self.error_count > 10 or self.started:
  92                         raise HTMLParser.HTMLParseError(message, self.getpos())
  93                 self.rawdata = '\n'.join(self.html.split('\n')[self.getpos()[0]:]) # skip one line
  94                 self.error_count += 1
  95                 self.goahead(1)
  96
  97         def loads(self, html):
  98                 self.html = html
  99                 self.feed(html)
 100                 self.close()
 101
 102         def handle_starttag(self, tag, attrs):
 103                 attrs = dict(attrs)
 104                 if self.started:
 105                         self.find_startpos(None)
 106                 if 'id' in attrs and attrs['id'] == self.id:
 107                         self.result = [tag]
 108                         self.started = True
 109                         self.watch_startpos = True
 110                 if self.started:
 111                         if not tag in self.depth: self.depth[tag] = 0
 112                         self.depth[tag] += 1
 113
 114         def handle_endtag(self, tag):
 115                 if self.started:
 116                         if tag in self.depth: self.depth[tag] -= 1
 117                         if self.depth[self.result[0]] == 0:
 118                                 self.started = False
 119                                 self.result.append(self.getpos())
 120
 121         def find_startpos(self, x):
 122                 """Needed to put the start position of the result (self.result[1])
 123                 after the opening tag with the requested id"""
 124                 if self.watch_startpos:
 125                         self.watch_startpos = False
 126                         self.result.append(self.getpos())
 127         handle_entityref = handle_charref = handle_data = handle_comment = \
 128         handle_decl = handle_pi = unknown_decl = find_startpos
 129
 130         def get_result(self):
 131                 if self.result == None: return None
 132                 if len(self.result) != 3: return None
 133                 lines = self.html.split('\n')
 134                 lines = lines[self.result[1][0]-1:self.result[2][0]]
 135                 lines[0] = lines[0][self.result[1][1]:]
 136                 if len(lines) == 1:
 137                         lines[-1] = lines[-1][:self.result[2][1]-self.result[1][1]]
 138                 lines[-1] = lines[-1][:self.result[2][1]]
 139                 return '\n'.join(lines).strip()
 140
 141 def get_element_by_id(id, html):
 142         """Return the content of the tag with the specified id in the passed HTML document"""
 143         parser = IDParser(id)
 144         try:
 145                 parser.loads(html)
 146         except HTMLParser.HTMLParseError:
 147                 pass
 148         return parser.get_result()
 149
 150
 151 def clean_html(html):
 152         """Clean an HTML snippet into a readable string"""
 153         # Newline vs <br />
 154         html = html.replace('\n', ' ')
 155         html = re.sub('\s*<\s*br\s*/?\s*>\s*', '\n', html)
 156         # Strip html tags
 157         html = re.sub('<.*?>', '', html)
 158         # Replace html entities
 159         html = unescapeHTML(html)
 160         return html
 161
 162
 163 def sanitize_open(filename, open_mode):
 164         """Try to open the given filename, and slightly tweak it if this fails.
 165
 166         Attempts to open the given filename. If this fails, it tries to change
 167         the filename slightly, step by step, until it's either able to open it
 168         or it fails and raises a final exception, like the standard open()
 169         function.
 170
 171         It returns the tuple (stream, definitive_file_name).
 172         """
 173         try:
 174                 if filename == u'-':
 175                         if sys.platform == 'win32':
 176                                 import msvcrt
 177                                 msvcrt.setmode(sys.stdout.fileno(), os.O_BINARY)
 178                         return (sys.stdout, filename)
 179                 stream = open(encodeFilename(filename), open_mode)
 180                 return (stream, filename)
 181         except (IOError, OSError), err:
 182                 # In case of error, try to remove win32 forbidden chars
 183                 filename = re.sub(ur'[/<>:"\|\?\*]', u'#', filename)
 184
 185                 # An exception here should be caught in the caller
 186                 stream = open(encodeFilename(filename), open_mode)
 187                 return (stream, filename)
 188
 189
 190 def timeconvert(timestr):
 191         """Convert RFC 2822 defined time string into system timestamp"""
 192         timestamp = None
 193         timetuple = email.utils.parsedate_tz(timestr)
 194         if timetuple is not None:
 195                 timestamp = email.utils.mktime_tz(timetuple)
 196         return timestamp
 197
 198 def sanitize_filename(s, restricted=False):
 199         """Sanitizes a string so it could be used as part of a filename.
 200         If restricted is set, use a stricter subset of allowed characters.
 201         """
 202         def replace_insane(char):
 203                 if char == '?' or ord(char) < 32 or ord(char) == 127:
 204                         return ''
 205                 elif char == '"':
 206                         return '' if restricted else '\''
 207                 elif char == ':':
 208                         return '_-' if restricted else ' -'
 209                 elif char in '\\/|*<>':
 210                         return '_'
 211                 if restricted and (char in '&\'' or char.isspace()):
 212                         return '_'
 213                 if restricted and ord(char) > 127:
 214                         return '_'
 215                 return char
 216
 217         result = u''.join(map(replace_insane, s))
 218         while '__' in result:
 219                 result = result.replace('__', '_')
 220         result = result.strip('_')
 221         # Common case of "Foreign band name - English song title"
 222         if restricted and result.startswith('-_'):
 223                 result = result[2:]
 224         if not result:
 225                 result = '_'
 226         return result
 227
 228 def orderedSet(iterable):
 229         """ Remove all duplicates from the input iterable """
 230         res = []
 231         for el in iterable:
 232                 if el not in res:
 233                         res.append(el)
 234         return res
 235
 236 def unescapeHTML(s):
 237         """
 238         @param s a string (of type unicode)
 239         """
 240         assert type(s) == type(u'')
 241
 242         result = re.sub(ur'(?u)&(.+?);', htmlentity_transform, s)
 243         return result
 244
 245 def encodeFilename(s):
 246         """
 247         @param s The name of the file (of type unicode)
 248         """
 249
 250         assert type(s) == type(u'')
 251
 252         if sys.platform == 'win32' and sys.getwindowsversion()[0] >= 5:
 253                 # Pass u'' directly to use Unicode APIs on Windows 2000 and up
 254                 # (Detecting Windows NT 4 is tricky because 'major >= 4' would
 255                 # match Windows 9x series as well. Besides, NT 4 is obsolete.)
 256                 return s
 257         else:
 258                 return s.encode(sys.getfilesystemencoding(), 'ignore')
 259
 260 class DownloadError(Exception):
 261         """Download Error exception.
 262
 263         This exception may be thrown by FileDownloader objects if they are not
 264         configured to continue on errors. They will contain the appropriate
 265         error message.
 266         """
 267         pass
 268
 269
 270 class SameFileError(Exception):
 271         """Same File exception.
 272
 273         This exception will be thrown by FileDownloader objects if they detect
 274         multiple files would have to be downloaded to the same file on disk.
 275         """
 276         pass
 277
 278
 279 class PostProcessingError(Exception):
 280         """Post Processing exception.
 281
 282         This exception may be raised by PostProcessor's .run() method to
 283         indicate an error in the postprocessing task.
 284         """
 285         pass
 286
 287 class MaxDownloadsReached(Exception):
 288         """ --max-downloads limit has been reached. """
 289         pass
 290
 291
 292 class UnavailableVideoError(Exception):
 293         """Unavailable Format exception.
 294
 295         This exception will be thrown when a video is requested
 296         in a format that is not available for that video.
 297         """
 298         pass
 299
 300
 301 class ContentTooShortError(Exception):
 302         """Content Too Short exception.
 303
 304         This exception may be raised by FileDownloader objects when a file they
 305         download is too small for what the server announced first, indicating
 306         the connection was probably interrupted.
 307         """
 308         # Both in bytes
 309         downloaded = None
 310         expected = None
 311
 312         def __init__(self, downloaded, expected):
 313                 self.downloaded = downloaded
 314                 self.expected = expected
 315
 316
 317 class Trouble(Exception):
 318         """Trouble helper exception
 319
 320         This is an exception to be handled with
 321         FileDownloader.trouble
 322         """
 323
 324 class YoutubeDLHandler(urllib2.HTTPHandler):
 325         """Handler for HTTP requests and responses.
 326
 327         This class, when installed with an OpenerDirector, automatically adds
 328         the standard headers to every HTTP request and handles gzipped and
 329         deflated responses from web servers. If compression is to be avoided in
 330         a particular request, the original request in the program code only has
 331         to include the HTTP header "Youtubedl-No-Compression", which will be
 332         removed before making the real request.
 333
 334         Part of this code was copied from:
 335
 336         http://techknack.net/python-urllib2-handlers/
 337
 338         Andrew Rowls, the author of that code, agreed to release it to the
 339         public domain.
 340         """
 341
 342         @staticmethod
 343         def deflate(data):
 344                 try:
 345                         return zlib.decompress(data, -zlib.MAX_WBITS)
 346                 except zlib.error:
 347                         return zlib.decompress(data)
 348
 349         @staticmethod
 350         def addinfourl_wrapper(stream, headers, url, code):
 351                 if hasattr(urllib2.addinfourl, 'getcode'):
 352                         return urllib2.addinfourl(stream, headers, url, code)
 353                 ret = urllib2.addinfourl(stream, headers, url)
 354                 ret.code = code
 355                 return ret
 356
 357         def http_request(self, req):
 358                 for h in std_headers:
 359                         if h in req.headers:
 360                                 del req.headers[h]
 361                         req.add_header(h, std_headers[h])
 362                 if 'Youtubedl-no-compression' in req.headers:
 363                         if 'Accept-encoding' in req.headers:
 364                                 del req.headers['Accept-encoding']
 365                         del req.headers['Youtubedl-no-compression']
 366                 return req
 367
 368         def http_response(self, req, resp):
 369                 old_resp = resp
 370                 # gzip
 371                 if resp.headers.get('Content-encoding', '') == 'gzip':
 372                         gz = gzip.GzipFile(fileobj=StringIO.StringIO(resp.read()), mode='r')
 373                         resp = self.addinfourl_wrapper(gz, old_resp.headers, old_resp.url, old_resp.code)
 374                         resp.msg = old_resp.msg
 375                 # deflate
 376                 if resp.headers.get('Content-encoding', '') == 'deflate':
 377                         gz = StringIO.StringIO(self.deflate(resp.read()))
 378                         resp = self.addinfourl_wrapper(gz, old_resp.headers, old_resp.url, old_resp.code)
 379                         resp.msg = old_resp.msg
 380                 return resp