<commit_before>#!/usr/bin/python

import argparse
import os.path
import cPickle
import tempfile
import random

from features import mfcc
from features import logfbank
import scipy.io.wavfile as wav
from scikits.samplerate import resample
import numpy as n
from pydub import AudioSegment
from itertools import tee, izip
from annoy import AnnoyIndex

# We can't work with files that don't have this desired_samplerate
desired_samplerate = 44100
desired_channels = 2            # We expect this for corpus files
FORCE_RESAMPLE = False          # This can be really slow

#ANN_NTREES = 100
ANN_NTREES = 10
ANN_CANDIDATES = 10

IGNORE_SAME_FRAME = True

SECONDS_OUTPUT = 30.0

SEED = 0
random.seed(SEED)

def filename_to_mfcc_frames(filename, winlen, winstep):
    samplerate = desired_samplerate
    opts = {"samplerate": samplerate,
            "winlen": winlen,
            "winstep": winstep,
            "numcep": 13,
            "nfilt": 26,
            "nfft": 512
            }
    cache_filename = filename + "." + "_".join("%s=%s" % (k, v) for k, v in sorted(opts.items())) + ".pkl"

    if not os.path.exists(cache_filename):
        print "No cached version for %s" % filename
        mfcc_feat = perform_mfcc_on_filename(filename, opts)
        cPickle.dump(mfcc_feat, open(cache_filename, "wb"))
        print "Wrote cache to %s" % cache_filename
    else:
        print "Reading cache from %s" % cache_filename
        mfcc_feat = cPickle.load(open(cache_filename, "rb"))
        if mfcc_feat is None:
            print "No MFCC for %s, perhaps has wrong samplerate" % filename
    if mfcc_feat is not None:
        print "%s has MFCC with shape %s" % (filename, repr(mfcc_feat.shape))
    return mfcc_feat

def perform_mfcc_on_filename(filename, opts):
    (samplerate, sig) = read_audio_to_numpy(filename)
    opts['samplerate'] = samplerate
    if sig.ndim > 1:
        # Mix to mono
        # TODO: Multi-channel
        nchannels = sig.shape[1]
        sig = n.mean(sig, axis=1)
    else:
        nchannels = 1
    print "Read %s with sample rate %s, #channels = %d" % (filename, samplerate, nchannels)

    if (samplerate != desired_samplerate and not FORCE_RESAMPLE):
        print "%s has the wrong samplerate, ignoring" % filename
        return None

    if (samplerate != desired_samplerate and FORCE_RESAMPLE):
        origsig = sig
        sig = resample(origsig, 1.0 * desired_samplerate/samplerate, 'sinc_best')
        print "Resampled file from rate %d to rate %d, shape %s to %s" % (samplerate, desired_samplerate, origsig.shape, sig.shape)

    mfcc_feat = mfcc(sig, **opts)
    print "Computed mfccs on %s" % filename
    return mfcc_feat

full_audiosegment_wave_cache = {}
def read_audio_to_numpy(filename):
    global full_audiosegment_wave_cache
    if filename not in full_audiosegment_wave_cache:
        (samplerate,signal) = read_audio_to_numpy_uncached(filename)
        assert samplerate == desired_samplerate
        full_audiosegment_wave_cache[filename] = (samplerate, signal)
        print "Read audio from %s" % filename
    return full_audiosegment_wave_cache[filename]

def read_audio_to_numpy_uncached(filename):
    if filename.endswith(".mp3"):
        song = AudioSegment.from_mp3(filename)
        filename = filename.replace(".mp3", ".wav")
        tmp = tempfile.NamedTemporaryFile(suffix=".wav")
        song.export(tmp.name, format="wav")
        print "Temporary export to %s" % tmp.name

        (samplerate,signal) = wav.read(tmp.name)
        tmp.close()
    else:
        assert filename.endswith(".wav")
        (samplerate,signal) = wav.read(filename)
    return (samplerate,signal)

# For the input file, find frames that are nearest in the corpus.
# Return a list of the following format:
#   (input frame start sec, input frame end sec, corpus filename, corpus frame start sec, corpus frame end sec)
annoy_mfcc_index, annoy_mfcc_list = None, None
def find_nearest_frames(input_filename, corpus_filenames, winlen, winstep):
    global annoy_mfcc_index, annoy_mfcc_list
    input_mfcc = filename_to_mfcc_frames(input_filename, winlen, winstep)
    input_nframes = input_mfcc.shape[0]
    dimension = input_mfcc.shape[1]

    corpus = []
    for corpus_filename in corpus_filenames:
        corpus_mfcc = filename_to_mfcc_frames(corpus_filename, winlen, winstep)
        if corpus_mfcc is not None:
            corpus.append((corpus_filename, corpus_mfcc))

    #Build an AnnoyIndex
    annoy_mfcc_index, annoy_mfcc_list = build_annoy_index(corpus, dimension)

    # For each frame, find the nearest frame
    dists = []
    near_frames = []
    approx_dist_error = []
    for frame_idx in range(min(int(SECONDS_OUTPUT / winstep), input_nframes)):
        this_frame = input_mfcc[frame_idx]
        (near_dist, corpus_filename, near_idx) = \
            find_nearest_frame_annoy(this_frame, input_filename, frame_idx, corpus)
        # Compute error against exhaustive
        if len(approx_dist_error) < 30 and random.random() < 0.01:
            (near_dist2, corpus_filename2, near_idx2) = \
                find_nearest_frame_exhaustive(this_frame, input_filename, frame_idx, corpus)
            assert near_dist2 <= near_dist
            approx_dist_error.append(near_dist - near_dist2)

        best_frame = (near_dist,
                      winstep * frame_idx,
                      winstep * frame_idx + winlen,
                      corpus_filename,
                      winstep * near_idx,
                      winstep * near_idx + winlen)
        dists.append(best_frame[0])
        near_frames.append(best_frame[1:])
    dists = n.array(dists)
    print "DISTANCE median=%.3f, mean=%.3f" % (n.median(dists), n.mean(dists))
    print "Average error (distance) because of approximate nearest neighbors: %.3f" % n.mean(approx_dist_error)
    return near_frames

# Exhaustive technique to find the nearest frame
# Return (distance, corpus file, corpus frame idx)
def find_nearest_frame_exhaustive(this_frame, input_filename, input_frame_idx, corpus):
    best_frames = []
    for (corpus_filename, corpus_mfcc) in corpus:
        # Don't allow it to use the same exact frame
        if input_filename == corpus_filename and IGNORE_SAME_FRAME: ignore_frame_idx = input_frame_idx
        else: ignore_frame_idx = None
        near_idx, near_dist = find_nearest_frame_for_one_with_one_corpus_file(this_frame, corpus_mfcc, ignore_frame_idx)
        best_frames.append((near_dist, corpus_filename, near_idx))
    best_frames.sort()
    return best_frames[0]

def find_nearest_frame_for_one_with_one_corpus_file(this_frame, corpus_mfcc, ignore_frame_idx):
    # Sum of squared distances (euclidean) against every frame:
    frame_dist = n.sqrt(n.square(corpus_mfcc - this_frame).sum(axis=1))
    dist_idx = [(dist, idx) for (idx, dist) in enumerate(frame_dist.tolist()) if idx != ignore_frame_idx]
    dist_idx.sort()

    near_frame_dist = dist_idx[0][0]
    near_frame_idx = dist_idx[0][1]
    return near_frame_idx, near_frame_dist

# Approx nearest neighbor technique to find the nearest frame
# Return (distance, corpus file, corpus frame idx)
def find_nearest_frame_annoy(this_frame, input_filename, input_frame_idx, corpus):
    nearest_neighbors = annoy_mfcc_index.get_nns_by_vector(this_frame.tolist(), ANN_CANDIDATES)
    candidates = []
    for nearest_neighbor in nearest_neighbors:
        corpus_filename, near_idx = annoy_mfcc_list[nearest_neighbor]
        if IGNORE_SAME_FRAME and corpus_filename == input_filename and near_idx == input_frame_idx: continue
        corpus_mfcc = None
        for (filename, mfcc) in corpus:
            if filename == corpus_filename: corpus_mfcc = mfcc
        near_dist = n.sqrt(n.square(corpus_mfcc[near_idx] - this_frame).sum())
        candidates.append((near_dist, corpus_filename, near_idx))
    candidates.sort()
    # If nothing is returned, try ANN_CANDIDATES > 1
    # Otherwise, we post filter and disallow you to return the frame being used to search
    return candidates[0]

def build_annoy_index(corpus, dimension):
    print "Adding to Annoy index"
    index = AnnoyIndex(dimension, "euclidean")
    mfcc_list = []
    i = 0
    for filename, frames in corpus:
        print filename, frames.shape
        for index_in_file, mfcc in enumerate(frames):
            mfcc_list.append((filename, index_in_file))
            index.add_item(i, mfcc.tolist())
            assert mfcc_list[i] == (filename, index_in_file)
            i += 1
    print "Building Annoy index with %d trees" % ANN_NTREES
#    index.build(-1)
    index.build(ANN_NTREES)
    return index, mfcc_list

def get_audiosegment_wave(filename, start, end):
    return read_audio_to_numpy(filename)[1][start:end]

def sec2sample(sec):
    return int(sec * desired_samplerate + 0.5)

