<commit_before>#!/usr/bin/python3
import signal
import argparse
import sys
import csv
import re
from pkg_resources import get_distribution
from bedparse import bedline
from bedparse import gtf2bed
# This allows using the program in a pipe
# The program is killed when it receives a sigpipe
signal.signal(signal.SIGPIPE, signal.SIG_DFL)
__version__ = get_distribution('bedparse').version

def introns(args):
    with args.bedfile as tsvfile:
        for line in tsvfile:
            introns=bedline(line.split('\t')).introns()
            if(introns): introns.print()
    tsvfile.close()

def threeP(args):
    with args.bedfile as tsvfile:
        for line in tsvfile:
            utr=bedline(line.split('\t')).utr(which=3)
            if(utr): utr.print()
    tsvfile.close()

def fiveP(args):
    with args.bedfile as tsvfile:
        for line in tsvfile:
            utr=bedline(line.split('\t')).utr(which=5)
            if(utr): utr.print()
    tsvfile.close()

def cds(args):
    with args.bedfile as tsvfile:
        for line in tsvfile:
            utr=bedline(line.split('\t')).cds(ignoreCDSonly=args.ignoreCDSonly)
            if(utr): utr.print()
    tsvfile.close()

def prom(args):
    with args.bedfile as tsvfile:
        for line in tsvfile:
            bedline(line.split('\t')).promoter(up=args.up, down=args.down, strand=(not args.unstranded)).print()
    tsvfile.close()

def bed12tobed6(args):
    with args.bedfile as tsvfile:
        for line in tsvfile:
            for el in bedline(line.split('\t')).bed12tobed6(appendExN=args.appendExN):
                el.print()
            if(args.keepIntrons):
                nameSub=re.compile("_Exon([0-9]+)")
                for el in bedline(line.split('\t')).introns().bed12tobed6(appendExN=args.appendExN):
                    el.name=nameSub.sub(r"_Intron\1", el.name)
                    el.print()
    tsvfile.close()

def filter(args):
    col=args.column-1
    filterset=set()
    try:
        annotation=open(args.annotation)
    except:
        raise BEDexception("Annotation file not valid")
    annotationReader = csv.reader(annotation, delimiter="\t")
    for line in annotationReader:
        filterset.add(line[col])
    annotation.close()
    with args.bedfile as tsvfile:
        for line in tsvfile:
            if(line.split('\t')[3] in filterset):
                print(line.rstrip())
    tsvfile.close()

def join(args):
    col=args.column-1
    annot=dict()
    try:
        annotation=open(args.annotation)
    except:
        raise BEDexception("Annotation file not valid")
    annotationReader = csv.reader(annotation, delimiter=args.separator)
    for line in annotationReader:
        annot.setdefault(line[col], []).append(line[0:col]+line[col+1:])
    annotation.close()
    with args.bedfile as tsvfile:
        for line in tsvfile:
            line=line.split('\t')
            if(args.noUnmatched==False or line[3] in annot.keys()):
                record=bedline(line)
                if(record):
                        nrec=len(annot.setdefault(record.name, []))
                        if(nrec==0):
                            record.print(end='')
                            print('',args.empty,sep="\t")
                        else:
                            for i in range(0,nrec):
                                record.print(end='')
                                print('',*annot[record.name][i], sep='\t')
    tsvfile.close()

def main(args=None):
    if args is None:
        args = sys.argv[1:]

    parser = argparse.ArgumentParser(
            description="""Perform various simple operations on BED files.""")

    parser.add_argument('--version', '-v', action='version', version='v'+__version__)
    subparsers = parser.add_subparsers(help='sub-command help', dest='sub-command')
    subparsers.required = True

    parser_3pUTR = subparsers.add_parser('3pUTR', help="Prints the 3' of coding genes.")
    parser_3pUTR.add_argument("bedfile", type=argparse.FileType('r'), nargs='?', default=sys.stdin, help="Path to the BED file.")
    parser_3pUTR.set_defaults(func=threeP)

    parser_5pUTR = subparsers.add_parser('5pUTR', help="Prints the 5' of coding genes.")
    parser_5pUTR.add_argument("bedfile", type=argparse.FileType('r'), nargs='?', default=sys.stdin, help="Path to the BED file.")
    parser_5pUTR.set_defaults(func=fiveP)

    parser_cds = subparsers.add_parser('cds', help="Prints the CDS of coding genes.")
    parser_cds.add_argument("--ignoreCDSonly",action="store_true", help="Ignore transcripts that only consist of CDS.")
    parser_cds.add_argument("bedfile", type=argparse.FileType('r'), nargs='?', default=sys.stdin, help="Path to the BED file.")
    parser_cds.set_defaults(func=cds)

    parser_prom = subparsers.add_parser('promoter', help="Prints the promoters of coding genes.")
    parser_prom.add_argument("--up",type=int, default=500, help="Get this many nt upstream of each feature.")
    parser_prom.add_argument("--down",type=int, default=500, help="Get this many nt downstream of each feature.")
    parser_prom.add_argument("--unstranded",action="store_true", help="Do not consider strands.")
    parser_prom.add_argument("bedfile", type=argparse.FileType('r'), nargs='?', default=sys.stdin, help="Path to the BED file.")
    parser_prom.set_defaults(func=prom)

    parser_introns = subparsers.add_parser('introns', help="Prints BED records corresponding to the introns of each transcript in the original file.")
    parser_introns.add_argument("bedfile", type=argparse.FileType('r'), nargs='?', default=sys.stdin, help="Path to the BED file.")
    parser_introns.set_defaults(func=introns)

    parser_filter = subparsers.add_parser('filter',
            help="""Filters a BED file based on an annotation.
            BED entries with a name (i.e. col4) that appears
            in the specified column of the annotation are
            printed to stdout. For efficiency reasons this
            command doesn't perform BED validation.""")
    parser_filter.add_argument("--annotation", "-a", type=str, help="Path to the annotation file.", required=True)
    parser_filter.add_argument("--column","-c",type=int, default=1, help="Column of the annotation file (1-based, default=1).")
    parser_filter.set_defaults(func=filter)
    parser_filter.add_argument("bedfile", type=argparse.FileType('r'), nargs='?', default=sys.stdin,
    help="Path to the BED file.")

    parser_join = subparsers.add_parser('join',
            help="""Joins a BED file with an annotation file using
            the BED name (col4) as the joining key.""")
    parser_join.add_argument("--annotation", "-a", type=str, help="Path to the annotation file.", required=True)
    parser_join.add_argument("--column","-c",type=int, default=1, help="Column of the annotation file (1-based, default=1).")
    parser_join.add_argument("--separator","-s",type=str, default='\t', help="Field separator for the annotation file (default tab)")
    parser_join.add_argument("--empty","-e",type=str, default='.', help="String to append to empty records (default '.').")
    parser_join.add_argument("--noUnmatched", "-n" ,action="store_true", help="Do not print unmatched lines.")
    parser_join.set_defaults(func=join)
    parser_join.add_argument("bedfile", type=argparse.FileType('r'), nargs='?', default=sys.stdin,
    help="Path to the BED file.")

    parser_gtf2bed = subparsers.add_parser('gtf2bed',
            help="""Converts a GTF file to BED12 format.
            This tool supports the Ensembl GTF format.
            The GTF file must contain 'transcript' and 'exon'
            features in field 3. If the GTF file also annotates
            'CDS' 'start_codon' or 'stop_codon' these are used
            to annotate the thickStart and thickEnd in the BED
            file.""")
    parser_gtf2bed.add_argument("gtf", type=argparse.FileType('r'), nargs='?', default=sys.stdin, help="Path to the GTF file.")
    parser_gtf2bed.add_argument("--extraFields",type=str, default='', help="Comma separated list of extra GTF fields to be added after col 12 (e.g. gene_id,gene_name).")
    parser_gtf2bed.add_argument("--filterType",type=str, default='', help="Comma separated list of 'transcript_type' types to retain.")
    parser_gtf2bed.set_defaults(func=lambda args: gtf2bed(args.gtf, args.extraFields.split(','), args.filterType.split(',')))

    parser_bed12tobed6 = subparsers.add_parser('bed12tobed6',
            help="""Converts a BED12 file to BED6 format.
            """)
    parser_bed12tobed6.add_argument("bedfile", type=argparse.FileType('r'), nargs='?', default=sys.stdin, help="Path to the GTF file.")
    parser_bed12tobed6.add_argument("--appendExN", action="store_true", help="Appends the exon number to the transcript name.")
    parser_bed12tobed6.add_argument("--keepIntrons", action="store_true", help="Add records for introns as well.")
    parser_bed12tobed6.set_defaults(func=bed12tobed6)

    args = parser.parse_args()
    args.func(args)

if __name__ == "__main__":
    main()
<commit_msg>Improved error handling in join<commit_after>#!/usr/bin/python3
import signal
import argparse
import sys
import csv
import re
from pkg_resources import get_distribution
from bedparse import bedline
from bedparse import gtf2bed
from bedparse import BEDexception
# This allows using the program in a pipe
# The program is killed when it receives a sigpipe
signal.signal(signal.SIGPIPE, signal.SIG_DFL)
__version__ = get_distribution('bedparse').version

