Switch back to underline for invalid characters, and make restricted ASCII-only
[youtube-dl] / 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         if not result:
222                 result = '_'
223         return result
224
225 def orderedSet(iterable):
226         """ Remove all duplicates from the input iterable """
227         res = []
228         for el in iterable:
229                 if el not in res:
230                         res.append(el)
231         return res
232
233 def unescapeHTML(s):
234         """
235         @param s a string (of type unicode)
236         """
237         assert type(s) == type(u'')
238
239         result = re.sub(ur'(?u)&(.+?);', htmlentity_transform, s)
240         return result
241
242 def encodeFilename(s):
243         """
244         @param s The name of the file (of type unicode)
245         """
246
247         assert type(s) == type(u'')
248
249         if sys.platform == 'win32' and sys.getwindowsversion()[0] >= 5:
250                 # Pass u'' directly to use Unicode APIs on Windows 2000 and up
251                 # (Detecting Windows NT 4 is tricky because 'major >= 4' would
252                 # match Windows 9x series as well. Besides, NT 4 is obsolete.)
253                 return s
254         else:
255                 return s.encode(sys.getfilesystemencoding(), 'ignore')
256
257 class DownloadError(Exception):
258         """Download Error exception.
259
260         This exception may be thrown by FileDownloader objects if they are not
261         configured to continue on errors. They will contain the appropriate
262         error message.
263         """
264         pass
265
266
267 class SameFileError(Exception):
268         """Same File exception.
269
270         This exception will be thrown by FileDownloader objects if they detect
271         multiple files would have to be downloaded to the same file on disk.
272         """
273         pass
274
275
276 class PostProcessingError(Exception):
277         """Post Processing exception.
278
279         This exception may be raised by PostProcessor's .run() method to
280         indicate an error in the postprocessing task.
281         """
282         pass
283
284 class MaxDownloadsReached(Exception):
285         """ --max-downloads limit has been reached. """
286         pass
287
288
289 class UnavailableVideoError(Exception):
290         """Unavailable Format exception.
291
292         This exception will be thrown when a video is requested
293         in a format that is not available for that video.
294         """
295         pass
296
297
298 class ContentTooShortError(Exception):
299         """Content Too Short exception.
300
301         This exception may be raised by FileDownloader objects when a file they
302         download is too small for what the server announced first, indicating
303         the connection was probably interrupted.
304         """
305         # Both in bytes
306         downloaded = None
307         expected = None
308
309         def __init__(self, downloaded, expected):
310                 self.downloaded = downloaded
311                 self.expected = expected
312
313
314 class Trouble(Exception):
315         """Trouble helper exception
316         
317         This is an exception to be handled with
318         FileDownloader.trouble
319         """
320
321 class YoutubeDLHandler(urllib2.HTTPHandler):
322         """Handler for HTTP requests and responses.
323
324         This class, when installed with an OpenerDirector, automatically adds
325         the standard headers to every HTTP request and handles gzipped and
326         deflated responses from web servers. If compression is to be avoided in
327         a particular request, the original request in the program code only has
328         to include the HTTP header "Youtubedl-No-Compression", which will be
329         removed before making the real request.
330
331         Part of this code was copied from:
332
333         http://techknack.net/python-urllib2-handlers/
334
335         Andrew Rowls, the author of that code, agreed to release it to the
336         public domain.
337         """
338
339         @staticmethod
340         def deflate(data):
341                 try:
342                         return zlib.decompress(data, -zlib.MAX_WBITS)
343                 except zlib.error:
344                         return zlib.decompress(data)
345
346         @staticmethod
347         def addinfourl_wrapper(stream, headers, url, code):
348                 if hasattr(urllib2.addinfourl, 'getcode'):
349                         return urllib2.addinfourl(stream, headers, url, code)
350                 ret = urllib2.addinfourl(stream, headers, url)
351                 ret.code = code
352                 return ret
353
354         def http_request(self, req):
355                 for h in std_headers:
356                         if h in req.headers:
357                                 del req.headers[h]
358                         req.add_header(h, std_headers[h])
359                 if 'Youtubedl-no-compression' in req.headers:
360                         if 'Accept-encoding' in req.headers:
361                                 del req.headers['Accept-encoding']
362                         del req.headers['Youtubedl-no-compression']
363                 return req
364
365         def http_response(self, req, resp):
366                 old_resp = resp
367                 # gzip
368                 if resp.headers.get('Content-encoding', '') == 'gzip':
369                         gz = gzip.GzipFile(fileobj=StringIO.StringIO(resp.read()), mode='r')
370                         resp = self.addinfourl_wrapper(gz, old_resp.headers, old_resp.url, old_resp.code)
371                         resp.msg = old_resp.msg
372                 # deflate
373                 if resp.headers.get('Content-encoding', '') == 'deflate':
374                         gz = StringIO.StringIO(self.deflate(resp.read()))
375                         resp = self.addinfourl_wrapper(gz, old_resp.headers, old_resp.url, old_resp.code)
376                         resp.msg = old_resp.msg
377                 return resp