def redub_overlay_wave(frame_locations, output_filename):
    print ""
    print "Beginning redub_overlay_wave to %s with %d frames" % (output_filename, len(frame_locations))
    start_points = set(sec2sample(frame[0]) for frame in frame_locations)
    end_points = set(sec2sample(frame[1]) for frame in frame_locations)
    cut_points = sorted(start_points.union(end_points))
    cuts = [x for x in window(cut_points, 2)]

    updated_frame_locations = []
    for (write_start_sec, write_end_sec, corpus_filename, corpus_start_sec, corpus_end_sec) in frame_locations:
        write_start = sec2sample(write_start_sec)
        write_end = sec2sample(write_end_sec)
        corpus_start = sec2sample(corpus_start_sec)
        corpus_end = sec2sample(corpus_end_sec)
        updated_frame_locations.append((write_start, write_end, corpus_filename, corpus_start, corpus_end))
    updated_frame_locations.sort()
    assert(len(updated_frame_locations) == len(frame_locations))
    frame_locations = updated_frame_locations

    fragments = []
    for (j, (cut_start, cut_end)) in enumerate(cuts):
        if j % 100 == 0:
            print "Making cut %d of %d" % (j, len(cuts))
        cut_length = cut_end - cut_start
        this_fragments = []
        # TODO: this nested loop can be a bit slow, but we're always searching in
        #       one direction. We could speed this up with some trickery.
        for (write_start, write_end, corpus_filename, corpus_start, corpus_end) in frame_locations:
            if write_start >= cut_end or write_end <= cut_start:
                continue

            desired_cut_start = max(write_start, cut_start)
            desired_cut_end   = min(write_end, cut_end)
            assert desired_cut_end >= desired_cut_start

            if (desired_cut_end - desired_cut_start) != cut_length:
                print "Weird. Skipping this cut, can't get the right size"
                continue

            actual_start = corpus_start + (desired_cut_start - write_start)
            actual_end   = actual_start + desired_cut_end - desired_cut_start
            assert actual_end >= actual_start

            segment = get_audiosegment_wave(corpus_filename, actual_start, actual_end)
            this_fragments.append(segment)

        to_avg = [n.zeros((cut_length, desired_channels))]
        for f in this_fragments:
            assert f.shape[1] == desired_channels
            if f.shape[0] != cut_length:
                print "Weird. Extracted cut of the wrong size. Expected = %d, received = %d" % (cut_length, f.shape[0])
                continue
            else:
                to_avg.append(f)
        total = to_avg[0].astype("float")
        for t in to_avg[1:]:
            assert(total.shape == t.shape)
            total += t.astype("float")
        total /= len(to_avg)
        assert(total.shape[0] == cut_length)
        fragments.append(total.astype("int16"))

    newsong = n.vstack(fragments)
    print "Composed %d fragments into %s" % (len(fragments), newsong.shape)
    wav.write(filename = output_filename, rate = desired_samplerate, data = newsong)
    print "Wrote new song to %s" % output_filename

def window(iterable, size):
    iters = tee(iterable, size)
    for i in xrange(1, size):
        for each in iters[i:]:
            next(each, None)
    return izip(*iters)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Aucode a sound.')
    parser.add_argument('-i', '--input', help='Input audio signal to be covered (mp3)')
    parser.add_argument('-o', '--output', help='Output filename (wav)')
    parser.add_argument('--winlen', default=250, help='Frame length, in ms')
    parser.add_argument('--winstep', help='Frame step, in ms (= frame length by default)')
    parser.add_argument('-c', '--corpus', help='Audio file(s) to use as samples (mp3)', nargs='*')

    args = parser.parse_args()
    winlen = float(args.winlen) / 1000.0
    winstep = float(args.winstep or args.winlen) / 1000.0

    assert args.input.endswith(".mp3")
    for c in args.corpus:
        assert c.endswith(".mp3")
    assert args.output.endswith(".wav")

    frame_locations = find_nearest_frames(args.input, args.corpus, winlen, winstep)
    redub_overlay_wave(frame_locations, args.output)
<commit_msg>More info<commit_after>#!/usr/bin/python

import argparse
import os.path
import cPickle
import tempfile
import random

from features import mfcc
from features import logfbank
import scipy.io.wavfile as wav
from scikits.samplerate import resample
import numpy as n
from pydub import AudioSegment
from itertools import tee, izip
from annoy import AnnoyIndex

# We can't work with files that don't have this desired_samplerate
desired_samplerate = 44100
desired_channels = 2            # We expect this for corpus files
FORCE_RESAMPLE = False          # This can be really slow

#ANN_NTREES = 100
ANN_NTREES = 10
ANN_CANDIDATES = 10

IGNORE_SAME_FRAME = True

SECONDS_OUTPUT = 30.0

SEED = 0
random.seed(SEED)

def filename_to_mfcc_frames(filename, winlen, winstep):
    samplerate = desired_samplerate
    opts = {"samplerate": samplerate,
            "winlen": winlen,
            "winstep": winstep,
            "numcep": 13,
            "nfilt": 26,
            "nfft": 512
            }
    cache_filename = filename + "." + "_".join("%s=%s" % (k, v) for k, v in sorted(opts.items())) + ".pkl"

    if not os.path.exists(cache_filename):
        print "No cached version for %s" % filename
        mfcc_feat = perform_mfcc_on_filename(filename, opts)
        cPickle.dump(mfcc_feat, open(cache_filename, "wb"))
        print "Wrote cache to %s" % cache_filename
    else:
        print "Reading cache from %s" % cache_filename
        mfcc_feat = cPickle.load(open(cache_filename, "rb"))
        if mfcc_feat is None:
            print "No MFCC for %s, perhaps has wrong samplerate" % filename
    if mfcc_feat is not None:
        print "%s has MFCC with shape %s" % (filename, repr(mfcc_feat.shape))
    return mfcc_feat

def perform_mfcc_on_filename(filename, opts):
    (samplerate, sig) = read_audio_to_numpy(filename)
    opts['samplerate'] = samplerate
    if sig.ndim > 1:
        # Mix to mono
        # TODO: Multi-channel
        nchannels = sig.shape[1]
        sig = n.mean(sig, axis=1)
    else:
        nchannels = 1
    print "Read %s with sample rate %s, #channels = %d" % (filename, samplerate, nchannels)

    if (samplerate != desired_samplerate and not FORCE_RESAMPLE):
        print "%s has the wrong samplerate, ignoring" % filename
        return None

    if (samplerate != desired_samplerate and FORCE_RESAMPLE):
        origsig = sig
        sig = resample(origsig, 1.0 * desired_samplerate/samplerate, 'sinc_best')
        print "Resampled file from rate %d to rate %d, shape %s to %s" % (samplerate, desired_samplerate, origsig.shape, sig.shape)

    mfcc_feat = mfcc(sig, **opts)
    print "Computed mfccs on %s" % filename
    return mfcc_feat

full_audiosegment_wave_cache = {}
def read_audio_to_numpy(filename):
    global full_audiosegment_wave_cache
    if filename not in full_audiosegment_wave_cache:
        (samplerate,signal) = read_audio_to_numpy_uncached(filename)
        assert samplerate == desired_samplerate
        full_audiosegment_wave_cache[filename] = (samplerate, signal)
        print "Read audio from %s" % filename
    return full_audiosegment_wave_cache[filename]

def read_audio_to_numpy_uncached(filename):
    if filename.endswith(".mp3"):
        song = AudioSegment.from_mp3(filename)
        filename = filename.replace(".mp3", ".wav")
        tmp = tempfile.NamedTemporaryFile(suffix=".wav")
        song.export(tmp.name, format="wav")
        print "Temporary export to %s" % tmp.name

        (samplerate,signal) = wav.read(tmp.name)
        tmp.close()
    else:
        assert filename.endswith(".wav")
        (samplerate,signal) = wav.read(filename)
    return (samplerate,signal)

# For the input file, find frames that are nearest in the corpus.
# Return a list of the following format:
#   (input frame start sec, input frame end sec, corpus filename, corpus frame start sec, corpus frame end sec)
annoy_mfcc_index, annoy_mfcc_list = None, None
def find_nearest_frames(input_filename, corpus_filenames, winlen, winstep):
    global annoy_mfcc_index, annoy_mfcc_list
    input_mfcc = filename_to_mfcc_frames(input_filename, winlen, winstep)
    input_nframes = input_mfcc.shape[0]
    dimension = input_mfcc.shape[1]

    corpus = []
    total_length = 0
    for corpus_filename in corpus_filenames:
        corpus_mfcc = filename_to_mfcc_frames(corpus_filename, winlen, winstep)
        total_length += corpus_mfcc.shape[0]
        if corpus_mfcc is not None:
            corpus.append((corpus_filename, corpus_mfcc))
    print "Corpus has %.2f hr, %d frames" % (total_length * winstep / 60 / 60, total_length)

    #Build an AnnoyIndex
    annoy_mfcc_index, annoy_mfcc_list = build_annoy_index(corpus, dimension)

    # For each frame, find the nearest frame
    dists = []
    near_frames = []
    approx_dist_error = []
    for frame_idx in range(min(int(SECONDS_OUTPUT / winstep), input_nframes)):
        this_frame = input_mfcc[frame_idx]
        (near_dist, corpus_filename, near_idx) = \
            find_nearest_frame_annoy(this_frame, input_filename, frame_idx, corpus)
        # Compute error against exhaustive
        if len(approx_dist_error) < 30 and random.random() < 0.01:
            (near_dist2, corpus_filename2, near_idx2) = \
                find_nearest_frame_exhaustive(this_frame, input_filename, frame_idx, corpus)
            assert near_dist2 <= near_dist
            approx_dist_error.append(near_dist - near_dist2)

        best_frame = (near_dist,
                      winstep * frame_idx,
                      winstep * frame_idx + winlen,
                      corpus_filename,
                      winstep * near_idx,
                      winstep * near_idx + winlen)
        dists.append(best_frame[0])
        near_frames.append(best_frame[1:])
    dists = n.array(dists)
    print "DISTANCE median=%.3f, mean=%.3f" % (n.median(dists), n.mean(dists))
    print "Average error (distance) because of approximate nearest neighbors: %.3f" % n.mean(approx_dist_error)
    return near_frames