def introns(args):
    with args.bedfile as tsvfile:
        for line in tsvfile:
            introns=bedline(line.split('\t')).introns()
            if(introns): introns.print()
    tsvfile.close()

def threeP(args):
    with args.bedfile as tsvfile:
        for line in tsvfile:
            utr=bedline(line.split('\t')).utr(which=3)
            if(utr): utr.print()
    tsvfile.close()

def fiveP(args):
    with args.bedfile as tsvfile:
        for line in tsvfile:
            utr=bedline(line.split('\t')).utr(which=5)
            if(utr): utr.print()
    tsvfile.close()

def cds(args):
    with args.bedfile as tsvfile:
        for line in tsvfile:
            utr=bedline(line.split('\t')).cds(ignoreCDSonly=args.ignoreCDSonly)
            if(utr): utr.print()
    tsvfile.close()

def prom(args):
    with args.bedfile as tsvfile:
        for line in tsvfile:
            bedline(line.split('\t')).promoter(up=args.up, down=args.down, strand=(not args.unstranded)).print()
    tsvfile.close()

def bed12tobed6(args):
    with args.bedfile as tsvfile:
        for line in tsvfile:
            for el in bedline(line.split('\t')).bed12tobed6(appendExN=args.appendExN):
                el.print()
            if(args.keepIntrons):
                nameSub=re.compile("_Exon([0-9]+)")
                for el in bedline(line.split('\t')).introns().bed12tobed6(appendExN=args.appendExN):
                    el.name=nameSub.sub(r"_Intron\1", el.name)
                    el.print()
    tsvfile.close()

def filter(args):
    col=args.column-1
    filterset=set()
    try:
        annotation=open(args.annotation)
    except:
        raise BEDexception("Annotation file not valid")
    annotationReader = csv.reader(annotation, delimiter="\t")
    for line in annotationReader:
        filterset.add(line[col])
    annotation.close()
    with args.bedfile as tsvfile:
        for line in tsvfile:
            if(line.split('\t')[3] in filterset):
                print(line.rstrip())
    tsvfile.close()

def join(args):
    col=args.column-1
    annot=dict()
    try:
        annotation=open(args.annotation)
    except:
        raise BEDexception("Annotation file not valid")
    annotationReader = csv.reader(annotation, delimiter=args.separator)
    for line in annotationReader:
        if(len(line)<=col):
            raise BEDexception("Some lines don't contain the annotation column")
        annot.setdefault(line[col], []).append(line[0:col]+line[col+1:])
    annotation.close()
    with args.bedfile as tsvfile:
        for line in tsvfile:
            line=line.split('\t')
            if(args.noUnmatched==False or line[3] in annot.keys()):
                record=bedline(line)
                if(record):
                        nrec=len(annot.setdefault(record.name, []))
                        if(nrec==0):
                            if(args.empty==''):
                                record.print()
                            else:
                                record.print(end='')
                                print('',args.empty,sep="\t")
                        else:
                            for i in range(0,nrec):
                                record.print(end='')
                                print('',*annot[record.name][i], sep='\t')
    tsvfile.close()

def main(args=None):
    if args is None:
        args = sys.argv[1:]

    parser = argparse.ArgumentParser(
            description="""Perform various simple operations on BED files.""")

    parser.add_argument('--version', '-v', action='version', version='v'+__version__)
    subparsers = parser.add_subparsers(help='sub-command help', dest='sub-command')
    subparsers.required = True

    parser_3pUTR = subparsers.add_parser('3pUTR', help="Prints the 3' of coding genes.")
    parser_3pUTR.add_argument("bedfile", type=argparse.FileType('r'), nargs='?', default=sys.stdin, help="Path to the BED file.")
    parser_3pUTR.set_defaults(func=threeP)

    parser_5pUTR = subparsers.add_parser('5pUTR', help="Prints the 5' of coding genes.")
    parser_5pUTR.add_argument("bedfile", type=argparse.FileType('r'), nargs='?', default=sys.stdin, help="Path to the BED file.")
    parser_5pUTR.set_defaults(func=fiveP)

    parser_cds = subparsers.add_parser('cds', help="Prints the CDS of coding genes.")
    parser_cds.add_argument("--ignoreCDSonly",action="store_true", help="Ignore transcripts that only consist of CDS.")
    parser_cds.add_argument("bedfile", type=argparse.FileType('r'), nargs='?', default=sys.stdin, help="Path to the BED file.")
    parser_cds.set_defaults(func=cds)

    parser_prom = subparsers.add_parser('promoter', help="Prints the promoters of coding genes.")
    parser_prom.add_argument("--up",type=int, default=500, help="Get this many nt upstream of each feature.")
    parser_prom.add_argument("--down",type=int, default=500, help="Get this many nt downstream of each feature.")
    parser_prom.add_argument("--unstranded",action="store_true", help="Do not consider strands.")
    parser_prom.add_argument("bedfile", type=argparse.FileType('r'), nargs='?', default=sys.stdin, help="Path to the BED file.")
    parser_prom.set_defaults(func=prom)

    parser_introns = subparsers.add_parser('introns', help="Prints BED records corresponding to the introns of each transcript in the original file.")
    parser_introns.add_argument("bedfile", type=argparse.FileType('r'), nargs='?', default=sys.stdin, help="Path to the BED file.")
    parser_introns.set_defaults(func=introns)

    parser_filter = subparsers.add_parser('filter',
            help="""Filters a BED file based on an annotation.
            BED entries with a name (i.e. col4) that appears
            in the specified column of the annotation are
            printed to stdout. For efficiency reasons this
            command doesn't perform BED validation.""")
    parser_filter.add_argument("--annotation", "-a", type=str, help="Path to the annotation file.", required=True)
    parser_filter.add_argument("--column","-c",type=int, default=1, help="Column of the annotation file (1-based, default=1).")
    parser_filter.set_defaults(func=filter)
    parser_filter.add_argument("bedfile", type=argparse.FileType('r'), nargs='?', default=sys.stdin,
    help="Path to the BED file.")

    parser_join = subparsers.add_parser('join',
            help="""Joins a BED file with an annotation file using
            the BED name (col4) as the joining key.""")
    parser_join.add_argument("--annotation", "-a", type=str, help="Path to the annotation file.", required=True)
    parser_join.add_argument("--column","-c",type=int, default=1, help="Column of the annotation file (1-based, default=1).")
    parser_join.add_argument("--separator","-s",type=str, default='\t', help="Field separator for the annotation file (default tab)")
    parser_join.add_argument("--empty","-e",type=str, default='.', help="String to append to empty records (default '.').")
    parser_join.add_argument("--noUnmatched", "-n" ,action="store_true", help="Do not print unmatched lines.")
    parser_join.set_defaults(func=join)
    parser_join.add_argument("bedfile", type=argparse.FileType('r'), nargs='?', default=sys.stdin,
    help="Path to the BED file.")

    parser_gtf2bed = subparsers.add_parser('gtf2bed',
            help="""Converts a GTF file to BED12 format.
            This tool supports the Ensembl GTF format.
            The GTF file must contain 'transcript' and 'exon'
            features in field 3. If the GTF file also annotates
            'CDS' 'start_codon' or 'stop_codon' these are used
            to annotate the thickStart and thickEnd in the BED
            file.""")
    parser_gtf2bed.add_argument("gtf", type=argparse.FileType('r'), nargs='?', default=sys.stdin, help="Path to the GTF file.")
    parser_gtf2bed.add_argument("--extraFields",type=str, default='', help="Comma separated list of extra GTF fields to be added after col 12 (e.g. gene_id,gene_name).")
    parser_gtf2bed.add_argument("--filterType",type=str, default='', help="Comma separated list of 'transcript_type' types to retain.")
    parser_gtf2bed.set_defaults(func=lambda args: gtf2bed(args.gtf, args.extraFields.split(','), args.filterType.split(',')))

    parser_bed12tobed6 = subparsers.add_parser('bed12tobed6',
            help="""Converts a BED12 file to BED6 format.
            """)
    parser_bed12tobed6.add_argument("bedfile", type=argparse.FileType('r'), nargs='?', default=sys.stdin, help="Path to the GTF file.")
    parser_bed12tobed6.add_argument("--appendExN", action="store_true", help="Appends the exon number to the transcript name.")
    parser_bed12tobed6.add_argument("--keepIntrons", action="store_true", help="Add records for introns as well.")
    parser_bed12tobed6.set_defaults(func=bed12tobed6)

    args = parser.parse_args()
    args.func(args)

if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before>from beer_search.models import Beer, Style, Region
from django.db.models import Max, Min, Q, Count

def get_update_date():
    return Beer.objects.aggregate(Min("updated_at"))["updated_at__min"]

