<commit_before>import hashlib
import logging
import mimetypes
import os
import subprocess
import time

from dropbox.client import DropboxClient

from app import analytics
from app import celery
from app import db
from app import emailer
from app import redis
from app.models import User, Book

log = logging.getLogger()

# Lock expires in 30 minutes, in case there are lots of epubs to convert.
LOCK_EXPIRE = 60 * 30

BASE_DIR = '/tmp/kindlebox'
try:
    os.makedirs(BASE_DIR)
except OSError:
    pass

# And can only email 25 books at a time. Sendgrid only allows 20MB at a time,
# after encoding to email text, so more like 15.
BOOK_ATTACHMENTS_LIMIT = 25
ATTACHMENTS_SIZE_LIMIT = 15 * (10**6)

# Supported filetypes.
# According to:
# http://www.amazon.com/gp/help/customer/display.html?nodeId=200375630
mimetypes.add_type('application/x-mobipocket-ebook', '.mobi')
mimetypes.add_type('application/x-mobipocket-ebook', '.prc')
mimetypes.add_type('application/vnd.amazon.ebook', '.azw')
mimetypes.add_type('application/vnd.amazon.ebook', '.azw1')

# Amazon doesn't support epub, but we do!
mimetypes.add_type('application/epub+zip', '.epub')
EPUB_MIMETYPE = 'application/epub+zip'

BOOK_MIMETYPES = {
    'application/epub+zip',
    'application/vnd.amazon.ebook',
    'text/plain',
    'application/x-mobipocket-ebook',
    'application/pdf',

    'application/msword',
    'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
    'application/rtf',
    'text/html',

    'image/jpeg',
    'image/gif',
    'image/x-ms-bmp',
    'image/png',
    }
BOOK_CHUNK = 5

@celery.task(ignore_result=True)
def upload_welcome_pdf(dropbox_id):
    user = User.query.filter_by(dropbox_id=dropbox_id,
                                active=True).first()
    if user is None:
        return False

    # If we've already sent the welcome PDF before, Dropbox webhook went
    # trigger, so do it here.
    if user.uploaded_welcome_pdf:
        return kindlebox(dropbox_id)

    analytics.track(str(user.id), 'Sent welcome pdf')

    client = DropboxClient(user.access_token)
    try:
        with open('app/static/kindlebox_welcome.pdf', 'rb') as f:
            response = client.put_file('Welcome to Kindlebox.pdf', f)
            if response:
                log.info("Welcome PDF sent to Dropbox ID {0}.".format(dropbox_id))
            else:
                raise Exception("No response received after sending welcome PDF")

        user.set_uploaded_welcome_pdf()
        db.session.commit()

    except:
        log.error(("Welcome PDF failed for Dropbox ID "
                   "{0}.").format(dropbox_id), exc_info=True)
        return False

    return True

def _kindlebox(dropbox_id, user, client):
    delta = client.delta(user.cursor)

    added_book_sizes = dict(get_added_book_sizes(delta['entries'], client))
    removed_books = get_removed_books(delta['entries'])
    log.debug("Delta contains {0} added books, {1} removed "
              "books".format(len(added_book_sizes), len(removed_books)))

    if len(added_book_sizes) == 0 and len(removed_books) == 0:
        return True

    # Download the changed files and get the hashes.
    # Also record the hashes/paths of any newly added books.
    new_books, duplicate_books = get_new_and_duplicates(client,
                                                        added_book_sizes,
                                                        user.id)

    # Email ze books.
    # Add all books that were supposed to be emailed. If email was
    # unsuccessful, mark book as unsent.
    email_from = user.emailer
    email_to = [row.kindle_name + '@kindle.com' for row in user.kindle_names.all()]
    attachments = {}
    attachment_size = 0
    for book_hash, book_path in new_books.iteritems():
        if (attachment_size + added_book_sizes[book_path] >
                ATTACHMENTS_SIZE_LIMIT or len(attachments) ==
                BOOK_ATTACHMENTS_LIMIT):
            _email_attachments(email_from, email_to, attachments, user.id)
            attachments = {}
            attachment_size = 0

        attachments[book_hash] = book_path
        attachment_size += added_book_sizes[book_path]

    if len(attachments) > 0:
        _email_attachments(email_from, email_to, attachments, user.id)

    # Commit all new books that should've been emailed.
    db.session.commit()

    # Add duplicates and remove the deleted books from the database.
    # books from and add the duplicates to the database.
    add_duplicate_books(user.id, duplicate_books)
    remove_books(user.id, removed_books)

    # Update the Dropbox delta cursor in database.
    user.cursor = delta['cursor']

    db.session.commit()

    # Clean up the temporary files.
    if len(added_book_sizes) > 0:
        clear_tmp_directory()

    return False

@celery.task(ignore_result=True)
def kindlebox(dropbox_id):
    # Lock per user.
    lock_id = '{0}-lock-{1}'.format(kindlebox.__name__, dropbox_id)
    lock = redis.lock(lock_id, timeout=LOCK_EXPIRE)

    # If unable to acquire lock, discard the task and hope that another worker
    # finishes (see note below).
    if not lock.acquire(blocking=False):
        log.debug("Couldn't acquire lock {0}.".format(lock_id))
        return False

    log.debug("Lock {0} acquired.".format(lock_id))

    # Only process Dropbox changes for active users.
    user = User.query.filter_by(dropbox_id=dropbox_id, active=True).first()
    if user is None:
        lock.release()
        return

    log.info("Processing dropbox webhook for dropbox id {0}".format(dropbox_id))
    # Loop until there is no delta.
    # NOTE: There is a slight chance of a race condition between dropbox
    # webhook and two celery workers that would result in a delta getting
    # dropped, but hopefully this is better than cluttering the task queues.
    client = DropboxClient(user.access_token)
    while True:
        log.debug("Processing one kindlebox iteration for dropbox id "
                  "{0}".format(dropbox_id))
        try:
            done = _kindlebox(dropbox_id, user, client)
            if done:
                break
        except:
            log.error(("Failed to process dropbox webhook for dropbox id "
                       "{0}.").format(dropbox_id), exc_info=True)
            break

    lock.release()

def mimetypes_filter(path):
    return mimetypes.guess_type(path)[0] in BOOK_MIMETYPES

def get_added_book_sizes(delta_entries, client):
    """
    Return a list of tuples of (book path, book size) representing books added
    during this delta. Book path must be of one of the accepted mimetypes and
    book size is under the BOOK_SIZE_LIMIT.
    """
    added_entries = [(canonicalize(entry[0]), entry[1]['bytes']) for entry in delta_entries if
                     entry[1] is not None and not entry[1]['is_dir']]
    return [entry for entry in added_entries if (mimetypes_filter(entry[0]) and
            entry[1] < ATTACHMENTS_SIZE_LIMIT)]

def get_removed_books(delta_entries):
    """
    Return a list of book paths that were deleted during this delta.
    """
    removed_entries = [canonicalize(entry[0]) for entry in delta_entries if
                       entry[1] is None]
    return [entry for entry in removed_entries if mimetypes_filter(entry)]

def get_new_and_duplicates(client, added_book_sizes, user_id):
    duplicate_books = {}
    new_books = {}
    for book_path, book_byte_size in added_book_sizes.iteritems():
        book_hash = download_book(client, book_path)

        # Make sure that the book is not a duplicate of either a previously
        # added book or a book added on this round.
        if (Book.query.filter_by(user_id=user_id)
                .filter_by(book_hash=book_hash).count() == 0 and
                book_hash not in new_books):
            new_books[book_hash] = book_path
        else:
            duplicate_books[book_hash] = book_path

    return new_books, duplicate_books

def download_book(client, book_path):
    # Make all the necessary nested directories in the temporary directory.
    tmp_path = get_tmp_path(book_path)
    try:
        book_dir = os.path.dirname(tmp_path)
        if not os.path.exists(book_dir):
            os.makedirs(book_dir)
    except OSError:
        log.error("Error creating directories for book {0}".format(book_path),
                  exc_info=True)

    md5 = hashlib.md5()
    with open(tmp_path, 'w') as tmp_book:
        with client.get_file(book_path) as book:
            data = book.read()
            tmp_book.write(data)
            md5.update(data)

    mobi_tmp_path = epub_to_mobi_path(tmp_path)
    if mobi_tmp_path is not None:
        subprocess.check_call(['ebook-convert', tmp_path, mobi_tmp_path])

    book_hash = md5.hexdigest()

    return book_hash

def remove_books(user_id, removed_books):
    for book_path in removed_books:
        book = Book.query.filter_by(user_id=user_id).filter_by(pathname=book_path).first()
        if book is not None:
            db.session.delete(book)

def _add_book(user_id, book_path, book_hash, unsent):
    book = Book(user_id, book_path, book_hash, unsent=unsent)
    db.session.add(book)

def add_duplicate_books(user_id, duplicate_books):
    for book_hash, book_path in duplicate_books.iteritems():
        duplicate = (Book.query.filter_by(user_id=user_id)
                         .filter_by(book_hash=book_hash).first())
        # There should always be a duplicate?
        if duplicate is None:
            continue
        _add_book(user_id, book_path, book_hash, duplicate.unsent)

def _email_attachments(email_from, email_to, attachments, user_id):
    attachment_paths = get_attachment_paths(attachments.values())
    log.debug("Sending email to " + ' '.join(email_to) + " " + ' '.join(attachment_paths))
    try:
        status, message = emailer.send_mail(email_from, email_to,
                                            attachment_paths)
        if status != 200:
            log.error("Failed to email books {books} for user id {id}, "
                      "message: {message}".format(books=' '.join(attachment_paths),
                                                  id=user_id,
                                                  message=msg))
        for book_hash, book_path in attachments.iteritems():
            _add_book(user_id, book_path, book_hash, status != 200)
    except:
        log.error("Failed to email books {books} for user id "
                  "{id}".format(books=' '.join(attachment_paths),
                                id=user_id),
                                exc_info=True)

