1
0
mirror of https://github.com/moparisthebest/SickRage synced 2024-12-25 01:08:50 -05:00
SickRage/lib/subliminal/videos.py

301 lines
10 KiB
Python

# -*- coding: utf-8 -*-
# Copyright 2011-2012 Antoine Bertin <diaoulael@gmail.com>
#
# This file is part of subliminal.
#
# subliminal is free software; you can redistribute it and/or modify it under
# the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# subliminal is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with subliminal. If not, see <http://www.gnu.org/licenses/>.
from . import subtitles
from .language import Language
from .utils import to_unicode
import enzyme.core
import guessit
import hashlib
import logging
import mimetypes
import os
import struct
from sickbeard import encodingKludge as ek
import sickbeard
__all__ = ['EXTENSIONS', 'MIMETYPES', 'Video', 'Episode', 'Movie', 'UnknownVideo',
'scan', 'hash_opensubtitles', 'hash_thesubdb']
logger = logging.getLogger("subliminal")
#: Video extensions
EXTENSIONS = ['.avi', '.mkv', '.mpg', '.mp4', '.m4v', '.mov', '.ogm', '.ogv', '.wmv',
'.divx', '.asf']
#: Video mimetypes
MIMETYPES = ['video/mpeg', 'video/mp4', 'video/quicktime', 'video/x-ms-wmv', 'video/x-msvideo',
'video/x-flv', 'video/x-matroska', 'video/x-matroska-3d']
class Video(object):
"""Base class for videos
:param string path: path
:param guess: guessed informations
:type guess: :class:`~guessit.guess.Guess`
:param string imdbid: imdbid
"""
def __init__(self, path, guess, imdbid=None):
self.release = path
self.guess = guess
self.imdbid = imdbid
self._path = None
self.hashes = {}
if isinstance(path, unicode):
path = path.encode('utf-8')
if os.path.exists(path):
self._path = path
self.size = os.path.getsize(self._path)
self._compute_hashes()
@classmethod
def from_path(cls, path):
"""Create a :class:`Video` subclass guessing all informations from the given path
:param string path: path
:return: video object
:rtype: :class:`Episode` or :class:`Movie` or :class:`UnknownVideo`
"""
guess = guessit.guess_file_info(path, 'autodetect')
result = None
if guess['type'] == 'episode' and 'series' in guess and 'season' in guess and 'episodeNumber' in guess:
title = None
if 'title' in guess:
title = guess['title']
result = Episode(path, guess['series'], guess['season'], guess['episodeNumber'], title, guess)
if guess['type'] == 'movie' and 'title' in guess:
year = None
if 'year' in guess:
year = guess['year']
result = Movie(path, guess['title'], year, guess)
if not result:
result = UnknownVideo(path, guess)
if not isinstance(result, cls):
raise ValueError('Video is not of requested type')
return result
@property
def exists(self):
"""Whether the video exists or not"""
if self._path:
return os.path.exists(self._path)
return False
@property
def path(self):
"""Path to the video"""
return self._path
@path.setter
def path(self, value):
if not os.path.exists(value):
raise ValueError('Path does not exists')
self._path = value
self.size = os.path.getsize(self._path)
self._compute_hashes()
def _compute_hashes(self):
"""Compute different hashes"""
self.hashes['OpenSubtitles'] = hash_opensubtitles(self.path)
self.hashes['TheSubDB'] = hash_thesubdb(self.path)
def scan(self):
"""Scan and return associated subtitles
:return: associated subtitles
:rtype: list of :class:`~subliminal.subtitles.Subtitle`
"""
if not self.exists:
return []
basepath = os.path.splitext(self.path)[0]
results = []
video_infos = None
try:
video_infos = enzyme.parse(self.path)
logger.debug(u'Succeeded parsing %s with enzyme: %r' % (self.path, video_infos))
except:
logger.debug(u'Failed parsing %s with enzyme' % self.path)
if isinstance(video_infos, enzyme.core.AVContainer):
results.extend([subtitles.EmbeddedSubtitle.from_enzyme(self.path, s) for s in video_infos.subtitles])
# cannot use glob here because it chokes if there are any square
# brackets inside the filename, so we have to use basic string
# startswith/endswith comparisons
folder, basename = os.path.split(basepath)
if folder == '':
folder = '.'
existing = [f for f in os.listdir(folder) if f.startswith(basename)]
if sickbeard.SUBTITLES_DIR:
subsDir = ek.ek(os.path.join, folder, sickbeard.SUBTITLES_DIR)
if ek.ek(os.path.isdir, subsDir):
existing.extend([f for f in os.listdir(subsDir) if f.startswith(basename)])
for path in existing:
for ext in subtitles.EXTENSIONS:
if path.endswith(ext):
language = Language(path[len(basename) + 1:-len(ext)], strict=False)
results.append(subtitles.ExternalSubtitle(path, language))
return results
def __unicode__(self):
return to_unicode(self.path or self.release)
def __str__(self):
return unicode(self).encode('utf-8')
def __repr__(self):
return '%s(%s)' % (self.__class__.__name__, self)
def __hash__(self):
return hash(self.path or self.release)
class Episode(Video):
"""Episode :class:`Video`
:param string path: path
:param string series: series
:param int season: season number
:param int episode: episode number
:param string title: title
:param guess: guessed informations
:type guess: :class:`~guessit.guess.Guess`
:param string tvdbid: tvdbid
:param string imdbid: imdbid
"""
def __init__(self, path, series, season, episode, title=None, guess=None, tvdbid=None, imdbid=None):
super(Episode, self).__init__(path, guess, imdbid)
self.series = series
self.title = title
self.season = season
self.episode = episode
self.tvdbid = tvdbid
class Movie(Video):
"""Movie :class:`Video`
:param string path: path
:param string title: title
:param int year: year
:param guess: guessed informations
:type guess: :class:`~guessit.guess.Guess`
:param string imdbid: imdbid
"""
def __init__(self, path, title, year=None, guess=None, imdbid=None):
super(Movie, self).__init__(path, guess, imdbid)
self.title = title
self.year = year
class UnknownVideo(Video):
"""Unknown video"""
pass
def scan(entry, max_depth=3, scan_filter=None, depth=0):
"""Scan a path for videos and subtitles
:param string entry: path
:param int max_depth: maximum folder depth
:param function scan_filter: filter function that takes a path as argument and returns a boolean indicating whether it has to be filtered out (``True``) or not (``False``)
:param int depth: starting depth
:return: found videos and subtitles
:rtype: list of (:class:`Video`, [:class:`~subliminal.subtitles.Subtitle`])
"""
if isinstance(entry, unicode):
entry = entry.encode('utf-8')
if depth > max_depth and max_depth != 0: # we do not want to search the whole file system except if max_depth = 0
return []
if os.path.isdir(entry): # a dir? recurse
logger.debug(u'Scanning directory %s with depth %d/%d' % (entry, depth, max_depth))
result = []
for e in os.listdir(entry):
result.extend(scan(os.path.join(entry, e), max_depth, scan_filter, depth + 1))
return result
if os.path.isfile(entry) or depth == 0:
logger.debug(u'Scanning file %s with depth %d/%d' % (entry, depth, max_depth))
if depth != 0: # trust the user: only check for valid format if recursing
if mimetypes.guess_type(entry)[0] not in MIMETYPES and os.path.splitext(entry)[1] not in EXTENSIONS:
return []
if scan_filter is not None and scan_filter(entry):
return []
video = Video.from_path(entry)
return [(video, video.scan())]
logger.warning(u'Scanning entry %s failed with depth %d/%d' % (entry, depth, max_depth))
return [] # anything else
def hash_opensubtitles(path):
"""Compute a hash using OpenSubtitles' algorithm
:param string path: path
:return: hash
:rtype: string
"""
longlongformat = 'q' # long long
bytesize = struct.calcsize(longlongformat)
with open(path, 'rb') as f:
filesize = os.path.getsize(path)
filehash = filesize
if filesize < 65536 * 2:
return None
for _ in range(65536 / bytesize):
filebuffer = f.read(bytesize)
(l_value,) = struct.unpack(longlongformat, filebuffer)
filehash += l_value
filehash = filehash & 0xFFFFFFFFFFFFFFFF # to remain as 64bit number
f.seek(max(0, filesize - 65536), 0)
for _ in range(65536 / bytesize):
filebuffer = f.read(bytesize)
(l_value,) = struct.unpack(longlongformat, filebuffer)
filehash += l_value
filehash = filehash & 0xFFFFFFFFFFFFFFFF
returnedhash = '%016x' % filehash
logger.debug(u'Computed OpenSubtitle hash %s for %s' % (returnedhash, path))
return returnedhash
def hash_thesubdb(path):
"""Compute a hash using TheSubDB's algorithm
:param string path: path
:return: hash
:rtype: string
"""
readsize = 64 * 1024
if os.path.getsize(path) < readsize:
return None
with open(path, 'rb') as f:
data = f.read(readsize)
f.seek(-readsize, os.SEEK_END)
data += f.read(readsize)
returnedhash = hashlib.md5(data).hexdigest()
logger.debug(u'Computed TheSubDB hash %s for %s' % (returnedhash, path))
return returnedhash