<commit_before>#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Make statistics on score files (stored in JSON files).
"""

import common_functions as common

import argparse
import numpy as np
import math

def hist_ratio(json_file_path_list,
               metric,
               min_npe=None,
               max_npe=None,
               exclude_aborted=False,
               aborted_only=False,
               tel_id=None,
               notebook=False,
               delta_angle_degrees=0.2):

    if len(json_file_path_list) != 2:
        raise Exception('"json_file_path_list" should have exactly 2 elements')

    if exclude_aborted and aborted_only:
        raise Exception('"exclude-aborted" and "aborted-only" options are not compatible"')

    # FETCH SCORE AND COMPUTE HISTOGRAM #######################################

    data_list = []
    label_list = []
    hist_list = []
    bins_list = []

    for json_file_path in json_file_path_list:
        if not notebook:
            print("Parsing {}...".format(json_file_path))

        json_dict = common.parse_json_file(json_file_path)

        if tel_id is not None:
            json_dict = common.image_filter_equals(json_dict, "tel_id", tel_id)

        if min_npe is not None:
            json_dict = common.image_filter_range(json_dict, "npe", min_value=min_npe)
            #json_dict = common.image_filter_range(json_dict, "img_cleaned_sum_pe", min_value=min_npe)

        if max_npe is not None:
            json_dict = common.image_filter_range(json_dict, "npe", max_value=max_npe)
            #json_dict = common.image_filter_range(json_dict, "img_cleaned_sum_pe", max_value=max_npe)

        if not notebook:
            print(len(json_dict["io"]), "images")

        score_array = common.extract_score_array(json_dict, metric)

        data_list.append(score_array)

        label_list.append(json_dict["label"])

        min_range, max_range = 0., math.sin(math.radians(delta_angle_degrees))

        hist, bin_edges = np.histogram(score_array,
                                       range=(min_range, max_range),
                                       bins=1)
                                       #bins=[0., math.sin(math.radians(0.2)), math.sin(math.radians(0.4)), math.sin(math.radians(0.6))] )

        hist_list.append(hist)
        bins_list.append(bin_edges)

    # COMPUTE RATIO ###########################################################

    #assert bins_list[0] == bins_list[1]

    edges_of_bins = bins_list[0]
    val_of_bins_data_1 = hist_list[0]
    val_of_bins_data_2 = hist_list[1]

    print(val_of_bins_data_1)
    print(val_of_bins_data_2)

    # Set ratio where val_of_bins_data_2 is not zero
    ratio = np.divide(val_of_bins_data_1,
                      val_of_bins_data_2,
                      where=(val_of_bins_data_2 != 0))

    # Compute error on ratio (null if cannot be computed)
    error = np.divide(val_of_bins_data_1 * np.sqrt(val_of_bins_data_2) + val_of_bins_data_2 * np.sqrt(val_of_bins_data_1),
                      np.power(val_of_bins_data_2, 2),
                      where=(val_of_bins_data_2 != 0))

    return ratio, error, edges_of_bins

if __name__ == '__main__':

    # PARSE OPTIONS ###########################################################

    parser = argparse.ArgumentParser(description="Make statistics on score files (JSON files).")

    parser.add_argument("--exclude-aborted", action="store_true", default=False,
                        help="Ignore values from aborted images")

    parser.add_argument("--aborted-only", action="store_true", default=False,
                        help="Only consider aborted images")

    parser.add_argument("--angle", type=float, default=None, metavar="FLOAT",
                        help="Delta angle (in degrees) of the first bin")

    parser.add_argument("--min-npe", type=float, default=None, metavar="FLOAT",
                        help="Only considere images having more than the specified total number of photo electrons")

    parser.add_argument("--max-npe", type=float, default=None, metavar="FLOAT",
                        help="Only considere images having less than the specified total number of photo electrons")

    parser.add_argument("--metric", "-m", required=True,
                        metavar="STRING",
                        help="The metric name to plot")

    parser.add_argument("--telid", type=int, default=None,
                        metavar="INTEGER",
                        help="Only plot results for this telescope")

    parser.add_argument("--notebook", action="store_true",
                        help="Notebook mode")

    parser.add_argument("fileargs", nargs=2, metavar="FILE",
                        help="The JSON file to process")

    args = parser.parse_args()

    exclude_aborted = args.exclude_aborted
    aborted_only = args.aborted_only
    angle = args.angle
    min_npe = args.min_npe
    max_npe = args.max_npe
    metric = args.metric
    tel_id = args.telid
    notebook = args.notebook
    json_file_path_list = args.fileargs

    ratio, error, bins = hist_ratio(json_file_path_list,
                                    metric,
                                    min_npe=min_npe,
                                    max_npe=max_npe,
                                    exclude_aborted=exclude_aborted,
                                    aborted_only=aborted_only,
                                    tel_id=tel_id,
                                    notebook=notebook,
                                    delta_angle_degrees=angle)

    print("ratio:", ratio)
    print("error:", error)
    print("bins:", bins)

<commit_msg>Use cleaned npe instead ref's npe.<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Make statistics on score files (stored in JSON files).
"""

import common_functions as common

import argparse
import numpy as np
import math

def hist_ratio(json_file_path_list,
               metric,
               min_npe=None,
               max_npe=None,
               exclude_aborted=False,
               aborted_only=False,
               tel_id=None,
               notebook=False,
               delta_angle_degrees=0.2):

    if len(json_file_path_list) != 2:
        raise Exception('"json_file_path_list" should have exactly 2 elements')

    if exclude_aborted and aborted_only:
        raise Exception('"exclude-aborted" and "aborted-only" options are not compatible"')

    # FETCH SCORE AND COMPUTE HISTOGRAM #######################################

    data_list = []
    label_list = []
    hist_list = []
    bins_list = []

    for json_file_path in json_file_path_list:
        if not notebook:
            print("Parsing {}...".format(json_file_path))

        json_dict = common.parse_json_file(json_file_path)

        if tel_id is not None:
            json_dict = common.image_filter_equals(json_dict, "tel_id", tel_id)

        if min_npe is not None:
            #json_dict = common.image_filter_range(json_dict, "npe", min_value=min_npe)
            json_dict = common.image_filter_range(json_dict, "img_cleaned_sum_pe", min_value=min_npe)

        if max_npe is not None:
            #json_dict = common.image_filter_range(json_dict, "npe", max_value=max_npe)
            json_dict = common.image_filter_range(json_dict, "img_cleaned_sum_pe", max_value=max_npe)

        if not notebook:
            print(len(json_dict["io"]), "images")

        score_array = common.extract_score_array(json_dict, metric)

        data_list.append(score_array)

        label_list.append(json_dict["label"])

        min_range, max_range = 0., math.sin(math.radians(delta_angle_degrees))

        hist, bin_edges = np.histogram(score_array,
                                       range=(min_range, max_range),
                                       bins=1)
                                       #bins=[0., math.sin(math.radians(0.2)), math.sin(math.radians(0.4)), math.sin(math.radians(0.6))] )

        hist_list.append(hist)
        bins_list.append(bin_edges)

    # COMPUTE RATIO ###########################################################

    #assert bins_list[0] == bins_list[1]

    edges_of_bins = bins_list[0]
    val_of_bins_data_1 = hist_list[0]
    val_of_bins_data_2 = hist_list[1]

    print(val_of_bins_data_1)
    print(val_of_bins_data_2)

    # Set ratio where val_of_bins_data_2 is not zero
    ratio = np.divide(val_of_bins_data_1,
                      val_of_bins_data_2,
                      where=(val_of_bins_data_2 != 0))

    # Compute error on ratio (null if cannot be computed)
    error = np.divide(val_of_bins_data_1 * np.sqrt(val_of_bins_data_2) + val_of_bins_data_2 * np.sqrt(val_of_bins_data_1),
                      np.power(val_of_bins_data_2, 2),
                      where=(val_of_bins_data_2 != 0))

    return ratio, error, edges_of_bins

if __name__ == '__main__':

    # PARSE OPTIONS ###########################################################

    parser = argparse.ArgumentParser(description="Make statistics on score files (JSON files).")

    parser.add_argument("--exclude-aborted", action="store_true", default=False,
                        help="Ignore values from aborted images")

    parser.add_argument("--aborted-only", action="store_true", default=False,
                        help="Only consider aborted images")

    parser.add_argument("--angle", type=float, default=None, metavar="FLOAT",
                        help="Delta angle (in degrees) of the first bin")

    parser.add_argument("--min-npe", type=float, default=None, metavar="FLOAT",
                        help="Only considere images having more than the specified total number of photo electrons")

    parser.add_argument("--max-npe", type=float, default=None, metavar="FLOAT",
                        help="Only considere images having less than the specified total number of photo electrons")

    parser.add_argument("--metric", "-m", required=True,
                        metavar="STRING",
                        help="The metric name to plot")

    parser.add_argument("--telid", type=int, default=None,
                        metavar="INTEGER",
                        help="Only plot results for this telescope")

    parser.add_argument("--notebook", action="store_true",
                        help="Notebook mode")

    parser.add_argument("fileargs", nargs=2, metavar="FILE",
                        help="The JSON file to process")

    args = parser.parse_args()

    exclude_aborted = args.exclude_aborted
    aborted_only = args.aborted_only
    angle = args.angle
    min_npe = args.min_npe
    max_npe = args.max_npe
    metric = args.metric
    tel_id = args.telid
    notebook = args.notebook
    json_file_path_list = args.fileargs

    ratio, error, bins = hist_ratio(json_file_path_list,
                                    metric,
                                    min_npe=min_npe,
                                    max_npe=max_npe,
                                    exclude_aborted=exclude_aborted,
                                    aborted_only=aborted_only,
                                    tel_id=tel_id,
                                    notebook=notebook,
                                    delta_angle_degrees=angle)

    print("ratio:", ratio)
    print("error:", error)
    print("bins:", bins)

<|endoftext|>


<commit_before>import amqp
from contextlib import closing

def __get_channel(connection):
    return connection.channel()

def __declare_exchange(channel, exchange, type):
    channel.exchange_declare(exchange=exchange, type=type, durable=True, auto_delete=False)

def __publish_message_to_exchange(channel, message, exchange, routing_key):
    channel.basic_publish(message, exchange=exchange, routing_key=routing_key)

def publish_message(message, exchange, type, routing_key):
    """ Publish a message to an exchange with exchange type and routing key specified.

    TODO The message format is to be confirmed. The message is sent to an exchange of specified type with the provided routing key. The exchange is durable and is not automatically deleted when the last consumer detaches.

    :param message: The message to be sent.
    :param exchange: The name of the exchange the message should be sent to.
    :param routing_key: The routing key to be sent with the message.

    Usage::

    >>> from utils import publish_message
    >>> publish_message(message, 'exchange', 'type', 'routing_key')

    """
    with closing(amqp.Connection()) as connection:
        channel = __get_channel(connection)
        __declare_exchange(channel, exchange, type)
        __publish_message_to_exchange(channel, message, exchange, routing_key)