def epub_to_mobi_path(epub_path):
    if mimetypes.guess_type(epub_path)[0] == EPUB_MIMETYPE:
        return epub_path[:-len('epub')] + 'mobi'

def _clear_directory(directory):
    """
    Remove all possible directories and files from a given directory.
    """
    for path in os.listdir(directory):
        subdirectory = os.path.join(directory, path)
        if os.path.isdir(subdirectory):
            _clear_directory(subdirectory)
            os.rmdir(subdirectory)
        else:
            os.unlink(subdirectory)

def clear_tmp_directory():
    """
    Remove all possible directories and files from the temporary directory.
    """
    try:
        _clear_directory(os.path.join(BASE_DIR, str(os.getpid())))
        os.rmdir(os.path.join(BASE_DIR, str(os.getpid())))
    except OSError:
        log.error("Failed to clear tmp directory", exc_info=True)

def get_tmp_path(book_path):
    return os.path.join(BASE_DIR, str(os.getpid()), book_path.strip('/'))

def get_attachment_paths(book_paths):
    attachment_paths = []
    for book_path in book_paths:
        tmp_path = get_tmp_path(book_path)
        mobi_tmp_path = epub_to_mobi_path(tmp_path)
        if mobi_tmp_path is not None:
            tmp_path = mobi_tmp_path
        attachment_paths.append(tmp_path)
    return attachment_paths

def canonicalize(pathname):
    return pathname.lower()
<commit_msg>oh dear<commit_after>import hashlib
import logging
import mimetypes
import os
import subprocess
import time

from dropbox.client import DropboxClient

from app import analytics
from app import celery
from app import db
from app import emailer
from app import redis
from app.models import User, Book

log = logging.getLogger()

# Lock expires in 30 minutes, in case there are lots of epubs to convert.
LOCK_EXPIRE = 60 * 30

BASE_DIR = '/tmp/kindlebox'
try:
    os.makedirs(BASE_DIR)
except OSError:
    pass

# And can only email 25 books at a time. Sendgrid only allows 20MB at a time,
# after encoding to email text, so more like 15.
BOOK_ATTACHMENTS_LIMIT = 25
ATTACHMENTS_SIZE_LIMIT = 15 * (10**6)

# Supported filetypes.
# According to:
# http://www.amazon.com/gp/help/customer/display.html?nodeId=200375630
mimetypes.add_type('application/x-mobipocket-ebook', '.mobi')
mimetypes.add_type('application/x-mobipocket-ebook', '.prc')
mimetypes.add_type('application/vnd.amazon.ebook', '.azw')
mimetypes.add_type('application/vnd.amazon.ebook', '.azw1')

# Amazon doesn't support epub, but we do!
mimetypes.add_type('application/epub+zip', '.epub')
EPUB_MIMETYPE = 'application/epub+zip'

BOOK_MIMETYPES = {
    'application/epub+zip',
    'application/vnd.amazon.ebook',
    'text/plain',
    'application/x-mobipocket-ebook',
    'application/pdf',

    'application/msword',
    'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
    'application/rtf',
    'text/html',

    'image/jpeg',
    'image/gif',
    'image/x-ms-bmp',
    'image/png',
    }
BOOK_CHUNK = 5

@celery.task(ignore_result=True)
def upload_welcome_pdf(dropbox_id):
    user = User.query.filter_by(dropbox_id=dropbox_id,
                                active=True).first()
    if user is None:
        return False

    # If we've already sent the welcome PDF before, Dropbox webhook went
    # trigger, so do it here.
    if user.uploaded_welcome_pdf:
        return kindlebox(dropbox_id)

    analytics.track(str(user.id), 'Sent welcome pdf')

    client = DropboxClient(user.access_token)
    try:
        with open('app/static/kindlebox_welcome.pdf', 'rb') as f:
            response = client.put_file('Welcome to Kindlebox.pdf', f)
            if response:
                log.info("Welcome PDF sent to Dropbox ID {0}.".format(dropbox_id))
            else:
                raise Exception("No response received after sending welcome PDF")

        user.set_uploaded_welcome_pdf()
        db.session.commit()

    except:
        log.error(("Welcome PDF failed for Dropbox ID "
                   "{0}.").format(dropbox_id), exc_info=True)
        return False

    return True

def _kindlebox(dropbox_id, user, client):
    delta = client.delta(user.cursor)

    added_book_sizes = dict(get_added_book_sizes(delta['entries'], client))
    removed_books = get_removed_books(delta['entries'])
    log.debug("Delta contains {0} added books, {1} removed "
              "books".format(len(added_book_sizes), len(removed_books)))

    if len(added_book_sizes) == 0 and len(removed_books) == 0:
        return True

    # Download the changed files and get the hashes.
    # Also record the hashes/paths of any newly added books.
    new_books, duplicate_books = get_new_and_duplicates(client,
                                                        added_book_sizes,
                                                        user.id)

    # Email ze books.
    # Add all books that were supposed to be emailed. If email was
    # unsuccessful, mark book as unsent.
    email_from = user.emailer
    email_to = [row.kindle_name + '@kindle.com' for row in user.kindle_names.all()]
    attachments = {}
    attachment_size = 0
    for book_hash, book_path in new_books.iteritems():
        if (attachment_size + added_book_sizes[book_path] >
                ATTACHMENTS_SIZE_LIMIT or len(attachments) ==
                BOOK_ATTACHMENTS_LIMIT):
            _email_attachments(email_from, email_to, attachments, user.id)
            attachments = {}
            attachment_size = 0

        attachments[book_hash] = book_path
        attachment_size += added_book_sizes[book_path]

    if len(attachments) > 0:
        _email_attachments(email_from, email_to, attachments, user.id)

    # Commit all new books that should've been emailed.
    db.session.commit()

    # Add duplicates and remove the deleted books from the database.
    # books from and add the duplicates to the database.
    add_duplicate_books(user.id, duplicate_books)
    remove_books(user.id, removed_books)

    # Update the Dropbox delta cursor in database.
    user.cursor = delta['cursor']

    db.session.commit()

    # Clean up the temporary files.
    if len(added_book_sizes) > 0:
        clear_tmp_directory()

    return False

@celery.task(ignore_result=True)
def kindlebox(dropbox_id):
    # Lock per user.
    lock_id = '{0}-lock-{1}'.format(kindlebox.__name__, dropbox_id)
    lock = redis.lock(lock_id, timeout=LOCK_EXPIRE)

    # If unable to acquire lock, discard the task and hope that another worker
    # finishes (see note below).
    if not lock.acquire(blocking=False):
        log.debug("Couldn't acquire lock {0}.".format(lock_id))
        return False

    log.debug("Lock {0} acquired.".format(lock_id))

    # Only process Dropbox changes for active users.
    user = User.query.filter_by(dropbox_id=dropbox_id, active=True).first()
    if user is None:
        lock.release()
        return

    log.info("Processing dropbox webhook for dropbox id {0}".format(dropbox_id))
    # Loop until there is no delta.
    # NOTE: There is a slight chance of a race condition between dropbox
    # webhook and two celery workers that would result in a delta getting
    # dropped, but hopefully this is better than cluttering the task queues.
    client = DropboxClient(user.access_token)
    while True:
        log.debug("Processing one kindlebox iteration for dropbox id "
                  "{0}".format(dropbox_id))
        try:
            done = _kindlebox(dropbox_id, user, client)
            if done:
                break
        except:
            log.error(("Failed to process dropbox webhook for dropbox id "
                       "{0}.").format(dropbox_id), exc_info=True)
            break

    lock.release()

def mimetypes_filter(path):
    return mimetypes.guess_type(path)[0] in BOOK_MIMETYPES

def get_added_book_sizes(delta_entries, client):
    """
    Return a list of tuples of (book path, book size) representing books added
    during this delta. Book path must be of one of the accepted mimetypes and
    book size is under the BOOK_SIZE_LIMIT.
    """
    added_entries = [(canonicalize(entry[0]), entry[1]['bytes']) for entry in delta_entries if
                     entry[1] is not None and not entry[1]['is_dir']]
    return [entry for entry in added_entries if (mimetypes_filter(entry[0]) and
            entry[1] < ATTACHMENTS_SIZE_LIMIT)]

def get_removed_books(delta_entries):
    """
    Return a list of book paths that were deleted during this delta.
    """
    removed_entries = [canonicalize(entry[0]) for entry in delta_entries if
                       entry[1] is None]
    return [entry for entry in removed_entries if mimetypes_filter(entry)]

def get_new_and_duplicates(client, added_book_sizes, user_id):
    duplicate_books = {}
    new_books = {}
    for book_path, book_byte_size in added_book_sizes.iteritems():
        book_hash = download_book(client, book_path)

        # Make sure that the book is not a duplicate of either a previously
        # added book or a book added on this round.
        if (Book.query.filter_by(user_id=user_id)
                .filter_by(book_hash=book_hash).count() == 0 and
                book_hash not in new_books):
            new_books[book_hash] = book_path
        else:
            duplicate_books[book_hash] = book_path

    return new_books, duplicate_books

def download_book(client, book_path):
    # Make all the necessary nested directories in the temporary directory.
    tmp_path = get_tmp_path(book_path)
    try:
        book_dir = os.path.dirname(tmp_path)
        if not os.path.exists(book_dir):
            os.makedirs(book_dir)
    except OSError:
        log.error("Error creating directories for book {0}".format(book_path),
                  exc_info=True)

    md5 = hashlib.md5()
    with open(tmp_path, 'w') as tmp_book:
        with client.get_file(book_path) as book:
            data = book.read()
            tmp_book.write(data)
            md5.update(data)

    mobi_tmp_path = epub_to_mobi_path(tmp_path)
    if mobi_tmp_path is not None:
        subprocess.check_call(['ebook-convert', tmp_path, mobi_tmp_path])

    book_hash = md5.hexdigest()

    return book_hash