# Exhaustive technique to find the nearest frame
# Return (distance, corpus file, corpus frame idx)
def find_nearest_frame_exhaustive(this_frame, input_filename, input_frame_idx, corpus):
    best_frames = []
    for (corpus_filename, corpus_mfcc) in corpus:
        # Don't allow it to use the same exact frame
        if input_filename == corpus_filename and IGNORE_SAME_FRAME: ignore_frame_idx = input_frame_idx
        else: ignore_frame_idx = None
        near_idx, near_dist = find_nearest_frame_for_one_with_one_corpus_file(this_frame, corpus_mfcc, ignore_frame_idx)
        best_frames.append((near_dist, corpus_filename, near_idx))
    best_frames.sort()
    return best_frames[0]

def find_nearest_frame_for_one_with_one_corpus_file(this_frame, corpus_mfcc, ignore_frame_idx):
    # Sum of squared distances (euclidean) against every frame:
    frame_dist = n.sqrt(n.square(corpus_mfcc - this_frame).sum(axis=1))
    dist_idx = [(dist, idx) for (idx, dist) in enumerate(frame_dist.tolist()) if idx != ignore_frame_idx]
    dist_idx.sort()

    near_frame_dist = dist_idx[0][0]
    near_frame_idx = dist_idx[0][1]
    return near_frame_idx, near_frame_dist

# Approx nearest neighbor technique to find the nearest frame
# Return (distance, corpus file, corpus frame idx)
def find_nearest_frame_annoy(this_frame, input_filename, input_frame_idx, corpus):
    nearest_neighbors = annoy_mfcc_index.get_nns_by_vector(this_frame.tolist(), ANN_CANDIDATES)
    candidates = []
    for nearest_neighbor in nearest_neighbors:
        corpus_filename, near_idx = annoy_mfcc_list[nearest_neighbor]
        if IGNORE_SAME_FRAME and corpus_filename == input_filename and near_idx == input_frame_idx: continue
        corpus_mfcc = None
        for (filename, mfcc) in corpus:
            if filename == corpus_filename: corpus_mfcc = mfcc
        near_dist = n.sqrt(n.square(corpus_mfcc[near_idx] - this_frame).sum())
        candidates.append((near_dist, corpus_filename, near_idx))
    candidates.sort()
    # If nothing is returned, try ANN_CANDIDATES > 1
    # Otherwise, we post filter and disallow you to return the frame being used to search
    return candidates[0]

def build_annoy_index(corpus, dimension):
    print "Adding to Annoy index"
    index = AnnoyIndex(dimension, "euclidean")
    mfcc_list = []
    i = 0
    for filename, frames in corpus:
        print filename, frames.shape
        for index_in_file, mfcc in enumerate(frames):
            mfcc_list.append((filename, index_in_file))
            index.add_item(i, mfcc.tolist())
            assert mfcc_list[i] == (filename, index_in_file)
            i += 1
    print "Building Annoy index with %d trees" % ANN_NTREES
#    index.build(-1)
    index.build(ANN_NTREES)
    return index, mfcc_list

def get_audiosegment_wave(filename, start, end):
    return read_audio_to_numpy(filename)[1][start:end]

def sec2sample(sec):
    return int(sec * desired_samplerate + 0.5)

def redub_overlay_wave(frame_locations, output_filename):
    print ""
    print "Beginning redub_overlay_wave to %s with %d frames" % (output_filename, len(frame_locations))
    start_points = set(sec2sample(frame[0]) for frame in frame_locations)
    end_points = set(sec2sample(frame[1]) for frame in frame_locations)
    cut_points = sorted(start_points.union(end_points))
    cuts = [x for x in window(cut_points, 2)]

    updated_frame_locations = []
    for (write_start_sec, write_end_sec, corpus_filename, corpus_start_sec, corpus_end_sec) in frame_locations:
        write_start = sec2sample(write_start_sec)
        write_end = sec2sample(write_end_sec)
        corpus_start = sec2sample(corpus_start_sec)
        corpus_end = sec2sample(corpus_end_sec)
        updated_frame_locations.append((write_start, write_end, corpus_filename, corpus_start, corpus_end))
    updated_frame_locations.sort()
    assert(len(updated_frame_locations) == len(frame_locations))
    frame_locations = updated_frame_locations

    fragments = []
    for (j, (cut_start, cut_end)) in enumerate(cuts):
        if j % 100 == 0:
            print "Making cut %d of %d" % (j, len(cuts))
        cut_length = cut_end - cut_start
        this_fragments = []
        # TODO: this nested loop can be a bit slow, but we're always searching in
        #       one direction. We could speed this up with some trickery.
        for (write_start, write_end, corpus_filename, corpus_start, corpus_end) in frame_locations:
            if write_start >= cut_end or write_end <= cut_start:
                continue

            desired_cut_start = max(write_start, cut_start)
            desired_cut_end   = min(write_end, cut_end)
            assert desired_cut_end >= desired_cut_start

            if (desired_cut_end - desired_cut_start) != cut_length:
                print "Weird. Skipping this cut, can't get the right size"
                continue

            actual_start = corpus_start + (desired_cut_start - write_start)
            actual_end   = actual_start + desired_cut_end - desired_cut_start
            assert actual_end >= actual_start

            segment = get_audiosegment_wave(corpus_filename, actual_start, actual_end)
            this_fragments.append(segment)

        to_avg = [n.zeros((cut_length, desired_channels))]
        for f in this_fragments:
            assert f.shape[1] == desired_channels
            if f.shape[0] != cut_length:
                print "Weird. Extracted cut of the wrong size. Expected = %d, received = %d" % (cut_length, f.shape[0])
                continue
            else:
                to_avg.append(f)
        total = to_avg[0].astype("float")
        for t in to_avg[1:]:
            assert(total.shape == t.shape)
            total += t.astype("float")
        total /= len(to_avg)
        assert(total.shape[0] == cut_length)
        fragments.append(total.astype("int16"))

    newsong = n.vstack(fragments)
    print "Composed %d fragments into %s" % (len(fragments), newsong.shape)
    wav.write(filename = output_filename, rate = desired_samplerate, data = newsong)
    print "Wrote new song to %s" % output_filename

def window(iterable, size):
    iters = tee(iterable, size)
    for i in xrange(1, size):
        for each in iters[i:]:
            next(each, None)
    return izip(*iters)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Aucode a sound.')
    parser.add_argument('-i', '--input', help='Input audio signal to be covered (mp3)')
    parser.add_argument('-o', '--output', help='Output filename (wav)')
    parser.add_argument('--winlen', default=250, help='Frame length, in ms')
    parser.add_argument('--winstep', help='Frame step, in ms (= frame length by default)')
    parser.add_argument('-c', '--corpus', help='Audio file(s) to use as samples (mp3)', nargs='*')

    args = parser.parse_args()
    winlen = float(args.winlen) / 1000.0
    winstep = float(args.winstep or args.winlen) / 1000.0

    assert args.input.endswith(".mp3")
    for c in args.corpus:
        assert c.endswith(".mp3")
    assert args.output.endswith(".wav")

    frame_locations = find_nearest_frames(args.input, args.corpus, winlen, winstep)
    redub_overlay_wave(frame_locations, args.output)
<|endoftext|>


<commit_before>#!/usr/bin/env python

"""
.. See the NOTICE file distributed with this work for additional information
   regarding copyright ownership.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
"""

from __future__ import print_function

# Required for ReadTheDocs
from functools import wraps # pylint: disable=unused-import

import argparse
import os.path

from basic_modules.workflow import Workflow
from basic_modules.metadata import Metadata
from utils import remap

from tool.bwa_aligner import bwaAlignerTool
from tool.biobambam_filter import biobambam
from tool.macs2 import macs2

# ------------------------------------------------------------------------------