def perform_filtering(beer_q, request_body):
    """

    :param beer_q: An un-evaluated query for Beer objects.
    :param request_body: A dict-like object of parameters by which beer_q
    should be filtered.
    :return: A filtered queryset.
    """

    # Filtering is performed in succession, on defined parameters only.
    # If the parameter count is zero, no filtering is performed.
    if len(request_body) > 0:

        # Case-insensitive substring check.
        if "beer_name" in request_body:
            name = request_body["beer_name"]
            if len(name) > 0:
                beer_q = beer_q.filter(beer_type__name__icontains=name)

        # Membership check for styles and containers.
        if "styles" in request_body:
            styles = request_body.getlist("styles")
            if len(styles) > 0:
                beer_q = beer_q.filter(beer_type__style__id__in=styles)
        if "containers" in request_body:
            containers = request_body.getlist("containers")
            if len(containers) > 0:
                beer_q = beer_q.filter(container__id__in=containers)

        # Numerical filtering, in order:
        # * minimum and maximum volume
        # * minimum and maximum price
        # * minimum and maximum abv
        if "min_volume" in request_body:
            min_volume = request_body["min_volume"]
            if not len(min_volume) > 0:
                min_volume = 0
            beer_q = beer_q.filter(volume__gte=min_volume)
        if "max_volume" in request_body:
            max_volume = request_body["max_volume"]
            if not len(max_volume) > 0:
                # Setting default in case an empty string is sent
                max_volume = Beer.objects.aggregate(Max("volume"))[
                    "volume__max"]
            beer_q = beer_q.filter(volume__lte=max_volume)

        if "min_price" in request_body:
            min_price = request_body["min_price"]
            if not len(min_price) > 0:
                min_price = 0
            beer_q = beer_q.filter(price__gte=min_price)
        if "max_price" in request_body:
            max_price = request_body["max_price"]
            if not len(max_price) > 0:
                # Setting default in case an empty string is sent
                max_price = Beer.objects.aggregate(Max("price"))[
                    "price__max"]
            beer_q = beer_q.filter(price__lte=max_price)

        if "max_abv" in request_body:
            max_abv = request_body["max_abv"]
            if not len(max_abv) > 0:
                # Setting default in case an empty string is sent
                max_abv = Beer.objects.aggregate(Max("abv"))[
                    "abv__max"]
            beer_q = beer_q.filter(beer_type__abv__lte=max_abv)
        if "min_abv" in request_body:
            min_abv = request_body["min_abv"]
            if not len(min_abv) > 0:
                min_abv = 0
            beer_q = beer_q.filter(beer_type__abv__gte=min_abv)

        # Filter for new and seasonal beers.
        if "noteworthy" in request_body:
            properties = request_body.getlist("noteworthy")
            if "new" in properties and "seasonal" in properties:
                beer_q = beer_q.filter(Q(new=True) | Q(seasonal=True))
            else:
                if "new" in properties:
                    beer_q = beer_q.filter(new=True)
                if "seasonal" in properties:
                    beer_q = beer_q.filter(seasonal=True)

        # Filter by stores.
        if "stores" in request_body:
            stores = request_body.getlist("stores")
            beer_q = beer_q.filter(store__in=stores)

    return beer_q

def num_per_style():
    """

    :return: A dict with style names as keys and the number of beers of
    that style as values.
    """
    style_qs = Style.objects.filter(beer__available=True). \
        annotate(num_beers=Count("beer"))
    return_set = {}
    for style in style_qs.all():
        return_set[style.name] = style.num_beers
    return return_set

def num_per_store():
    """

    :return: A dict with region names as keys, and as values
    a dict that has store names as keys and their beer count as a value.
    """

    regions = Region.objects.all()

    return_set = {}

    for region in regions:
        stores = {}
        for store in region.store_set.all():
            stores[store.location] = store.beers_available.count()
        return_set[region.name] = stores

    return return_set<commit_msg>Reverted reference to beer types in the filtering, those beers that don't have a type were left out in the cold...<commit_after>from beer_search.models import Beer, Style, Region
from django.db.models import Max, Min, Q, Count

def get_update_date():
    return Beer.objects.aggregate(Min("updated_at"))["updated_at__min"]

def perform_filtering(beer_q, request_body):
    """

    :param beer_q: An un-evaluated query for Beer objects.
    :param request_body: A dict-like object of parameters by which beer_q
    should be filtered.
    :return: A filtered queryset.
    """

    # Filtering is performed in succession, on defined parameters only.
    # If the parameter count is zero, no filtering is performed.
    if len(request_body) > 0:

        # Case-insensitive substring check.
        if "beer_name" in request_body:
            name = request_body["beer_name"]
            if len(name) > 0:
                beer_q = beer_q.filter(name__icontains=name)

        # Membership check for styles and containers.
        if "styles" in request_body:
            styles = request_body.getlist("styles")
            if len(styles) > 0:
                beer_q = beer_q.filter(style__id__in=styles)
        if "containers" in request_body:
            containers = request_body.getlist("containers")
            if len(containers) > 0:
                beer_q = beer_q.filter(container__id__in=containers)

        # Numerical filtering, in order:
        # * minimum and maximum volume
        # * minimum and maximum price
        # * minimum and maximum abv
        if "min_volume" in request_body:
            min_volume = request_body["min_volume"]
            if not len(min_volume) > 0:
                min_volume = 0
            beer_q = beer_q.filter(volume__gte=min_volume)
        if "max_volume" in request_body:
            max_volume = request_body["max_volume"]
            if not len(max_volume) > 0:
                # Setting default in case an empty string is sent
                max_volume = Beer.objects.aggregate(Max("volume"))[
                    "volume__max"]
            beer_q = beer_q.filter(volume__lte=max_volume)

        if "min_price" in request_body:
            min_price = request_body["min_price"]
            if not len(min_price) > 0:
                min_price = 0
            beer_q = beer_q.filter(price__gte=min_price)
        if "max_price" in request_body:
            max_price = request_body["max_price"]
            if not len(max_price) > 0:
                # Setting default in case an empty string is sent
                max_price = Beer.objects.aggregate(Max("price"))[
                    "price__max"]
            beer_q = beer_q.filter(price__lte=max_price)

        if "max_abv" in request_body:
            max_abv = request_body["max_abv"]
            if not len(max_abv) > 0:
                # Setting default in case an empty string is sent
                max_abv = Beer.objects.aggregate(Max("abv"))[
                    "abv__max"]
            beer_q = beer_q.filter(abv__lte=max_abv)
        if "min_abv" in request_body:
            min_abv = request_body["min_abv"]
            if not len(min_abv) > 0:
                min_abv = 0
            beer_q = beer_q.filter(abv__gte=min_abv)

        # Filter for new and seasonal beers.
        if "noteworthy" in request_body:
            properties = request_body.getlist("noteworthy")
            if "new" in properties and "seasonal" in properties:
                beer_q = beer_q.filter(Q(new=True) | Q(seasonal=True))
            else:
                if "new" in properties:
                    beer_q = beer_q.filter(new=True)
                if "seasonal" in properties:
                    beer_q = beer_q.filter(seasonal=True)

        # Filter by stores.
        if "stores" in request_body:
            stores = request_body.getlist("stores")
            beer_q = beer_q.filter(store__in=stores)

    return beer_q

def num_per_style():
    """

    :return: A dict with style names as keys and the number of beers of
    that style as values.
    """
    style_qs = Style.objects.filter(beer__available=True). \
        annotate(num_beers=Count("beer"))
    return_set = {}
    for style in style_qs.all():
        return_set[style.name] = style.num_beers
    return return_set

def num_per_store():
    """

    :return: A dict with region names as keys, and as values
    a dict that has store names as keys and their beer count as a value.
    """

    regions = Region.objects.all()

    return_set = {}

    for region in regions:
        stores = {}
        for store in region.store_set.all():
            stores[store.location] = store.beers_available.count()
        return_set[region.name] = stores

    return return_set<|endoftext|>


<commit_before># This file is part of beets.
# Copyright 2013, Jakob Schnitzer.
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.

"""Converts tracks or albums to external directory
"""
import logging
import os
import threading
from subprocess import Popen
import tempfile
from string import Template

from beets.plugins import BeetsPlugin
from beets import ui, util
from beetsplug.embedart import _embed
from beets import config

log = logging.getLogger('beets')
DEVNULL = open(os.devnull, 'wb')
_fs_lock = threading.Lock()
_temp_files = []  # Keep track of temporary transcoded files for deletion.

def _destination(lib, dest_dir, item, keep_new, path_formats):
    """Return the path under `dest_dir` where the file should be placed
    (possibly after conversion).
    """
    dest = lib.destination(item, basedir=dest_dir, path_formats=path_formats)
    if keep_new:
        # When we're keeping the converted file, no extension munging
        # occurs.
        return dest
    else:
        # Otherwise, replace the extension.
        return os.path.splitext(dest)[0] + '.' + config['convert']['extension'].get(unicode)

def encode(source, dest):
    log.info(u'Started encoding {0}'.format(util.displayable_path(source)))

    command = config['transport']['command'].get(unicode).split(u' ')
    opts = []

    for arg in command:
        opts.append(Template(arg).substitute({
            'source':   source,
            'dest':     dest
        }))

    encode = Popen(opts, close_fds=True, stderr=DEVNULL)

    #opts = config['convert']['opts'].get(unicode).split(u' ')
    #encode = Popen([config['convert']['ffmpeg'].get(unicode), '-i',
    #                source, '-y'] + opts + [dest],
    #               close_fds=True, stderr=DEVNULL)

    encode.wait()
    if encode.returncode != 0:
        # Something went wrong (probably Ctrl+C), remove temporary files
        log.info(u'Encoding {0} failed. Cleaning up...'
                 .format(util.displayable_path(source)))
        util.remove(dest)
        util.prune_dirs(os.path.dirname(dest))
        return
    log.info(u'Finished encoding {0}'.format(util.displayable_path(source)))

