trackverify

#!/usr/bin/python

#Audio Tools, a module and set of tools for manipulating audio data
#Copyright (C) 2007-2013  Brian Langenberger

#This program is free software; you can redistribute it and/or modify
#it under the terms of the GNU General Public License as published by
#the Free Software Foundation; either version 2 of the License, or
#(at your option) any later version.

#This program 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 General Public License for more details.

#You should have received a copy of the GNU General Public License
#along with this program; if not, write to the Free Software
#Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

import sys
import os.path
import audiotools
import audiotools.text as _
from operator import or_

MAX_CPUS = audiotools.MAX_JOBS


class Counter:
    def __init__(self):
        self.value = 0

    def __int__(self):
        return self.value

    def increment(self):
        self.value += 1


class FailedAudioFile:
    def __init__(self, class_name, filename, err):
        self.NAME = class_name
        self.filename = filename
        self.err = err

    def verify(self):
        raise self.err


def open_file(filename):
    #this is much like audiotools.open
    #except that file init errors fail differently

    f = open(filename, "rb")
    try:
        audio_class = audiotools.file_type(f)
        if (((audio_class is not None) and
             (audio_class.available(audiotools.BIN)))):
            class_name = audio_class.NAME
            try:
                return audio_class(filename)
            except audiotools.InvalidFile, err:
                return FailedAudioFile(class_name, filename, err)
        else:
            raise audiotools.UnsupportedFile(filename)
    finally:
        f.close()


def get_tracks(args, queued_files, accept_list=None):
    if (accept_list is not None):
        accept_list = set(accept_list)

    for path in args:
        if (os.path.isfile(path)):
            try:
                filename = audiotools.Filename(path)
                if (filename not in queued_files):
                    queued_files.add(filename)
                    track = open_file(str(filename))
                    if ((accept_list is None) or (track.NAME in accept_list)):
                        yield track
            except (audiotools.UnsupportedFile, IOError, OSError):
                continue
        elif (os.path.isdir(path)):
            for (d, ds, fs) in os.walk(path):
                for track in get_tracks([os.path.join(d, f) for f in fs],
                                        queued_files,
                                        accept_list=accept_list):
                    yield track


def track_number(track, default):
    metadata = track.get_metadata()
    if (metadata is not None):
        if (metadata.track_number is not None):
            return metadata.track_number
        else:
            return default
    else:
        return default


def verify(progress, track):
    try:
        track.verify(progress)
        return (unicode(audiotools.Filename(track.filename)),
                track.NAME,
                None)
    except audiotools.InvalidFile, err:
        return (unicode(audiotools.Filename(track.filename)),
                track.NAME,
                unicode(err))


def display_results(result):
    (filename, track_type, error) = result
    if (error is None):
        return _.LAB_TRACKVERIFY_RESULT % {
            "path": filename,
            "result": _.LAB_TRACKVERIFY_OK}
    else:
        return _.LAB_TRACKVERIFY_RESULT % {
            "path": filename,
            "result": error}


def display_results_tty(result):
    (filename, track_type, error) = result
    if (error is None):
        return _.LAB_TRACKVERIFY_RESULT % {
            "path": filename,
            "result": audiotools.output_text(_.LAB_TRACKVERIFY_OK,
                                             fg_color="green").format(True)}
    else:
        return _.LAB_TRACKVERIFY_RESULT % {
            "path": filename,
            "result": audiotools.output_text(error,
                                             fg_color="red").format(True)}


#returned if the track isn't found in the AccurateRip database
AR_NOT_FOUND = -1

#returned if the track has one or more matches in the AccurateRip database
#but our track doesn't match any of the checksums
AR_MISMATCH = -2