class process_chipseq(Workflow):
    """
    Functions for processing Chip-Seq FastQ files. Files are the aligned,
    filtered and analysed for peak calling
    """

    def __init__(self, configuration=None):
        """
        Initialise the tool with its configuration.

        Parameters
        ----------
        configuration : dict
            a dictionary containing parameters that define how the operation
            should be carried out, which are specific to each Tool.
        """
        if configuration is None:
            configuration = {}

        self.configuration.update(configuration)

    def run(self, input_files, metadata, output_files):
        """
        Main run function for processing ChIP-seq FastQ data. Pipeline aligns
        the FASTQ files to the genome using BWA. MACS 2 is then used for peak
        calling to identify transcription factor binding sites within the
        genome.

        Currently this can only handle a single data file and a single
        background file.

        Parameters
        ----------
        input_files : dict
            Location of the initial input files required by the workflow
            genome : str
                Genome FASTA file
            amb : str
                BWA index file
            ann : str
                BWA index file
            bwt : str
                BWA index file
            pac : str
                BWA index file
            sa : str
                BWA index file
            loc : str
                Location of the FASTQ reads files
            bg_loc : str
                Location of the background FASTQ reads files [OPTIONAL]
        metadata : dict
            Input file meta data associated with their roles

            genome : str
            amb : str
            ann : str
            bwt : str
            pac : str
            sa : str
            loc : str
            bg_loc : str
                [OPTIONAL]
        output_files : dict
            Output file locations

            bam [, "bam_bg"] : str
            filtered [, "filtered_bg"] : str
            narrow_peak : str
            summits : str
            broad_peak : str
            gapped_peak : str

        Returns
        -------
        output_files : dict
            Output file locations associated with their roles, for the output

            bam [, "bam_bg"] : str
                Aligned FASTQ short read file [ and aligned background file]
                locations
            filtered [, "filtered_bg"] : str
                Filtered versions of the respective bam files
            narrow_peak : str
                Results files in bed4+1 format
            summits : str
                Results files in bed6+4 format
            broad_peak : str
                Results files in bed6+3 format
            gapped_peak : str
                Results files in bed12+3 format
        output_metadata : dict
            Output metadata for the associated files in output_files

            bam [, "bam_bg"] : Metadata
            filtered [, "filtered_bg"] : Metadata
            narrow_peak : Metadata
            summits : Metadata
            broad_peak : Metadata
            gapped_peak : Metadata
        """
        print("INPUT RUN METADATA:", metadata)

        bwa = bwaAlignerTool(self.configuration)

        print(
            "PROCESS CHIPSEQ - FILES PASSED TO TOOLS:",
            remap(input_files, "genome", "loc", "index")
        )

        print("PROCESS CHIPSEQ - DEFINED OUTPUT:", output_files["bam"])
        bwa_files, bwa_meta = bwa.run(
            # ideally parameter "roles" don't change
            remap(input_files,
                  "genome", "loc", "index"),
            remap(metadata,
                  "genome", "loc", "index"),
            {"output": output_files["bam"]}
        )

        if "bg_loc" in input_files:
            bwa_bg_files, bwa_bg_meta = bwa.run(
                # Small changes can be handled easily using "remap"
                remap(input_files,
                      "genome", "index", loc="bg_loc"),
                remap(metadata,
                      "genome", "index", loc="bg_loc"),
                # Intermediate outputs should be created via tempfile?
                {"output": output_files["bam_bg"]}
            )

        # For multiple files there will need to be merging into a single bam file

        # Filter the bams
        b3f = biobambam(self.configuration)
        b3f_files, b3f_meta = b3f.run(
            # Alternatively, we can rebuild the dict
            {"input": bwa_files['bam']},
            {"input": bwa_meta['bam']},
            # Intermediate outputs should be created via tempfile?
            {"output": output_files["filtered"]}
        )

        if "bg_loc" in input_files:
            b3f_bg_files, b3f_bg_meta = b3f.run(
                {"input": bwa_bg_files['bam']},
                {"input": bwa_bg_meta['bam']},
                # Intermediate outputs should be created via tempfile?
                {"output": output_files["filtered_bg"]}
            )

        ## MACS2 to call peaks
        macs_caller = macs2(self.configuration)
        macs_inputs = {"input": b3f_files['bam']}
        macs_metadt = {"input": b3f_meta['bam']}

        if "bg_loc" in input_files:
            # The dicts can be built incrementally
            macs_inputs["background"] = b3f_bg_files['bam']
            macs_metadt["background"] = b3f_bg_meta['bam']

        m_results_files, m_results_meta = macs_caller.run(
            macs_inputs, macs_metadt,
            # Outputs of the final step may match workflow outputs;
            # Extra entries in output_files will be disregarded.
            remap(output_files, 'narrow_peak', 'summits', 'broad_peak', 'gapped_peak'))

        m_results_meta['narrow_peak'].meta_data['tool_description'] = m_results_meta['narrow_peak'].meta_data['tool']
        m_results_meta['narrow_peak'].meta_data['tool'] = "process_chipseq"
        m_results_meta['summits'].meta_data['tool_description'] = m_results_meta['summits'].meta_data['tool']
        m_results_meta['summits'].meta_data['tool'] = "process_chipseq"
        m_results_meta['broad_peak'].meta_data['tool_description'] = m_results_meta['broad_peak'].meta_data['tool']
        m_results_meta['broad_peak'].meta_data['tool'] = "process_chipseq"
        m_results_meta['gapped_peak'].meta_data['tool_description'] = m_results_meta['gapped_peak'].meta_data['tool']
        m_results_meta['gapped_peak'].meta_data['tool'] = "process_chipseq"

        # Outputs are collected with some name changes
        m_results_files["bam"] = bwa_files["bam"]
        m_results_files["filtered"] = b3f_files["bam"]

        m_results_meta['bam'].meta_data['tool_description'] = m_results_meta['bam'].meta_data['tool']
        m_results_meta['bam'].meta_data['tool'] = "process_chipseq"
        m_results_meta['filtered'].meta_data['tool_description'] = m_results_meta['filtered'].meta_data['tool']
        m_results_meta['filtered'].meta_data['tool'] = "process_chipseq"

        # Equivalent meta data is collected
        m_results_meta["bam"] = bwa_meta["bam"]
        m_results_meta["filtered"] = b3f_meta["bam"]

        if "bg_loc" in input_files:
            m_results_files["bam_bg"] = bwa_bg_files["bam"]
            m_results_files["filtered_bg"] = b3f_bg_files["bam"]

            m_results_meta["bam_bg"] = bwa_bg_meta["bam"]
            m_results_meta["filtered_bg"] = b3f_bg_meta["bam"]

            m_results_meta['bam_bg'].meta_data['tool_description'] = m_results_meta['bam_bg'].meta_data['tool']
            m_results_meta['bam_bg'].meta_data['tool'] = "process_chipseq"
            m_results_meta['filtered_bg'].meta_data['tool_description'] = m_results_meta['filtered_bg'].meta_data['tool']
            m_results_meta['filtered_bg'].meta_data['tool'] = "process_chipseq"

        print("CHIPSEQ RESULTS:", m_results_meta)
        return m_results_files, m_results_meta

# ------------------------------------------------------------------------------

def main_json(config, in_metadata, out_metadata):
    """
    Alternative main function
    -------------

    This function launches the app using configuration written in
    two json files: config.json and input_metadata.json.
    """
    # 1. Instantiate and launch the App
    print("1. Instantiate and launch the App")
    from apps.jsonapp import JSONApp
    app = JSONApp()
    result = app.launch(process_chipseq,
                        config,
                        in_metadata,
                        out_metadata)

    # 2. The App has finished
    print("2. Execution finished; see " + out_metadata)
    print(result)

    return result

# ------------------------------------------------------------------------------

if __name__ == "__main__":
    import sys
    sys._run_from_cmdl = True  # pylint: disable=protected-access

    # Set up the command line parameters
    PARSER = argparse.ArgumentParser(description="ChIP-seq peak calling")
    PARSER.add_argument("--config", help="Configuration file")
    PARSER.add_argument("--in_metadata", help="Location of input metadata file")
    PARSER.add_argument("--out_metadata", help="Location of output metadata file")

    # Get the matching parameters from the command line
    ARGS = PARSER.parse_args()

    CONFIG = ARGS.config
    IN_METADATA = ARGS.in_metadata
    OUT_METADATA = ARGS.out_metadata

    RESULTS = main_json(CONFIG, IN_METADATA, OUT_METADATA)

    print(RESULTS)
<commit_msg>Fix for the missing files<commit_after>#!/usr/bin/env python

"""
.. See the NOTICE file distributed with this work for additional information
   regarding copyright ownership.

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
"""

from __future__ import print_function

# Required for ReadTheDocs
from functools import wraps # pylint: disable=unused-import

import argparse
import os.path

from basic_modules.workflow import Workflow
from basic_modules.metadata import Metadata
from utils import remap

from tool.bwa_aligner import bwaAlignerTool
from tool.biobambam_filter import biobambam
from tool.macs2 import macs2

# ------------------------------------------------------------------------------