def should_transcode(item):
    """Determine whether the item should be transcoded as part of
    conversion (i.e., its bitrate is high or it has the wrong format).
    """
    maxbr = config['convert']['max_bitrate'].get(int)
    return item.format != 'MP3' or item.bitrate >= 1000 * maxbr

def convert_item(lib, dest_dir, keep_new, path_formats):
    while True:
        item = yield
        dest = _destination(lib, dest_dir, item, keep_new, path_formats)

        if os.path.exists(util.syspath(dest)):
            log.info(u'Skipping {0} (target file exists)'.format(
                util.displayable_path(item.path)
            ))
            continue

        # Ensure that only one thread tries to create directories at a
        # time. (The existence check is not atomic with the directory
        # creation inside this function.)
        with _fs_lock:
            util.mkdirall(dest)

        # When keeping the new file in the library, we first move the
        # current (pristine) file to the destination. We'll then copy it
        # back to its old path or transcode it to a new path.
        if keep_new:
            log.info(u'Moving to {0}'.
                     format(util.displayable_path(dest)))
            util.move(item.path, dest)

        if not should_transcode(item):
            # No transcoding necessary.
            log.info(u'Copying {0}'.format(util.displayable_path(item.path)))
            if keep_new:
                util.copy(dest, item.path)
            else:
                util.copy(item.path, dest)

        else:
            if keep_new:
                item.path = os.path.splitext(item.path)[0] + '.' + config['convert']['extension'].get(unicode)
                encode(dest, item.path)
            else:
                encode(item.path, dest)

        # Write tags from the database to the converted file.
        if not keep_new:
            item.path = dest
        item.write()

        # If we're keeping the transcoded file, read it again (after
        # writing) to get new bitrate, duration, etc.
        if keep_new:
            item.read()
            item.store()  # Store new path and audio data.

        if config['convert']['embed']:
            album = lib.get_album(item)
            if album:
                artpath = album.artpath
                if artpath:
                    _embed(artpath, [item])

def convert_on_import(lib, item):
    """Transcode a file automatically after it is imported into the
    library.
    """
    if should_transcode(item):
        fd, dest = tempfile.mkstemp('.' + config['convert']['extension'].get(unicode))
        os.close(fd)
        _temp_files.append(dest)  # Delete the transcode later.
        encode(item.path, dest)
        item.path = dest
        item.write()
        item.read()  # Load new audio information data.
        item.store()

def convert_func(lib, opts, args):
    dest = opts.dest if opts.dest is not None else \
            config['convert']['dest'].get()
    if not dest:
        raise ui.UserError('no convert destination set')
    dest = util.bytestring_path(dest)
    threads = opts.threads if opts.threads is not None else \
            config['convert']['threads'].get(int)
    keep_new = opts.keep_new

    if not config['convert']['paths']:
        path_formats = ui.get_path_formats()
    else:
        path_formats = ui.get_path_formats(config['convert']['paths'])

    ui.commands.list_items(lib, ui.decargs(args), opts.album, None)

    if not ui.input_yn("Convert? (Y/n)"):
        return

    if opts.album:
        items = (i for a in lib.albums(ui.decargs(args)) for i in a.items())
    else:
        items = iter(lib.items(ui.decargs(args)))
    convert = [convert_item(lib, dest, keep_new, path_formats) for i in range(threads)]
    pipe = util.pipeline.Pipeline([items, convert])
    pipe.run_parallel()

class ConvertPlugin(BeetsPlugin):
    def __init__(self):
        super(ConvertPlugin, self).__init__()
        self.config.add({
            u'dest': None,
            u'threads': util.cpu_count(),
            u'command': u'ffmpeg -i $source -y -aq 2 $dest',
            u'extension': u'mp3',
            u'max_bitrate': 500,
            u'embed': True,
            u'auto': False,
            u'paths': {},
        })
        self.import_stages = [self.auto_convert]

    def commands(self):
        cmd = ui.Subcommand('convert', help='convert to external location')
        cmd.parser.add_option('-a', '--album', action='store_true',
                              help='choose albums instead of tracks')
        cmd.parser.add_option('-t', '--threads', action='store', type='int',
                              help='change the number of threads, \
                              defaults to maximum available processors')
        cmd.parser.add_option('-k', '--keep-new', action='store_true',
                              dest='keep_new', help='keep only the converted \
                              and move the old files')
        cmd.parser.add_option('-d', '--dest', action='store',
                              help='set the destination directory')
        cmd.func = convert_func
        return [cmd]

    def auto_convert(self, config, task):
        if self.config['auto'].get():
            if not task.is_album:
                convert_on_import(config.lib, task.item)
            else:
                for item in task.items:
                    convert_on_import(config.lib, item)

@ConvertPlugin.listen('import_task_files')
def _cleanup(task, session):
    for path in task.old_paths:
        if path in _temp_files:
            if os.path.isfile(path):
                util.remove(path)
            _temp_files.remove(path)
<commit_msg>Removed reference to my testing plugin, oops.<commit_after># This file is part of beets.
# Copyright 2013, Jakob Schnitzer.
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.

"""Converts tracks or albums to external directory
"""
import logging
import os
import threading
from subprocess import Popen
import tempfile
from string import Template

from beets.plugins import BeetsPlugin
from beets import ui, util
from beetsplug.embedart import _embed
from beets import config

log = logging.getLogger('beets')
DEVNULL = open(os.devnull, 'wb')
_fs_lock = threading.Lock()
_temp_files = []  # Keep track of temporary transcoded files for deletion.

def _destination(lib, dest_dir, item, keep_new, path_formats):
    """Return the path under `dest_dir` where the file should be placed
    (possibly after conversion).
    """
    dest = lib.destination(item, basedir=dest_dir, path_formats=path_formats)
    if keep_new:
        # When we're keeping the converted file, no extension munging
        # occurs.
        return dest
    else:
        # Otherwise, replace the extension.
        return os.path.splitext(dest)[0] + '.' + config['convert']['extension'].get(unicode)

def encode(source, dest):
    log.info(u'Started encoding {0}'.format(util.displayable_path(source)))

    command = config['convert']['command'].get(unicode).split(u' ')
    opts = []

    for arg in command:
        opts.append(Template(arg).substitute({
            'source':   source,
            'dest':     dest
        }))

    encode = Popen(opts, close_fds=True, stderr=DEVNULL)

    #opts = config['convert']['opts'].get(unicode).split(u' ')
    #encode = Popen([config['convert']['ffmpeg'].get(unicode), '-i',
    #                source, '-y'] + opts + [dest],
    #               close_fds=True, stderr=DEVNULL)

    encode.wait()
    if encode.returncode != 0:
        # Something went wrong (probably Ctrl+C), remove temporary files
        log.info(u'Encoding {0} failed. Cleaning up...'
                 .format(util.displayable_path(source)))
        util.remove(dest)
        util.prune_dirs(os.path.dirname(dest))
        return
    log.info(u'Finished encoding {0}'.format(util.displayable_path(source)))

def should_transcode(item):
    """Determine whether the item should be transcoded as part of
    conversion (i.e., its bitrate is high or it has the wrong format).
    """
    maxbr = config['convert']['max_bitrate'].get(int)
    return item.format != 'MP3' or item.bitrate >= 1000 * maxbr

def convert_item(lib, dest_dir, keep_new, path_formats):
    while True:
        item = yield
        dest = _destination(lib, dest_dir, item, keep_new, path_formats)

        if os.path.exists(util.syspath(dest)):
            log.info(u'Skipping {0} (target file exists)'.format(
                util.displayable_path(item.path)
            ))
            continue

        # Ensure that only one thread tries to create directories at a
        # time. (The existence check is not atomic with the directory
        # creation inside this function.)
        with _fs_lock:
            util.mkdirall(dest)

        # When keeping the new file in the library, we first move the
        # current (pristine) file to the destination. We'll then copy it
        # back to its old path or transcode it to a new path.
        if keep_new:
            log.info(u'Moving to {0}'.
                     format(util.displayable_path(dest)))
            util.move(item.path, dest)

        if not should_transcode(item):
            # No transcoding necessary.
            log.info(u'Copying {0}'.format(util.displayable_path(item.path)))
            if keep_new:
                util.copy(dest, item.path)
            else:
                util.copy(item.path, dest)

        else:
            if keep_new:
                item.path = os.path.splitext(item.path)[0] + '.' + config['convert']['extension'].get(unicode)
                encode(dest, item.path)
            else:
                encode(item.path, dest)

        # Write tags from the database to the converted file.
        if not keep_new:
            item.path = dest
        item.write()

        # If we're keeping the transcoded file, read it again (after
        # writing) to get new bitrate, duration, etc.
        if keep_new:
            item.read()
            item.store()  # Store new path and audio data.

        if config['convert']['embed']:
            album = lib.get_album(item)
            if album:
                artpath = album.artpath
                if artpath:
                    _embed(artpath, [item])

