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