class process_chipseq(Workflow):
    """
    Functions for processing Chip-Seq FastQ files. Files are the aligned,
    filtered and analysed for peak calling
    """

    def __init__(self, configuration=None):
        """
        Initialise the tool with its configuration.

        Parameters
        ----------
        configuration : dict
            a dictionary containing parameters that define how the operation
            should be carried out, which are specific to each Tool.
        """
        if configuration is None:
            configuration = {}

        self.configuration.update(configuration)

    def run(self, input_files, metadata, output_files):
        """
        Main run function for processing ChIP-seq FastQ data. Pipeline aligns
        the FASTQ files to the genome using BWA. MACS 2 is then used for peak
        calling to identify transcription factor binding sites within the
        genome.

        Currently this can only handle a single data file and a single
        background file.

        Parameters
        ----------
        input_files : dict
            Location of the initial input files required by the workflow
            genome : str
                Genome FASTA file
            amb : str
                BWA index file
            ann : str
                BWA index file
            bwt : str
                BWA index file
            pac : str
                BWA index file
            sa : str
                BWA index file
            loc : str
                Location of the FASTQ reads files
            bg_loc : str
                Location of the background FASTQ reads files [OPTIONAL]
        metadata : dict
            Input file meta data associated with their roles

            genome : str
            amb : str
            ann : str
            bwt : str
            pac : str
            sa : str
            loc : str
            bg_loc : str
                [OPTIONAL]
        output_files : dict
            Output file locations

            bam [, "bam_bg"] : str
            filtered [, "filtered_bg"] : str
            narrow_peak : str
            summits : str
            broad_peak : str
            gapped_peak : str

        Returns
        -------
        output_files : dict
            Output file locations associated with their roles, for the output

            bam [, "bam_bg"] : str
                Aligned FASTQ short read file [ and aligned background file]
                locations
            filtered [, "filtered_bg"] : str
                Filtered versions of the respective bam files
            narrow_peak : str
                Results files in bed4+1 format
            summits : str
                Results files in bed6+4 format
            broad_peak : str
                Results files in bed6+3 format
            gapped_peak : str
                Results files in bed12+3 format
        output_metadata : dict
            Output metadata for the associated files in output_files

            bam [, "bam_bg"] : Metadata
            filtered [, "filtered_bg"] : Metadata
            narrow_peak : Metadata
            summits : Metadata
            broad_peak : Metadata
            gapped_peak : Metadata
        """
        print("INPUT RUN METADATA:", metadata)

        bwa = bwaAlignerTool(self.configuration)

        print(
            "PROCESS CHIPSEQ - FILES PASSED TO TOOLS:",
            remap(input_files, "genome", "loc", "index")
        )

        print("PROCESS CHIPSEQ - DEFINED OUTPUT:", output_files["bam"])
        bwa_files, bwa_meta = bwa.run(
            # ideally parameter "roles" don't change
            remap(input_files,
                  "genome", "loc", "index"),
            remap(metadata,
                  "genome", "loc", "index"),
            {"output": output_files["bam"]}
        )

        if "bg_loc" in input_files:
            bwa_bg_files, bwa_bg_meta = bwa.run(
                # Small changes can be handled easily using "remap"
                remap(input_files,
                      "genome", "index", loc="bg_loc"),
                remap(metadata,
                      "genome", "index", loc="bg_loc"),
                # Intermediate outputs should be created via tempfile?
                {"output": output_files["bam_bg"]}
            )

        # For multiple files there will need to be merging into a single bam file

        # Filter the bams
        b3f = biobambam(self.configuration)
        b3f_files, b3f_meta = b3f.run(
            # Alternatively, we can rebuild the dict
            {"input": bwa_files['bam']},
            {"input": bwa_meta['bam']},
            # Intermediate outputs should be created via tempfile?
            {"output": output_files["filtered"]}
        )

        if "bg_loc" in input_files:
            b3f_bg_files, b3f_bg_meta = b3f.run(
                {"input": bwa_bg_files['bam']},
                {"input": bwa_bg_meta['bam']},
                # Intermediate outputs should be created via tempfile?
                {"output": output_files["filtered_bg"]}
            )

        ## MACS2 to call peaks
        macs_caller = macs2(self.configuration)
        macs_inputs = {"input": b3f_files['bam']}
        macs_metadt = {"input": b3f_meta['bam']}

        if "bg_loc" in input_files:
            # The dicts can be built incrementally
            macs_inputs["background"] = b3f_bg_files['bam']
            macs_metadt["background"] = b3f_bg_meta['bam']

        m_results_files, m_results_meta = macs_caller.run(
            macs_inputs, macs_metadt,
            # Outputs of the final step may match workflow outputs;
            # Extra entries in output_files will be disregarded.
            remap(output_files, 'narrow_peak', 'summits', 'broad_peak', 'gapped_peak'))

        if 'narrow_peak' in m_results_meta:
            m_results_meta['narrow_peak'].meta_data['tool_description'] = m_results_meta['narrow_peak'].meta_data['tool']
            m_results_meta['narrow_peak'].meta_data['tool'] = "process_chipseq"
        if 'summits' in m_results_meta:
            m_results_meta['summits'].meta_data['tool_description'] = m_results_meta['summits'].meta_data['tool']
            m_results_meta['summits'].meta_data['tool'] = "process_chipseq"
        if 'broad_peak' in m_results_meta:
            m_results_meta['broad_peak'].meta_data['tool_description'] = m_results_meta['broad_peak'].meta_data['tool']
            m_results_meta['broad_peak'].meta_data['tool'] = "process_chipseq"
        if 'gapped_peak' in m_results_meta:
            m_results_meta['gapped_peak'].meta_data['tool_description'] = m_results_meta['gapped_peak'].meta_data['tool']
            m_results_meta['gapped_peak'].meta_data['tool'] = "process_chipseq"

        # Outputs are collected with some name changes
        m_results_files["bam"] = bwa_files["bam"]
        m_results_files["filtered"] = b3f_files["bam"]

        m_results_meta['bam'].meta_data['tool_description'] = m_results_meta['bam'].meta_data['tool']
        m_results_meta['bam'].meta_data['tool'] = "process_chipseq"
        m_results_meta['filtered'].meta_data['tool_description'] = m_results_meta['filtered'].meta_data['tool']
        m_results_meta['filtered'].meta_data['tool'] = "process_chipseq"

        # Equivalent meta data is collected
        m_results_meta["bam"] = bwa_meta["bam"]
        m_results_meta["filtered"] = b3f_meta["bam"]

        if "bg_loc" in input_files:
            m_results_files["bam_bg"] = bwa_bg_files["bam"]
            m_results_files["filtered_bg"] = b3f_bg_files["bam"]

            m_results_meta["bam_bg"] = bwa_bg_meta["bam"]
            m_results_meta["filtered_bg"] = b3f_bg_meta["bam"]

            m_results_meta['bam_bg'].meta_data['tool_description'] = m_results_meta['bam_bg'].meta_data['tool']
            m_results_meta['bam_bg'].meta_data['tool'] = "process_chipseq"
            m_results_meta['filtered_bg'].meta_data['tool_description'] = m_results_meta['filtered_bg'].meta_data['tool']
            m_results_meta['filtered_bg'].meta_data['tool'] = "process_chipseq"

        print("CHIPSEQ RESULTS:", m_results_meta)
        return m_results_files, m_results_meta

# ------------------------------------------------------------------------------

def main_json(config, in_metadata, out_metadata):
    """
    Alternative main function
    -------------

    This function launches the app using configuration written in
    two json files: config.json and input_metadata.json.
    """
    # 1. Instantiate and launch the App
    print("1. Instantiate and launch the App")
    from apps.jsonapp import JSONApp
    app = JSONApp()
    result = app.launch(process_chipseq,
                        config,
                        in_metadata,
                        out_metadata)

    # 2. The App has finished
    print("2. Execution finished; see " + out_metadata)
    print(result)

    return result

# ------------------------------------------------------------------------------

if __name__ == "__main__":
    import sys
    sys._run_from_cmdl = True  # pylint: disable=protected-access

    # Set up the command line parameters
    PARSER = argparse.ArgumentParser(description="ChIP-seq peak calling")
    PARSER.add_argument("--config", help="Configuration file")
    PARSER.add_argument("--in_metadata", help="Location of input metadata file")
    PARSER.add_argument("--out_metadata", help="Location of output metadata file")

    # Get the matching parameters from the command line
    ARGS = PARSER.parse_args()

    CONFIG = ARGS.config
    IN_METADATA = ARGS.in_metadata
    OUT_METADATA = ARGS.out_metadata

    RESULTS = main_json(CONFIG, IN_METADATA, OUT_METADATA)

    print(RESULTS)
<|endoftext|>


<commit_before>import re

def strip_end_spaces(s):
    end = s.find(r"\275")
    count = s.count(r"\275")
    return s[:end] + r" (...plus \275 * %d)" % count

def lookup_function (val):
    "Look-up and return a pretty-printer that can print val."

    # Get the type.
    type = val.type

    # If it points to a reference, get the reference.
    if type.code == gdb.TYPE_CODE_REF:
        type = type.target ()

    # Get the unqualified type, stripped of typedefs.
    type = type.unqualified ().strip_typedefs ()

    # Get the type name.
    typename = type.tag

    if typename == None:
        return None

    # Iterate over local dictionary of types to determine
    # if a printer is registered for that type.  Return an
    # instantiation of the printer if found.
    for function in pretty_printers_dict:
       if function.match (typename):
           return pretty_printers_dict[function] (val)

    # Cannot find a pretty printer.  Return None.
    return None