def convert_on_import(lib, item):
    """Transcode a file automatically after it is imported into the
    library.
    """
    if should_transcode(item):
        fd, dest = tempfile.mkstemp('.' + config['convert']['extension'].get(unicode))
        os.close(fd)
        _temp_files.append(dest)  # Delete the transcode later.
        encode(item.path, dest)
        item.path = dest
        item.write()
        item.read()  # Load new audio information data.
        item.store()

def convert_func(lib, opts, args):
    dest = opts.dest if opts.dest is not None else \
            config['convert']['dest'].get()
    if not dest:
        raise ui.UserError('no convert destination set')
    dest = util.bytestring_path(dest)
    threads = opts.threads if opts.threads is not None else \
            config['convert']['threads'].get(int)
    keep_new = opts.keep_new

    if not config['convert']['paths']:
        path_formats = ui.get_path_formats()
    else:
        path_formats = ui.get_path_formats(config['convert']['paths'])

    ui.commands.list_items(lib, ui.decargs(args), opts.album, None)

    if not ui.input_yn("Convert? (Y/n)"):
        return

    if opts.album:
        items = (i for a in lib.albums(ui.decargs(args)) for i in a.items())
    else:
        items = iter(lib.items(ui.decargs(args)))
    convert = [convert_item(lib, dest, keep_new, path_formats) for i in range(threads)]
    pipe = util.pipeline.Pipeline([items, convert])
    pipe.run_parallel()

class ConvertPlugin(BeetsPlugin):
    def __init__(self):
        super(ConvertPlugin, self).__init__()
        self.config.add({
            u'dest': None,
            u'threads': util.cpu_count(),
            u'command': u'ffmpeg -i $source -y -aq 2 $dest',
            u'extension': u'mp3',
            u'max_bitrate': 500,
            u'embed': True,
            u'auto': False,
            u'paths': {},
        })
        self.import_stages = [self.auto_convert]

    def commands(self):
        cmd = ui.Subcommand('convert', help='convert to external location')
        cmd.parser.add_option('-a', '--album', action='store_true',
                              help='choose albums instead of tracks')
        cmd.parser.add_option('-t', '--threads', action='store', type='int',
                              help='change the number of threads, \
                              defaults to maximum available processors')
        cmd.parser.add_option('-k', '--keep-new', action='store_true',
                              dest='keep_new', help='keep only the converted \
                              and move the old files')
        cmd.parser.add_option('-d', '--dest', action='store',
                              help='set the destination directory')
        cmd.func = convert_func
        return [cmd]

    def auto_convert(self, config, task):
        if self.config['auto'].get():
            if not task.is_album:
                convert_on_import(config.lib, task.item)
            else:
                for item in task.items:
                    convert_on_import(config.lib, item)

@ConvertPlugin.listen('import_task_files')
def _cleanup(task, session):
    for path in task.old_paths:
        if path in _temp_files:
            if os.path.isfile(path):
                util.remove(path)
            _temp_files.remove(path)
<|endoftext|>


<commit_before>import binascii
import FT245RL_ftdi
import array

class SASEBO:
    def __init__(self):
        self.ADDR_CONT = 0x0002
        self.ADDR_IPSEL = 0x0004
        self.ADDR_OUTSEL = 0x0008
        self.ADDR_MODE = 0x000C
        self.ADDR_RSEL = 0x000E
        self.ADDR_KEY0 = 0x0100
        self.ADDR_ITEXT0 = 0x0140
        self.ADDR_OTEXT0 = 0x0180
        self.ADDR_VERSION = 0xFFFC

        self.MODE_ENC = 0x0000
        self.MODE_DEC = 0x0001

        # idVendor = 0x0403
        # idProduct = 0x6001
        self.port = FT245RL_ftdi.FTDI_USB()  #
        # device with idVendor =
        # 0x0403, idProduct = 0x6001

    # ------------------------------------------------------------------ open()
    def open(self):
        self.port.open()

        # Reset
        self.write(self.ADDR_CONT, 0x0004)
        self.write(self.ADDR_CONT, 0x0004)

    # ----------------------------------------------------------------- close()
    def close(self):
        self.port.close()

    # ---------------------------------------------------------------- setKey()
    def setKey(self, key, len):
        self.writeBurst(self.ADDR_KEY0, key, len)
        self.write(self.ADDR_CONT, 0x0002)  # Execute key generation
        while self.read(self.ADDR_CONT) != 0:
            pass

    # ---------------------------------------------------------------- setEnc()
    def setEnc(self):
        self.write(self.ADDR_MODE, self.MODE_ENC)

    # ---------------------------------------------------------------- setDec()
    def setDec(self):
        self.write(self.ADDR_MODE, self.MODE_DEC)

    # ------------------------------------------------------------- writeText()
    def writeText(self, text, len):
        self.writeBurst(self.ADDR_ITEXT0, text, len)

    # -------------------------------------------------------------- readText()
    def readText(self, text, len):
        # print text
        self.readBurst(self.ADDR_OTEXT0, text, len)

    # ----------------------------------------------------------------execute()
    def execute(self):
        self.write(self.ADDR_CONT, 0x0001)  # Execute cipher processing
        while self.read(self.ADDR_CONT) != 0:
            pass

    # ----------------------------------------------------------------- write()
    def write(self, addr, data):  # write function from SASEBO
        # buf = bytearray(5)
        buf = array.array('B', [0] * 5)
        buf[0] = 0x01
        buf[1] = ((addr >> 8) & 0xFF)
        buf[2] = ((addr) & 0xFF)
        buf[3] = ((data >> 8) & 0xFF)
        buf[4] = ((data) & 0xFF)
        return buf  # for testing purpose
        # port.write(buf,5)   # write function from FTDI

    # ------------------------------------------------------------------ read()
    def read(self, addr):
        buf = array.array('B', [0] * 3)
        buf[0] = 0x00
        buf[1] = ((addr >> 8) & 0xFF)
        buf[2] = ((addr) & 0xFF)
        self.port.write(buf, 3)
        self.port.read(buf, 2)
        return (buf[0] << 8) + buf[1]

    def writeBurst(self, addr, data, len):
        buf = array.array('B', [0] * (5 * len / 2))
        for i in range(len / 2):
            buf[i * 5 + 0] = 0x01
            buf[i * 5 + 1] = (((addr + i * 2) >> 8) & 0xFF)
            buf[i * 5 + 2] = (((addr + i * 2)) & 0xFF)
            buf[i * 5 + 3] = data[i * 2]
            buf[i * 5 + 4] = data[i * 2 + 1]
        self.port.write(buf, 5 * (len / 2))

    # ------------------------------------------------------------- readBurst()
    def readBurst(self, addr, data, len):
        buf = array.array('B', [0] * (3 * len / 2))
        for i in range(len / 2):
            buf[i * 3 + 0] = 0x00
            buf[i * 3 + 1] = (((addr + i * 2) >> 8) & 0xFF)
            buf[i * 3 + 2] = (((addr + i * 2)) & 0xFF)
        print binascii.b2a_hex(buf).upper()  # testing
        print binascii.b2a_hex(data).upper()  # testing
        self.port.write(buf, 3 * (len / 2))
        self.port.read(data, len)

# testing purpose

sasebo = SASEBO()
# write = sasebo.write(0x00C, 45)
# print binascii.hexlify(write).upper()
text_out = array.array('B', [0] * 16)
sasebo.open()
sasebo.setEnc()
sasebo.readText(text_out, 16)
sasebo.close()
<commit_msg>More functions added<commit_after>#!/usr/bin/env python
"""Sasebo_ftdi.py: Class for Sasebo-G device."""

__author__ = 'Ionut Gorgos'
__copyright__ = 'Copyright (c) 2016 Ionut Gorgos'

import binascii
import FT245RL_ftdi
import array