def remove_books(user_id, removed_books):
    for book_path in removed_books:
        book = Book.query.filter_by(user_id=user_id).filter_by(pathname=book_path).first()
        if book is not None:
            db.session.delete(book)

def _add_book(user_id, book_path, book_hash, unsent):
    book = Book(user_id, book_path, book_hash, unsent=unsent)
    db.session.add(book)

def add_duplicate_books(user_id, duplicate_books):
    for book_hash, book_path in duplicate_books.iteritems():
        duplicate = (Book.query.filter_by(user_id=user_id)
                         .filter_by(book_hash=book_hash).first())
        # There should always be a duplicate?
        if duplicate is None:
            continue
        _add_book(user_id, book_path, book_hash, duplicate.unsent)

def _email_attachments(email_from, email_to, attachments, user_id):
    attachment_paths = get_attachment_paths(attachments.values())
    log.debug("Sending email to " + ' '.join(email_to) + " " + ' '.join(attachment_paths))
    try:
        status, message = emailer.send_mail(email_from, email_to,
                                            attachment_paths)
        if status != 200:
            log.error("Failed to email books {books} for user id {id}, "
                      "message: {message}".format(books=' '.join(attachment_paths),
                                                  id=user_id,
                                                  message=message))
        for book_hash, book_path in attachments.iteritems():
            _add_book(user_id, book_path, book_hash, status != 200)
    except:
        log.error("Failed to email books {books} for user id "
                  "{id}".format(books=' '.join(attachment_paths),
                                id=user_id),
                                exc_info=True)

def epub_to_mobi_path(epub_path):
    if mimetypes.guess_type(epub_path)[0] == EPUB_MIMETYPE:
        return epub_path[:-len('epub')] + 'mobi'

def _clear_directory(directory):
    """
    Remove all possible directories and files from a given directory.
    """
    for path in os.listdir(directory):
        subdirectory = os.path.join(directory, path)
        if os.path.isdir(subdirectory):
            _clear_directory(subdirectory)
            os.rmdir(subdirectory)
        else:
            os.unlink(subdirectory)

def clear_tmp_directory():
    """
    Remove all possible directories and files from the temporary directory.
    """
    try:
        _clear_directory(os.path.join(BASE_DIR, str(os.getpid())))
        os.rmdir(os.path.join(BASE_DIR, str(os.getpid())))
    except OSError:
        log.error("Failed to clear tmp directory", exc_info=True)

def get_tmp_path(book_path):
    return os.path.join(BASE_DIR, str(os.getpid()), book_path.strip('/'))

def get_attachment_paths(book_paths):
    attachment_paths = []
    for book_path in book_paths:
        tmp_path = get_tmp_path(book_path)
        mobi_tmp_path = epub_to_mobi_path(tmp_path)
        if mobi_tmp_path is not None:
            tmp_path = mobi_tmp_path
        attachment_paths.append(tmp_path)
    return attachment_paths

def canonicalize(pathname):
    return pathname.lower()
<|endoftext|>


<commit_before>#!/usr/bin/env python
from optparse import OptionParser
from collections import OrderedDict
import subprocess
import sys

import numpy as np
import pyBigWig

################################################################################
# preprocess_chromatin.py
#
# Annotate the sites in a BED file with a set of Wig/BigWig tracks.
################################################################################

################################################################################
# main
################################################################################
def main():
    usage = 'usage: %prog [options] <sites_bed> <sample_wigs_file> <out_file>'
    parser = OptionParser(usage)
    parser.add_option('-b', dest='bin_size', default=None, type='int', help='Bin size to take the mean track value [Default: %default]')
    parser.add_option('-f', dest='function', default='mean', help='Function to compute in each bin [Default: %default]')
    parser.add_option('-l', dest='log2', default=False, action='store_true', help='Take log2 [Default: %default')
    parser.add_option('-n', dest='normalize', default=False, action='store_true', help='Normalize [Default: %default]')
    parser.add_option('-s', dest='span', default=200, type='int', help='Span of sequence to consider around each site Default: %default]')
    (options,args) = parser.parse_args()

    if len(args) != 3:
        parser.error('Must provide sites BED and file with Wig/BigWig labels and paths')
    else:
        sites_bed_file = args[0]
        sample_wigs_file = args[1]
        out_file = args[2]

    if options.bin_size is None:
        options.bin_size = options.span

    if options.span % options.bin_size != 0:
        parser.error('Bin size must evenly divide span')

    #################################################################
    # setup
    #################################################################
    # determine span on each side
    span_left = int(round(options.span/2.0))
    span_right = options.span/2

    # read in sites
    sites = []
    for line in open(sites_bed_file):
        a = line.split()
        start = int(a[1])
        end = int(a[2])
        mid = (start + end)/2

        sites.append(Site(a[0], mid-span_left, mid+span_right))

    # count bins
    num_bins = options.span / options.bin_size

    # get wig files and labels
    target_wigs = OrderedDict()
    for line in open(sample_wigs_file):
        a = line.split()
        target_wigs[a[0]] = a[1]

    # label features
    feature_labels = []
    for target in target_wigs.keys():
        if num_bins == 1:
            feature_labels.append(target)
        else:
            bin_mid = -span_left + options.bin_size/2
            for bi in range(num_bins):
                feature_labels.append('%s_pos%d' % (target,bin_mid))
                bin_mid += options.bin_size

    # initialize features array
    site_features = np.zeros((len(sites),len(feature_labels)), dtype='float16')

    #################################################################
    # process wigs
    #################################################################
    wi = 0
    for sample in target_wigs:
        wig_file = target_wigs[sample]
        print wig_file

        # open wig
        wig_in = pyBigWig.open(wig_file)

        for si in range(len(sites)):
            s = sites[si]

            # pull stats from wig
            bin_stats = wig_in.stats(s.chrom, s.start, s.end, type=options.function, nBins=num_bins)

            # copy into primary data structure
            for bi in range(num_bins):
                site_features[si,wi+bi] = bin_stats[bi]

        wi += num_bins

    #################################################################
    # normalize
    #################################################################
    # replace nan with zero
    site_features = np.nan_to_num(site_features)

    # acknowledge the max of float16
    site_features = site_features.clip(0, 65504.0)

    if options.log2:
        site_features = np.log2(site_features+1)

    # normalize
    if options.normalize:
        site_features = site_features - np.nanmean(site_features, axis=0)
        site_features = site_features / np.nanstd(site_features, axis=0)

    #################################################################
    # output
    #################################################################
    sv_out = open(out_file, 'w')

    cols = [''] + feature_labels
    print >> sv_out, '\t'.join(cols)

    si = 0
    for line in open(sites_bed_file):
        a = line.split()
        site_id = '%s:%s-%s(%s)' % (a[0], a[1], a[2], a[5])
        cols = [site_id] + [str(v) for v in site_features[si]]
        print >> sv_out, '\t'.join(cols)
        si += 1

    sv_out.close()

class Site:
    def __init__(self, chrom, start, end):
        self.chrom = chrom
        self.start = start
        self.end = end

################################################################################
# __main__
################################################################################
if __name__ == '__main__':
    main()
<commit_msg>clip values and add pseudocounts<commit_after>#!/usr/bin/env python
from optparse import OptionParser
from collections import OrderedDict
import subprocess
import sys

import numpy as np
import pyBigWig

################################################################################
# preprocess_chromatin.py
#
# Annotate the sites in a BED file with a set of Wig/BigWig tracks.
################################################################################