class Btree_FsmPrinter:
    def __init__(self, val):
        self.val = val

    def to_string(self):
        return "========================================\nbtree object with the following members:"

    def children(self):
        yield 'fsm_type', Fsm_TypePrinter(int(self.val['fsm_type'])).to_string()
        yield 'noreply', "false" if self.val['noreply']==False else "true"

        try: yield 'request', "a pointer to a " + str(self.val['request'].dereference())
        except: yield "request", "no request object"

        try: yield 'cache', "a pointer to a " + str(self.val['cache'].dereference())
        except RuntimeError: yield "cache", "no cache object"

        try: yield 'on_complete', "a pointer to a " + str(self.val['on_complete'].dereference())
        except RuntimeError: yield "on_complete", "no on_complete object"

        try: yield "transaction", "a pointer to a " + str(self.val['transaction'].dereference())
        except RuntimeError: yield "transaction", "no transaction object"

        try: yield 'prev', "a pointer to a " + str(self.val['prev'].dereference())
        except: yield "prev", "no prev object"

        try: yield 'next', "a pointer to a " + str(self.val['next'].dereference())
        except: yield "next", "no next object"

        yield 'type', self.val['type']
        yield 'return_cpu', self.val['return_cpu']
        yield 'key_memory', strip_end_spaces(str(self.val['key_memory']))
        yield 'value_memory', strip_end_spaces(str(self.val['value_memory']))

        fsm_type = Fsm_TypePrinter(int(self.val['fsm_type'])).to_string()
        if fsm_type == "btree_set_fsm":
            yield "state", self.val['state']
            try: yield "key", "a pointer to a " + str(self.val['key'].dereference())
            except RuntimeError: yield "key", "no key object"

            try: yield "value", "a pointer to a " + str(self.val['value'].dereference())
            except RuntimeError: yield "value", "no value object"

            try: yield "sb_buf", "a pointer to a " + str(self.val['sb_buf'].dereference())
            except RuntimeError: yield "sb_buf", "no sb_buf object"

            try: yield "buf", "a pointer to a " + str(self.val['buf'].dereference())
            except RuntimeError: yield "buf", "no buf object"

            try: yield "last_buf", "a pointer to a " + str(self.val['last_buf'].dereference())
            except RuntimeError: yield "last_buf", "no last_buf object"

            yield "node_id", self.val['node_id']
            yield "last_node_id", self.val['last_node_id']
            yield "set_type", self.val['set_type']
            yield "set_was_successful", self.val['set_was_successful']
        elif fsm_type == "btree_get_fsm":
            yield "state", self.val['state']
            yield "key_memory", strip_end_spaces(str(self.val['key_memory']))
            yield "value_memory", strip_end_spaces(str(self.val['value_memory']))

            try: yield "buf", "a pointer to a " + str(self.val['buf'].dereference())
            except RuntimeError: yield "buf", "no buf object"

            try: yield "last_buf", "a pointer to a " + str(self.val['last_buf'].dereference())
            except RuntimeError: yield "last_buf", "no last_buf object"

            yield "node_id", self.val['node_id']
        elif fsm_type == "btree_delete_fsm":
            yield "state", self.val['state']
            yield "key_memory", strip_end_spaces(str(self.val['key_memory']))
            yield "node_id", self.val['node_id']
            yield "last_node_id", self.val['last_node_id']
            yield "sib_node_id", self.val['sib_node_id']

            try: yield "key", "a pointer to a " + str(self.val['key'].dereference())
            except RuntimeError: yield "key", "no key object"

            try: yield "sb_buf", "a pointer to a " + str(self.val['sb_buf'].dereference())
            except RuntimeError: yield "sb_buf", "no sb_buf object"

            try: yield "sib_buf", "a pointer to a " + str(self.val['sib_buf'].dereference())
            except RuntimeError: yield "sib_buf", "no sib_buf object"

            try: yield "last_buf", "a pointer to a " + str(self.val['last_buf'].dereference())
            except RuntimeError: yield "last_buf", "no last_buf object"

class Fsm_TypePrinter:
    def __init__(self, val):
        self.val = int(val)

    def to_string(self):
        a = ["btree_mock_fsm", "btree_get_fsm", "btree_set_fsm", "btree_delete_fsm"]
        return a[int(self.val)]

class RequestPrinter:
    def __init__(self, val=None):
        self.val = val

    def to_string(self):
        return "request object with the following members:"

    def children(self):
        yield 'nstarted', 1# int(self.val['nstarted'])
        yield 'ncompleted',2# int(self.val['ncompleted'])

class CachePrinter:
    def __init__(self, val):
        self.val = val

    def to_string(self):
        return "cache object with the following members:"

    def children(self):
        yield 'n_trans_created', int(self.val['n_trans_created'])
        yield 'n_trans_freed', int(self.val['n_trans_freed'])
        yield 'n_blocks_acquired', int(self.val['n_blocks_acquired'])
        yield 'n_blocks_released', int(self.val['n_blocks_released'])

class TransactionPrinter:
    def __init__(self, val):
        self.val = val

    def to_string(self):
        return "transaction object with the following members:"

    def children(self):
        try: yield 'cache', "a pointer to a " + str(self.val['cache'].dereference())
        except RuntimeError: yield "cache", "no cache object"

        try: yield 'access', "a pointer to a " + str(self.val['access'].dereference())
        except RuntimeError: yield "access", "no access object"

        try: yield "begin_callback", "a pointer to a" + str(self.val['begin_callback'].dereference())
        except RuntimeError: yield "begin_callback", "no transaction_begin_callback object"

        try: yield "commit_callback", "a pointer to a" + str(self.val['commit_callback'].dereference())
        except RuntimeError: yield "commit_callback", "no transaction_commit_callback object"

class AccessPrinter:
    def __init__(self, val):
        self.val = val

    def to_string(self):
        return "access object with the following members:"

    def children(self):
        pass

class Transaction_Begin_CallbackPrinter:
    def __init__(self, val):
        self.val = val
    def to_string(self):
        return "transaction_begin_callback object"

class Transaction_Commit_CallbackPrinter:
    def __init__(self, val):
        self.val = val
    def to_string(self):
        return "transaction_commit_callback object"

class Cpu_MessagePrinter:
    def __init__(self, val):
        self.val = val
    def to_string(self):
        return "cpu_message object"

class Btree_KeyPrinter:
    def __init__(self, val):
        self.val = val
    def to_string(self):
        return "btree_key object with the following members:"
    def children(self):
        yield "size", self.val['size']
        yield "contents", strip_end_spaces(str(self.val['contents']))

class Btree_ValuePrinter:
    def __init__(self, val):
        self.val = val
    def to_string(self):
        return "btree_value object with the following members:"
    def children(self):
        yield "size", self.val['size']
        yield "contents", strip_end_spaces(str(self.val['contents']))

class BufPrinter:
    def __init__(self, val):
        self.val = val
    def to_string(self):
        return "buf object"

class Msg_TypePrinter:
    def __init__(self, val):
        self.val = val

    def to_string(self):
        a = ["mt_btree", "mt_lock" , "mt_perfmon"]
        return a[int(self.val)]

class StatePrinter:
    def __init__(self, val):
        self.val = val

    def to_string(self):
        a = ["uninitialized", "start_transaction", "acquire_superblock", "acquire_root", "insert_root", "insert_root_on_split", "acquire_node", "update_complete", "committing"]
        return a[int(self.val)]

class Btree_Set_TypePrinter:
    def __init__(self, val):
        self.val = val

    def to_string(self):
        a = ["btree_set_type_set", "btree_set_type_add", "btree_set_type_replace", "btree_set_type_incr", "btree_set_type_decr"]
        return a[int(self.val)]

pretty_printers_dict = {}
pretty_printers_dict[re.compile ('.*state.*')]                           = StatePrinter
pretty_printers_dict[re.compile ('^btree_fsm.*')]                       = Btree_FsmPrinter
pretty_printers_dict[re.compile ('^btree_set_fsm.*')]                   = Btree_FsmPrinter
pretty_printers_dict[re.compile ('^btree_get_fsm.*')]                   = Btree_FsmPrinter
pretty_printers_dict[re.compile ('^btree_delete_fsm.*')]                = Btree_FsmPrinter
pretty_printers_dict[re.compile ('.*cpu_message.*')]                     = Cpu_MessagePrinter
pretty_printers_dict[re.compile ('.*fsm_type.*')]                        = Fsm_TypePrinter
pretty_printers_dict[re.compile ('.*request.*')]                         = RequestPrinter
pretty_printers_dict[re.compile ('.*transaction.*')]                     = TransactionPrinter
pretty_printers_dict[re.compile ('.*access.*')]                          = AccessPrinter
pretty_printers_dict[re.compile ('.*cache.*')]                           = CachePrinter
pretty_printers_dict[re.compile ('.*transaction_begin_callback.*')]      = Transaction_Begin_CallbackPrinter
pretty_printers_dict[re.compile ('.*transaction_commit_callback.*')]     = Transaction_Commit_CallbackPrinter
pretty_printers_dict[re.compile ('.*msg_type.*')]                        = Msg_TypePrinter
pretty_printers_dict[re.compile ('.*btree_key.*')]                       = Btree_KeyPrinter
pretty_printers_dict[re.compile ('.*btree_value.*')]                     = Btree_ValuePrinter
pretty_printers_dict[re.compile ('.*buf.*')]                             = BufPrinter
pretty_printers_dict[re.compile ('.*btree_set_type.*')]                  = Btree_Set_TypePrinter
gdb.pretty_printers.append (lookup_function)
<commit_msg>pretty-printing bug finally fixed.<commit_after>import re

def strip_end_spaces(s):
    end = s.find(r"\275")
    count = s.count(r"\275")
    return s[:end] + r" (...plus \275 * %d)" % count

def lookup_function (val):
    "Look-up and return a pretty-printer that can print val."

    # Get the type.
    type = val.type

    # If it points to a reference, get the reference.
    if type.code == gdb.TYPE_CODE_REF:
        type = type.target ()

    # Get the unqualified type, stripped of typedefs.
    type = type.unqualified ().strip_typedefs ()

    # Get the type name.
    typename = type.tag

    if typename == None:
        return None

    # Iterate over local dictionary of types to determine
    # if a printer is registered for that type.  Return an
    # instantiation of the printer if found.
    for function in sorted(pretty_printers_dict):
       if function.match (typename):
           return pretty_printers_dict[function] (val)

    # Cannot find a pretty printer.  Return None.
    return None