class SASEBO:
    def __init__(self):
        self.ADDR_CONT = 0x0002
        self.ADDR_IPSEL = 0x0004
        self.ADDR_OUTSEL = 0x0008
        self.ADDR_MODE = 0x000C
        self.ADDR_RSEL = 0x000E
        self.ADDR_KEY0 = 0x0100
        self.ADDR_ITEXT0 = 0x0140
        self.ADDR_OTEXT0 = 0x0180
        self.ADDR_VERSION = 0xFFFC

        # Cipher Mode
        self.MODE_ENC = 0x0000
        self.MODE_DEC = 0x0001

        # idVendor = 0x0403
        # idProduct = 0x6001
        self.port = FT245RL_ftdi.FTDI_USB()  # device with idVendor =0x0403,
                                                        # idProduct = 0x6001

    # ------------------------------------------------------------------ open()
    def open(self):
        self.port.open()

        # Reset
        self.write(self.ADDR_CONT, 0x0004)
        self.write(self.ADDR_CONT, 0x0004)

    # ----------------------------------------------------------------- close()
    def close(self):
        self.port.close()

    # ---------------------------------------------------------------- setKey()
    def setKey(self, key, len):
        self.writeBurst(self.ADDR_KEY0, key, len)
        self.write(self.ADDR_CONT, 0x0002)  # Execute key generation
        while self.read(self.ADDR_CONT) != 0:
            pass

    # ---------------------------------------------------------------- setEnc()
    def setEnc(self):
        self.write(self.ADDR_MODE, self.MODE_ENC)

    # ---------------------------------------------------------------- setDec()
    def setDec(self):
        self.write(self.ADDR_MODE, self.MODE_DEC)

    # ------------------------------------------------------------- writeText()
    def writeText(self, text, len):
        self.writeBurst(self.ADDR_ITEXT0, text, len)

    # -------------------------------------------------------------- readText()
    def readText(self, text, len):
        # print text
        self.readBurst(self.ADDR_OTEXT0, text, len)

    # ----------------------------------------------------------------execute()
    def execute(self):
        self.write(self.ADDR_CONT, 0x0001)  # Execute cipher processing
        while self.read(self.ADDR_CONT) != 0:
            pass

    # ----------------------------------------------------------------- write()
    def write(self, addr, data):  # write function from SASEBO
        # buf = bytearray(5)
        buf = array.array('B', [0] * 5)
        buf[0] = 0x01
        buf[1] = ((addr >> 8) & 0xFF)
        buf[2] = ((addr) & 0xFF)
        buf[3] = ((data >> 8) & 0xFF)
        buf[4] = ((data) & 0xFF)
        return buf  # for testing purpose
        # port.write(buf,5)   # write function from FTDI

    # ------------------------------------------------------------------ read()
    def read(self, addr):
        buf = array.array('B', [0] * 3)
        buf[0] = 0x00
        buf[1] = ((addr >> 8) & 0xFF)
        buf[2] = ((addr) & 0xFF)
        self.port.write(buf, 3)
        self.port.read(buf, 2)
        return (buf[0] << 8) + buf[1]

    def writeBurst(self, addr, data, len):
        buf = array.array('B', [0] * (5 * len / 2))
        for i in range(len / 2):
            buf[i * 5 + 0] = 0x01
            buf[i * 5 + 1] = (((addr + i * 2) >> 8) & 0xFF)
            buf[i * 5 + 2] = (((addr + i * 2)) & 0xFF)
            buf[i * 5 + 3] = data[i * 2]
            buf[i * 5 + 4] = data[i * 2 + 1]
        self.port.write(buf, 5 * (len / 2))

    # ------------------------------------------------------------- readBurst()
    def readBurst(self, addr, data, len):
        buf = array.array('B', [0] * (3 * len / 2))
        for i in range(len / 2):
            buf[i * 3 + 0] = 0x00
            buf[i * 3 + 1] = (((addr + i * 2) >> 8) & 0xFF)
            buf[i * 3 + 2] = (((addr + i * 2)) & 0xFF)
        print binascii.b2a_hex(buf).upper()  # testing
        print binascii.b2a_hex(data).upper()  # testing
        self.port.write(buf, 3 * (len / 2))
        self.port.read(data, len)

# testing purpose

sasebo = SASEBO()
# write = sasebo.write(0x00C, 45)
# print binascii.hexlify(write).upper()
text_out = array.array('B', [0] * 16)
sasebo.open()
sasebo.setEnc()
sasebo.readText(text_out, 16)
sasebo.close()
<|endoftext|>


<commit_before># This file is part of beets.
# Copyright 2015, Adrian Sampson.
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.

"""Adds Discogs album search support to the autotagger. Requires the
discogs-client library.
"""
from __future__ import (division, absolute_import, print_function,
                        unicode_literals)

import beets.ui
from beets import logging
from beets.autotag.hooks import AlbumInfo, TrackInfo, Distance
from beets.plugins import BeetsPlugin
from beets.util import confit
from discogs_client import Release, Client
from discogs_client.exceptions import DiscogsAPIError
from requests.exceptions import ConnectionError
import beets
import re
import time
import json
import socket
import httplib
import os

# Silence spurious INFO log lines generated by urllib3.
urllib3_logger = logging.getLogger('requests.packages.urllib3')
urllib3_logger.setLevel(logging.CRITICAL)

USER_AGENT = 'beets/{0} +http://beets.radbox.org/'.format(beets.__version__) \
    .encode('utf8')

# Exceptions that discogs_client should really handle but does not.
CONNECTION_ERRORS = (ConnectionError, socket.error, httplib.HTTPException,
                     ValueError,  # JSON decoding raises a ValueError.
                     DiscogsAPIError)