################################################################################
# main
################################################################################
def main():
    usage = 'usage: %prog [options] <sites_bed> <sample_wigs_file> <out_file>'
    parser = OptionParser(usage)
    parser.add_option('-b', dest='bin_size', default=None, type='int', help='Bin size to take the mean track value [Default: %default]')
    parser.add_option('--clip_min', dest='clip_min', default=None, type='float', help='Clip the distribution minimums at this proportion. [Default: %default; Suggested: 0.05]')
    parser.add_option('--clip_max', dest='clip_max', default=None, type='float', help='Clip the distribution maximums at this proportion [Default: %default; Suggested: .001]')
    parser.add_option('-f', dest='function', default='mean', help='Function to compute in each bin [Default: %default]')
    parser.add_option('-l', dest='log2', default=False, action='store_true', help='Take log2 [Default: %default')
    parser.add_option('-n', dest='normalize', default=False, action='store_true', help='Normalize [Default: %default]')
    parser.add_option('-p', dest='pseudocount', default=0, type='float', help='Pseudocount added before transformations [Default: %default]')
    parser.add_option('-s', dest='span', default=200, type='int', help='Span of sequence to consider around each site Default: %default]')
    (options,args) = parser.parse_args()

    if len(args) != 3:
        parser.error('Must provide sites BED and file with Wig/BigWig labels and paths')
    else:
        sites_bed_file = args[0]
        sample_wigs_file = args[1]
        out_file = args[2]

    if options.bin_size is None:
        options.bin_size = options.span

    if options.span % options.bin_size != 0:
        parser.error('Bin size must evenly divide span')

    #################################################################
    # setup
    #################################################################
    # determine span on each side
    span_left = int(round(options.span/2.0))
    span_right = options.span/2

    # read in sites
    sites = []
    for line in open(sites_bed_file):
        a = line.split()
        start = int(a[1])
        end = int(a[2])
        mid = (start + end)/2

        sites.append(Site(a[0], mid-span_left, mid+span_right))

    # count bins
    num_bins = options.span / options.bin_size

    # get wig files and labels
    target_wigs = OrderedDict()
    for line in open(sample_wigs_file):
        a = line.split()
        target_wigs[a[0]] = a[1]

    # label features
    feature_labels = []
    for target in target_wigs.keys():
        if num_bins == 1:
            feature_labels.append(target)
        else:
            bin_mid = -span_left + options.bin_size/2
            for bi in range(num_bins):
                feature_labels.append('%s_pos%d' % (target,bin_mid))
                bin_mid += options.bin_size

    # initialize features array
    site_features = np.zeros((len(sites),len(feature_labels)), dtype='float16')

    #################################################################
    # process wigs
    #################################################################
    wi = 0
    for sample in target_wigs:
        wig_file = target_wigs[sample]
        print wig_file

        # open wig
        wig_in = pyBigWig.open(wig_file)

        for si in range(len(sites)):
            s = sites[si]

            # pull stats from wig
            bin_stats = wig_in.stats(s.chrom, s.start, s.end, type=options.function, nBins=num_bins)

            # copy into primary data structure
            for bi in range(num_bins):
                site_features[si,wi+bi] = bin_stats[bi]

        wi += num_bins

    #################################################################
    # normalize
    #################################################################
    # replace nan with zero
    site_features = np.nan_to_num(site_features)

    # add pseudocount
    if options.pseudocount:
        site_features += options.pseudocount

    # acknowledge the max of float16
    site_features = site_features.clip(0, 65504.0)

    if options.log2:
        site_features = np.log2(site_features)

    # normalize
    if options.normalize:
        site_features = site_features - np.nanmean(site_features, axis=0)
        site_features = site_features / np.nanstd(site_features, axis=0)

    # clip minimum
    if options.clip_min is not None:
        wmins = np.percentile(site_features, 100*options.clip_min, axis=0)
        site_features = site_features.clip(min=wmins)

    # clip maximum
    if options.clip_max is not None:
        wmaxs = np.percentile(site_features, 100*(1-options.clip_max), axis=0)
        site_features = site_features.clip(max=wmaxs)

    #################################################################
    # output
    #################################################################
    sv_out = open(out_file, 'w')

    cols = [''] + feature_labels
    print >> sv_out, '\t'.join(cols)

    si = 0
    for line in open(sites_bed_file):
        a = line.split()
        site_id = '%s:%s-%s(%s)' % (a[0], a[1], a[2], a[5])
        cols = [site_id] + [str(v) for v in site_features[si]]
        print >> sv_out, '\t'.join(cols)
        si += 1

    sv_out.close()

class Site:
    def __init__(self, chrom, start, end):
        self.chrom = chrom
        self.start = start
        self.end = end

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


<commit_before>import sys
import re

#from ctypeslib.codegen.gccxmlparser import parse
from gccxmlparser import parse
from ctypeslib.codegen import typedesc
from codegenlib import Func, parse_type
from funcs import generic_as_arg
from cytypes import generic_decl

def query_items(xml, filter=None):
    # XXX: support filter
    xml_items = parse(xml)
    #items = {}
    keep = set()
    named = {}
    locations = {}
    for it in xml_items:
        #items[it] = it
        keep.add(it)

        if hasattr(it, 'name'):
            named[it] = it.name

        if hasattr(it, 'location'):
            locations[it] = it.location

    return keep, named, locations

def classify(items, locations, ifilter=None):
    # Dictionaries name -> typedesc instances
    funcs = {}
    tpdefs = {}
    enumvals = {}
    enums = {}
    structs = {}
    vars = {}

    if ifilter is None:
        ifilter = lambda x : True

    for it in items:
        try:
            origin = locations[it][0]
            if ifilter(origin):
                if isinstance(it, typedesc.Function):
                    funcs[it.name] = it
                elif isinstance(it, typedesc.EnumValue):
                    enumvals[it.name] = it
                elif isinstance(it, typedesc.Enumeration):
                    enums[it.name] = it
                elif isinstance(it, typedesc.Typedef):
                    tpdefs[it.name] = it
                elif isinstance(it, typedesc.Structure):
                    structs[it.name] = it
                elif isinstance(it, typedesc.Variable):
                    vars[it.name] = it
                else:
                    print "Do not know how to handle", str(it)
        except KeyError:
            if isinstance(it, typedesc.EnumValue):
                enumvals[it.name] = it
            #print "No location for item %s, ignoring" % str(it)

    return funcs, tpdefs, enumvals, enums, structs, vars

def find_named_type(tp):
    if hasattr(tp, 'name'):
        return tp.name
    elif isinstance(tp, typedesc.CvQualifiedType) or \
         isinstance(tp, typedesc.PointerType):
        return find_named_type(tp.typ)
    elif isinstance(tp, typedesc.FunctionType):
        return None
    else:
        raise ValueError("Unhandled type %s" % str(tp))

def find_unqualified_type(tp):
    if isinstance(tp, typedesc.FundamentalType) or \
            isinstance(tp, typedesc.Structure) or \
            isinstance(tp, typedesc.Enumeration) or \
            isinstance(tp, typedesc.Typedef):
        return tp
    elif isinstance(tp, typedesc.CvQualifiedType) or \
         isinstance(tp, typedesc.PointerType):
        return find_unqualified_type(tp.typ)
    elif isinstance(tp, typedesc.FunctionType):
        return None
    else:
        raise ValueError("Unhandled type %s" % str(tp))

def signatures_types(funcs):
    """Given a sequence of typedesc.Function instances, generate a set of all
    typedesc instances used in function declarations."""
    arguments = set()

    for f in funcs:
        for t in f.iterArgTypes():
            ut = find_unqualified_type(t)
            if ut in items:
                arguments.add(ut)

    return arguments

class TypePuller:
    def __init__(self, all):
        self._items = []
        #self._all = sorted(all, cmpitems)
        self._all = all

    def pull_fundamental(self, item):
        pass

    def pull_cv_qualified_type(self, item):
        self.pull(item.typ)
        #names.add(item.name)
        self._items.append(item)

    def pull_typedef(self, item):
        # XXX: Generate the typdef itself
        self.pull(item.typ)
        #names.add(item.name)
        self._items.append(item)

    def pull_function(self, item):
        # XXX: fix signatures_type for single item
        types = signatures_types([item])
        #names.append(item.name)
        for t in types:
            ut = find_unqualified_type(t)
            if ut in self._all:
                self.pull(ut)
        self._items.append(item)

    def pull_function_type(self, item):
        # XXX: fix signatures_type for single item
        types = signatures_types([item])
        #self._items.append(item)
        for t in types:
            ut = find_unqualified_type(t)
            if ut in self._all:
                self.pull(ut)

    def pull_structure(self, item):
        #names.append(item.name)
        for m in item.members:
            g = self.pull(m)
            if g:
                self._items.append(pull(m))
        self._items.append(item)

    def pull_union(self, item):
        #names.append(item.name)
        for m in item.members:
            g = self.pull(m)
            if g:
                self._items.append(pull(m))
        self._items.append(item)

    def pull_array_type(self, item):
        #names.append(item.name)
        self.pull(item.typ)
        self._items.append(item)

    def pull_enumeration(self, item):
        #names.append(item.name)
        for v in item.values:
            self.pull(v)
        self._items.append(item)

    def pull_enum_value(self, item):
        #names.append(item.name)
        self._items.append(item)

    def pull(self, item):
        if isinstance(item, typedesc.FundamentalType):
            #print "Fund Pulling", item, item.name
            self.pull_fundamental(item)
            return
        elif isinstance(item, typedesc.Enumeration):
            #print "Enumeration Pulling", item, item.name
            self.pull_enumeration(item)
            return
        elif isinstance(item, typedesc.EnumValue):
            #print "Enumeration Pulling", item, item.name
            self.pull_enum_value(item)
            return
        elif isinstance(item, typedesc.Typedef):
            #print "Typedef Pulling", item, item.name
            self.pull_typedef(item)
            return
        elif isinstance(item, typedesc.Structure):
            #print "Struct Pulling", item, item.name
            self.pull_structure(item)
            return
        elif isinstance(item, typedesc.Union):
            #print "FunctionType Pulling", item
            self.pull_union(item)
            return
        elif isinstance(item, typedesc.Function):
            #print "Func Pulling", item
            self.pull_function(item)
            return
        elif isinstance(item, typedesc.Field):
            #print "Field Pulling", item
            self.pull(item.typ)
            return
        elif isinstance(item, typedesc.PointerType):
            #print "Pointer Pulling", item
            self.pull(item.typ)
            return
        elif isinstance(item, typedesc.FunctionType):
            #print "FunctionType Pulling", item
            self.pull_function_type(item)
            return
        elif isinstance(item, typedesc.CvQualifiedType):
            #print "FunctionType Pulling", item
            self.pull_cv_qualified_type(item)
            return
        elif isinstance(item, typedesc.ArrayType):
            #print "FunctionType Pulling", item
            self.pull_array_type(item)
            return
        else:
            raise ValueError, ("item not handled:", item)

    def values(self):
        return self._items

def instance_puller(tp, all):
    p = TypePuller(all)
    p.pull(tp)
    return p.values()

def cmpitems(a, b):
    aloc = getattr(a, "location", None)
    bloc = getattr(b, "location", None)
    if aloc is None:
        return -1
    if bloc is None:
        return 1

    st = cmp(aloc[0],bloc[0]) or cmp(int(aloc[1]),int(bloc[1]))
    if st == 0:
        # Two items as the same location: if it is a typedef'd structure with
        # different name and tag, we make sure the structure is defined before
        # the typedef. If it is a different case, just do nothing for now
        if isinstance(a, typedesc.Structure):
            if isinstance(b, typedesc.Typedef):
                return -1
            else:
                # XXX
                print "Hm, not sure what to do here"
                return 0
        if isinstance(b, typedesc.Structure):
            if isinstance(a, typedesc.Typedef):
                return 1
            else:
                print "Hm, not sure what to do here"
                return 0
    else:
        return st