<commit_msg>Add message body parsing<commit_after>import amqp
from contextlib import closing

def __get_channel(connection):
    return connection.channel()

def __get_message(message_body):
    return amqp.Message(message_body)

def __declare_exchange(channel, exchange, type):
    channel.exchange_declare(exchange=exchange, type=type, durable=True, auto_delete=False)

def __publish_message_to_exchange(channel, message, exchange, routing_key):
    channel.basic_publish(message, exchange=exchange, routing_key=routing_key)

def publish_message(message_body, exchange, type, routing_key):
    """ Publish a message to an exchange with exchange type and routing key specified.

    A message is sent with the given body. It is sent to an exchange of specified type with the provided routing key. The exchange is durable and is not automatically deleted when the last consumer detaches.

    :param message_body: The body of the message to be sent.
    :param exchange: The name of the exchange the message is sent to.
    :param type: The type of the exchange the message is sent to.
    :param routing_key: The routing key to be sent with the message.

    Usage::

    >>> from utils import publish_message
    >>> publish_message(message, 'exchange', 'type', 'routing_key')

    """
    with closing(amqp.Connection()) as connection:
        channel = __get_channel(connection)
        message = __get_message(message_body)
        __declare_exchange(channel, exchange, type)
        __publish_message_to_exchange(channel, message, exchange, routing_key)
<|endoftext|>


<commit_before>from __future__ import absolute_import, unicode_literals, print_function

import collections
import copy
import json
import logging
import re
import zlib

from .compat import OrderedDict, parse_qsl, quote
from .filters import fallback_elider

logger = logging.getLogger(__name__)

def wrap_before_record_response(wrapped, **kwargs):
    before_record_response = make_before_record_response(**kwargs)
    def wrapper(response):
        response = before_record_response(response)
        response = wrapped(response)
        return response
    return wrapper

def make_before_record_response(elide_access_token,
                                elider_prefix):

    def before_record_response(response):
        if 'facebook-api-version' not in response['headers']:
            return response

        # The response at this point is both (1) what will be recorded, and
        # (2) what will be returned to the application. The problem is that a
        # test needs to run through with the original responses the first time,
        # so that paging links with embedded access tokens are still valid.
        # Later when the test replays in full, the elided tokens will match
        # from the response and the request for the next page, but the first
        # time we have to be careful not to break the paging links.
        response = copy.deepcopy(response)

        response = ungzip(response)
        response = filter_access_tokens(
            response, elide_access_token, elider_prefix)
        response = update_content_length(response)
        return response
    return before_record_response

def ungzip(response):
    headers, body = response['headers'], response['body']

    if 'gzip' in headers.get('content-encoding', []):
        body['string'] = zlib.decompress(body['string'], 16 + zlib.MAX_WBITS)
        headers['content-encoding'].remove('gzip')
        if not headers['content-encoding']:
            del headers['content-encoding']

    return response

def filter_access_tokens(response,
                         elide_access_token,
                         elider_prefix):

    # One approach would be to detect batch responses, then branch
    # accordingly. Either know in advance all the places that access tokens
    # can appear, or use something like jsonpath-rw to find them. But in
    # fact access tokens only appear in two forms in the JSON response:
    #
    #   "access_token":"..."  in JSON
    #   access_token=...      in paging URLs
    #
    # so it's much easier to use regular expression matching.

    # Decode so unicode patterns can operate on the body.
    body = response['body']['string'].decode('utf-8')

    body = re.sub(
        r'"access_token":\s*"([^"]+)"',
        lambda m: '"access_token":"{0}"'.format(
            elide(m.group(1), elide_access_token, elider_prefix)),
        body,
    )

    body = re.sub(
        r'access_token=([^&"]+)',
        lambda m: 'access_token={0}'.format(
            elide(m.group(1), elide_access_token, elider_prefix)),
        body,
    )

    response['body']['string'] = body.encode('utf-8')
    return response

def elide(orig, fun, prefix):
    if prefix and orig.startswith(prefix):
        return orig
    value = None
    if fun:
        value = fun(orig)
    if not value:
        value = fallback_elider(orig)
    return prefix + value

def update_content_length(response):
    headers, body = response['headers'], response['body']
    if 'content-length' in headers:
        headers['content-length'] = [str(len(body['string']))]
    return response
<commit_msg>Don't swallow backslash when eliding tokens in nested JSON.<commit_after>from __future__ import absolute_import, unicode_literals, print_function

import collections
import copy
import json
import logging
import re
import zlib

from .compat import OrderedDict, parse_qsl, quote
from .filters import fallback_elider

logger = logging.getLogger(__name__)

def wrap_before_record_response(wrapped, **kwargs):
    before_record_response = make_before_record_response(**kwargs)
    def wrapper(response):
        response = before_record_response(response)
        response = wrapped(response)
        return response
    return wrapper

def make_before_record_response(elide_access_token,
                                elider_prefix):

    def before_record_response(response):
        if 'facebook-api-version' not in response['headers']:
            return response

        # The response at this point is both (1) what will be recorded, and
        # (2) what will be returned to the application. The problem is that a
        # test needs to run through with the original responses the first time,
        # so that paging links with embedded access tokens are still valid.
        # Later when the test replays in full, the elided tokens will match
        # from the response and the request for the next page, but the first
        # time we have to be careful not to break the paging links.
        response = copy.deepcopy(response)

        response = ungzip(response)
        response = filter_access_tokens(
            response, elide_access_token, elider_prefix)
        response = update_content_length(response)
        return response
    return before_record_response

def ungzip(response):
    headers, body = response['headers'], response['body']

    if 'gzip' in headers.get('content-encoding', []):
        body['string'] = zlib.decompress(body['string'], 16 + zlib.MAX_WBITS)
        headers['content-encoding'].remove('gzip')
        if not headers['content-encoding']:
            del headers['content-encoding']

    return response

def filter_access_tokens(response,
                         elide_access_token,
                         elider_prefix):

    # One approach would be to detect batch responses, then branch
    # accordingly. Either know in advance all the places that access tokens
    # can appear, or use something like jsonpath-rw to find them. But in
    # fact access tokens only appear in two forms in the JSON response:
    #
    #   "access_token":"..."  in JSON
    #   access_token=...      in paging URLs
    #
    # so it's much easier to use regular expression matching.

    # Decode so unicode patterns can operate on the body.
    body = response['body']['string'].decode('utf-8')

    # Variable-length escapes here handle nested JSON for batch responses.
    body = re.sub(
        r'((\\*")access_token\2:\s*\2)([^\\"]+)\2',
        lambda m: ''.join([
            m.group(1),
            elide(m.group(3), elide_access_token, elider_prefix),
            m.group(2),
        ]),
        body,
    )

    body = re.sub(
        r'(access_token=)([^\\&"]+)',
        lambda m: ''.join([
            m.group(1),
            elide(m.group(2), elide_access_token, elider_prefix),
        ]),
        body,
    )

    response['body']['string'] = body.encode('utf-8')
    return response

def elide(orig, fun, prefix):
    if prefix and orig.startswith(prefix):
        return orig
    value = None
    if fun:
        value = fun(orig)
    if not value:
        value = fallback_elider(orig)
    return prefix + value

def update_content_length(response):
    headers, body = response['headers'], response['body']
    if 'content-length' in headers:
        headers['content-length'] = [str(len(body['string']))]
    return response
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""
For a given motif annotated vcf file (already run through motifs.py) and a gene expression file, remove all motif
matches for TFs that are not expressed in at least one sample above the specified threshold.

Usage: tf_expression.py -i <input.vcf> -e <expression.bed> -o <output.vcf> [OPTIONS]

Args:
    -i (required) <input.vcf>: Name of sorted variant file to process.
    -o (required) <output.vcf>: Name of output file to be created.
    -e (required) <expression.bed>: An expression 'bed' file.
    -m (optional) <motif.txt>: Tab-delimited key file containing a frequency
        matrix with each row corresponding to a base and each column
        corresponding to a position (JASPAR format).
        If specified ignores input.vcf and output.vcf
    -mo (optional) <motif_output.txt>: Name of output motif file to be created.
        if blank and -m then creates .tf_filtered version of motif.txt file
    -th (optional) <5>: TFs are considered expressed if they are above this threshold.