def accuraterip_checksum(progress, track, track_number, track_total,
                         ar_matches):
    from audiotools._accuraterip import ChecksumV1

    reader = audiotools.PCMReaderProgress(track.to_pcm(),
                                          track.total_frames(),
                                          progress)

    metadata = track.get_metadata()
    if (metadata is not None):
        if (metadata.track_number is not None):
            is_first = (metadata.track_number == 1)
            if (metadata.track_total is not None):
                is_last = (metadata.track_number == metadata.track_total)
            else:
                is_last = (metadata.track_number == track_total)
        else:
            is_first = (track_number == 1)
            is_last = (track_number == track_total)
    else:
        is_first = (track_number == 1)
        is_last = (track_number == track_total)

    checksum = ChecksumV1(is_first,
                          is_last,
                          track.sample_rate(),
                          track.total_frames())

    try:
        audiotools.transfer_data(reader.read, checksum.update)
    except (IOError, ValueError), err:
        return (track.filename, None, None, str(err))

    checksum = checksum.checksum()

    if (len(ar_matches) > 0):
        for (confidence, ar_checksum, ar_crc2) in ar_matches:
            if (checksum == ar_checksum):
                return (track.filename, checksum, confidence, None)
        else:
            return (track.filename, checksum, AR_MISMATCH, None)
    else:
        return (track.filename, checksum, AR_NOT_FOUND, None)


def accuraterip_display_result(result):
    (filename, checksum, confidence, error) = result

    if (error is None):
        if (confidence == AR_NOT_FOUND):
            return _.LAB_TRACKVERIFY_RESULT % \
                {"path": audiotools.Filename(filename),
                 "result": _.LAB_ACCURATERIP_NOT_FOUND}
        elif (confidence == AR_MISMATCH):
            return _.LAB_TRACKVERIFY_RESULT % \
                {"path": audiotools.Filename(filename),
                 "result": _.LAB_ACCURATERIP_MISMATCH}
        else:
            return _.LAB_TRACKVERIFY_RESULT % \
                {"path": audiotools.Filename(filename),
                 "result": u"%s (%s)" %
                 (_.LAB_ACCURATERIP_FOUND,
                  _.LAB_ACCURATERIP_CONFIDENCE %
                  (confidence))}
    else:
        #FIXME - display actual error
        return u"read error"


def accuraterip_image_checksum(progress, track, track_number, track_total,
                               ar_matches,
                               displayed_filename,
                               pcm_frames_offset, total_pcm_frames):
    from audiotools._accuraterip import ChecksumV1

    reader = track.to_pcm()

    #if PCMReader has seek(), use it to reduce the amount of frames to skip
    if (hasattr(reader, "seek") and callable(reader.seek)):
        pcm_frames_offset -= reader.seek(pcm_frames_offset)

    reader = audiotools.PCMReaderProgress(
        audiotools.PCMReaderWindow(reader,
                                   pcm_frames_offset,
                                   total_pcm_frames),
        total_pcm_frames,
        progress)

    checksum = ChecksumV1(track_number == 1,
                          track_number == track_total,
                          track.sample_rate(),
                          track.total_frames())

    try:
        audiotools.transfer_data(reader.read, checksum.update)
    except (IOError, ValueError), err:
        return (displayed_filename,
                None,
                None,
                str(err))

    checksum = checksum.checksum()

    if (len(ar_matches) > 0):
        for (confidence, ar_checksum, ar_crc2) in ar_matches:
            if (checksum == ar_checksum):
                return (displayed_filename,
                        checksum,
                        confidence,
                        None)
        else:
            return (displayed_filename,
                    checksum,
                    AR_MISMATCH,
                    None)
    else:
        return (track.filename,
                None,
                AR_NOT_FOUND,
                None)