def named_pointer_decl(tp):
    if isinstance(tp.typ, typedesc.FunctionType):
        args = [generic_decl(arg) for arg in tp.typ.iterArgTypes()]
        return generic_decl(tp.typ.returns) + '(*%s)' + '(%s)' % ", ".join(args)
    else:
        return '%s *' % generic_decl(tp.typ)

def cy_generate_typedef(item):
    if not isinstance(item.typ, typedesc.PointerType):
        return ["ctypedef %s %s" % (item.typ.name, item.name)]
    else:
        return ["ctypedef %s" % (named_pointer_decl(item.typ) % item.name)]

def cy_generate_structure(tp):
    output = ['cdef struct %s:' % tp.name]
    for f in tp.members:
        if isinstance(f, typedesc.Field):
            output.append("\t%s %s" % (generic_decl(f.typ), f.name))
        elif isinstance(f, typedesc.Structure):
            output.append("\t%s" % generic_decl(f))
        else:
            print "Struct member not handled:", f
    if not tp.members:
        output.append("\tpass")

    return output

def cy_generate_function(func):
    args = [generic_as_arg(a) for a in func.iterArgTypes()]
    return ["%s %s(%s)" % (generic_as_arg(func.returns),
            func.name, ", ".join(args))]

def cy_generate(item):
    if isinstance(item, typedesc.Typedef):
        #print "Typedef Generating", item
        return cy_generate_typedef(item)
    elif isinstance(item, typedesc.Structure):
        #print "Struct Generating", item
        return cy_generate_structure(item)
    elif isinstance(item, typedesc.Function):
        #print "FunctionType Generating", item
        return cy_generate_function(item)
    else:
        print "Item not handled", item
    #    raise ValueError, ("item not handled:", item)

#root = 'asoundlib'
root = 'foo'
header_name = '%s.h' % root
#header_matcher = re.compile('alsa')
header_matcher = re.compile(header_name)
xml_name = '%s.xml' % root
pyx_name = '_%s.pyx' % root
if sys.platform[:7] == 'darwin':
    so_name = root
else:
    so_name = 'lib%s.so' % root

items, named, locations = query_items(xml_name)
funcs, tpdefs, enumvals, enums, structs, vars = \
        classify(items, locations, ifilter=header_matcher.search)

arguments = signatures_types(funcs.values())
print "Need to pull out arguments", [named[i] for i in arguments]

puller = TypePuller(items)
for a in arguments:
    puller.pull(a)

needed = puller.values()
print "Pulled out items:", [named[i] for i in needed]

# List of items to generate code for
gen = list(needed) + funcs.values()
gen_names = [named[i] for i in gen]

cython_code = [cy_generate(i) for i in gen]

output = open(pyx_name, 'w')
output.write("cdef extern from '%s':\n" % header_name)
for i in cython_code:
    if len(i) > 1:
        output.write("\t%s\n" % i[0])
        for j in i[1:]:
            output.write("\t%s\n" % j)
    else:
        output.write("\t%s\n" % i[0])
output.close()
<commit_msg>Handle enum values in cython code generator.<commit_after>import sys
import re

#from ctypeslib.codegen.gccxmlparser import parse
from gccxmlparser import parse
from ctypeslib.codegen import typedesc
from codegenlib import Func, parse_type
from funcs import generic_as_arg
from cytypes import generic_decl

def query_items(xml, filter=None):
    # XXX: support filter
    xml_items = parse(xml)
    #items = {}
    keep = set()
    named = {}
    locations = {}
    for it in xml_items:
        #items[it] = it
        keep.add(it)

        if hasattr(it, 'name'):
            named[it] = it.name

        if hasattr(it, 'location'):
            locations[it] = it.location

    return keep, named, locations

def classify(items, locations, ifilter=None):
    # Dictionaries name -> typedesc instances
    funcs = {}
    tpdefs = {}
    enumvals = {}
    enums = {}
    structs = {}
    vars = {}

    if ifilter is None:
        ifilter = lambda x : True

    for it in items:
        try:
            origin = locations[it][0]
            if ifilter(origin):
                if isinstance(it, typedesc.Function):
                    funcs[it.name] = it
                elif isinstance(it, typedesc.EnumValue):
                    enumvals[it.name] = it
                elif isinstance(it, typedesc.Enumeration):
                    enums[it.name] = it
                elif isinstance(it, typedesc.Typedef):
                    tpdefs[it.name] = it
                elif isinstance(it, typedesc.Structure):
                    structs[it.name] = it
                elif isinstance(it, typedesc.Variable):
                    vars[it.name] = it
                else:
                    print "Do not know how to handle", str(it)
        except KeyError:
            if isinstance(it, typedesc.EnumValue):
                enumvals[it.name] = it
            #print "No location for item %s, ignoring" % str(it)

    return funcs, tpdefs, enumvals, enums, structs, vars

def find_named_type(tp):
    if hasattr(tp, 'name'):
        return tp.name
    elif isinstance(tp, typedesc.CvQualifiedType) or \
         isinstance(tp, typedesc.PointerType):
        return find_named_type(tp.typ)
    elif isinstance(tp, typedesc.FunctionType):
        return None
    else:
        raise ValueError("Unhandled type %s" % str(tp))

def find_unqualified_type(tp):
    if isinstance(tp, typedesc.FundamentalType) or \
            isinstance(tp, typedesc.Structure) or \
            isinstance(tp, typedesc.Enumeration) or \
            isinstance(tp, typedesc.Typedef):
        return tp
    elif isinstance(tp, typedesc.CvQualifiedType) or \
         isinstance(tp, typedesc.PointerType):
        return find_unqualified_type(tp.typ)
    elif isinstance(tp, typedesc.FunctionType):
        return None
    else:
        raise ValueError("Unhandled type %s" % str(tp))

def signatures_types(funcs):
    """Given a sequence of typedesc.Function instances, generate a set of all
    typedesc instances used in function declarations."""
    arguments = set()

    for f in funcs:
        for t in f.iterArgTypes():
            ut = find_unqualified_type(t)
            if ut in items:
                arguments.add(ut)

    return arguments

class TypePuller:
    def __init__(self, all):
        self._items = []
        #self._all = sorted(all, cmpitems)
        self._all = all

    def pull_fundamental(self, item):
        pass

    def pull_cv_qualified_type(self, item):
        self.pull(item.typ)
        #names.add(item.name)
        self._items.append(item)

    def pull_typedef(self, item):
        # XXX: Generate the typdef itself
        self.pull(item.typ)
        #names.add(item.name)
        self._items.append(item)

    def pull_function(self, item):
        # XXX: fix signatures_type for single item
        types = signatures_types([item])
        #names.append(item.name)
        for t in types:
            ut = find_unqualified_type(t)
            if ut in self._all:
                self.pull(ut)
        self._items.append(item)

    def pull_function_type(self, item):
        # XXX: fix signatures_type for single item
        types = signatures_types([item])
        #self._items.append(item)
        for t in types:
            ut = find_unqualified_type(t)
            if ut in self._all:
                self.pull(ut)

    def pull_structure(self, item):
        #names.append(item.name)
        for m in item.members:
            g = self.pull(m)
            if g:
                self._items.append(pull(m))
        self._items.append(item)

    def pull_union(self, item):
        #names.append(item.name)
        for m in item.members:
            g = self.pull(m)
            if g:
                self._items.append(pull(m))
        self._items.append(item)

    def pull_array_type(self, item):
        #names.append(item.name)
        self.pull(item.typ)
        self._items.append(item)

    def pull_enumeration(self, item):
        #names.append(item.name)
        for v in item.values:
            self.pull(v)
        self._items.append(item)

    def pull_enum_value(self, item):
        #names.append(item.name)
        self._items.append(item)

    def pull(self, item):
        if isinstance(item, typedesc.FundamentalType):
            #print "Fund Pulling", item, item.name
            self.pull_fundamental(item)
            return
        elif isinstance(item, typedesc.Enumeration):
            #print "Enumeration Pulling", item, item.name
            self.pull_enumeration(item)
            return
        elif isinstance(item, typedesc.EnumValue):
            #print "Enumeration Pulling", item, item.name
            self.pull_enum_value(item)
            return
        elif isinstance(item, typedesc.Typedef):
            #print "Typedef Pulling", item, item.name
            self.pull_typedef(item)
            return
        elif isinstance(item, typedesc.Structure):
            #print "Struct Pulling", item, item.name
            self.pull_structure(item)
            return
        elif isinstance(item, typedesc.Union):
            #print "FunctionType Pulling", item
            self.pull_union(item)
            return
        elif isinstance(item, typedesc.Function):
            #print "Func Pulling", item
            self.pull_function(item)
            return
        elif isinstance(item, typedesc.Field):
            #print "Field Pulling", item
            self.pull(item.typ)
            return
        elif isinstance(item, typedesc.PointerType):
            #print "Pointer Pulling", item
            self.pull(item.typ)
            return
        elif isinstance(item, typedesc.FunctionType):
            #print "FunctionType Pulling", item
            self.pull_function_type(item)
            return
        elif isinstance(item, typedesc.CvQualifiedType):
            #print "FunctionType Pulling", item
            self.pull_cv_qualified_type(item)
            return
        elif isinstance(item, typedesc.ArrayType):
            #print "FunctionType Pulling", item
            self.pull_array_type(item)
            return
        else:
            raise ValueError, ("item not handled:", item)

    def values(self):
        return self._items

def instance_puller(tp, all):
    p = TypePuller(all)
    p.pull(tp)
    return p.values()