"""
import sys
import argparse
import time
import motif

def parse_header(line):
    """
    Get all sample names present in the vcf header.

    Assume each column is in format 'extra_info.samplename'.

    Args:
        line (str): Header line.

    Returns:
        samples (list of str): List of sample names.
    """
    line = line.strip().split("\t")[9:]
    samples = []

    for x in line:
        samp = x.split('.')[-1]
        if samp not in samples:
            samples.append(samp)

    return samples

def filter_motifs(motifs, gene_dict, thresh):
    """
    Get lists of motifs above threshold and below for a list of motif names & the gene dict.

    Args:
        motifs (list of str): Motif names for a given variant.
        gene_dict (dict): Dictionary containing gene names and expression values.
        thresh (float): Expression threshold that TFs must meet to be included in output.

    Returns:
        passed (list of int): List of motif indices that pass threshold.
        failed (list of int): List of motif indices that fail threshold.
    """

    passed = []
    failed = []

    for item in motifs:
        pass_th = False

        try:
            # TODO - Try to handle complexes and such - 'ATX::TCF3', etc.
            exp_vals = gene_dict[item]    # this is the line that must be in try

            # Check if any of the expression values meet the threshold.
            for x in exp_vals:
                if float(x) >= thresh:
                    pass_th = True
                    indices = [i for i, x in enumerate(motifs) if x == item]  # Find multiple motifs for same TF.
                    for x in indices:
                        if x not in passed:
                            passed.append(x)
                    break    # breaks out of x in exp_vals

        except:
            # TODO - Add option to retain motifs that don't match a gene in the expression file.
            pass_th = False

        if pass_th is False:
            indices = [i for i, x in enumerate(motifs) if x == item]  # Find multiple motifs for same TF.
            for x in indices:
                if x not in failed:
                    failed.append(x)

    return (passed, failed)

def process_line(line, gene_dict, thresh):
    """
    Process a vcf record and return modified version of input line for output
    Returning line rather than direct print enable later integration with
    process stack class elements.

    Args:
        line (str): Line to parse.
        gene_dict (dict): Dict with gene names as keys and expression values for samples in vcf as values.
        thresh (float): Expression threshold that TFs must meet to be included in output.
    Returns:
        modified version of input line
    """
    line = line.strip()
    line_list = line.split("\t")
    info_fields = line_list[7].split(";")
    (motifns, motif_other, other_info) = ([], {}, [])

    motifn_present = False  # Skip line if no motif was matched.

    for field in info_fields:
        if field.startswith("MOTIFN="):
            motifn_present = True
            motifns = field[7:].split(',')    # the set of names

            # Get motifs that pass expression threshold.
            passed_idx, failed_idx = filter_motifs(motifns, gene_dict, thresh)

            # Delete by high index first so that lower indices aren't changed.
            for i in sorted(failed_idx, reverse=True):
                del motifns[i]

        # Create dict from other motif fields.
        elif (field.startswith("MOTIFV=") or field.startswith("MOTIFR=") or field.startswith("MOTIFC=")):
            name = field[:7]
            values = field[7:].split(',')
            motif_other[name] = values
        elif (field.startswith("MOTIFVH=") or field.startswith("MOTIFRH=") or field.startswith("MOTIFVG=") or
              field.startswith("MOTIFRG=")):
            name = field[:8]
            values = field[8:].split(',')
            motif_other[name] = values
        else:
            other_info.append(field)

    # Just go to next line if motif info isn't present.
    if motifn_present is False:
        return

    new_info_fields = "MOTIFN=" + ",".join(motifns)

    # Delete matches for the motif info fields that didn't pass the threshold.
    for f in motif_other:
        vals = motif_other[f]
        for i in sorted(failed_idx, reverse=True):
            del vals[i]
        motif_other[f] = vals
        new_info_fields += ";" + f + ",".join(vals)

    for x in other_info:
        new_info_fields += ";" + x

    line_list[7] = new_info_fields
    new_line = "\t".join(line_list)  # New, filtered output line.

    return new_line

def get_genes(exp_file, samples, threshold, max_only):
    """
    Reads in and parses the .bed expression file.
    File format expected to be:
        Whose format is tab seperated columns with header line:
        CHR  START  STOP  GENE  <sample 1>  <sample 2>  ...  <sample n>

    Args:
        exp_file (str): Name of expression file.
        samples (list): Names of the samples in the vcf file.
        threshold (float): Expression threshold to filter lowly/unexpressed genes.
        max_only (bool): if true, gene_dict value is 1 value = max expression
            if false gene_dict value is list of expression values
            YYY: WARNING: if want list to have meaning
                then values needs to be tied to header sample names

    Returns:
        gene_dict (dict): {gene_name: [expression_vals]}.
            Only include values for samples in the vcf.
    """
    data_cols = []
    gene_dict = {}

    print('start read exp_file:' + format(exp_file))

    if max_only:
        # read and only return max exp value in gene_dict
        with open(exp_file) as f:
            header = f.readline().strip().split('\t')
            for samp in header[4:]:
                if samp in samples:
                    data_idx = header.index(samp)
                    data_cols.append(data_idx)

            # Read in expression levels for each gene.
            for line in f:
                line = line.strip().split('\t')
                gene_name = line[3].upper()
                exp_val = -1e1000
                for idx in data_cols:
                    if float(line[idx]) > exp_val:
                        exp_val = float(line[idx])

                gene_dict[gene_name] = exp_val

    else:
        # read and return exp value list in gene_dict
        with open(exp_file) as f:
            header = f.readline().strip().split('\t')
            for samp in header[4:]:
                if samp in samples:
                    data_idx = header.index(samp)
                    data_cols.append(data_idx)

            # Read in expression levels for each gene.
            for line in f:
                line = line.strip().split('\t')
                gene_name = line[3].upper()
                exp_vals = []
                for idx in data_cols:
                    exp_vals.append(line[idx])

                gene_dict[gene_name] = exp_vals

    return gene_dict

def main(inp_file, exp_file, out_file, th=5, motif_file="", motifout_file="", use_vcf=True):
    """
        If use_vcf true then:
            For a given motif annotated vcf file (already run through motifs.py),
            remove all motif matches for TFs that are
            not expressed in at least one sample above the threshold.
        Else:
            For a given vcf file read the samples,
            then filter the motif_file, only outputing to motifout_file those items
            not expressed in at least one sample above the threshold

        Args:
            -i (str): Name of sorted variant file to process.
            -o (str): Name of output file to be created.
            -e (str): Name of expression file.
            -th (float): TFs are considered expressed if they are above this threshold.
    """

    if use_vcf:
        output_f = open(out_file, "w")

        with open(inp_file) as vcf:

            line = vcf.readline().strip()

            # Skip info lines.
            while line.startswith("##"):
                print(line, file=output_f)
                line = vcf.readline().strip()

            # First non-## line is the header. Get sample names and print to output.
            samples = parse_header(line)
            print(line, file=output_f)

            print("Creating gene dictionary for expression data.")
            gene_dict = get_genes(exp_file, samples, th, True)

            if len(gene_dict) == 0:
                print("Error, no genes above threshold found in expression file.",
                      "\nTry lowering the threshold and ensure the expression fil",
                      "e has values in the range that you expect.")
                sys.exit()

            print("Filtering motif info for TFs that don't meet the expression threshold of " +
                str(th) + ". Found " + format(len(gene_dict)) + " genes. Start processing vcf file.")
            for line in vcf:
                new_line = process_line(line, gene_dict, th)
                if new_line is not None:
                    print(new_line, file=output_f)

        output_f.close()
    else:
        # this version processes the motif file

        with open(inp_file) as vcf:

            line = vcf.readline().strip()

            # Skip info lines.
            while line.startswith("##"):
                line = vcf.readline().strip()

            # First non-## line is the header. Get sample names and print to output.
            samples = parse_header(line)

        # done with vcf file; only used to get samples

        print("Creating gene dictionary for expression data.")
        gene_dict = get_genes(exp_file, samples, th, True)

        if len(gene_dict) == 0:
            print("Error, no genes above threshold found in expression file.",
                  "\nTry lowering the threshold and ensure the expression fil",
                  "e has values in the range that you expect.")
            sys.exit()

        print("Filtering motif info for TFs that don't meet the expression threshold of " +
            str(th) + ". Found " + format(len(gene_dict)) + " genes. Start filtering motifs.")

        motif.get_filterbygene_put_motifs(motif_file, motifout_file, th, gene_dict)

    print("COMPLETE.")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(usage=__doc__)

    parser.add_argument("-i", "--input", dest="input_file", required=True)
    parser.add_argument("-e", "--expression", dest="exp_file", required=True)
    parser.add_argument("-m", "--motif_input", dest="motif_file", required=False)
    parser.add_argument("-mo", "--motif_output", dest="motif_out_file", required=False)
    parser.add_argument("-o", "--output", dest="output_file", required=True)
    parser.add_argument("-th", "--threshold", dest="threshold", required=False, default=5, type=float)

    args = parser.parse_args()

    inp_file = args.input_file
    exp_file = args.exp_file
    out_file = args.output_file
    th = args.threshold

    if args.motif_file is not None:
        motif_file = args.motif_file
        if args.motif_out_file is not None:
            motifout_file = args.motif_out_file
        else:
            motifout_file = motif_file.replace('.txt', '.tf_filtered.txt')
        main(inp_file, exp_file, out_file, th, motif_file, motifout_file, False)
    else:
        main(inp_file, exp_file, out_file, th, None, None, True)
<commit_msg>clarify comments and other small changes<commit_after>#!/usr/bin/env python
"""
Input type 1: (-i, -o): For a given motif annotated vcf file (already run through motifs.py)
Input type 2: (-i, -m, -mo): For a given vcf file (used for samples), filter the motif file

With motifs from either input type and a gene expression file (-e), process the motif set to
remove all motif matches for TFs that are not expressed in at least one sample above the
specified threshold.

Usage: tf_expression.py -i <input.vcf> -e <expression.bed> -o <output.vcf> [OPTIONS]

Args:
    -i (required) <input.vcf>: Name of sorted variant file to process.
    -o (required) <output.vcf>: Name of output file to be created.
        Not created if using -m and -mo
    -e (required) <expression.bed>: An expression 'bed' file.
    -m (optional) <motif.txt>: Tab-delimited key file containing a frequency
        matrix with each row corresponding to a base and each column
        corresponding to a position (JASPAR format).
        If specified ignores input.vcf and output.vcf
    -mo (optional) <motif_output.txt>: Name of output motif file to be created.
        if blank and -m then creates .tf_filtered version of motif.txt file
    -th (optional) <5>: TFs are considered expressed if they are above this threshold.