class DiscogsPlugin(BeetsPlugin):

    def __init__(self):
        super(DiscogsPlugin, self).__init__()
        self.config.add({
            'apikey': 'rAzVUQYRaoFjeBjyWuWZ',
            'apisecret': 'plxtUTqoCzwxZpqdPysCwGuBSmZNdZVy',
            'tokenfile': 'discogs_token.json',
            'source_weight': 0.5,
        })
        self.config['apikey'].redact = True
        self.config['apisecret'].redact = True
        self.discogs_client = None
        self.register_listener('import_begin', self.setup)

    def setup(self, session=None):
        """Create the `discogs_client` field. Authenticate if necessary.
        """
        c_key = self.config['apikey'].get(unicode).encode('utf8')
        c_secret = self.config['apisecret'].get(unicode).encode('utf8')

        # Get the OAuth token from a file or log in.
        try:
            with open(self._tokenfile()) as f:
                tokendata = json.load(f)
        except IOError:
            # No token yet. Generate one.
            token, secret = self.authenticate(c_key, c_secret)
        else:
            token = tokendata['token'].encode('utf8')
            secret = tokendata['secret'].encode('utf8')

        self.discogs_client = Client(USER_AGENT, c_key, c_secret,
                                     token, secret)

    def reset_auth(self):
        """Delete toke file & redo the auth steps.
        """
        os.remove(self._tokenfile())
        self.setup()

    def _tokenfile(self):
        """Get the path to the JSON file for storing the OAuth token.
        """
        return self.config['tokenfile'].get(confit.Filename(in_app_dir=True))

    def authenticate(self, c_key, c_secret):
        # Get the link for the OAuth page.
        auth_client = Client(USER_AGENT, c_key, c_secret)
        try:
            _, _, url = auth_client.get_authorize_url()
        except CONNECTION_ERRORS as e:
            self._log.debug('connection error: {0}', e)
            raise beets.ui.UserError('communication with Discogs failed')

        beets.ui.print_("To authenticate with Discogs, visit:")
        beets.ui.print_(url)

        # Ask for the code and validate it.
        code = beets.ui.input_("Enter the code:")
        try:
            token, secret = auth_client.get_access_token(code)
        except DiscogsAPIError:
            raise beets.ui.UserError('Discogs authorization failed')
        except CONNECTION_ERRORS as e:
            self._log.debug(u'connection error: {0}', e)
            raise beets.ui.UserError('Discogs token request failed')

        # Save the token for later use.
        self._log.debug('Discogs token {0}, secret {1}', token, secret)
        with open(self._tokenfile(), 'w') as f:
            json.dump({'token': token, 'secret': secret}, f)

        return token.encode('utf8'), secret.encode('utf8')

    def album_distance(self, items, album_info, mapping):
        """Returns the album distance.
        """
        dist = Distance()
        if album_info.data_source == 'Discogs':
            dist.add('source', self.config['source_weight'].as_number())
        return dist

    def candidates(self, items, artist, album, va_likely):
        """Returns a list of AlbumInfo objects for discogs search results
        matching an album and artist (if not various).
        """
        if not self.discogs_client:
            return

        if va_likely:
            query = album
        else:
            query = '%s %s' % (artist, album)
        try:
            return self.get_albums(query)
        except DiscogsAPIError as e:
            self._log.debug(u'API Error: {0} (query: {1})', e, query)
            if e.status_code == 401:
                self.reset_auth()
                return self.candidates(items, artist, album, va_likely)
            else:
                return []
        except CONNECTION_ERRORS as e:
            self._log.debug(u'HTTP Connection Error: {0}', e)
            return []

    def album_for_id(self, album_id):
        """Fetches an album by its Discogs ID and returns an AlbumInfo object
        or None if the album is not found.
        """
        if not self.discogs_client:
            return

        self._log.debug(u'Searching for release {0}', album_id)
        # Discogs-IDs are simple integers. We only look for those at the end
        # of an input string as to avoid confusion with other metadata plugins.
        # An optional bracket can follow the integer, as this is how discogs
        # displays the release ID on its webpage.
        match = re.search(r'(^|\[*r|discogs\.com/.+/release/)(\d+)($|\])',
                          album_id)
        if not match:
            return None
        result = Release(self.discogs_client, {'id': int(match.group(2))})
        # Try to obtain title to verify that we indeed have a valid Release
        try:
            getattr(result, 'title')
        except DiscogsAPIError as e:
            if e.status_code != 404:
                self._log.debug(u'API Error: {0} (query: {1})', e, result._uri)
                if e.status_code == 401:
                    self.reset_auth()
                    return self.album_for_id(album_id)
            return None
        except CONNECTION_ERRORS as e:
            self._log.debug(u'HTTP Connection Error: {0}', e)
            return None
        return self.get_album_info(result)

    def get_albums(self, query):
        """Returns a list of AlbumInfo objects for a discogs search query.
        """
        # Strip non-word characters from query. Things like "!" and "-" can
        # cause a query to return no results, even if they match the artist or
        # album title. Use `re.UNICODE` flag to avoid stripping non-english
        # word characters.
        # TEMPORARY: Encode as ASCII to work around a bug:
        # https://github.com/sampsyo/beets/issues/1051
        # When the library is fixed, we should encode as UTF-8.
        query = re.sub(r'(?u)\W+', ' ', query).encode('ascii', "replace")
        # Strip medium information from query, Things like "CD1" and "disk 1"
        # can also negate an otherwise positive result.
        query = re.sub(r'(?i)\b(CD|disc)\s*\d+', '', query)
        try:
            releases = self.discogs_client.search(query,
                                                  type='release').page(1)
        except CONNECTION_ERRORS as exc:
            self._log.debug("Communication error while searching for {0!r}: "
                            "{1}".format(query, exc))
            return []
        return [self.get_album_info(release) for release in releases[:5]]

    def get_album_info(self, result):
        """Returns an AlbumInfo object for a discogs Release object.
        """
        artist, artist_id = self.get_artist([a.data for a in result.artists])
        album = re.sub(r' +', ' ', result.title)
        album_id = result.data['id']
        # Use `.data` to access the tracklist directly instead of the
        # convenient `.tracklist` property, which will strip out useful artist
        # information and leave us with skeleton `Artist` objects that will
        # each make an API call just to get the same data back.
        tracks = self.get_tracks(result.data['tracklist'])
        albumtype = ', '.join(
            result.data['formats'][0].get('descriptions', [])) or None
        va = result.data['artists'][0]['name'].lower() == 'various'
        year = result.data['year']
        label = result.data['labels'][0]['name']
        mediums = len(set(t.medium for t in tracks))
        catalogno = result.data['labels'][0]['catno']
        if catalogno == 'none':
            catalogno = None
        country = result.data.get('country')
        media = result.data['formats'][0]['name']
        data_url = result.data['uri']
        return AlbumInfo(album, album_id, artist, artist_id, tracks, asin=None,
                         albumtype=albumtype, va=va, year=year, month=None,
                         day=None, label=label, mediums=mediums,
                         artist_sort=None, releasegroup_id=None,
                         catalognum=catalogno, script=None, language=None,
                         country=country, albumstatus=None, media=media,
                         albumdisambig=None, artist_credit=None,
                         original_year=None, original_month=None,
                         original_day=None, data_source='Discogs',
                         data_url=data_url)

    def get_artist(self, artists):
        """Returns an artist string (all artists) and an artist_id (the main
        artist) for a list of discogs album or track artists.
        """
        artist_id = None
        bits = []
        for i, artist in enumerate(artists):
            if not artist_id:
                artist_id = artist['id']
            name = artist['name']
            # Strip disambiguation number.
            name = re.sub(r' \(\d+\)$', '', name)
            # Move articles to the front.
            name = re.sub(r'(?i)^(.*?), (a|an|the)$', r'\2 \1', name)
            bits.append(name)
            if artist['join'] and i < len(artists) - 1:
                bits.append(artist['join'])
        artist = ' '.join(bits).replace(' ,', ',') or None
        return artist, artist_id

    def get_tracks(self, tracklist):
        """Returns a list of TrackInfo objects for a discogs tracklist.
        """
        tracks = []
        index_tracks = {}
        index = 0
        for track in tracklist:
            # Only real tracks have `position`. Otherwise, it's an index track.
            if track['position']:
                index += 1
                tracks.append(self.get_track_info(track, index))
            else:
                index_tracks[index + 1] = track['title']

        # Fix up medium and medium_index for each track. Discogs position is
        # unreliable, but tracks are in order.
        medium = None
        medium_count, index_count = 0, 0
        for track in tracks:
            # Handle special case where a different medium does not indicate a
            # new disc, when there is no medium_index and the ordinal of medium
            # is not sequential. For example, I, II, III, IV, V. Assume these
            # are the track index, not the medium.
            medium_is_index = track.medium and not track.medium_index and (
                len(track.medium) != 1 or
                ord(track.medium) - 64 != medium_count + 1
            )

            if not medium_is_index and medium != track.medium:
                # Increment medium_count and reset index_count when medium
                # changes.
                medium = track.medium
                medium_count += 1
                index_count = 0
            index_count += 1
            track.medium, track.medium_index = medium_count, index_count

        # Get `disctitle` from Discogs index tracks. Assume that an index track
        # before the first track of each medium is a disc title.
        for track in tracks:
            if track.medium_index == 1:
                if track.index in index_tracks:
                    disctitle = index_tracks[track.index]
                else:
                    disctitle = None
            track.disctitle = disctitle

        return tracks

    def get_track_info(self, track, index):
        """Returns a TrackInfo object for a discogs track.
        """
        title = track['title']
        track_id = None
        medium, medium_index = self.get_track_index(track['position'])
        artist, artist_id = self.get_artist(track.get('artists', []))
        length = self.get_track_length(track['duration'])
        return TrackInfo(title, track_id, artist, artist_id, length, index,
                         medium, medium_index, artist_sort=None,
                         disctitle=None, artist_credit=None)

    def get_track_index(self, position):
        """Returns the medium and medium index for a discogs track position.
        """
        # medium_index is a number at the end of position. medium is everything
        # else. E.g. (A)(1), (Side A, Track )(1), (A)(), ()(1), etc.
        match = re.match(r'^(.*?)(\d*)$', position.upper())
        if match:
            medium, index = match.groups()
        else:
            self._log.debug(u'Invalid position: {0}', position)
            medium = index = None
        return medium or None, index or None

    def get_track_length(self, duration):
        """Returns the track length in seconds for a discogs duration.
        """
        try:
            length = time.strptime(duration, '%M:%S')
        except ValueError:
            return None
        return length.tm_min * 60 + length.tm_sec
<commit_msg>Discogs: option to change 'various' album artist to 'Various Artists' (Musicbrainz naming)<commit_after># This file is part of beets.
# Copyright 2015, Adrian Sampson.
#
# Permission is hereby granted, free of charge, to any person obtaining
# a copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.

"""Adds Discogs album search support to the autotagger. Requires the
discogs-client library.
"""
from __future__ import (division, absolute_import, print_function,
                        unicode_literals)

import beets.ui
from beets import logging
from beets.autotag.hooks import AlbumInfo, TrackInfo, Distance
from beets.plugins import BeetsPlugin
from beets.util import confit
from discogs_client import Release, Client
from discogs_client.exceptions import DiscogsAPIError
from requests.exceptions import ConnectionError
import beets
import re
import time
import json
import socket
import httplib
import os

# Silence spurious INFO log lines generated by urllib3.
urllib3_logger = logging.getLogger('requests.packages.urllib3')
urllib3_logger.setLevel(logging.CRITICAL)

USER_AGENT = 'beets/{0} +http://beets.radbox.org/'.format(beets.__version__) \
    .encode('utf8')

# Exceptions that discogs_client should really handle but does not.
CONNECTION_ERRORS = (ConnectionError, socket.error, httplib.HTTPException,
                     ValueError,  # JSON decoding raises a ValueError.
                     DiscogsAPIError)