def cmpitems(a, b):
    aloc = getattr(a, "location", None)
    bloc = getattr(b, "location", None)
    if aloc is None:
        return -1
    if bloc is None:
        return 1

    st = cmp(aloc[0],bloc[0]) or cmp(int(aloc[1]),int(bloc[1]))
    if st == 0:
        # Two items as the same location: if it is a typedef'd structure with
        # different name and tag, we make sure the structure is defined before
        # the typedef. If it is a different case, just do nothing for now
        if isinstance(a, typedesc.Structure):
            if isinstance(b, typedesc.Typedef):
                return -1
            else:
                # XXX
                print "Hm, not sure what to do here"
                return 0
        if isinstance(b, typedesc.Structure):
            if isinstance(a, typedesc.Typedef):
                return 1
            else:
                print "Hm, not sure what to do here"
                return 0
    else:
        return st

def named_pointer_decl(tp):
    if isinstance(tp.typ, typedesc.FunctionType):
        args = [generic_decl(arg) for arg in tp.typ.iterArgTypes()]
        return generic_decl(tp.typ.returns) + '(*%s)' + '(%s)' % ", ".join(args)
    else:
        return '%s *' % generic_decl(tp.typ)

def cy_generate_typedef(item):
    if not isinstance(item.typ, typedesc.PointerType):
        return ["ctypedef %s %s" % (item.typ.name, item.name)]
    else:
        return ["ctypedef %s" % (named_pointer_decl(item.typ) % item.name)]

def cy_generate_structure(tp):
    output = ['cdef struct %s:' % tp.name]
    for f in tp.members:
        if isinstance(f, typedesc.Field):
            output.append("\t%s %s" % (generic_decl(f.typ), f.name))
        elif isinstance(f, typedesc.Structure):
            output.append("\t%s" % generic_decl(f))
        else:
            print "Struct member not handled:", f
    if not tp.members:
        output.append("\tpass")

    return output

def cy_generate_function(func):
    args = [generic_as_arg(a) for a in func.iterArgTypes()]
    return ["%s %s(%s)" % (generic_as_arg(func.returns),
            func.name, ", ".join(args))]

def cy_generate_enum_value(tp):
    output = ['cdef enum:']
    output.append("\t%s = %d" % (tp.name, int(tp.value)))
    return output

def cy_generate(item):
    if isinstance(item, typedesc.Typedef):
        #print "Typedef Generating", item
        return cy_generate_typedef(item)
    elif isinstance(item, typedesc.Structure):
        #print "Struct Generating", item
        return cy_generate_structure(item)
    elif isinstance(item, typedesc.Function):
        #print "FunctionType Generating", item
        return cy_generate_function(item)
    elif isinstance(item, typedesc.EnumValue):
        #print "FunctionType Generating", item
        return cy_generate_enum_value(item)
    else:
        print "Item not handled", item
    #    raise ValueError, ("item not handled:", item)

#root = 'asoundlib'
root = 'foo'
header_name = '%s.h' % root
#header_matcher = re.compile('alsa')
header_matcher = re.compile(header_name)
xml_name = '%s.xml' % root
pyx_name = '_%s.pyx' % root
if sys.platform[:7] == 'darwin':
    so_name = root
else:
    so_name = 'lib%s.so' % root

items, named, locations = query_items(xml_name)
funcs, tpdefs, enumvals, enums, structs, vars = \
        classify(items, locations, ifilter=header_matcher.search)

arguments = signatures_types(funcs.values())
print "Need to pull out arguments", [named[i] for i in arguments]

puller = TypePuller(items)
for a in arguments:
    puller.pull(a)

needed = puller.values()
print "Pulled out items:", [named[i] for i in needed]

# List of items to generate code for
gen = enumvals.values() + list(needed) + funcs.values()
gen_names = [named[i] for i in gen]

cython_code = [cy_generate(i) for i in gen]

output = open(pyx_name, 'w')
output.write("cdef extern from '%s':\n" % header_name)
for i in cython_code:
    if len(i) > 1:
        output.write("\t%s\n" % i[0])
        for j in i[1:]:
            output.write("\t%s\n" % j)
    else:
        output.write("\t%s\n" % i[0])
output.close()
<|endoftext|>


<commit_before># -*- coding: utf8 -*-
# This file is part of PYBOSSA.
#
# Copyright (C) 2015 Scifabric LTD.
#
# PYBOSSA is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# PYBOSSA is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with PYBOSSA.  If not, see <http://www.gnu.org/licenses/>.
import json
from default import with_context, Test
from nose.tools import assert_raises
from werkzeug.exceptions import MethodNotAllowed
from pybossa.api.user import UserAPI
from test_api import TestAPI
from pybossa.core import db

from factories import UserFactory

