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