class Btree_FsmPrinter:
    def __init__(self, val):
        self.val = val

    def to_string(self):
        return "========================================\nbtree object with the following members:"

    def children(self):
        yield 'fsm_type', Fsm_TypePrinter(int(self.val['fsm_type'])).to_string()
        yield 'noreply', "false" if self.val['noreply']==False else "true"

        try: yield 'request', "a pointer to a " + str(self.val['request'].dereference())
        except: yield "request", "no request object"

        try: yield 'cache', "a pointer to a " + str(self.val['cache'].dereference())
        except RuntimeError: yield "cache", "no cache object"

        try: yield 'on_complete', "a pointer to a " + str(self.val['on_complete'].dereference())
        except RuntimeError: yield "on_complete", "no on_complete object"

        try: yield "transaction", "a pointer to a " + str(self.val['transaction'].dereference())
        except RuntimeError: yield "transaction", "no transaction object"

        try: yield 'prev', "a pointer to a " + str(self.val['prev'].dereference())
        except: yield "prev", "no prev object"

        try: yield 'next', "a pointer to a " + str(self.val['next'].dereference())
        except: yield "next", "no next object"

        yield 'type', self.val['type']
        yield 'return_cpu', self.val['return_cpu']
        yield 'key_memory', strip_end_spaces(str(self.val['key_memory']))
        yield 'value_memory', strip_end_spaces(str(self.val['value_memory']))

        fsm_type = Fsm_TypePrinter(int(self.val['fsm_type'])).to_string()
        if fsm_type == "btree_set_fsm":
            yield "state", self.val['state']
            try: yield "key", "a pointer to a " + str(self.val['key'].dereference())
            except RuntimeError: yield "key", "no key object"

            try: yield "value", "a pointer to a " + str(self.val['value'].dereference())
            except RuntimeError: yield "value", "no value object"

            try: yield "sb_buf", "a pointer to a " + str(self.val['sb_buf'].dereference())
            except RuntimeError: yield "sb_buf", "no sb_buf object"

            try: yield "buf", "a pointer to a " + str(self.val['buf'].dereference())
            except RuntimeError: yield "buf", "no buf object"

            try: yield "last_buf", "a pointer to a " + str(self.val['last_buf'].dereference())
            except RuntimeError: yield "last_buf", "no last_buf object"

            yield "node_id", self.val['node_id']
            yield "last_node_id", self.val['last_node_id']
            yield "set_type", self.val['set_type']
            yield "set_was_successful", self.val['set_was_successful']
        elif fsm_type == "btree_get_fsm":
            yield "state", self.val['state']
            yield "key_memory", strip_end_spaces(str(self.val['key_memory']))
            yield "value_memory", strip_end_spaces(str(self.val['value_memory']))

            try: yield "buf", "a pointer to a " + str(self.val['buf'].dereference())
            except RuntimeError: yield "buf", "no buf object"

            try: yield "last_buf", "a pointer to a " + str(self.val['last_buf'].dereference())
            except RuntimeError: yield "last_buf", "no last_buf object"

            yield "node_id", self.val['node_id']
        elif fsm_type == "btree_delete_fsm":
            yield "state", self.val['state']
            yield "key_memory", strip_end_spaces(str(self.val['key_memory']))
            yield "node_id", self.val['node_id']
            yield "last_node_id", self.val['last_node_id']
            yield "sib_node_id", self.val['sib_node_id']

            try: yield "key", "a pointer to a " + str(self.val['key'].dereference())
            except RuntimeError: yield "key", "no key object"

            try: yield "sb_buf", "a pointer to a " + str(self.val['sb_buf'].dereference())
            except RuntimeError: yield "sb_buf", "no sb_buf object"

            try: yield "sib_buf", "a pointer to a " + str(self.val['sib_buf'].dereference())
            except RuntimeError: yield "sib_buf", "no sib_buf object"

            try: yield "last_buf", "a pointer to a " + str(self.val['last_buf'].dereference())
            except RuntimeError: yield "last_buf", "no last_buf object"

class Fsm_TypePrinter:
    def __init__(self, val):
        self.val = int(val)

    def to_string(self):
        a = ["btree_mock_fsm", "btree_get_fsm", "btree_set_fsm", "btree_delete_fsm"]
        return a[int(self.val)]

class RequestPrinter:
    def __init__(self, val=None):
        self.val = val

    def to_string(self):
        return "request object with the following members:"

    def children(self):
        yield 'nstarted', 1# int(self.val['nstarted'])
        yield 'ncompleted',2# int(self.val['ncompleted'])

class CachePrinter:
    def __init__(self, val):
        self.val = val

    def to_string(self):
        return "cache object with the following members:"

    def children(self):
        yield 'n_trans_created', int(self.val['n_trans_created'])
        yield 'n_trans_freed', int(self.val['n_trans_freed'])
        yield 'n_blocks_acquired', int(self.val['n_blocks_acquired'])
        yield 'n_blocks_released', int(self.val['n_blocks_released'])

class TransactionPrinter:
    def __init__(self, val):
        self.val = val

    def to_string(self):
        return "transaction object with the following members:"

    def children(self):
        try: yield 'cache', "a pointer to a " + str(self.val['cache'].dereference())
        except RuntimeError: yield "cache", "no cache object"

        try: yield 'access', "a pointer to a " + str(self.val['access'].dereference())
        except RuntimeError: yield "access", "no access object"

        try: yield "begin_callback", "a pointer to a" + str(self.val['begin_callback'].dereference())
        except RuntimeError: yield "begin_callback", "no transaction_begin_callback object"

        try: yield "commit_callback", "a pointer to a" + str(self.val['commit_callback'].dereference())
        except RuntimeError: yield "commit_callback", "no transaction_commit_callback object"

class AccessPrinter:
    def __init__(self, val):
        self.val = val

    def to_string(self):
        return "access object with the following members:"

    def children(self):
        pass

class Transaction_Begin_CallbackPrinter:
    def __init__(self, val):
        self.val = val
    def to_string(self):
        return "transaction_begin_callback object"

class Transaction_Commit_CallbackPrinter:
    def __init__(self, val):
        self.val = val
    def to_string(self):
        return "transaction_commit_callback object"

class Cpu_MessagePrinter:
    def __init__(self, val):
        self.val = val
    def to_string(self):
        return "cpu_message object"

class Btree_KeyPrinter:
    def __init__(self, val):
        self.val = val
    def to_string(self):
        return "btree_key object with the following members:"
    def children(self):
        yield "size", self.val['size']
        yield "contents", strip_end_spaces(str(self.val['contents']))

class Btree_ValuePrinter:
    def __init__(self, val):
        self.val = val
    def to_string(self):
        return "btree_value object with the following members:"
    def children(self):
        yield "size", self.val['size']
        yield "contents", strip_end_spaces(str(self.val['contents']))

class BufPrinter:
    def __init__(self, val):
        self.val = val
    def to_string(self):
        return "buf object"

class Msg_TypePrinter:
    def __init__(self, val):
        self.val = val

    def to_string(self):
        a = ["mt_btree", "mt_lock" , "mt_perfmon"]
        return a[int(self.val)]

class StatePrinter:
    def __init__(self, val):
        self.val = val

    def to_string(self):
        a = ["uninitialized", "start_transaction", "acquire_superblock", "acquire_root", "insert_root", "insert_root_on_split", "acquire_node", "update_complete", "committing"]
        return a[int(self.val)]

class Btree_Set_TypePrinter:
    def __init__(self, val):
        self.val = val

    def to_string(self):
        a = ["btree_set_type_set", "btree_set_type_add", "btree_set_type_replace", "btree_set_type_incr", "btree_set_type_decr"]
        return a[int(self.val)]

pretty_printers_dict = {}
pretty_printers_dict[re.compile ('.*state.*')]                           = StatePrinter
pretty_printers_dict[re.compile ('^btree_fsm[^:].*')]                       = Btree_FsmPrinter
pretty_printers_dict[re.compile ('^btree_set_fsm[^:].*')]                   = Btree_FsmPrinter
pretty_printers_dict[re.compile ('^btree_get_fsm[^:].*')]                   = Btree_FsmPrinter
pretty_printers_dict[re.compile ('^btree_delete_fsm[^:].*')]                = Btree_FsmPrinter
pretty_printers_dict[re.compile ('.*cpu_message.*')]                     = Cpu_MessagePrinter
pretty_printers_dict[re.compile ('.*fsm_type.*')]                        = Fsm_TypePrinter
pretty_printers_dict[re.compile ('.*request.*')]                         = RequestPrinter
pretty_printers_dict[re.compile ('.*transaction.*')]                     = TransactionPrinter
pretty_printers_dict[re.compile ('.*access.*')]                          = AccessPrinter
pretty_printers_dict[re.compile ('.*cache.*')]                           = CachePrinter
pretty_printers_dict[re.compile ('.*transaction_begin_callback.*')]      = Transaction_Begin_CallbackPrinter
pretty_printers_dict[re.compile ('.*transaction_commit_callback.*')]     = Transaction_Commit_CallbackPrinter
pretty_printers_dict[re.compile ('.*msg_type.*')]                        = Msg_TypePrinter
pretty_printers_dict[re.compile ('.*btree_key.*')]                       = Btree_KeyPrinter
pretty_printers_dict[re.compile ('.*btree_value.*')]                     = Btree_ValuePrinter
pretty_printers_dict[re.compile ('.*buf.*')]                             = BufPrinter
pretty_printers_dict[re.compile ('.*btree_set_type.*')]                  = Btree_Set_TypePrinter
gdb.pretty_printers.append (lookup_function)
<|endoftext|>