if (__name__ == '__main__'):
    import argparse

    parser = argparse.ArgumentParser(description=_.DESCRIPTION_TRACKVERIFY)

    parser.add_argument("--version",
                        action="version",
                        version="Python Audio Tools %s" % (audiotools.VERSION))

    parser.add_argument("-V", "--verbose",
                        dest="verbosity",
                        choices=audiotools.VERBOSITY_LEVELS,
                        default=audiotools.DEFAULT_VERBOSITY,
                        help=_.OPT_VERBOSE)

    parser.add_argument("-t", "--type",
                        action="append",
                        dest="accept_list",
                        metavar="type",
                        choices=audiotools.TYPE_MAP.keys(),
                        help=_.OPT_TYPE_TRACKVERIFY)

    parser.add_argument("-S", "--no-summary",
                        action="store_true",
                        dest="no_summary",
                        help=_.OPT_NO_SUMMARY)

    parser.add_argument("-R", "--accuraterip",
                        action="store_true",
                        dest="accuraterip",
                        default=False,
                        help=_.OPT_ACCURATERIP)

    parser.add_argument("-j", "--joint",
                        type=int,
                        default=MAX_CPUS,
                        dest="max_processes",
                        help=_.OPT_JOINT)

    parser.add_argument("filenames",
                        metavar="PATH",
                        nargs="+",
                        help=_.OPT_INPUT_FILENAME_OR_DIR)

    options = parser.parse_args()
    msg = audiotools.Messenger("trackverify", options)

    if (not options.accuraterip):
        queued_files = set([])  # a set of Filename objects already encountered
        queue = audiotools.ExecProgressQueue(audiotools.ProgressDisplay(msg))
        for track in get_tracks(options.filenames,
                                queued_files,
                                options.accept_list):
            queue.execute(
                function=verify,
                progress_text=unicode(audiotools.Filename(track.filename)),
                completion_output=(display_results_tty
                                   if msg.info_isatty() else
                                   display_results),
                track=track)

        msg.ansi_clearline()
        results = queue.run(options.max_processes)

        formats = sorted(list(set([r[1] for r in results])))
        success_total = len([r for r in results if r[2] is None])
        failure_total = len(results) - success_total

        if ((len(formats) > 0) and (not options.no_summary)):
            msg.output(_.LAB_TRACKVERIFY_RESULTS)
            msg.output(u"")

            table = audiotools.output_table()

            row = table.row()
            row.add_column(_.LAB_TRACKVERIFY_RESULT_FORMAT, "right")
            row.add_column(u" ")
            row.add_column(_.LAB_TRACKVERIFY_RESULT_SUCCESS, "right")
            row.add_column(u" ")
            row.add_column(_.LAB_TRACKVERIFY_RESULT_FAILURE, "right")
            row.add_column(u" ")
            row.add_column(_.LAB_TRACKVERIFY_RESULT_TOTAL, "right")

            table.divider_row([_.DIV, u" ", _.DIV, u" ", _.DIV, u" ", _.DIV])

            for format in formats:
                success = len([r for r in results if ((r[1] == format) and
                                                      (r[2] is None))])
                failure = len([r for r in results if ((r[1] == format) and
                                                      (r[2] is not None))])
                row = table.row()
                row.add_column(format.decode('ascii'), "right")
                row.add_column(u" ")
                row.add_column(unicode(success), "right")
                row.add_column(u" ")
                row.add_column(unicode(failure), "right")
                row.add_column(u" ")
                row.add_column(unicode(success + failure), "right")

            table.divider_row([_.DIV, u" ", _.DIV, u" ", _.DIV, u" ", _.DIV])
            row = table.row()
            row.add_column(_.LAB_TRACKVERIFY_SUMMARY, "right")
            row.add_column(u" ")
            row.add_column(unicode(success_total), "right")
            row.add_column(u" ")
            row.add_column(unicode(failure_total), "right")
            row.add_column(u" ")
            row.add_column(unicode(success_total + failure_total), "right")

            for row in table.format(msg.output_isatty()):
                msg.output(row)

        if (failure_total > 0):
            sys.exit(1)
    else:
        queued_files = set([])  # a set of Filename objects already encountered
        queue = audiotools.ExecProgressQueue(audiotools.ProgressDisplay(msg))

        for tracks in audiotools.group_tracks(
                get_tracks(options.filenames,
                           queued_files,
                           options.accept_list)):
            #perform AccurateRip lookup on album's worth of tracks
            #if tracks are CD formatted
            if (((set([t.channels() for t in tracks]) == set([2])) and
                 (set([t.sample_rate() for t in tracks]) == set([44100])) and
                 (set([t.bits_per_sample() for t in tracks]) == set([16])))):
                if (((len(tracks) == 1) and
                     (tracks[0].get_cuesheet() is not None))):
                    filename = audiotools.Filename(tracks[0].filename)
                    sheet = tracks[0].get_cuesheet()
                    total_frames = tracks[0].total_frames()
                    sample_rate = tracks[0].sample_rate()

                    #process tracks in CD image individually
                    ar_results = audiotools.accuraterip_sheet_lookup(
                        sheet, total_frames, sample_rate)

                    pcm_frames_offset = 0

                    for (i, pcm_frames) in enumerate(sheet.pcm_lengths(
                            total_frames, sample_rate), 1):
                        filename = "%s - track %2.2d" % \
                            (tracks[0].filename, i)
                        queue.execute(
                            function=accuraterip_image_checksum,
                            progress_text=filename,
                            completion_output=accuraterip_display_result,
                            track=tracks[0],
                            track_number=i,
                            track_total=len(tracks),
                            ar_matches=ar_results.get(i, []),
                            displayed_filename=filename,
                            pcm_frames_offset=pcm_frames_offset,
                            total_pcm_frames=pcm_frames)

                        pcm_frames_offset += pcm_frames
                else:
                    #process each track as if it were a CD
                    tracks = audiotools.sorted_tracks(tracks)
                    ar_results = audiotools.accuraterip_lookup(tracks)

                    for (i, track) in enumerate(tracks, 1):
                        filename = audiotools.Filename(track.filename)

                        queue.execute(
                            function=accuraterip_checksum,
                            progress_text=unicode(filename),
                            completion_output=accuraterip_display_result,
                            track=track,
                            track_number=i,
                            track_total=len(tracks),
                            ar_matches=ar_results.get(
                                track_number(track, i), []))
            else:
                for track in tracks:
                    msg.error(u"\"%(path)s\" %(result)s" %
                              {"path": audiotools.Filename(track.filename),
                               "result": _.ERR_TRACKVERIFY})

        msg.ansi_clearline()

        results = queue.run(options.max_processes)

        table = audiotools.output_table()

        row = table.row()
        row.add_column(_.LAB_TRACKVERIFY_AR_TRACK)
        row.add_column(u" ")
        row.add_column(_.LAB_TRACKVERIFY_AR_CHECKSUM1)
        row.add_column(u" ")
        row.add_column(_.LAB_TRACKVERIFY_AR_CONFIDENCE)

        filename_max = (msg.terminal_size(sys.stdout)[1] -
                        (sum(map(len, [u" ",
                                       audiotools.output_text(
                                _.LAB_TRACKVERIFY_AR_CHECKSUM1),
                                       u" ",
                                       audiotools.output_text(
                                _.LAB_TRACKVERIFY_AR_CONFIDENCE)]))))

        table.divider_row([_.DIV, u" ", _.DIV, u" ", _.DIV])

        for (filename, checksum, confidence, err) in results:
            row = table.row()
            #truncate long filenames if necessary
            filename = audiotools.output_text(
                audiotools.Filename(filename).basename())
            if (len(filename) <= filename_max):
                row.add_column(filename)
            else:
                row.add_column(unicode(filename.head(filename_max - 1)) +
                               u"\u2026")
            row.add_column(u"")
            if (err is None):
                row.add_column(u"%8.8X" % (checksum), "right")
                row.add_column(u"")
                row.add_column(unicode(confidence), "right")
            else:
                row.add_column(_.ERR_READ_ERROR, "right")
                row.add_column(u"")
                row.add_column(u"")

        for row in table.format(msg.output_isatty()):
            msg.output(row)

        if (len([r for r in results if r[3] is not None]) > 0):
            sys.exit(1)