"""
import sys
import argparse
import time
import motif

def parse_header(line):
    """
    Get all sample names present in the vcf header.

    Assume each column is in format 'extra_info.samplename'.

    Args:
        line (str): Header line.

    Returns:
        samples (list of str): List of sample names.
    """
    line = line.strip().split("\t")[9:]
    samples = []

    for x in line:
        samp = x.split('.')[-1]
        if samp not in samples:
            samples.append(samp)

    return samples

def filter_motifs(motifs, gene_dict, thresh):
    """
    Get lists of motifs above threshold and below for a list of motif names & the gene dict.

    Args:
        motifs (list of str): Motif names for a given variant.
        gene_dict (dict): Dictionary containing gene names and expression values.
        thresh (float): Expression threshold that TFs must meet to be included in output.

    Returns:
        passed (list of int): List of motif indices that pass threshold.
        failed (list of int): List of motif indices that fail threshold.
    """

    passed = []
    failed = []

    for item in motifs:
        pass_th = False

        try:
            # TODO - Try to handle complexes and such - 'ATX::TCF3', etc.
            exp_vals = gene_dict[item]    # this is the line that must be in try

            # Check if any of the expression values meet the threshold.
            for x in exp_vals:
                if float(x) >= thresh:
                    pass_th = True
                    indices = [i for i, x in enumerate(motifs) if x == item]  # Find multiple motifs for same TF.
                    for x in indices:
                        if x not in passed:
                            passed.append(x)
                    break    # breaks out of x in exp_vals

        except:
            # TODO - Add option to retain motifs that don't match a gene in the expression file.
            pass_th = False

        if pass_th is False:
            indices = [i for i, x in enumerate(motifs) if x == item]  # Find multiple motifs for same TF.
            for x in indices:
                if x not in failed:
                    failed.append(x)

    return (passed, failed)

def process_line(line, gene_dict, thresh):
    """
    Process a vcf record and return modified version of input line for output
    Returning line rather than direct print enable later integration with
    process stack class elements.

    Args:
        line (str): Line to parse.
        gene_dict (dict): Dict with gene names as keys and expression values for samples in vcf as values.
        thresh (float): Expression threshold that TFs must meet to be included in output.
    Returns:
        modified version of input line
    """
    line = line.strip()
    line_list = line.split("\t")
    info_fields = line_list[7].split(";")
    (motifns, motif_other, other_info) = ([], {}, [])

    motifn_present = False  # Skip line if no motif was matched.

    for field in info_fields:
        if field.startswith("MOTIFN="):
            motifn_present = True
            motifns = field[7:].split(',')    # the set of names

            # Get motifs that pass expression threshold.
            passed_idx, failed_idx = filter_motifs(motifns, gene_dict, thresh)

            # Delete by high index first so that lower indices aren't changed.
            for i in sorted(failed_idx, reverse=True):
                del motifns[i]

        # Create dict from other motif fields.
        elif (field.startswith("MOTIFV=") or field.startswith("MOTIFR=") or field.startswith("MOTIFC=")):
            name = field[:7]
            values = field[7:].split(',')
            motif_other[name] = values
        elif (field.startswith("MOTIFVH=") or field.startswith("MOTIFRH=") or field.startswith("MOTIFVG=") or
              field.startswith("MOTIFRG=")):
            name = field[:8]
            values = field[8:].split(',')
            motif_other[name] = values
        else:
            other_info.append(field)

    # Just go to next line if motif info isn't present.
    if motifn_present is False:
        return

    new_info_fields = "MOTIFN=" + ",".join(motifns)

    # Delete matches for the motif info fields that didn't pass the threshold.
    for f in motif_other:
        vals = motif_other[f]
        for i in sorted(failed_idx, reverse=True):
            del vals[i]
        motif_other[f] = vals
        new_info_fields += ";" + f + ",".join(vals)

    for x in other_info:
        new_info_fields += ";" + x

    line_list[7] = new_info_fields
    new_line = "\t".join(line_list)  # New, filtered output line.

    return new_line

def get_genes(exp_file, samples, threshold, max_only):
    """
    Reads in and parses the .bed expression file.
    File format expected to be:
        Whose format is tab seperated columns with header line:
        CHR  START  STOP  GENE  <sample 1>  <sample 2>  ...  <sample n>

    Args:
        exp_file (str): Name of expression file.
        samples (list): Names of the samples in the vcf file.
        threshold (float): Expression threshold to filter lowly/unexpressed genes.
        max_only (bool): if true, gene_dict value is 1 value = max expression
            if false gene_dict value is list of expression values
            YYY: WARNING: if want list to have meaning
                then values needs to be tied to header sample names

    Returns:
        gene_dict (dict): {gene_name: [expression_vals]}.
            Only include values for samples in the vcf.
    """
    data_cols = []
    gene_dict = {}

    print('start read exp_file:' + format(exp_file))

    if max_only:
        # read and only return max exp value in gene_dict
        with open(exp_file) as f:
            header = f.readline().strip().split('\t')
            for samp in header[4:]:
                if samp in samples:
                    data_idx = header.index(samp)
                    data_cols.append(data_idx)

            # Read in expression levels for each gene.
            for line in f:
                line = line.strip().split('\t')
                gene_name = line[3].upper()
                exp_val = -1e1000
                for idx in data_cols:
                    if float(line[idx]) > exp_val:
                        exp_val = float(line[idx])

                gene_dict[gene_name] = exp_val

    else:
        # read and return exp value list in gene_dict
        with open(exp_file) as f:
            header = f.readline().strip().split('\t')
            for samp in header[4:]:
                if samp in samples:
                    data_idx = header.index(samp)
                    data_cols.append(data_idx)

            # Read in expression levels for each gene.
            for line in f:
                line = line.strip().split('\t')
                gene_name = line[3].upper()
                exp_vals = []
                for idx in data_cols:
                    exp_vals.append(line[idx])

                gene_dict[gene_name] = exp_vals

    return gene_dict

def main(inp_file, exp_file, out_file, th=5, motif_file="", motifout_file="", use_vcf=True):
    """
        If use_vcf true then:
            For a given motif annotated vcf file (already run through motifs.py),
            remove all motif matches for TFs that are
            not expressed in at least one sample above the threshold.
        Else:
            For a given vcf file read the samples,
            then filter the motif_file, only outputing to motifout_file those items
            not expressed in at least one sample above the threshold

        Args:
            -i (str): Name of sorted variant file to process.
            -o (str): Name of output file to be created.
            -e (str): Name of expression file.
            -th (float): TFs are considered expressed if they are above this threshold.
    """

    if use_vcf:
        output_f = open(out_file, "w")

        with open(inp_file) as vcf:

            line = vcf.readline().strip()

            # Skip info lines.
            while line.startswith("##"):
                print(line, file=output_f)
                line = vcf.readline().strip()

            # First non-## line is the header. Get sample names and print to output.
            samples = parse_header(line)
            print(line, file=output_f)

            print("Creating gene dictionary for expression data.")
            gene_dict = get_genes(exp_file, samples, th, True)

            if len(gene_dict) == 0:
                print("Error, no genes above threshold found in expression file.",
                      "\nTry lowering the threshold and ensure the expression fil",
                      "e has values in the range that you expect.")
                sys.exit()

            print("Filtering motif info for TFs that don't meet the expression threshold of " +
                str(th) + ". Found " + format(len(gene_dict)) + " genes. Start processing vcf file.")
            for line in vcf:
                new_line = process_line(line, gene_dict, th)
                if new_line is not None:
                    print(new_line, file=output_f)

        output_f.close()
    else:
        # this version processes the motif file

        with open(inp_file) as vcf:

            line = vcf.readline().strip()

            # Skip info lines.
            while line.startswith("##"):
                line = vcf.readline().strip()

            # First non-## line is the header. Get sample names and print to output.
            samples = parse_header(line)

        # done with vcf file; only used to get samples

        print("Creating gene dictionary for expression data.")
        gene_dict = get_genes(exp_file, samples, th, True)

        if len(gene_dict) == 0:
            print("Error, no genes above threshold found in expression file.",
                  "\nTry lowering the threshold and ensure the expression fil",
                  "e has values in the range that you expect.")
            sys.exit()

        print("Filtering motif info for TFs that don't meet the expression threshold of " +
            str(th) + ". Found " + format(len(gene_dict)) + " genes. Start filtering motifs.")

        motif.get_filterbygene_put_motifs(motif_file, motifout_file, th, gene_dict)

    print("COMPLETE.")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(usage=__doc__)

    parser.add_argument("-i", "--input", dest="input_file", required=True)
    parser.add_argument("-e", "--expression", dest="exp_file", required=True)
    parser.add_argument("-m", "--motif_input", dest="motif_file", required=False)
    parser.add_argument("-mo", "--motif_output", dest="motif_out_file", required=False)
    parser.add_argument("-o", "--output", dest="output_file", required=True)
    parser.add_argument("-th", "--threshold", dest="threshold", required=False, default=5, type=float)

    args = parser.parse_args()

    inp_file = args.input_file
    exp_file = args.exp_file
    out_file = args.output_file
    th = args.threshold

    if args.motif_file is not None:
        motif_file = args.motif_file
        if args.motif_out_file is not None:
            motifout_file = args.motif_out_file
        else:
            motifout_file = motif_file.replace('.txt', '.tf_filtered.txt')
        main(inp_file, exp_file, out_file, th, motif_file, motifout_file, False)
    else:
        main(inp_file, exp_file, out_file, th, None, None, True)
<|endoftext|>


<commit_before>"""
Provides request handlers for file downloads.

"""
import gzip
import json

from aiohttp import web

import virtool.bio
import virtool.db.downloads
import virtool.db.history
import virtool.db.otus
import virtool.db.references
import virtool.errors
import virtool.http.routes
import virtool.otus
from virtool.api.utils import CustomEncoder, not_found

routes = virtool.http.routes.Routes()

@routes.get("/download/otus/{otu_id}/isolates/{isolate_id}")
async def download_isolate(req):
    """
    Download a FASTA file containing the sequences for a single Virtool isolate.

    """
    db = req.app["db"]

    otu_id = req.match_info["otu_id"]
    isolate_id = req.match_info["isolate_id"]

    try:
        filename, fasta = await virtool.db.downloads.generate_isolate_fasta(db, otu_id, isolate_id)
    except virtool.errors.DatabaseError as err:
        if "OTU does not exist" in str(err):
            return not_found("OTU not found")

        if "Isolate does not exist" in str(err):
            return not_found("Isolate not found")

        raise

    return web.Response(text=fasta, headers={
        "Content-Disposition": f"attachment; filename='{filename}'"
    })

@routes.get("/download/otus/{otu_id}")
async def download_otu(req):
    """
    Download a FASTA file containing the sequences for all isolates in a single Virtool otu.

    """
    db = req.app["db"]

    otu_id = req.match_info["otu_id"]

    try:
        filename, fasta = await virtool.db.downloads.generate_otu_fasta(db, otu_id)
    except virtool.errors.DatabaseError as err:
        if "Sequence does not exist" in str(err):
            return not_found("Sequence not found")

        if "OTU does not exist" in str(err):
            return not_found("OTU not found")

        raise

    if not fasta:
        return web.Response(status=404)

    return web.Response(text=fasta, headers={
        "Content-Disposition": f"attachment; filename='{filename}'"
    })

@routes.get("/download/refs/{ref_id}")
async def download_reference(req):
    """
    Export all otus and sequences for a given reference as a gzipped JSON string. Made available as a downloadable file
    named ``reference.json.gz``.

    """
    db = req.app["db"]

    ref_id = req.match_info["ref_id"]

    document = await db.references.find_one(ref_id, ["data_type", "organism"])

    if document is None:
        return not_found()

    scope = req.query.get("scope", "built")

    if scope not in ["built", "remote", "unbuilt", "unverified"]:
        scope = "built"

    otu_list = await virtool.db.references.export(db, ref_id, scope)

    data = {
        "otus": otu_list,
        "data_type": document["data_type"],
        "organism": document["organism"]
    }

    # Convert the list of OTUs to a JSON-formatted string.
    json_string = json.dumps(data, cls=CustomEncoder)

    # Compress the JSON string with gzip.
    body = await req.app["run_in_process"](gzip.compress, bytes(json_string, "utf-8"))

    return web.Response(
        headers={"Content-Disposition": f"attachment; filename='reference.{scope}.json.gz'"},
        content_type="application/gzip",
        body=body
    )

@routes.get("/download/sequences/{sequence_id}")
async def download_sequence(req):
    """
    Download a FASTA file containing a single Virtool sequence.

    """
    db = req.app["db"]

    sequence_id = req.match_info["sequence_id"]

    try:
        filename, fasta = await virtool.db.downloads.generate_sequence_fasta(db, sequence_id)
    except virtool.errors.DatabaseError as err:
        if "Sequence does not exist" in str(err):
            return not_found("Sequence not found")

        if "Isolate does not exist" in str(err):
            return not_found("Isolate not found")

        if "OTU does not exist" in str(err):
            return not_found("OTU not found")

        raise

    if fasta is None:
        return web.Response(status=404)

    return web.Response(text=fasta, headers={
        "Content-Disposition": f"attachment; filename='{filename}'"
    })
<commit_msg>Add sample file download handler<commit_after>"""
Provides request handlers for file downloads.