<commit_before>#!/usr/bin/python

# Imports
import argparse
import os
import glob
import subprocess
import sys

class colors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

def print_warnings(msg):
    print colors.WARNING + msg

def print_errors(msg):
    print colors.FAIL + msg

def print_errors_or_warnings(msg):
    if ERROR:
        print_errors(msg)
        return True
    else:
        print_warnings(msg)

def which(pgm):
    path = os.getenv('PATH')
    for p in path.split(os.path.pathsep):
        p = os.path.join(p, pgm)
        if os.path.exists(p) and os.access(p, os.X_OK):
            return p

# Verbose output
def debug(string):
    if VERBOSE:
        print string

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Run a program with multiple input files')
    parser.add_argument('-i', '--in', help='Input file prefix', default='in')
    parser.add_argument('-o', '--out', help='Output file prefix', default='out')
    parser.add_argument('-p', '--path', help='Path to input/output files', default='.')
    parser.add_argument('-v', '--verbose', help='Increase output verbosity', action='store_true')
    parser.add_argument('-e', '--error', help='Change warnings to errors', action='store_true')
    parser.add_argument('command', help='Command to run with input files')
    args = vars(parser.parse_args())

    # Set Constants
    COMMAND = args['command']
    IN = args['in']
    OUT = args['out']
    VERBOSE = args['verbose']
    PATH = args['path']
    ERROR = args['error']

    debug(args)

    # Check if command is valid
    command = COMMAND.split(" ")[0]
    debug("Checking if " + command + " is available...")
    path_file = which(command)  # Command is somewhere in path
    local_file = os.path.isfile(command)  # Command is a local file
    if not (path_file or local_file):
        print_errors("Command \"" + command + "\" does not exist")
        debug("Terminating Program.")
        exit()

    # Find all input/output files
    input_files = glob.glob(os.path.join(PATH, IN + '*'))
    output_files = glob.glob(os.path.join(PATH, OUT + '*'))

    debug("Checking path: " + os.path.join(PATH, IN + '*'))
    debug("Detected input Files:")
    debug(input_files)

    valid_pairs = []
    terminate_program = False

    # Make
    for test in input_files:
        debug("Checking " + test + "'s output file")
        suffix = test.rsplit(IN)[-1]
        outpath = os.path.join(PATH, OUT + suffix)
        if not os.path.isfile(outpath):
            terminate_program = terminate_program or print_errors_or_warnings("Output file does not exist: " + outpath)
        else:
            valid_pairs.append((test, outpath))
            debug("Output file found: " + outpath)

    if terminate_program:
        debug("Terminating Program.")
        exit()

    for test, outpath in valid_pairs:
        input_file = open(test)
        process = subprocess.Popen(COMMAND, stdin=input_file, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, stderr = process.communicate()

        # Read in expected output
        output = open(outpath).read()

        for i, (char1, char2) in enumerate(zip(output, stdout)):
            if char1 != char2:
                # Output differs from expected
                print_errors_or_warnings("Output of " + test + " differs from the expected output of " + outpath)
                break
<commit_msg>added error codes<commit_after>#!/usr/bin/python

# Imports
import argparse
import os
import glob
import subprocess
import sys

class colors:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKGREEN = '\033[92m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

def print_warnings(msg):
    print colors.WARNING + msg

def print_errors(msg):
    print colors.FAIL + msg

def print_errors_or_warnings(msg):
    if ERROR:
        print_errors(msg)
        return True
    else:
        print_warnings(msg)

def which(pgm):
    path = os.getenv('PATH')
    for p in path.split(os.path.pathsep):
        p = os.path.join(p, pgm)
        if os.path.exists(p) and os.access(p, os.X_OK):
            return p

# Verbose output
def debug(string):
    if VERBOSE:
        print string

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Run a program with multiple input files')
    parser.add_argument('-i', '--in', help='Input file prefix', default='in')
    parser.add_argument('-o', '--out', help='Output file prefix', default='out')
    parser.add_argument('-p', '--path', help='Path to input/output files', default='.')
    parser.add_argument('-v', '--verbose', help='Increase output verbosity', action='store_true')
    parser.add_argument('-e', '--error', help='Change warnings to errors', action='store_true')
    parser.add_argument('command', help='Command to run with input files')
    args = vars(parser.parse_args())

    # Set Constants
    COMMAND = args['command']
    IN = args['in']
    OUT = args['out']
    VERBOSE = args['verbose']
    PATH = args['path']
    ERROR = args['error']

    debug(args)

    # Check if command is valid
    command = COMMAND.split(" ")[0]
    debug("Checking if " + command + " is available...")
    path_file = which(command)  # Command is somewhere in path
    local_file = os.path.isfile(command)  # Command is a local file
    if not (path_file or local_file):
        print_errors("Command \"" + command + "\" does not exist")
        debug("Terminating Program.")
        sys.exit(100)

    # Find all input/output files
    input_files = glob.glob(os.path.join(PATH, IN + '*'))
    output_files = glob.glob(os.path.join(PATH, OUT + '*'))

    debug("Checking path: " + os.path.join(PATH, IN + '*'))
    debug("Detected input Files:")
    debug(input_files)

    valid_pairs = []
    terminate_program = False

    # Make
    for test in input_files:
        debug("Checking " + test + "'s output file")
        suffix = test.rsplit(IN)[-1]
        outpath = os.path.join(PATH, OUT + suffix)
        if not os.path.isfile(outpath):
            terminate_program = terminate_program or print_errors_or_warnings("Output file does not exist: " + outpath)
        else:
            valid_pairs.append((test, outpath))
            debug("Output file found: " + outpath)

    if terminate_program:
        debug("Terminating Program.")
        sys.exit(101)

    for test, outpath in valid_pairs:
        input_file = open(test)
        process = subprocess.Popen(COMMAND, stdin=input_file, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, stderr = process.communicate()

        # Read in expected output
        output = open(outpath).read()

        for i, (char1, char2) in enumerate(zip(output, stdout)):
            if char1 != char2:
                # Output differs from expected
                print_errors_or_warnings("Output of " + test + " differs from the expected output of " + outpath)
                break
<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import unicode_literals
from abc import ABCMeta, abstractmethod, abstractproperty
from functools import total_ordering

import six

from corehq.apps.aggregate_ucrs.date_utils import Month

AGGREGATION_UNIT_CHOICE_MONTH = 'month'

AGG_WINDOW_START_PARAM = 'agg_window_start'
AGG_WINDOW_END_PARAM = 'agg_window_end'

class TimeAggregation(six.with_metaclass(ABCMeta, object)):

    def __init__(self, datetime):
        self._datetime = datetime

    @classmethod
    def from_aggregation_unit(cls, unit):
        adapter_classes = {
            AGGREGATION_UNIT_CHOICE_MONTH: MonthAggregation
        }
        return adapter_classes[unit]

    @abstractmethod
    def next_window(self):
        pass

    @abstractproperty
    def start(self):
        pass

    @abstractproperty
    def end(self):
        pass

    def __str__(self):
        return '{}: {}-{}'.format(type(self).__name__, self.start, self.end)

@total_ordering
class MonthAggregation(six.with_metaclass(ABCMeta, TimeAggregation)):

    def next_window(self):
        return MonthAggregation(self._month.end)

    @property
    def start(self):
        return self._month.start

    @property
    def end(self):
        return self._month.end

    def __init__(self, datetime):
        super(MonthAggregation, self).__init__(datetime)
        self._month = Month.datetime_to_month(self._datetime)

    def __eq__(self, other):
        return self._month == other._month

    def __lt__(self, other):
        return self._month < other._month
<commit_msg>implement hash and improve eq<commit_after>from __future__ import absolute_import
from __future__ import unicode_literals
from abc import ABCMeta, abstractmethod, abstractproperty
from functools import total_ordering

import six

from corehq.apps.aggregate_ucrs.date_utils import Month

AGGREGATION_UNIT_CHOICE_MONTH = 'month'

AGG_WINDOW_START_PARAM = 'agg_window_start'
AGG_WINDOW_END_PARAM = 'agg_window_end'

class TimeAggregation(six.with_metaclass(ABCMeta, object)):

    def __init__(self, datetime):
        self._datetime = datetime

    @classmethod
    def from_aggregation_unit(cls, unit):
        adapter_classes = {
            AGGREGATION_UNIT_CHOICE_MONTH: MonthAggregation
        }
        return adapter_classes[unit]

    @abstractmethod
    def next_window(self):
        pass

    @abstractproperty
    def start(self):
        pass

    @abstractproperty
    def end(self):
        pass

    def __str__(self):
        return '{}: {}-{}'.format(type(self).__name__, self.start, self.end)

@total_ordering
class MonthAggregation(six.with_metaclass(ABCMeta, TimeAggregation)):

    def next_window(self):
        return MonthAggregation(self._month.end)

    @property
    def start(self):
        return self._month.start

    @property
    def end(self):
        return self._month.end

    def __init__(self, datetime):
        super(MonthAggregation, self).__init__(datetime)
        self._month = Month.datetime_to_month(self._datetime)

    def __eq__(self, other):
        return isinstance(other, MonthAggregation) and self._month == other._month

    def __hash__(self):
        return hash((type(self), self.start))

    def __lt__(self, other):
        return self._month < other._month
<|endoftext|>