class TestUserAPI(Test):

    def setUp(self):
        super(TestUserAPI, self).setUp()
        with self.flask_app.app_context():
            db.create_all()
            self.redis_flushall()

    def tearDown(self):
        with self.flask_app.app_context():
            db.drop_all()
            db.session.remove()
            self.redis_flushall()
            #reset_all_pk_sequences()

    @with_context
    def test_user_get(self):
        """Test API User GET"""
        expected_user = UserFactory.create()
        # Test GET all users
        res = self.app.get('/api/user')
        data = json.loads(res.data)
        user = data[0]
        assert len(data) == 1, data
        assert user['name'] == expected_user.name, data

        # The output should have a mime-type: application/json
        assert res.mimetype == 'application/json', res

        # Test GETting a specific user by ID
        res = self.app.get('/api/user/1')
        data = json.loads(res.data)
        user = data
        assert user['name'] == expected_user.name, data

        # Test a non-existant ID
        res = self.app.get('/api/user/3434209')
        err = json.loads(res.data)
        assert res.status_code == 404, err
        assert err['status'] == 'failed', err
        assert err['target'] == 'user', err
        assert err['exception_cls'] == 'NotFound', err
        assert err['action'] == 'GET', err

    @with_context
    def test_query_user(self):
        """Test API query for user endpoint works"""
        expected_user = UserFactory.create_batch(2)[0]
        # When querying with a valid existing field which is unique
        # It should return one correct result if exists
        res = self.app.get('/api/user?name=%s' % expected_user.name)
        data = json.loads(res.data)
        assert len(data) == 1, data
        assert data[0]['name'] == expected_user.name, data
        # And it should return no results if there are no matches
        res = self.app.get('/api/user?name=Godzilla')
        data = json.loads(res.data)
        assert len(data) == 0, data

        # When querying with a valid existing non-unique field
        res = self.app.get("/api/user?locale=en")
        data = json.loads(res.data)
        # It should return 3 results, as every registered user has locale=en by default
        assert len(data) == 2, data
        # And they should be the correct ones
        assert (data[0]['locale'] == data[1]['locale'] == 'en'
               and data[0] != data[1]), data

        # When querying with multiple valid fields
        res = self.app.get('/api/user?name=%s&locale=en' % expected_user.name)
        data = json.loads(res.data)
        # It should find and return one correct result
        assert len(data) == 1, data
        assert data[0]['name'] == expected_user.name, data
        assert data[0]['locale'] == 'en', data

        # When querying with non-valid fields -- Errors
        res = self.app.get('/api/user?something_invalid=whatever')
        err = json.loads(res.data)
        err_msg = "AttributeError exception should be raised"
        assert res.status_code == 415, err_msg
        assert err['action'] == 'GET', err_msg
        assert err['status'] == 'failed', err_msg
        assert err['exception_cls'] == 'AttributeError', err_msg

    @with_context
    def test_user_not_allowed_actions_anon(self):
        """Test POST, PUT and DELETE for ANON actions are not allowed for user
        in the API"""
        user = UserFactory.create()
        url = 'api/user'
        res = self.app.post(url, data=json.dumps(user.to_public_json()))
        data = json.loads(res.data)
        assert res.status_code == 405, res.status_code
        assert data['status_code'] == 405, data

        url += '/%s' % user.id
        res = self.app.put(url, data=json.dumps(dict(name='new')))
        assert res.status_code == 401, res.data

        res = self.app.delete(url)
        assert res.status_code == 405, res.status_code

    @with_context
    def test_user_not_allowed_actions_user(self):
        """Test POST, PUT and DELETE for USER actions are not allowed for user
        in the API"""
        admin = UserFactory.create()
        auth = UserFactory.create()
        user = UserFactory.create()
        url = 'api/user'
        res = self.app.post(url + '?api_key=%s' % auth.api_key,
                            data=json.dumps(user.to_public_json()))
        data = json.loads(res.data)
        assert res.status_code == 405, res.status_code
        assert data['status_code'] == 405, data

        # Update another user
        url += '/%s' % user.id
        res = self.app.put(url + '?api_key=%s' % auth.api_key,
                           data=json.dumps(dict(name='new')))
        assert res.status_code == 403, res.data

        res = self.app.delete(url + '?apikey=%s' % auth.api_key)
        assert res.status_code == 405, res.status_code

        # Update the same user
        url = 'api/user/%s' % auth.id
        res = self.app.put(url + '?api_key=%s' % auth.api_key,
                           data=json.dumps(dict(name='new')))
        data = json.loads(res.data)
        assert res.status_code == 200, res.data
        assert data['name'] == 'new', data

    @with_context
    def test_user_not_allowed_actions_admin (self):
        """Test POST, PUT and DELETE for ADMIN actions are not allowed for user
        in the API"""
        admin = UserFactory.create()
        auth = UserFactory.create()
        user = UserFactory.create()
        url = 'api/user'
        res = self.app.post(url + '?api_key=%s' % admin.api_key,
                            data=json.dumps(user.to_public_json()))
        data = json.loads(res.data)
        assert res.status_code == 405, res.status_code
        assert data['status_code'] == 405, data

        # Update another user
        url += '/%s' % user.id
        new_info = user.info
        new_info['foo'] = 'bar'
        res = self.app.put(url + '?api_key=%s' % admin.api_key,
                           data=json.dumps(dict(name='new', info=new_info)))
        data = json.loads(res.data)
        assert res.status_code == 200, res.data
        assert data['name'] == 'new', data
        assert data['info']['foo'] == 'bar', data

        res = self.app.delete(url + '?apikey=%s' % auth.api_key)
        assert res.status_code == 405, res.status_code

        # Update the same user
        url = 'api/user/%s' % admin.id
        res = self.app.put(url + '?api_key=%s' % admin.api_key,
                           data=json.dumps(dict(name='newadmin')))
        data = json.loads(res.data)
        assert res.status_code == 200, res.data
        assert data['name'] == 'newadmin', data

    @with_context
    def test_privacy_mode_user_get(self):
        """Test API user queries for privacy mode"""
        admin = UserFactory.create()
        user = UserFactory.create()
        # Add user with fullname 'Public user', privacy mode disabled
        user_with_privacy_disabled = UserFactory.create(email_addr='public@user.com',
                                    name='publicUser', fullname='Public user',
                                    privacy_mode=False)
        # Add user with fullname 'Private user', privacy mode enabled
        user_with_privacy_enabled = UserFactory.create(email_addr='private@user.com',
                                    name='privateUser', fullname='Private user',
                                    privacy_mode=True)

        # With no API-KEY
        # User with privacy disabled
        res = self.app.get('/api/user/3')
        data = json.loads(res.data)
        user_with_privacy_disabled = data
        # When checking a public field it should be returned
        assert user_with_privacy_disabled['locale'] == 'en', data
        # When checking a private field it should be returned too
        assert user_with_privacy_disabled['fullname'] == 'Public user', data
        # User with privacy enabled
        res = self.app.get('/api/user/4')
        data = json.loads(res.data)
        user_with_privacy_enabled = data
        # When checking a public field it should be returned
        assert user_with_privacy_enabled['locale'] == 'en', data
        # When checking a private field it should not be returned
        assert 'fullname' not in user_with_privacy_enabled, data
        # Users with privacy enabled and disabled, mixed together
        res = self.app.get('/api/user')
        data = json.loads(res.data)
        user_with_privacy_disabled = data[2]
        user_with_privacy_enabled = data[3]
        assert user_with_privacy_disabled['locale'] == 'en', data
        assert user_with_privacy_disabled['fullname'] == 'Public user', data
        assert user_with_privacy_enabled['locale'] == 'en', data
        assert 'fullname' not in user_with_privacy_enabled, data

        # With a non-admin API-KEY
        # User with privacy disabled
        res = self.app.get('/api/user/3?api_key=' + user.api_key)
        data = json.loads(res.data)
        user_with_privacy_disabled = data
        # When checking a public field it should be returned
        assert user_with_privacy_disabled['locale'] == 'en', data
        # When checking a private field it should be returned too
        assert user_with_privacy_disabled['fullname'] == 'Public user', data
        # User with privacy enabled
        res = self.app.get('/api/user/4?api_key=' + user.api_key)
        data = json.loads(res.data)
        user_with_privacy_enabled = data
        # When checking a public field it should be returned
        assert user_with_privacy_enabled['locale'] == 'en', data
        # When checking a private field it should not be returned
        assert 'fullname' not in user_with_privacy_enabled, data
        # Users with privacy enabled and disabled, mixed together
        res = self.app.get('/api/user?api_key=' + user.api_key)
        data = json.loads(res.data)
        user_with_privacy_disabled = data[2]
        user_with_privacy_enabled = data[3]
        assert user_with_privacy_disabled['locale'] == 'en', data
        assert user_with_privacy_disabled['fullname'] == 'Public user', data
        assert user_with_privacy_enabled['locale'] == 'en', data
        assert 'fullname' not in user_with_privacy_enabled, data

        # Admin API-KEY should be able to retrieve every field in user
        res = self.app.get('/api/user/3?api_key=' + admin.api_key)
        data = json.loads(res.data)
        user_with_privacy_disabled = data
        # When checking a public field it should be returned
        assert user_with_privacy_disabled['locale'] == 'en', data
        # When checking a private field it should be returned too
        assert user_with_privacy_disabled['fullname'] == 'Public user', data
        # User with privacy enabled
        res = self.app.get('/api/user/4?api_key=' + admin.api_key)
        data = json.loads(res.data)
        user_with_privacy_enabled = data
        # When checking a public field it should be returned
        assert user_with_privacy_enabled['locale'] == 'en', data
        # When checking a private field it should be returned too
        assert user_with_privacy_enabled['fullname'] == 'Private user', data
        # Users with privacy enabled and disabled, mixed together
        res = self.app.get('/api/user?api_key=' + admin.api_key)
        data = json.loads(res.data)
        user_with_privacy_disabled = data[2]
        user_with_privacy_enabled = data[3]
        assert user_with_privacy_disabled['locale'] == 'en', data
        assert user_with_privacy_disabled['fullname'] == 'Public user', data
        assert user_with_privacy_enabled['locale'] == 'en', data
        assert user_with_privacy_enabled['fullname'] == 'Private user', data

    @with_context
    def test_privacy_mode_user_queries(self):
        """Test API user queries for privacy mode with private fields in query
        """
        admin = UserFactory.create()
        user = UserFactory.create()
        # Add user with fullname 'Public user', privacy mode disabled
        user_with_privacy_disabled = UserFactory(email_addr='public@user.com',
                                    name='publicUser', fullname='User',
                                    privacy_mode=False)
        # Add user with fullname 'Private user', privacy mode enabled
        user_with_privacy_enabled = UserFactory(email_addr='private@user.com',
                                    name='privateUser', fullname='User',
                                    privacy_mode=True)

        # When querying with private fields
        query = 'api/user?fullname=User'
        # with no API-KEY, no user with privacy enabled should be returned,
        # even if it matches the query
        res = self.app.get(query)
        data = json.loads(res.data)
        assert len(data) == 1, data
        public_user = data[0]
        assert public_user['name'] == 'publicUser', public_user

        # with a non-admin API-KEY, the result should be the same
        res = self.app.get(query + '&api_key=' + user.api_key)
        data = json.loads(res.data)
        assert len(data) == 1, data
        public_user = data[0]
        assert public_user['name'] == 'publicUser', public_user

        # with an admin API-KEY, all the matching results should be returned
        res = self.app.get(query + '&api_key=' + admin.api_key)
        data = json.loads(res.data)
        assert len(data) == 2, data
        public_user = data[0]
        assert public_user['name'] == 'publicUser', public_user
        private_user = data[1]
        assert private_user['name'] == 'privateUser', private_user
<commit_msg>Test for private/public data.<commit_after># -*- coding: utf8 -*-
# This file is part of PYBOSSA.
#
# Copyright (C) 2015 Scifabric LTD.
#
# PYBOSSA is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# PYBOSSA is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with PYBOSSA.  If not, see <http://www.gnu.org/licenses/>.
import json
from default import with_context, Test
from nose.tools import assert_raises
from werkzeug.exceptions import MethodNotAllowed
from pybossa.api.user import UserAPI
from test_api import TestAPI
from pybossa.core import db

from factories import UserFactory