"""
import os
import gzip
import json

from aiohttp import web

import virtool.bio
import virtool.db.downloads
import virtool.db.history
import virtool.db.otus
import virtool.db.references
import virtool.errors
import virtool.http.routes
import virtool.otus
import virtool.utils
from virtool.api.utils import CustomEncoder, bad_request, not_found

routes = virtool.http.routes.Routes()

@routes.get("/download/samples/{sample_id}/{prefix}_{suffix}.fq")
async def download_uncompressed_sample_reads(req):
    db = req.app["db"]
    data_path = req.app["settings"]["data_path"]

    sample_id = req.match_info["sample_id"]

    document = await db.samples.find_one(sample_id, ["paired", "files"])

    if document is None or not document.get("files", False):
        return not_found()

    suffix = req.match_info["suffix"]

    path = os.path.join(data_path, "samples", sample_id, f"reads_{suffix}.fastq")

    if not os.path.isfile(path):
        return not_found()

    file_stats = virtool.utils.file_stats(path)

    headers = {
        "Content-Length": file_stats["size"],
        "Content-Type": "text"
    }

    return web.FileResponse(path, chunk_size=1024*1024, headers=headers)

@routes.get("/download/samples/{sample_id}/{prefix}_{suffix}.fq.gz")
async def download_sample_reads(req):
    db = req.app["db"]
    data_path = req.app["settings"]["data_path"]

    sample_id = req.match_info["sample_id"]

    document = await db.samples.find_one(sample_id, ["paired", "files"])

    if document is None or not document.get("files", False):
        return not_found()

    suffix = req.match_info["suffix"]

    path = os.path.join(data_path, "samples", sample_id, f"reads_{suffix}.fq.gz")

    if not os.path.isfile(path):
        return not_found()

    file_stats = virtool.utils.file_stats(path)

    headers = {
        "Content-Length": file_stats["size"],
        "Content-Type": "application/gzip"
    }

    return web.FileResponse(path, chunk_size=1024*1024, headers=headers)

@routes.get("/download/otus/{otu_id}/isolates/{isolate_id}")
async def download_isolate(req):
    """
    Download a FASTA file containing the sequences for a single Virtool isolate.

    """
    db = req.app["db"]

    otu_id = req.match_info["otu_id"]
    isolate_id = req.match_info["isolate_id"]

    try:
        filename, fasta = await virtool.db.downloads.generate_isolate_fasta(db, otu_id, isolate_id)
    except virtool.errors.DatabaseError as err:
        if "OTU does not exist" in str(err):
            return not_found("OTU not found")

        if "Isolate does not exist" in str(err):
            return not_found("Isolate not found")

        raise

    return web.Response(text=fasta, headers={
        "Content-Disposition": f"attachment; filename='{filename}'"
    })

@routes.get("/download/otus/{otu_id}")
async def download_otu(req):
    """
    Download a FASTA file containing the sequences for all isolates in a single Virtool otu.

    """
    db = req.app["db"]

    otu_id = req.match_info["otu_id"]

    try:
        filename, fasta = await virtool.db.downloads.generate_otu_fasta(db, otu_id)
    except virtool.errors.DatabaseError as err:
        if "Sequence does not exist" in str(err):
            return not_found("Sequence not found")

        if "OTU does not exist" in str(err):
            return not_found("OTU not found")

        raise

    if not fasta:
        return web.Response(status=404)

    return web.Response(text=fasta, headers={
        "Content-Disposition": f"attachment; filename='{filename}'"
    })

@routes.get("/download/refs/{ref_id}")
async def download_reference(req):
    """
    Export all otus and sequences for a given reference as a gzipped JSON string. Made available as a downloadable file
    named ``reference.json.gz``.

    """
    db = req.app["db"]

    ref_id = req.match_info["ref_id"]

    document = await db.references.find_one(ref_id, ["data_type", "organism"])

    if document is None:
        return not_found()

    scope = req.query.get("scope", "built")

    if scope not in ["built", "remote", "unbuilt", "unverified"]:
        scope = "built"

    otu_list = await virtool.db.references.export(db, ref_id, scope)

    data = {
        "otus": otu_list,
        "data_type": document["data_type"],
        "organism": document["organism"]
    }

    # Convert the list of OTUs to a JSON-formatted string.
    json_string = json.dumps(data, cls=CustomEncoder)

    # Compress the JSON string with gzip.
    body = await req.app["run_in_process"](gzip.compress, bytes(json_string, "utf-8"))

    return web.Response(
        headers={"Content-Disposition": f"attachment; filename='reference.{scope}.json.gz'"},
        content_type="application/gzip",
        body=body
    )

@routes.get("/download/sequences/{sequence_id}")
async def download_sequence(req):
    """
    Download a FASTA file containing a single Virtool sequence.

    """
    db = req.app["db"]

    sequence_id = req.match_info["sequence_id"]

    try:
        filename, fasta = await virtool.db.downloads.generate_sequence_fasta(db, sequence_id)
    except virtool.errors.DatabaseError as err:
        if "Sequence does not exist" in str(err):
            return not_found("Sequence not found")

        if "Isolate does not exist" in str(err):
            return not_found("Isolate not found")

        if "OTU does not exist" in str(err):
            return not_found("OTU not found")

        raise

    if fasta is None:
        return web.Response(status=404)

    return web.Response(text=fasta, headers={
        "Content-Disposition": f"attachment; filename='{filename}'"
    })
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from django.test import TestCase
from django.utils import timezone
import json

import mock
from vkontakte_groups.factories import GroupFactory

from factories import AlbumFactory, VideoFactory
from models import Album, Video, Comment
from vkontakte_users.factories import UserFactory, User
#from datetime import datetime
#import mock
GROUP_ID = 16297716  # https://vk.com/cocacola
ALBUM_ID = 50850761  # 9 videos
VIDEO_ID = 166742757  # 12 comments

GROUP_CRUD_ID = 59154616  # https://vk.com/club59154616 # django-vkontakte-wall crud operations
ALBUM_CRUD_ID = 55964907
VIDEO_CRUD_ID = 170947024
#USER_AUTHOR_ID = 201164356

class AlbumTest(TestCase):

    def test_fetch_owner_albums(self):

        owner = GroupFactory(remote_id=GROUP_ID)

        self.assertEqual(Album.objects.count(), 0)

        albums = Album.remote.fetch(owner=owner)

        self.assertTrue(len(albums) > 0)
        self.assertEqual(Album.objects.count(), len(albums))
        self.assertEqual(albums[0].owner, owner)

    def test_fetch_with_count_and_offset(self):
        # testing `count` parameter, count is the same as limit
        owner = GroupFactory(remote_id=GROUP_ID)

        self.assertEqual(Album.objects.count(), 0)

        albums = Album.remote.fetch(owner=owner, count=5)

        self.assertEqual(len(albums), 5)
        self.assertEqual(Album.objects.count(), 5)

        # testing `offset` parameter
        albums2 = Album.remote.fetch(owner=owner, count=2, offset=4)

        self.assertEqual(len(albums2), 2)
        self.assertEqual(Album.objects.count(), 6)

        self.assertEqual(albums[4].remote_id, albums2[0].remote_id)

    def test_parse_album(self):

        owner = GroupFactory(remote_id=GROUP_ID)

        d = {u'count': 16, u'photo_320': u'http://cs619722.vk.me/u8704019/video/l_6369beb6.jpg', u'title': u'Coca-Cola Football',
             u'photo_160': u'http://cs619722.vk.me/u8704019/video/m_ef3493e1.jpg', u'id': 54387280, u'owner_id': -16297716}

        instance = Album()
        instance.parse(d)
        instance.save()

        self.assertEqual(instance.owner, owner)

        self.assertEqual(instance.pk, d['id'])
        self.assertEqual(instance.title, d['title'])
        self.assertEqual(instance.videos_count, d['count'])
        self.assertEqual(instance.photo_160, d['photo_160'])

class VideoTest(TestCase):

    def test_album_fetch_videos(self):

        owner = GroupFactory(remote_id=GROUP_ID)
        #album = AlbumFactory(remote_id=ALBUM_ID, owner=owner)

        albums = Album.remote.fetch(owner=owner)   # have to fetch for album.videos_count
        album = Album.objects.get(remote_id=ALBUM_ID)

        self.assertEqual(Video.objects.count(), 0)

        videos = album.fetch_videos()  # extended=True

        self.assertTrue(len(videos) > 0)
        self.assertEqual(album.videos_count, len(videos))
        self.assertEqual(Video.objects.count(), len(videos))
        self.assertEqual(videos[0].owner, owner)
        self.assertEqual(videos[0].album, album)
        #self.assertTrue(videos[0].likes_count > 0)
        self.assertTrue(videos[0].comments_count > 0)

        # testing `after` parameter
        after = Video.objects.order_by('-date')[4].date

        Video.objects.all().delete()
        self.assertEqual(Video.objects.count(), 0)

        videos = album.fetch_videos(after=after)
        self.assertEqual(len(videos), Video.objects.count())
        self.assertEqual(len(videos), 5)

        date = videos[0].date
        self.assertGreaterEqual(date, after)

        # testing `before` parameter
        before = Video.objects.order_by('-date')[2].date

        Video.objects.all().delete()
        self.assertEqual(Video.objects.count(), 0)

        videos = album.fetch_videos(before=before, after=after)
        self.assertEqual(len(videos), Video.objects.count())
        self.assertEqual(len(videos), 3)

        date = videos[0].date
        self.assertGreaterEqual(date, after)

        date = videos.last().date
        self.assertLessEqual(date, before)

    def test_fetch_with_count_and_offset(self):
        # testing `count` parameter, count is the same as limit
        owner = GroupFactory(remote_id=GROUP_ID)
        album = AlbumFactory(remote_id=ALBUM_ID, owner=owner)

        self.assertEqual(Video.objects.count(), 0)

        videos = album.fetch_videos(count=5)

        self.assertEqual(len(videos), 5)
        self.assertEqual(Video.objects.count(), 5)

        # testing `offset` parameter
        videos2 = album.fetch_videos(count=2, offset=4)

        self.assertEqual(len(videos2), 2)
        self.assertEqual(Video.objects.count(), 6)

        # TODO: to fix
        print videos[4].remote_id
        print videos2[0].remote_id
        print videos2[1].remote_id

        self.assertEqual(videos[4].remote_id, videos2[0].remote_id)

    def test_fetch_videos_by_ids(self):
        owner = GroupFactory(remote_id=GROUP_ID)
        album = AlbumFactory(remote_id=ALBUM_ID, owner=owner)

        self.assertEqual(Video.objects.count(), 0)

        videos = Video.remote.fetch(owner=owner, ids=[VIDEO_ID])

        self.assertEqual(len(videos), 1)
        #self.assertEqual(album.videos_count, 1)
        self.assertEqual(Video.objects.count(), 1)
        self.assertEqual(videos[0].owner, owner)
        self.assertEqual(videos[0].album, album)

        # fetch by album parameter
        videos = Video.remote.fetch(album=album, ids=[VIDEO_ID])
        self.assertEqual(len(videos), 1)

    def test_parse_video(self):

        owner = GroupFactory(remote_id=GROUP_ID)
        album = AlbumFactory(remote_id=ALBUM_ID, owner=owner)

        response = '''{"photo_130": "http://cs313422.vk.me/u163668241/video/s_6819a7d1.jpg",
            "repeat": 0,
            "photo_320": "http://cs313422.vk.me/u163668241/video/l_4cc8a38a.jpg",
            "description": "bla bla bla",
            "title": "Эстафета Олимпийского огня «Сочи 2014». Неделя 3-я",
            "can_repost": 1, "views": 928, "album_id": 50850761, "comments": 12, "player": "http://www.youtube.com/embed/UmDAmM53bU0", "date": 1386074580, "likes": {"count": 191, "user_likes": 0}, "duration": 206, "can_comment": 1, "id": 166742757, "owner_id": -16297716}
        '''
        d = json.loads(response)

        instance = Video()
        instance.parse(d)
        instance.save()

        self.assertEqual(instance.album, album)
        self.assertEqual(instance.owner, owner)

        self.assertEqual(instance.remote_id, d['id'])
        self.assertEqual(instance.title, d['title'])
        self.assertEqual(instance.description, d['description'])
        self.assertEqual(instance.photo_130, d['photo_130'])
        self.assertEqual(instance.player, d['player'])
        self.assertEqual(instance.views_count, d['views'])
        self.assertEqual(instance.comments_count, d['comments'])
        self.assertEqual(instance.duration, d['duration'])

        self.assertIsNotNone(instance.date)

class CommentTest(TestCase):

    def setUp(self):
        self.objects_to_delete = []

    def tearDown(self):
        for object in self.objects_to_delete:
            object.delete(commit_remote=True)

    @mock.patch('vkontakte_users.models.User.remote.fetch', side_effect=lambda ids, **kw: User.objects.filter(id__in=[user.id for user in [UserFactory.create(remote_id=i) for i in ids]]))
    def test_video_fetch_comments(self, *kwargs):

        owner = GroupFactory(remote_id=GROUP_ID)
        #album = AlbumFactory(remote_id=ALBUM_ID, owner=owner)
        # not factory coz we need video.comments_count later
        video = Video.remote.fetch(owner=owner, ids=[VIDEO_ID])[0]

        comments = video.fetch_comments(count=10, sort='desc')
        self.assertEqual(len(comments), video.comments.count())
        self.assertEqual(len(comments), 10)

        # testing `after` parameter
        after = Comment.objects.order_by('-date')[2].date

        Comment.objects.all().delete()
        self.assertEqual(Comment.objects.count(), 0)

        comments = video.fetch_comments(after=after, sort='desc')
        self.assertEqual(len(comments), Comment.objects.count())
        self.assertEqual(len(comments), video.comments.count())
        self.assertEqual(len(comments), 3)

        date = comments[0].date
        self.assertGreaterEqual(date, after)

        # testing `all` parameter
        Comment.objects.all().delete()
        self.assertEqual(Comment.objects.count(), 0)

        comments = video.fetch_comments(all=True)
        self.assertEqual(len(comments), Comment.objects.count())
        self.assertEqual(len(comments), video.comments.count())
        self.assertEqual(len(comments), video.comments_count)
        self.assertTrue(video.comments.count() > 10)

    def test_fetch_with_count_and_offset(self):
        # testing `count` parameter, count is the same as limit
        owner = GroupFactory(remote_id=GROUP_ID)
        album = AlbumFactory(remote_id=ALBUM_ID, owner=owner)
        video = VideoFactory(remote_id=VIDEO_ID, album=album, owner=owner)

        self.assertEqual(Comment.objects.count(), 0)

        comments = video.fetch_comments(count=5)

        self.assertEqual(len(comments), 5)
        self.assertEqual(Comment.objects.count(), 5)

        # testing `offset` parameter
        comments2 = video.fetch_comments(count=2, offset=4)

        self.assertEqual(len(comments2), 2)
        self.assertEqual(Comment.objects.count(), 6)

        self.assertEqual(comments[4].remote_id, comments2[0].remote_id)

    #@mock.patch('vkontakte_users.models.User.remote.fetch', side_effect=lambda ids, **kw: User.objects.filter(id__in=[user.id for user in [UserFactory.create(remote_id=i) for i in ids]]))
    def test_video_fetch_likes(self, *kwargs):

        owner = GroupFactory(remote_id=GROUP_ID)
        album = AlbumFactory(remote_id=ALBUM_ID, owner=owner)
        video = VideoFactory(remote_id=VIDEO_ID, album=album, owner=owner)

        self.assertEqual(video.likes_count, 0)
        users_initial = User.objects.count()

        users = video.fetch_likes(all=True)

        self.assertTrue(video.likes_count > 0)
        self.assertEqual(video.likes_count, len(users))
        self.assertEqual(video.likes_count, User.objects.count() - users_initial)
        self.assertEqual(video.likes_count, video.like_users.count())

class OtherTests(TestCase):

    def test_fetch_by_user_parameter(self):
        user = UserFactory(remote_id=13312307)

        # fetch albums
        albums = Album.remote.fetch(owner=user)
        self.assertGreater(len(albums), 0)
        self.assertEqual(Album.objects.count(), len(albums))
        self.assertEqual(albums[0].owner, user)

        # fetch album videos
        album = albums[0]
        videos = album.fetch_videos()
        self.assertGreater(len(videos), 0)
        self.assertEqual(Video.objects.count(), len(videos))
        self.assertEqual(videos[0].owner, user)

        # fetch user video comments
        video = videos[0]
        comments = video.fetch_comments()
        self.assertGreater(len(comments), 0)
        self.assertEqual(Comment.objects.count(), len(comments))
        self.assertEqual(comments[0].author, user)

        # fetch user video likes
        users = video.fetch_likes(all=True)
        self.assertTrue(video.likes_count > 0)
        self.assertEqual(video.likes_count, len(users))

        # fetch all user videos
        videos = Video.remote.fetch(owner=user)
        self.assertGreater(len(videos), 0)
        self.assertEqual(Video.objects.count(), len(videos))
        self.assertEqual(videos[0].owner, user)

    def test_link(self):
        owner = GroupFactory(remote_id=GROUP_ID)
        album = AlbumFactory(remote_id=ALBUM_ID, owner=owner)
        video = VideoFactory(remote_id=VIDEO_ID, album=album, owner=owner)

        self.assertEqual(album.link.count("-"), 1)
        self.assertEqual(video.link.count("-"), 1)

        user = UserFactory(remote_id=13312307)
        album = AlbumFactory(remote_id=55976289, owner=user)
        video = VideoFactory(remote_id=165144348, album=album, owner=user)

        self.assertEqual(album.link.count("-"), 0)
        self.assertEqual(video.link.count("-"), 0)

class OldTests():

    def test_fetch_photo_likes_parser(self):

        owner = GroupFactory(remote_id=GROUP_ID)
        album = AlbumFactory(remote_id=ALBUM_ID, owner=owner)
        photo = PhotoFactory(remote_id=PHOTO_ID, album=album)

        self.assertEqual(photo.likes_count, 0)
        photo.fetch_likes_parser()
        self.assertTrue(photo.likes_count > 0)

    def test_fetch_photo_comments_parser(self):

        owner = GroupFactory(remote_id=GROUP_ID)
        album = AlbumFactory(remote_id=ALBUM_ID, owner=owner)
        photo = PhotoFactory(remote_id=PHOTO_ID, album=album)

        self.assertEqual(photo.comments_count, 0)
        photo.fetch_comments_parser()
        self.assertTrue(photo.comments_count > 0)
<commit_msg>remove print<commit_after># -*- coding: utf-8 -*-
from django.test import TestCase
from django.utils import timezone
import json

import mock
from vkontakte_groups.factories import GroupFactory

from factories import AlbumFactory, VideoFactory
from models import Album, Video, Comment
from vkontakte_users.factories import UserFactory, User
#from datetime import datetime
#import mock
GROUP_ID = 16297716  # https://vk.com/cocacola
ALBUM_ID = 50850761  # 9 videos
VIDEO_ID = 166742757  # 12 comments

GROUP_CRUD_ID = 59154616  # https://vk.com/club59154616 # django-vkontakte-wall crud operations
ALBUM_CRUD_ID = 55964907
VIDEO_CRUD_ID = 170947024
#USER_AUTHOR_ID = 201164356

class AlbumTest(TestCase):

    def test_fetch_owner_albums(self):

        owner = GroupFactory(remote_id=GROUP_ID)

        self.assertEqual(Album.objects.count(), 0)

        albums = Album.remote.fetch(owner=owner)

        self.assertTrue(len(albums) > 0)
        self.assertEqual(Album.objects.count(), len(albums))
        self.assertEqual(albums[0].owner, owner)

    def test_fetch_with_count_and_offset(self):
        # testing `count` parameter, count is the same as limit
        owner = GroupFactory(remote_id=GROUP_ID)

        self.assertEqual(Album.objects.count(), 0)

        albums = Album.remote.fetch(owner=owner, count=5)

        self.assertEqual(len(albums), 5)
        self.assertEqual(Album.objects.count(), 5)

        # testing `offset` parameter
        albums2 = Album.remote.fetch(owner=owner, count=2, offset=4)

        self.assertEqual(len(albums2), 2)
        self.assertEqual(Album.objects.count(), 6)

        self.assertEqual(albums[4].remote_id, albums2[0].remote_id)

    def test_parse_album(self):

        owner = GroupFactory(remote_id=GROUP_ID)

        d = {u'count': 16, u'photo_320': u'http://cs619722.vk.me/u8704019/video/l_6369beb6.jpg', u'title': u'Coca-Cola Football',
             u'photo_160': u'http://cs619722.vk.me/u8704019/video/m_ef3493e1.jpg', u'id': 54387280, u'owner_id': -16297716}

        instance = Album()
        instance.parse(d)
        instance.save()

        self.assertEqual(instance.owner, owner)

        self.assertEqual(instance.pk, d['id'])
        self.assertEqual(instance.title, d['title'])
        self.assertEqual(instance.videos_count, d['count'])
        self.assertEqual(instance.photo_160, d['photo_160'])

class VideoTest(TestCase):

    def test_album_fetch_videos(self):

        owner = GroupFactory(remote_id=GROUP_ID)
        #album = AlbumFactory(remote_id=ALBUM_ID, owner=owner)

        albums = Album.remote.fetch(owner=owner)   # have to fetch for album.videos_count
        album = Album.objects.get(remote_id=ALBUM_ID)

        self.assertEqual(Video.objects.count(), 0)

        videos = album.fetch_videos()  # extended=True

        self.assertTrue(len(videos) > 0)
        self.assertEqual(album.videos_count, len(videos))
        self.assertEqual(Video.objects.count(), len(videos))
        self.assertEqual(videos[0].owner, owner)
        self.assertEqual(videos[0].album, album)
        #self.assertTrue(videos[0].likes_count > 0)
        self.assertTrue(videos[0].comments_count > 0)

        # testing `after` parameter
        after = Video.objects.order_by('-date')[4].date

        Video.objects.all().delete()
        self.assertEqual(Video.objects.count(), 0)

        videos = album.fetch_videos(after=after)
        self.assertEqual(len(videos), Video.objects.count())
        self.assertEqual(len(videos), 5)

        date = videos[0].date
        self.assertGreaterEqual(date, after)

        # testing `before` parameter
        before = Video.objects.order_by('-date')[2].date

        Video.objects.all().delete()
        self.assertEqual(Video.objects.count(), 0)

        videos = album.fetch_videos(before=before, after=after)
        self.assertEqual(len(videos), Video.objects.count())
        self.assertEqual(len(videos), 3)

        date = videos[0].date
        self.assertGreaterEqual(date, after)

        date = videos.last().date
        self.assertLessEqual(date, before)

    def test_fetch_with_count_and_offset(self):
        # testing `count` parameter, count is the same as limit
        owner = GroupFactory(remote_id=GROUP_ID)
        album = AlbumFactory(remote_id=ALBUM_ID, owner=owner)

        self.assertEqual(Video.objects.count(), 0)

        videos = album.fetch_videos(count=5)

        self.assertEqual(len(videos), 5)
        self.assertEqual(Video.objects.count(), 5)

        # testing `offset` parameter
        videos2 = album.fetch_videos(count=2, offset=4)

        self.assertEqual(len(videos2), 2)
        self.assertEqual(Video.objects.count(), 6)

        # print videos[4].remote_id
        # print videos2[0].remote_id
        # print videos2[1].remote_id

        self.assertEqual(videos[4].remote_id, videos2[0].remote_id)

    def test_fetch_videos_by_ids(self):
        owner = GroupFactory(remote_id=GROUP_ID)
        album = AlbumFactory(remote_id=ALBUM_ID, owner=owner)

        self.assertEqual(Video.objects.count(), 0)

        videos = Video.remote.fetch(owner=owner, ids=[VIDEO_ID])

        self.assertEqual(len(videos), 1)
        #self.assertEqual(album.videos_count, 1)
        self.assertEqual(Video.objects.count(), 1)
        self.assertEqual(videos[0].owner, owner)
        self.assertEqual(videos[0].album, album)

        # fetch by album parameter
        videos = Video.remote.fetch(album=album, ids=[VIDEO_ID])
        self.assertEqual(len(videos), 1)

    def test_parse_video(self):

        owner = GroupFactory(remote_id=GROUP_ID)
        album = AlbumFactory(remote_id=ALBUM_ID, owner=owner)

        response = '''{"photo_130": "http://cs313422.vk.me/u163668241/video/s_6819a7d1.jpg",
            "repeat": 0,
            "photo_320": "http://cs313422.vk.me/u163668241/video/l_4cc8a38a.jpg",
            "description": "bla bla bla",
            "title": "Эстафета Олимпийского огня «Сочи 2014». Неделя 3-я",
            "can_repost": 1, "views": 928, "album_id": 50850761, "comments": 12, "player": "http://www.youtube.com/embed/UmDAmM53bU0", "date": 1386074580, "likes": {"count": 191, "user_likes": 0}, "duration": 206, "can_comment": 1, "id": 166742757, "owner_id": -16297716}
        '''
        d = json.loads(response)

        instance = Video()
        instance.parse(d)
        instance.save()

        self.assertEqual(instance.album, album)
        self.assertEqual(instance.owner, owner)

        self.assertEqual(instance.remote_id, d['id'])
        self.assertEqual(instance.title, d['title'])
        self.assertEqual(instance.description, d['description'])
        self.assertEqual(instance.photo_130, d['photo_130'])
        self.assertEqual(instance.player, d['player'])
        self.assertEqual(instance.views_count, d['views'])
        self.assertEqual(instance.comments_count, d['comments'])
        self.assertEqual(instance.duration, d['duration'])

        self.assertIsNotNone(instance.date)

class CommentTest(TestCase):

    def setUp(self):
        self.objects_to_delete = []

    def tearDown(self):
        for object in self.objects_to_delete:
            object.delete(commit_remote=True)

    @mock.patch('vkontakte_users.models.User.remote.fetch', side_effect=lambda ids, **kw: User.objects.filter(id__in=[user.id for user in [UserFactory.create(remote_id=i) for i in ids]]))
    def test_video_fetch_comments(self, *kwargs):

        owner = GroupFactory(remote_id=GROUP_ID)
        #album = AlbumFactory(remote_id=ALBUM_ID, owner=owner)
        # not factory coz we need video.comments_count later
        video = Video.remote.fetch(owner=owner, ids=[VIDEO_ID])[0]

        comments = video.fetch_comments(count=10, sort='desc')
        self.assertEqual(len(comments), video.comments.count())
        self.assertEqual(len(comments), 10)

        # testing `after` parameter
        after = Comment.objects.order_by('-date')[2].date

        Comment.objects.all().delete()
        self.assertEqual(Comment.objects.count(), 0)

        comments = video.fetch_comments(after=after, sort='desc')
        self.assertEqual(len(comments), Comment.objects.count())
        self.assertEqual(len(comments), video.comments.count())
        self.assertEqual(len(comments), 3)

        date = comments[0].date
        self.assertGreaterEqual(date, after)

        # testing `all` parameter
        Comment.objects.all().delete()
        self.assertEqual(Comment.objects.count(), 0)

        comments = video.fetch_comments(all=True)
        self.assertEqual(len(comments), Comment.objects.count())
        self.assertEqual(len(comments), video.comments.count())
        self.assertEqual(len(comments), video.comments_count)
        self.assertTrue(video.comments.count() > 10)

    def test_fetch_with_count_and_offset(self):
        # testing `count` parameter, count is the same as limit
        owner = GroupFactory(remote_id=GROUP_ID)
        album = AlbumFactory(remote_id=ALBUM_ID, owner=owner)
        video = VideoFactory(remote_id=VIDEO_ID, album=album, owner=owner)

        self.assertEqual(Comment.objects.count(), 0)

        comments = video.fetch_comments(count=5)

        self.assertEqual(len(comments), 5)
        self.assertEqual(Comment.objects.count(), 5)

        # testing `offset` parameter
        comments2 = video.fetch_comments(count=2, offset=4)

        self.assertEqual(len(comments2), 2)
        self.assertEqual(Comment.objects.count(), 6)

        self.assertEqual(comments[4].remote_id, comments2[0].remote_id)

    #@mock.patch('vkontakte_users.models.User.remote.fetch', side_effect=lambda ids, **kw: User.objects.filter(id__in=[user.id for user in [UserFactory.create(remote_id=i) for i in ids]]))
    def test_video_fetch_likes(self, *kwargs):

        owner = GroupFactory(remote_id=GROUP_ID)
        album = AlbumFactory(remote_id=ALBUM_ID, owner=owner)
        video = VideoFactory(remote_id=VIDEO_ID, album=album, owner=owner)

        self.assertEqual(video.likes_count, 0)
        users_initial = User.objects.count()

        users = video.fetch_likes(all=True)

        self.assertTrue(video.likes_count > 0)
        self.assertEqual(video.likes_count, len(users))
        self.assertEqual(video.likes_count, User.objects.count() - users_initial)
        self.assertEqual(video.likes_count, video.like_users.count())

class OtherTests(TestCase):

    def test_fetch_by_user_parameter(self):
        user = UserFactory(remote_id=13312307)

        # fetch albums
        albums = Album.remote.fetch(owner=user)
        self.assertGreater(len(albums), 0)
        self.assertEqual(Album.objects.count(), len(albums))
        self.assertEqual(albums[0].owner, user)

        # fetch album videos
        album = albums[0]
        videos = album.fetch_videos()
        self.assertGreater(len(videos), 0)
        self.assertEqual(Video.objects.count(), len(videos))
        self.assertEqual(videos[0].owner, user)

        # fetch user video comments
        video = videos[0]
        comments = video.fetch_comments()
        self.assertGreater(len(comments), 0)
        self.assertEqual(Comment.objects.count(), len(comments))
        self.assertEqual(comments[0].author, user)

        # fetch user video likes
        users = video.fetch_likes(all=True)
        self.assertTrue(video.likes_count > 0)
        self.assertEqual(video.likes_count, len(users))

        # fetch all user videos
        videos = Video.remote.fetch(owner=user)
        self.assertGreater(len(videos), 0)
        self.assertEqual(Video.objects.count(), len(videos))
        self.assertEqual(videos[0].owner, user)

    def test_link(self):
        owner = GroupFactory(remote_id=GROUP_ID)
        album = AlbumFactory(remote_id=ALBUM_ID, owner=owner)
        video = VideoFactory(remote_id=VIDEO_ID, album=album, owner=owner)

        self.assertEqual(album.link.count("-"), 1)
        self.assertEqual(video.link.count("-"), 1)

        user = UserFactory(remote_id=13312307)
        album = AlbumFactory(remote_id=55976289, owner=user)
        video = VideoFactory(remote_id=165144348, album=album, owner=user)

        self.assertEqual(album.link.count("-"), 0)
        self.assertEqual(video.link.count("-"), 0)

class OldTests():

    def test_fetch_photo_likes_parser(self):

        owner = GroupFactory(remote_id=GROUP_ID)
        album = AlbumFactory(remote_id=ALBUM_ID, owner=owner)
        photo = PhotoFactory(remote_id=PHOTO_ID, album=album)

        self.assertEqual(photo.likes_count, 0)
        photo.fetch_likes_parser()
        self.assertTrue(photo.likes_count > 0)

    def test_fetch_photo_comments_parser(self):

        owner = GroupFactory(remote_id=GROUP_ID)
        album = AlbumFactory(remote_id=ALBUM_ID, owner=owner)
        photo = PhotoFactory(remote_id=PHOTO_ID, album=album)

        self.assertEqual(photo.comments_count, 0)
        photo.fetch_comments_parser()
        self.assertTrue(photo.comments_count > 0)
<|endoftext|>


<commit_before>from flask.ext.wtf import Form
from wtforms import StringField, SubmitField, \
    TextAreaField, SelectField
from wtforms.validators import Required

class NoteForm(Form):
    title = StringField('Title:', validators=[Required()])
    body = TextAreaField('Dump Your Brain:', validators=[Required()])
    body_html = TextAreaField()
    tags = StringField()
    notebook = SelectField(coerce=int)
    submit = SubmitField('Submit')

class ShareForm(Form):
    recipient_email = StringField('Recipient Email', validators=[Required()])
    submit = SubmitField('Share')

class NotebookForm(Form):
    title = StringField('Title:', validators=[Required()])
    submit = SubmitField('Submit')

class SearchForm(Form):
    search_field = StringField()
    submit = SubmitField('Search')
<commit_msg>Added validation for title length in Note and Notebook, as present in the Database.<commit_after>from flask.ext.wtf import Form
from wtforms import StringField, SubmitField, \
    TextAreaField, SelectField
from wtforms.validators import Required

class NoteForm(Form):
    title = StringField('Title:', validators=[Required(), Length(1, 200)])
    body = TextAreaField('Dump Your Brain:', validators=[Required()])
    body_html = TextAreaField()
    tags = StringField()
    notebook = SelectField(coerce=int)
    submit = SubmitField('Submit')

class ShareForm(Form):
    recipient_email = StringField('Recipient Email', validators=[Required()])
    submit = SubmitField('Share')

class NotebookForm(Form):
    title = StringField('Title:', validators=[Required(), Length(1, 200)])
    submit = SubmitField('Submit')

class SearchForm(Form):
    search_field = StringField()
    submit = SubmitField('Search')
<|endoftext|>


<commit_before>from flask.ext.wtf import Form
from wtforms import StringField, TextAreaField, BooleanField, SelectField,\
    SubmitField
from wtforms.validators import Required, Length, Email, Regexp
from wtforms import ValidationError
from flask.ext.pagedown.fields import PageDownField
from ..models import Role, User

class NameForm(Form):
    name = StringField('What is your name?', validators=[Required()])
    # name űЧ֣û룬ֵΪ None
    submit = SubmitField('Submit')

class EditProfileForm(Form):
    name = StringField('Real name', validators=[Length(0, 64)])
    location = StringField('Location', validators=[Length(0, 64)])
    about_me = TextAreaField('About me')
    submit = SubmitField('Submit')

class EditProfileAdminForm(Form):
    email = StringField('Email', validators=[Required(), Length(1, 64),
                                             Email()])
    username = StringField('Username', validators=[
        Required(), Length(1, 64), Regexp('^[A-Za-z][A-Za-z0-9_.]*$', 0,
                                          'Usernames must have only letters, '
                                          'numbers, dots or underscores')])
    confirmed = BooleanField('Confirmed')
    role = SelectField('Role', coerce=int)
    name = StringField('Real name', validators=[Length(0, 64)])
    location = StringField('Location', validators=[Length(0, 64)])
    about_me = TextAreaField('About me')
    submit = SubmitField('Submit')

    def __init__(self, user, *args, **kwargs):
        super(EditProfileAdminForm, self).__init__(*args, **kwargs)
        self.role.choices = [(role.id, role.name)
                             for role in Role.query.order_by(Role.name).all()]
        self.user = user

    def validate_email(self, field):
        if field.data != self.user.email and \
                User.query.filter_by(email=field.data).first():
            raise ValidationError('Email already registered.')

    def validate_username(self, field):
        if field.data != self.user.username and \
                User.query.filter_by(username=field.data).first():
            raise ValidationError('Username already in use.')

class PostForm(Form):
    body = PageDownField("What's on your mind?", validators=[Required()])
    submit = SubmitField('Submit')
###<commit_msg>nothing<commit_after>from flask.ext.wtf import Form
from wtforms import StringField, TextAreaField, BooleanField, SelectField,\
    SubmitField
from wtforms.validators import Required, Length, Email, Regexp
from wtforms import ValidationError
from flask.ext.pagedown.fields import PageDownField
from ..models import Role, User

class NameForm(Form):
    name = StringField('What is your name?', validators=[Required()])
    # name űЧ֣û룬ֵΪ None
    submit = SubmitField('Submit')

class EditProfileForm(Form):
    name = StringField('Real name', validators=[Length(0, 64)])
    location = StringField('Location', validators=[Length(0, 64)])
    about_me = TextAreaField('About me')
    submit = SubmitField('Submit')

class EditProfileAdminForm(Form):
    email = StringField('Email', validators=[Required(), Length(1, 64),
                                             Email()])
    username = StringField('Username', validators=[
        Required(), Length(1, 64), Regexp('^[A-Za-z][A-Za-z0-9_.]*$', 0,
                                          'Usernames must have only letters, '
                                          'numbers, dots or underscores')])
    confirmed = BooleanField('Confirmed')
    role = SelectField('Role', coerce=int)
    name = StringField('Real name', validators=[Length(0, 64)])
    location = StringField('Location', validators=[Length(0, 64)])
    about_me = TextAreaField('About me')
    submit = SubmitField('Submit')

    def __init__(self, user, *args, **kwargs):
        super(EditProfileAdminForm, self).__init__(*args, **kwargs)
        self.role.choices = [(role.id, role.name)
                             for role in Role.query.order_by(Role.name).all()]
        self.user = user

    def validate_email(self, field):
        if field.data != self.user.email and \
                User.query.filter_by(email=field.data).first():
            raise ValidationError('Email already registered.')

    def validate_username(self, field):
        if field.data != self.user.username and \
                User.query.filter_by(username=field.data).first():
            raise ValidationError('Username already in use.')

class PostForm(Form):
    body = PageDownField("What's on your mind?", validators=[Required()])
    submit = SubmitField('Submit')
###
###
###
####
###<|endoftext|>


<commit_before>import os
import tarfile
from flask import Blueprint, render_template, request, jsonify, current_app, g, url_for, redirect
from rq import Queue
from redis import Redis

from .forms import UploadForm
from .. import spfy

from werkzeug.utils import secure_filename

from datetime import datetime

bp = Blueprint('main', __name__)

def fetch_job(job_id):
    '''
    Iterates through all queues looking for the job.
    '''
    print 'received'
    print job_id
    for queue in current_app.config['QUEUES']:
        q = Queue(queue, connection=Redis())
        job = q.fetch_job(job_id)
        if job is not None:
            return job

@bp.route('/results/<job_id>')
def job_status(job_id):
    job = fetch_job(job_id)
    if job.is_finished:
        return jsonify(job.result)
    else:
        return "Nay!", 202

@bp.route('/upload', methods=['POST'])
def upload():
    if request.method == 'POST':
        form = request.form
        print "=== Form Data ==="
        for key, value in form.items():
            print key, "=>", value
        f = request.form
        for key in f.keys():
            for value in f.getlist(key):
                print key,":",value

        file = request.files['file']
        if file:
            # for saving file
            now = datetime.now()
            now = now.strftime("%Y-%m-%d-%H-%M-%S-%f")

            filename = os.path.join(current_app.config[
                                    'UPLOAD_FOLDER'], now + '-' + secure_filename(file.filename))
            file.save(filename)

            if tarfile.is_tarfile(filename):
                tar = tarfile.open(filename)
                extracted_dir = os.path.join(
                    current_app.config['UPLOAD_FOLDER'] + '/' + now)
                os.mkdir(extracted_dir)
                for member in tar.getmembers():
                    if not secure_filename(member.name):
                        return 'invalid upload', 500
                        # TODO: wipe temp data
                tar.extractall(path=extracted_dir)
                for fn in os.listdir(extracted_dir):
                    os.rename(extracted_dir +'/' + fn, extracted_dir +'/'+ now + '-' + fn)
                tar.close()

                # set filename to dir for spfy call
                filename = extracted_dir

            # for enqueing task
            jobs_dict = spfy.spfy(
                {'i': filename, 'disable_serotype': False, 'disable_amr': False, 'disable_vf': False})
            return jsonify(jobs_dict)
    return 500

@bp.route('/', methods=['GET', 'POST'])
def index():
    return render_template("index.html")

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1] in current_app.config['ALLOWED_EXTENSIONS']
<commit_msg>form options being posted and recv ok<commit_after>import os
import tarfile
from flask import Blueprint, render_template, request, jsonify, current_app, g, url_for, redirect
from rq import Queue
from redis import Redis

from .forms import UploadForm
from .. import spfy

from werkzeug.utils import secure_filename

from datetime import datetime

bp = Blueprint('main', __name__)

def fetch_job(job_id):
    '''
    Iterates through all queues looking for the job.
    '''
    print 'received'
    print job_id
    for queue in current_app.config['QUEUES']:
        q = Queue(queue, connection=Redis())
        job = q.fetch_job(job_id)
        if job is not None:
            return job

@bp.route('/results/<job_id>')
def job_status(job_id):
    job = fetch_job(job_id)
    if job.is_finished:
        return jsonify(job.result)
    else:
        return "Nay!", 202

@bp.route('/upload', methods=['POST'])
def upload():
    if request.method == 'POST':
        form = request.form
        print "=== Form Data ==="
        for key, value in form.items():
            print key, "=>", value

        file = request.files['file']
        if file:
            # for saving file
            now = datetime.now()
            now = now.strftime("%Y-%m-%d-%H-%M-%S-%f")

            filename = os.path.join(current_app.config[
                                    'UPLOAD_FOLDER'], now + '-' + secure_filename(file.filename))
            file.save(filename)

            if tarfile.is_tarfile(filename):
                tar = tarfile.open(filename)
                extracted_dir = os.path.join(
                    current_app.config['UPLOAD_FOLDER'] + '/' + now)
                os.mkdir(extracted_dir)
                for member in tar.getmembers():
                    if not secure_filename(member.name):
                        return 'invalid upload', 500
                        # TODO: wipe temp data
                tar.extractall(path=extracted_dir)
                for fn in os.listdir(extracted_dir):
                    os.rename(extracted_dir +'/' + fn, extracted_dir +'/'+ now + '-' + fn)
                tar.close()

                # set filename to dir for spfy call
                filename = extracted_dir

            # for enqueing task
            jobs_dict = spfy.spfy(
                {'i': filename, 'disable_serotype': False, 'disable_amr': False, 'disable_vf': False})
            return jsonify(jobs_dict)
    return 500

@bp.route('/', methods=['GET', 'POST'])
def index():
    return render_template("index.html")

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1] in current_app.config['ALLOWED_EXTENSIONS']
<|endoftext|>