audiomangler/codecs.py

   1 # -*- coding: utf-8 -*-
   2 ###########################################################################
   3 #    Copyright (C) 2008 by Andrew Mahone
   4 #    <andrew.mahone@gmail.com>
   5 #
   6 # Copyright: See COPYING file that comes with this distribution
   7 #
   8 ###########################################################################
   9 # pylint: disable=E1101
  10 import os, os.path
  11 import sys
  12 import re
  13 import Image
  14 try:
  15     from cStringIO import StringIO
  16 except ImportError:
  17     from StringIO import StringIO
  18 from tempfile import mkdtemp
  19 from threading import BoundedSemaphore, RLock
  20 from subprocess import Popen, PIPE
  21 from mutagen import FileType
  22 from audiomangler.config import Config
  23 from audiomangler.tag import NormMetaData
  24 from audiomangler.task import CLITask, generator_task
  25 from audiomangler.expression import Expr, Format
  26 from audiomangler import util
  27 from audiomangler.util import ClassInitMeta
  28 from mutagen import File
  29
  30 codec_map = {}
  31
  32 class Codec(object):
  33     __metaclass__ = ClassInitMeta
  34
  35     has_from_wav_multi = False
  36     has_from_wav_pipe = False
  37     has_to_wav_pipe = False
  38     has_replaygain = False
  39     lossless = False
  40     @classmethod
  41     def __classinit__(cls, name__, bases__, cls_dict):
  42         if 'type_' in cls_dict:
  43             codec_map[cls_dict['type_']] = cls
  44
  45     @classmethod
  46     def _conv_out_filename(cls, filename):
  47         return ''.join((filename.rsplit('.', 1)[0], '.', cls.ext))
  48
  49     @classmethod
  50     def from_wav_multi(cls, indir, infiles, outfiles):
  51         if not hasattr('from_wav_multi_cmd'):
  52             return None
  53         encopts = Config['encopts']
  54         if not encopts:
  55             encopts = ()
  56         else:
  57             encopts = tuple(encopts.split())
  58         args = cls.from_wav_multi_cmd.evaluate({
  59            'indir':indir,
  60            'infiles':tuple(infiles),
  61            'outfiles':tuple(outfiles),
  62            'ext':cls.ext,
  63            'type':cls.type_,
  64            'encopts':encopts,
  65            'encoder':cls.encoder
  66         })
  67         return (CLITask(args=args, stdin='/dev/null', stdout='/dev/null', stderr=sys.stderr, background=True))
  68
  69     @classmethod
  70     def from_wav_pipe(cls, infile, outfile):
  71         if not hasattr(cls, 'from_wav_pipe_cmd'):
  72             return None
  73         encopts = Config['encopts']
  74         if not encopts:
  75             encopts = ()
  76         else:
  77             encopts = tuple(encopts.split())
  78         outfile = cls._conv_out_filename(infile)
  79         env = {
  80             'infile': infile,
  81             'outfile': outfile,
  82             'ext': cls.ext,
  83             'type': cls.type_,
  84             'encopts': encopts,
  85             'encoder': cls.encoder
  86         }
  87         args = cls.from_wav_pipe_cmd.evaluate(env)
  88         stdin = '/dev/null'
  89         if hasattr(cls, '_from_wav_pipe_stdin'):
  90             stdin = cls._from_wav_pipe_stdin.evaluate(env)
  91         return CLITask(args=args, stdin=stdin, stdout='/dev/null', stderr=sys.stderr, background=True)
  92
  93     @classmethod
  94     def to_wav_pipe(cls, infile, outfile):
  95         if not hasattr(cls, 'to_wav_pipe_cmd'):
  96             return None
  97         env = {
  98            'infile':infile,
  99            'outfile':outfile,
 100            'ext':cls.ext,
 101            'type':cls.type_,
 102            'decoder':cls.decoder
 103         }
 104         args = cls.to_wav_pipe_cmd.evaluate(env)
 105         stdout = '/dev/null'
 106         if hasattr(cls, 'to_wav_pipe_stdout'):
 107             stdout = cls.to_wav_pipe_stdout.evaluate(env)
 108         return CLITask(args=args, stdin='/dev/null', stdout=stdout, stderr=sys.stderr, background=False)
 109
 110     @classmethod
 111     @generator_task
 112     def add_replaygain(cls, files, metas=None):
 113         env = {
 114             'replaygain':cls.replaygain,
 115             'files':tuple(files)
 116         }
 117         if metas and hasattr(cls, 'calc_replaygain_cmd'):
 118             task = CLITask(*cls.calc_replaygain_cmd.evaluate(env))
 119             output = yield task
 120             tracks, album = cls.calc_replaygain(output)
 121             if tracks:
 122                 for meta, track in zip(metas, tracks):
 123                     meta.update(track)
 124                     meta.update(album)
 125             yield metas
 126         elif hasattr(cls, 'replaygain_cmd'):
 127             task = CLITask(*cls.replaygain_cmd.evaluate(env))
 128             yield task
 129         elif hasattr(cls, 'calc_replaygain'):
 130             task = CLITask(*cls.calc_replaygain_cmd.evaluate(env))
 131             output = yield task
 132             tracks, album = cls.calc_replaygain(output)
 133             for trackfile, trackgain in zip(files, tracks):
 134                 f = File(trackfile)
 135                 m = NormMetaData(trackgain + album)
 136                 m.apply(f)
 137                 f.save()
 138
 139
 140 class MP3Codec(Codec):
 141     ext = 'mp3'
 142     type_ = 'mp3'
 143     encoder = 'lame'
 144     replaygain = 'mp3gain'
 145     has_from_wav_multi = True
 146     has_replaygain = True
 147     from_wav_multi_cmd = Expr("(encoder, '--quiet') + encopts + ('--noreplaygain', '--nogapout', indir, '--nogaptags', '--nogap') + infiles")
 148     calc_replaygain_cmd = Expr("(replaygain, '-q', '-o', '-s', 's')+files")
 149
 150     @staticmethod
 151     def calc_replaygain(out):
 152         (out, err__) = out
 153         out = [l.split('\t')[2:4] for l in out.splitlines()[1:]]
 154         tracks = []
 155         for i in out[:-1]:
 156             gain = ' '.join((i[0], 'dB'))
 157             peak = '%.8f'% (float(i[1]) / 32768)
 158             tracks.append((('replaygain_track_gain', gain), ('replaygain_track_peak', peak)))
 159         gain = ' '.join((out[-1][0], 'dB'))
 160         peak = '%.8f'% (float(out[-1][1]) / 32768)
 161         album = (('replaygain_album_gain', gain), ('replaygain_album_peak', peak))
 162         return tracks, album
 163
 164 class WavPackCodec(Codec):
 165     ext = 'wv'
 166     type_ = 'wavpack'
 167     encoder = 'wavpack'
 168     decoder = 'wvunpack'
 169     replaygain = 'wvgain'
 170     has_to_wav_pipe = True
 171     has_from_wav_pipe = True
 172     has_replaygain = True
 173     lossless = True
 174     to_wav_pipe_cmd = Expr("(decoder, '-q', '-w', infile, '-o', '-')")
 175     to_wav_pipe_stdout = Expr("outfile")
 176     from_wav_pipe_cmd = Expr("(encoder, '-q')+encopts+(infile, '-o', outfile)")
 177     replaygain_cmd = Expr("(replaygain, '-a')+files")
 178
 179 class FLACCodec(Codec):
 180     ext = 'flac'
 181     type_ = 'flac'
 182     encoder = 'flac'
 183     decoder = 'flac'
 184     replaygain = 'metaflac'
 185     has_to_wav_pipe = True
 186     has_from_wav_pipe = True
 187     has_replaygain = True
 188     lossless = True
 189     to_wav_pipe_cmd = Expr("(decoder, '-s', '-c', '-d', infile)")
 190     to_wav_pipe_stdout = Expr("outfile")
 191     from_wav_pipe_cmd = Expr("(encoder, '-s')+encopts+(infile, )")
 192     replaygain_cmd = Expr("(replaygain, '--add-replay-gain')+files")
 193
 194 class OggVorbisCodec(Codec):
 195     ext = 'ogg'
 196     type_ = 'oggvorbis'
 197     encoder = 'oggenc'
 198     decoder = 'oggdec'
 199     replaygain = 'vorbisgain'
 200     has_to_wav_pipe = True
 201     has_from_wav_pipe = True
 202     has_replaygain = True
 203     to_wav_pipe_cmd = Expr("(decoder, '-Q', '-o', '-', infile)")
 204     to_wav_pipe_stdout = Expr("outfile")
 205     from_wav_pipe_cmd = Expr("(encoder, '-Q')+encopts+('-o', outfile, infile)")
 206     replaygain_cmd = Expr("(replaygain, '-q', '-a')+files")
 207     calc_replaygain_cmd = Expr("(replaygain, '-a', '-n', '-d')+files")
 208
 209     @classmethod
 210     def calc_replaygain(cls, files):
 211         tracks = []
 212         args = [cls.replaygain, '-and']
 213         args.extend(files)
 214         p = Popen(args=args, stdout=PIPE, stderr=PIPE)
 215         (out, err) = p.communicate()
 216         apeak = 0.0
 217         for match in re.finditer('^\s*(\S+ dB)\s*\|\s*([0-9]+)\s*\|', out,
 218             re.M):
 219             gain = match.group(1)
 220             peak = float(match.group(2)) / 32768
 221             apeak = max(apeak, peak)
 222             peak = "%.8f" % peak
 223             tracks.append((('replaygain_track_gain', gain), ('replaygain_track_peak', peak)))
 224         again = re.search('^Recommended Album Gain:\s*(\S+ dB)', err, re.M)
 225         if again:
 226             album = (('replaygain_album_gain', again.group(1)), ('replaygain_album_peak', "%.8f" % apeak))
 227         else:
 228             album = (('replaygain_album_peak', apeak),)
 229         return tracks, album
 230
 231 def transcode_track(dtask, etask, sem):
 232     etask.run()
 233     dtask.run()
 234     etask.wait()
 235     if sem:
 236         sem.release()
 237
 238 def check_and_copy_cover(fileset, targetfiles):
 239     cover_sizes = Config['cover_sizes']
 240     if not cover_sizes:
 241         return
 242     cover_out_filename = Config['cover_out_filename']
 243     if not cover_out_filename:
 244         return
 245     cover_out_filename = Format(cover_out_filename)
 246     cover_sizes = cover_sizes.split(',')
 247     covers_loaded = {}
 248     covers_written = {}
 249     outdirs = set()
 250     cover_filenames = Config['cover_filenames']
 251     if cover_filenames:
 252         cover_filenames = cover_filenames.split(',')
 253     else:
 254         cover_filenames = ()
 255     cover_out_filenames = [cover_out_filename.evaluate({'size':s}) for s in cover_sizes]
 256     for (infile, targetfile) in zip(fileset, targetfiles):
 257         outdir = os.path.split(targetfile)[0]
 258         if outdir in outdirs: continue
 259         if all(os.path.isfile(os.path.join(outdir, filename) for filename in cover_out_filenames)):
 260             outdirs.add(outdir)
 261             continue
 262         i = None
 263         for filename in (os.path.join(infile.meta['dir'], file_) for file_ in cover_filenames):
 264             try:
 265                 d = open(filename).read()
 266                 i = Image.open(StringIO(d))
 267                 i.load()
 268             except Exception:
 269                 continue
 270             if i: break
 271         if not i:
 272             tags = [(value.type, value) for key, value in infile.tags.items()
 273                 if key.startswith('APIC') and hasattr(value, 'type')
 274                 and value.type in (0, 3)]
 275             tags.sort(None, None, True)
 276             for t, value in tags:
 277                 i = None
 278                 try:
 279                     d = value.data
 280                     i = Image.open(StringIO(d))
 281                     i.load()
 282                     break
 283                 except Exception:
 284                     continue
 285         if not i: continue
 286         for s in cover_sizes:
 287             try:
 288                 s = int(s)
 289             except Exception:
 290                 continue
 291             w, h = i.size
 292             sc = 1.0*s/max(w, h)
 293             w = int(w*sc+0.5)
 294             h = int(h*sc+0.5)
 295             iw = i.resize((w, h), Image.ADAPTIVE)
 296             filename = os.path.join(
 297                outdir, cover_out_filename.evaluate({'size':s})
 298             )
 299             print "save cover %s" % filename
 300             iw.save(filename)
 301         outdirs.add(outdir)
 302
 303 rg_keys = 'replaygain_track_gain', 'replaygain_track_peak', 'replaygain_album_gain', 'replaygain_album_peak'
 304 def transcode_set(targetcodec, fileset, targetfiles, alsem, trsem, workdirs, workdirs_l):
 305     try:
 306         if not fileset:
 307             workdirs_l = None
 308             return
 309         workdirs_l.acquire()
 310         workdir, pipefiles = workdirs.pop()
 311         workdirs_l.release()
 312         outfiles = map(targetcodec._conv_out_filename, pipefiles[:len(fileset)])
 313         if targetcodec._from_wav_pipe:
 314             for i, p, o in zip(fileset, pipefiles, outfiles):
 315                 bgprocs = set()
 316                 dtask = get_codec(i).to_wav_pipe(i.meta['path'], p)
 317                 etask = targetcodec.from_wav_pipe(p, o)
 318                 # FuncTask removed
 319                 #ttask = FuncTask(background=True, target=transcode_track,
 320                    #args=(dtask, etask, trsem)
 321                 #)
 322                 if trsem:
 323                     trsem.acquire()
 324                     bgprocs.add(ttask.run())
 325                 else:
 326                     ttask.runfg()
 327             for task in bgprocs:
 328                 task.wait()
 329         elif targetcodec._from_wav_multi:
 330             etask = targetcodec.from_wav_multi(
 331                workdir, pipefiles[:len(fileset)], outfiles
 332             )
 333             etask.run()
 334             for i, o in zip(fileset, pipefiles):
 335                 task = get_codec(i).to_wav_pipe(i.meta['path'], o)
 336                 task.run()
 337             etask.wait()
 338         dirs = set()
 339         metas = []
 340         newreplaygain = False
 341         for i, o in zip(fileset, outfiles):
 342             meta = i.meta.copy()
 343             if not (i.lossless and targetcodec.lossless):
 344                 for key in rg_keys:
 345                     if key in meta:
 346                         del meta[key]
 347                 newreplaygain = True
 348             if not newreplaygain:
 349                 for key in rg_keys:
 350                     if key not in meta:
 351                         newreplaygain=True
 352                         break
 353             metas.append(meta)
 354         if newreplaygain and targetcodec._replaygain:
 355             targetcodec.add_replaygain(outfiles, metas)
 356         for i, m, o, t in zip(fileset, metas, outfiles, targetfiles):
 357             o = File(o)
 358             m.apply(o)
 359             o.save()
 360             targetdir = os.path.split(t)[0]
 361             if targetdir not in dirs:
 362                 dirs.add(targetdir)
 363                 if not os.path.isdir(targetdir):
 364                     os.makedirs(targetdir)
 365             print "%s -> %s" %(i.filename, t)
 366             util.move(o.filename, t)
 367         check_and_copy_cover(fileset, targetfiles)
 368     finally:
 369         if workdirs_l:
 370             workdirs_l.acquire()
 371             workdirs.add((workdir, pipefiles))
 372             workdirs_l.release()
 373         if alsem:
 374             alsem.release()
 375
 376 def sync_sets(sets=[], targettids=()):
 377     try:
 378         semct = int(Config['jobs'])
 379     except (ValueError, TypeError):
 380         semct = 1
 381     bgtasks = set()
 382     targetcodec = Config['type']
 383     if ',' in targetcodec:
 384         allowedcodecs = targetcodec.split(',')
 385         targetcodec = allowedcodecs[0]
 386         allowedcodecs = set(allowedcodecs)
 387     else:
 388         allowedcodecs = set((targetcodec,))
 389     targetcodec = get_codec(targetcodec)
 390     workdir = Config['workdir'] or Config['base']
 391     workdir = mkdtemp(dir=workdir, prefix='audiomangler_work_')
 392     if targetcodec._from_wav_pipe:
 393         if len(sets) > semct * 2:
 394             alsem = BoundedSemaphore(semct)
 395             trsem = None
 396         else:
 397             trsem = BoundedSemaphore(semct)
 398             alsem = None
 399     elif targetcodec._from_wav_multi:
 400         trsem = None
 401         alsem = BoundedSemaphore(semct)
 402     numpipes = max(len(s) for s in sets)
 403     workdirs = set()
 404     workdirs_l = RLock()
 405     for n in range(semct):
 406         w = os.path.join(workdir, "%02d" % n)
 407         os.mkdir(w)
 408         pipes = []
 409         for m in range(numpipes):
 410             pipes.append(os.path.join(w, "%02d.wav"%m))
 411             os.mkfifo(pipes[-1])
 412         pipes = tuple(pipes)
 413         workdirs.add((w, pipes))
 414     for fileset in sets:
 415         if all(file_.type_ in allowedcodecs for file_ in fileset):
 416             targetfiles = [f.format() for f in fileset]
 417             if not all(file_.tid in targettids for file_ in fileset):
 418                 print "copying files"
 419                 dirs = set()
 420                 for i in fileset:
 421                     t = i.format()
 422                     targetdir = os.path.split(t)[0]
 423                     if targetdir not in dirs:
 424                         dirs.add(targetdir)
 425                         if not os.path.isdir(targetdir):
 426                             os.makedirs(targetdir)
 427                     print "%s -> %s" % (i.filename, t)
 428                     util.copy(i.filename, t)
 429                 codecs = set((get_codec(f) for f in fileset))
 430                 codec = codecs.pop()
 431                 if codec and not codecs and codec._replaygain:
 432                     codec.add_replaygain(targetfiles)
 433             check_and_copy_cover(fileset, targetfiles)
 434             continue
 435         postadd = {'type':targetcodec.type_, 'ext':targetcodec.ext}
 436         targetfiles = [f.format(postadd=postadd)for f in fileset]
 437         if all(file_.tid in targettids for file_ in fileset):
 438             check_and_copy_cover(fileset, targetfiles)
 439             continue
 440         if alsem:
 441             alsem.acquire()
 442             for task in list(bgtasks):
 443                 if task.poll():
 444                     bgtasks.remove(task)
 445         # FuncTask use needs rewrite
 446         #task = FuncTask(
 447            #background=True, target=transcode_set, args=(
 448               #targetcodec, fileset, targetfiles, alsem, trsem, workdirs,
 449               #workdirs_l
 450         #))
 451         if alsem:
 452             bgtasks.add(task.run())
 453         else:
 454             task.runfg()
 455     for task in bgtasks:
 456         task.wait()
 457     for w, ps in workdirs:
 458         for p in ps:
 459             os.unlink(p)
 460         os.rmdir(w)
 461     os.rmdir(workdir)
 462
 463 def get_codec(item):
 464     if isinstance(item, FileType):
 465         item = getattr(item, 'type_')
 466     return codec_map[item]
 467
 468 __all__ = ['sync_sets', 'get_codec']