class TestUserAPI(Test):

    def setUp(self):
        super(TestUserAPI, self).setUp()
        with self.flask_app.app_context():
            db.create_all()
            self.redis_flushall()

    def tearDown(self):
        with self.flask_app.app_context():
            db.drop_all()
            db.session.remove()
            self.redis_flushall()
            #reset_all_pk_sequences()

    @with_context
    def test_user_get(self):
        """Test API User GET"""
        expected_user = UserFactory.create()
        # Test GET all users
        res = self.app.get('/api/user')
        data = json.loads(res.data)
        user = data[0]
        assert len(data) == 1, data
        assert user['name'] == expected_user.name, data

        # The output should have a mime-type: application/json
        assert res.mimetype == 'application/json', res

        # Test GETting a specific user by ID
        res = self.app.get('/api/user/1')
        data = json.loads(res.data)
        user = data
        assert user['name'] == expected_user.name, data

        # Test a non-existant ID
        res = self.app.get('/api/user/3434209')
        err = json.loads(res.data)
        assert res.status_code == 404, err
        assert err['status'] == 'failed', err
        assert err['target'] == 'user', err
        assert err['exception_cls'] == 'NotFound', err
        assert err['action'] == 'GET', err

    @with_context
    def test_query_user(self):
        """Test API query for user endpoint works"""
        expected_user = UserFactory.create_batch(2)[0]
        # When querying with a valid existing field which is unique
        # It should return one correct result if exists
        res = self.app.get('/api/user?name=%s' % expected_user.name)
        data = json.loads(res.data)
        assert len(data) == 1, data
        assert data[0]['name'] == expected_user.name, data
        # And it should return no results if there are no matches
        res = self.app.get('/api/user?name=Godzilla')
        data = json.loads(res.data)
        assert len(data) == 0, data

        # When querying with a valid existing non-unique field
        res = self.app.get("/api/user?locale=en")
        data = json.loads(res.data)
        # It should return 3 results, as every registered user has locale=en by default
        assert len(data) == 2, data
        # And they should be the correct ones
        assert (data[0]['locale'] == data[1]['locale'] == 'en'
               and data[0] != data[1]), data

        # When querying with multiple valid fields
        res = self.app.get('/api/user?name=%s&locale=en' % expected_user.name)
        data = json.loads(res.data)
        # It should find and return one correct result
        assert len(data) == 1, data
        assert data[0]['name'] == expected_user.name, data
        assert data[0]['locale'] == 'en', data

        # When querying with non-valid fields -- Errors
        res = self.app.get('/api/user?something_invalid=whatever')
        err = json.loads(res.data)
        err_msg = "AttributeError exception should be raised"
        assert res.status_code == 415, err_msg
        assert err['action'] == 'GET', err_msg
        assert err['status'] == 'failed', err_msg
        assert err['exception_cls'] == 'AttributeError', err_msg

    @with_context
    def test_user_not_allowed_actions_anon(self):
        """Test POST, PUT and DELETE for ANON actions are not allowed for user
        in the API"""
        user = UserFactory.create()
        url = 'api/user'
        res = self.app.post(url, data=json.dumps(user.to_public_json()))
        data = json.loads(res.data)
        assert res.status_code == 405, res.status_code
        assert data['status_code'] == 405, data

        url += '/%s' % user.id
        res = self.app.put(url, data=json.dumps(dict(name='new')))
        assert res.status_code == 401, res.data

        res = self.app.delete(url)
        assert res.status_code == 405, res.status_code

    @with_context
    def test_user_not_allowed_actions_user(self):
        """Test POST, PUT and DELETE for USER actions are not allowed for user
        in the API"""
        admin = UserFactory.create()
        auth = UserFactory.create()
        user = UserFactory.create()
        url = 'api/user'
        res = self.app.post(url + '?api_key=%s' % auth.api_key,
                            data=json.dumps(user.to_public_json()))
        data = json.loads(res.data)
        assert res.status_code == 405, res.status_code
        assert data['status_code'] == 405, data

        # Update another user
        url += '/%s' % user.id
        res = self.app.put(url + '?api_key=%s' % auth.api_key,
                           data=json.dumps(dict(name='new')))
        assert res.status_code == 403, res.data

        res = self.app.delete(url + '?apikey=%s' % auth.api_key)
        assert res.status_code == 405, res.status_code

        # Update the same user
        url = 'api/user/%s' % auth.id
        res = self.app.put(url + '?api_key=%s' % auth.api_key,
                           data=json.dumps(dict(name='new')))
        data = json.loads(res.data)
        assert res.status_code == 200, res.data
        assert data['name'] == 'new', data

    @with_context
    def test_user_not_allowed_actions_admin (self):
        """Test POST, PUT and DELETE for ADMIN actions are not allowed for user
        in the API"""
        admin = UserFactory.create()
        auth = UserFactory.create()
        user = UserFactory.create()
        url = 'api/user'
        res = self.app.post(url + '?api_key=%s' % admin.api_key,
                            data=json.dumps(user.to_public_json()))
        data = json.loads(res.data)
        assert res.status_code == 405, res.status_code
        assert data['status_code'] == 405, data

        # Update another user
        url += '/%s' % user.id
        new_info = user.info
        new_info['foo'] = 'bar'
        res = self.app.put(url + '?api_key=%s' % admin.api_key,
                           data=json.dumps(dict(name='new', info=new_info)))
        data = json.loads(res.data)
        assert res.status_code == 200, res.data
        assert data['name'] == 'new', data
        assert data['info']['foo'] == 'bar', data

        res = self.app.delete(url + '?apikey=%s' % auth.api_key)
        assert res.status_code == 405, res.status_code

        # Update the same user
        url = 'api/user/%s' % admin.id
        res = self.app.put(url + '?api_key=%s' % admin.api_key,
                           data=json.dumps(dict(name='newadmin')))
        data = json.loads(res.data)
        assert res.status_code == 200, res.data
        assert data['name'] == 'newadmin', data

    @with_context
    def test_privacy_mode_user_get(self):
        """Test API user queries for privacy mode"""
        admin = UserFactory.create()
        user = UserFactory.create()
        # Add user with fullname 'Public user', privacy mode disabled
        user_with_privacy_disabled = UserFactory.create(email_addr='public@user.com',
                                    name='publicUser', fullname='Public user',
                                    privacy_mode=False)
        # Add user with fullname 'Private user', privacy mode enabled
        user_with_privacy_enabled = UserFactory.create(email_addr='private@user.com',
                                    name='privateUser', fullname='Private user',
                                    privacy_mode=True)

        # With no API-KEY
        # User with privacy disabled
        res = self.app.get('/api/user/3')
        data = json.loads(res.data)
        user_with_privacy_disabled = data
        # When checking a public field it should be returned
        assert user_with_privacy_disabled['locale'] == 'en', data
        # When checking a private field it should be returned too
        assert user_with_privacy_disabled['fullname'] == 'Public user', data
        # User with privacy enabled
        res = self.app.get('/api/user/4')
        data = json.loads(res.data)
        user_with_privacy_enabled = data
        # When checking a public field it should be returned
        assert user_with_privacy_enabled['locale'] == 'en', data
        # When checking a private field it should not be returned
        assert 'fullname' not in user_with_privacy_enabled, data
        # Users with privacy enabled and disabled, mixed together
        res = self.app.get('/api/user')
        data = json.loads(res.data)
        user_with_privacy_disabled = data[2]
        user_with_privacy_enabled = data[3]
        assert user_with_privacy_disabled['locale'] == 'en', data
        assert user_with_privacy_disabled['fullname'] == 'Public user', data
        assert user_with_privacy_enabled['locale'] == 'en', data
        assert 'fullname' not in user_with_privacy_enabled, data

        # With a non-admin API-KEY
        # User with privacy disabled
        res = self.app.get('/api/user/3?api_key=' + user.api_key)
        data = json.loads(res.data)
        user_with_privacy_disabled = data
        # When checking a public field it should be returned
        assert user_with_privacy_disabled['locale'] == 'en', data
        # When checking a private field it should be returned too
        assert user_with_privacy_disabled['fullname'] == 'Public user', data
        # User with privacy enabled
        res = self.app.get('/api/user/4?api_key=' + user.api_key)
        data = json.loads(res.data)
        user_with_privacy_enabled = data
        # When checking a public field it should be returned
        assert user_with_privacy_enabled['locale'] == 'en', data
        # When checking a private field it should not be returned
        assert 'fullname' not in user_with_privacy_enabled, data
        # Users with privacy enabled and disabled, mixed together
        res = self.app.get('/api/user?api_key=' + user.api_key)
        data = json.loads(res.data)
        user_with_privacy_disabled = data[2]
        user_with_privacy_enabled = data[3]
        assert user_with_privacy_disabled['locale'] == 'en', data
        assert user_with_privacy_disabled['fullname'] == 'Public user', data
        assert user_with_privacy_enabled['locale'] == 'en', data
        assert 'fullname' not in user_with_privacy_enabled, data

        # Admin API-KEY should be able to retrieve every field in user
        res = self.app.get('/api/user/3?api_key=' + admin.api_key)
        data = json.loads(res.data)
        user_with_privacy_disabled = data
        # When checking a public field it should be returned
        assert user_with_privacy_disabled['locale'] == 'en', data
        # When checking a private field it should be returned too
        assert user_with_privacy_disabled['fullname'] == 'Public user', data
        # User with privacy enabled
        res = self.app.get('/api/user/4?api_key=' + admin.api_key)
        data = json.loads(res.data)
        user_with_privacy_enabled = data
        # When checking a public field it should be returned
        assert user_with_privacy_enabled['locale'] == 'en', data
        # When checking a private field it should be returned too
        assert user_with_privacy_enabled['fullname'] == 'Private user', data
        # Users with privacy enabled and disabled, mixed together
        res = self.app.get('/api/user?api_key=' + admin.api_key)
        data = json.loads(res.data)
        user_with_privacy_disabled = data[2]
        user_with_privacy_enabled = data[3]
        assert user_with_privacy_disabled['locale'] == 'en', data
        assert user_with_privacy_disabled['fullname'] == 'Public user', data
        assert user_with_privacy_enabled['locale'] == 'en', data
        assert user_with_privacy_enabled['fullname'] == 'Private user', data

    @with_context
    def test_privacy_mode_user_queries(self):
        """Test API user queries for privacy mode with private fields in query
        """
        admin = UserFactory.create()
        user = UserFactory.create()
        # Add user with fullname 'Public user', privacy mode disabled
        user_with_privacy_disabled = UserFactory(email_addr='public@user.com',
                                    name='publicUser', fullname='User',
                                    privacy_mode=False)
        # Add user with fullname 'Private user', privacy mode enabled
        user_with_privacy_enabled = UserFactory(email_addr='private@user.com',
                                    name='privateUser', fullname='User',
                                    info=dict(container=1, avatar='png', extra='badge1.png'),
                                    privacy_mode=True)

        # When querying with private fields
        query = 'api/user?fullname=User'
        # with no API-KEY, no user with privacy enabled should be returned,
        # even if it matches the query
        res = self.app.get(query)
        data = json.loads(res.data)
        assert len(data) == 1, data
        public_user = data[0]
        assert public_user['name'] == 'publicUser', public_user
        assert 'email_addr' not in public_user.keys(), public_user
        assert 'id' not in public_user.keys(), public_user
        assert 'info' not in public_user.keys(), public_user

        # with a non-admin API-KEY, the result should be the same
        res = self.app.get(query + '&api_key=' + user.api_key)
        data = json.loads(res.data)
        assert len(data) == 1, data
        public_user = data[0]
        assert public_user['name'] == 'publicUser', public_user
        assert 'email_addr' not in public_user.keys(), public_user
        assert 'id' not in public_user.keys(), public_user
        assert 'info' not in public_user.keys(), public_user

        # with an admin API-KEY, all the matching results should be returned
        res = self.app.get(query + '&api_key=' + admin.api_key)
        data = json.loads(res.data)
        assert len(data) == 2, data
        public_user = data[0]
        assert public_user['name'] == 'publicUser', public_user
        private_user = data[1]
        assert private_user['name'] == 'privateUser', private_user
        assert private_user['email_addr'] == user_with_privacy_enabled.email_addr, private_user
        assert private_user['id'] == user_with_privacy_enabled.id, private_user
        assert private_user['info'] == user_with_privacy_enabled.info, private_user

<|endoftext|>