class DiscogsPlugin(BeetsPlugin):

    def __init__(self):
        super(DiscogsPlugin, self).__init__()
        self.config.add({
            'apikey': 'rAzVUQYRaoFjeBjyWuWZ',
            'apisecret': 'plxtUTqoCzwxZpqdPysCwGuBSmZNdZVy',
            'tokenfile': 'discogs_token.json',
            'source_weight': 0.5,
            'change_va': False,
        })
        self.config['apikey'].redact = True
        self.config['apisecret'].redact = True
        self.discogs_client = None
        self.register_listener('import_begin', self.setup)

    def setup(self, session=None):
        """Create the `discogs_client` field. Authenticate if necessary.
        """
        c_key = self.config['apikey'].get(unicode).encode('utf8')
        c_secret = self.config['apisecret'].get(unicode).encode('utf8')

        # Get the OAuth token from a file or log in.
        try:
            with open(self._tokenfile()) as f:
                tokendata = json.load(f)
        except IOError:
            # No token yet. Generate one.
            token, secret = self.authenticate(c_key, c_secret)
        else:
            token = tokendata['token'].encode('utf8')
            secret = tokendata['secret'].encode('utf8')

        self.discogs_client = Client(USER_AGENT, c_key, c_secret,
                                     token, secret)

    def reset_auth(self):
        """Delete toke file & redo the auth steps.
        """
        os.remove(self._tokenfile())
        self.setup()

    def _tokenfile(self):
        """Get the path to the JSON file for storing the OAuth token.
        """
        return self.config['tokenfile'].get(confit.Filename(in_app_dir=True))

    def authenticate(self, c_key, c_secret):
        # Get the link for the OAuth page.
        auth_client = Client(USER_AGENT, c_key, c_secret)
        try:
            _, _, url = auth_client.get_authorize_url()
        except CONNECTION_ERRORS as e:
            self._log.debug('connection error: {0}', e)
            raise beets.ui.UserError('communication with Discogs failed')

        beets.ui.print_("To authenticate with Discogs, visit:")
        beets.ui.print_(url)

        # Ask for the code and validate it.
        code = beets.ui.input_("Enter the code:")
        try:
            token, secret = auth_client.get_access_token(code)
        except DiscogsAPIError:
            raise beets.ui.UserError('Discogs authorization failed')
        except CONNECTION_ERRORS as e:
            self._log.debug(u'connection error: {0}', e)
            raise beets.ui.UserError('Discogs token request failed')

        # Save the token for later use.
        self._log.debug('Discogs token {0}, secret {1}', token, secret)
        with open(self._tokenfile(), 'w') as f:
            json.dump({'token': token, 'secret': secret}, f)

        return token.encode('utf8'), secret.encode('utf8')

    def album_distance(self, items, album_info, mapping):
        """Returns the album distance.
        """
        dist = Distance()
        if album_info.data_source == 'Discogs':
            dist.add('source', self.config['source_weight'].as_number())
        return dist

    def candidates(self, items, artist, album, va_likely):
        """Returns a list of AlbumInfo objects for discogs search results
        matching an album and artist (if not various).
        """
        if not self.discogs_client:
            return

        if va_likely:
            query = album
        else:
            query = '%s %s' % (artist, album)
        try:
            return self.get_albums(query)
        except DiscogsAPIError as e:
            self._log.debug(u'API Error: {0} (query: {1})', e, query)
            if e.status_code == 401:
                self.reset_auth()
                return self.candidates(items, artist, album, va_likely)
            else:
                return []
        except CONNECTION_ERRORS as e:
            self._log.debug(u'HTTP Connection Error: {0}', e)
            return []

    def album_for_id(self, album_id):
        """Fetches an album by its Discogs ID and returns an AlbumInfo object
        or None if the album is not found.
        """
        if not self.discogs_client:
            return

        self._log.debug(u'Searching for release {0}', album_id)
        # Discogs-IDs are simple integers. We only look for those at the end
        # of an input string as to avoid confusion with other metadata plugins.
        # An optional bracket can follow the integer, as this is how discogs
        # displays the release ID on its webpage.
        match = re.search(r'(^|\[*r|discogs\.com/.+/release/)(\d+)($|\])',
                          album_id)
        if not match:
            return None
        result = Release(self.discogs_client, {'id': int(match.group(2))})
        # Try to obtain title to verify that we indeed have a valid Release
        try:
            getattr(result, 'title')
        except DiscogsAPIError as e:
            if e.status_code != 404:
                self._log.debug(u'API Error: {0} (query: {1})', e, result._uri)
                if e.status_code == 401:
                    self.reset_auth()
                    return self.album_for_id(album_id)
            return None
        except CONNECTION_ERRORS as e:
            self._log.debug(u'HTTP Connection Error: {0}', e)
            return None
        return self.get_album_info(result)

    def get_albums(self, query):
        """Returns a list of AlbumInfo objects for a discogs search query.
        """
        # Strip non-word characters from query. Things like "!" and "-" can
        # cause a query to return no results, even if they match the artist or
        # album title. Use `re.UNICODE` flag to avoid stripping non-english
        # word characters.
        # TEMPORARY: Encode as ASCII to work around a bug:
        # https://github.com/sampsyo/beets/issues/1051
        # When the library is fixed, we should encode as UTF-8.
        query = re.sub(r'(?u)\W+', ' ', query).encode('ascii', "replace")
        # Strip medium information from query, Things like "CD1" and "disk 1"
        # can also negate an otherwise positive result.
        query = re.sub(r'(?i)\b(CD|disc)\s*\d+', '', query)
        try:
            releases = self.discogs_client.search(query,
                                                  type='release').page(1)
        except CONNECTION_ERRORS as exc:
            self._log.debug("Communication error while searching for {0!r}: "
                            "{1}".format(query, exc))
            return []
        return [self.get_album_info(release) for release in releases[:5]]

    def get_album_info(self, result):
        """Returns an AlbumInfo object for a discogs Release object.
        """
        artist, artist_id = self.get_artist([a.data for a in result.artists])
        album = re.sub(r' +', ' ', result.title)
        album_id = result.data['id']
        # Use `.data` to access the tracklist directly instead of the
        # convenient `.tracklist` property, which will strip out useful artist
        # information and leave us with skeleton `Artist` objects that will
        # each make an API call just to get the same data back.
        tracks = self.get_tracks(result.data['tracklist'])
        albumtype = ', '.join(
            result.data['formats'][0].get('descriptions', [])) or None
        va = result.data['artists'][0]['name'].lower() == 'various'
        if va and self.config['change_va']:
           artist = 'Various Artists'
        year = result.data['year']
        label = result.data['labels'][0]['name']
        mediums = len(set(t.medium for t in tracks))
        catalogno = result.data['labels'][0]['catno']
        if catalogno == 'none':
            catalogno = None
        country = result.data.get('country')
        media = result.data['formats'][0]['name']
        data_url = result.data['uri']
        return AlbumInfo(album, album_id, artist, artist_id, tracks, asin=None,
                         albumtype=albumtype, va=va, year=year, month=None,
                         day=None, label=label, mediums=mediums,
                         artist_sort=None, releasegroup_id=None,
                         catalognum=catalogno, script=None, language=None,
                         country=country, albumstatus=None, media=media,
                         albumdisambig=None, artist_credit=None,
                         original_year=None, original_month=None,
                         original_day=None, data_source='Discogs',
                         data_url=data_url)

    def get_artist(self, artists):
        """Returns an artist string (all artists) and an artist_id (the main
        artist) for a list of discogs album or track artists.
        """
        artist_id = None
        bits = []
        for i, artist in enumerate(artists):
            if not artist_id:
                artist_id = artist['id']
            name = artist['name']
            # Strip disambiguation number.
            name = re.sub(r' \(\d+\)$', '', name)
            # Move articles to the front.
            name = re.sub(r'(?i)^(.*?), (a|an|the)$', r'\2 \1', name)
            bits.append(name)
            if artist['join'] and i < len(artists) - 1:
                bits.append(artist['join'])
        artist = ' '.join(bits).replace(' ,', ',') or None
        return artist, artist_id

    def get_tracks(self, tracklist):
        """Returns a list of TrackInfo objects for a discogs tracklist.
        """
        tracks = []
        index_tracks = {}
        index = 0
        for track in tracklist:
            # Only real tracks have `position`. Otherwise, it's an index track.
            if track['position']:
                index += 1
                tracks.append(self.get_track_info(track, index))
            else:
                index_tracks[index + 1] = track['title']

        # Fix up medium and medium_index for each track. Discogs position is
        # unreliable, but tracks are in order.
        medium = None
        medium_count, index_count = 0, 0
        for track in tracks:
            # Handle special case where a different medium does not indicate a
            # new disc, when there is no medium_index and the ordinal of medium
            # is not sequential. For example, I, II, III, IV, V. Assume these
            # are the track index, not the medium.
            medium_is_index = track.medium and not track.medium_index and (
                len(track.medium) != 1 or
                ord(track.medium) - 64 != medium_count + 1
            )

            if not medium_is_index and medium != track.medium:
                # Increment medium_count and reset index_count when medium
                # changes.
                medium = track.medium
                medium_count += 1
                index_count = 0
            index_count += 1
            track.medium, track.medium_index = medium_count, index_count

        # Get `disctitle` from Discogs index tracks. Assume that an index track
        # before the first track of each medium is a disc title.
        for track in tracks:
            if track.medium_index == 1:
                if track.index in index_tracks:
                    disctitle = index_tracks[track.index]
                else:
                    disctitle = None
            track.disctitle = disctitle

        return tracks

    def get_track_info(self, track, index):
        """Returns a TrackInfo object for a discogs track.
        """
        title = track['title']
        track_id = None
        medium, medium_index = self.get_track_index(track['position'])
        artist, artist_id = self.get_artist(track.get('artists', []))
        length = self.get_track_length(track['duration'])
        return TrackInfo(title, track_id, artist, artist_id, length, index,
                         medium, medium_index, artist_sort=None,
                         disctitle=None, artist_credit=None)

    def get_track_index(self, position):
        """Returns the medium and medium index for a discogs track position.
        """
        # medium_index is a number at the end of position. medium is everything
        # else. E.g. (A)(1), (Side A, Track )(1), (A)(), ()(1), etc.
        match = re.match(r'^(.*?)(\d*)$', position.upper())
        if match:
            medium, index = match.groups()
        else:
            self._log.debug(u'Invalid position: {0}', position)
            medium = index = None
        return medium or None, index or None

    def get_track_length(self, duration):
        """Returns the track length in seconds for a discogs duration.
        """
        try:
            length = time.strptime(duration, '%M:%S')
        except ValueError:
            return None
        return length.tm_min * 60 + length.tm_sec
<|endoftext|>