<commit_before>#!/usr/bin/env python
"""
This script parses the SALAMI v2.0 dataset into the JAMS format.

Usage example:
    ./salami_parser.py salami-data-public/ OutputSalamiJAMS
"""

__author__ = "Oriol Nieto"
__copyright__ = "Copyright 2016, Music and Audio Research Lab (MARL)"
__license__ = "MIT"
__version__ = "1.1"
__email__ = "oriol@nyu.edu"

import argparse
import csv
from joblib import Parallel, delayed
import logging
import numpy as np
import os
import time

import jams

# Mapping of the labels to the actual labels allowed by the SALAMI guidelines
labels_map = {
    "a'/a''": "no_function",
    "b/c'": "no_function",
    "bagpipes": "instrumental",
    "banjo": "instrumental",
    "code": "coda",
    "dialog": "spoken",
    "female": "voice_female",
    "first": "main theme",
    "guitar": "instrumental",
    "hammond": "instrumental",
    "main_theme": "main theme",
    "male": "voice_male",
    "muted": "fade-out",
    "organ": "instrumental",
    "out": "outro",
    "piano": "instrumental",
    "recap": "recapitulation",
    "secondary_theme": "secondary theme",
    "spoken_voice": "spoken",
    "stage_speaking": "spoken",
    "steel": "instrumental",
    "tag": "head",
    "trumpet": "instrumental",
    "third": "third theme",
    "violin": "instrumental",
    "vocalizations": "voice",
    "vocals": "voice",
    "w/dialog": "spoken"
}

def fix_label(label):
    """Fixes the given label to comply the SALAMI guidelines and JAMS
    namespace."""
    label = label.lower()
    return labels_map.get(label, label)

def parse_annotation(jam, path, annotation_id, level, metadata):
    """Parses one annotation for the given level

    Parameters
    ----------
    jam: object
        The top-level JAMS object.
    path: str
        path to the track in the SALAMI dataset.
    annotation_id: int
        Whether to use the first or the second annotation.
    level: str
        Level of annotation.
    metadata: list
        List containing the information of the CSV file for the current track.
    """
    level_dict = {
        "function": "_functions",
        "large_scale": "_uppercase",
        "small_scale": "_lowercase"
    }
    namespace_dict = {
        "function": "segment_salami_function",
        "large_scale": "segment_salami_upper",
        "small_scale": "segment_salami_lower"
    }

    # File to open
    file_path = os.path.join(path, "parsed",
                             "textfile" + str(annotation_id) +
                             level_dict[level] + ".txt")

    # Open file
    try:
        f = open(file_path, "r")
    except IOError:
        logging.warning("Annotation missing in %s", file_path)
        return

    # Annotation Metadata
    curator = jams.Curator(name="Jordan Smith",
                           email="jblsmith@gmail.com")
    annotator = {
        "name": metadata[annotation_id + 1],
        "submission_date": metadata[annotation_id + 15]
    }
    ann_meta = jams.AnnotationMetadata(curator=curator,
                                       version="2.0",
                                       corpus="SALAMI",
                                       annotator=annotator,
                                       data_source=metadata[1],
                                       annotation_tools="Sonic Visualizer")

    # Create Annotation
    annot = jams.Annotation(namespace=namespace_dict[level],
                            annotation_metadata=ann_meta)

    # Actual Data
    lines = f.readlines()
    for i, line in enumerate(lines[:-1]):
        start_time, label = line.strip("\n").split("\t")
        end_time = lines[i + 1].split("\t")[0]
        start_time = float(start_time)
        end_time = float(end_time)
        dur = end_time - start_time
        if start_time - end_time == 0:
            continue

        if level == "function":
            label = fix_label(label)

        annot.data.add_observation(time=start_time, duration=dur, value=label)
    f.close()

    # Add annotation to the jams
    jam.annotations.append(annot)

def fill_global_metadata(jam, metadata, dur):
    """Fills the global metada into the JAMS jam."""
    meta = jams.FileMetadata(title=metadata[7],
                             artist=metadata[8],
                             duration=dur,
                             jams_version=jams.version.version)
    jam.file_metadata = meta

def create_annotations(jam, path, annotation_id, metadata):
    """Fills the JAMS annot annotation given a path to the original
    SALAMI annotations. The variable "annotator" let's you choose which
    SALAMI annotation to use.

    Parameters
    ----------
    jam: object
        The top-level JAMS object.
    path: str
        Patht to the file containing the references.
    annotation_id: int
        1 or 2 depending on which annotation to use
    metadata: list
        List containing the information of the CSV file for the current track.
    """
    # Parse all level annotations
    levels = ["function", "large_scale", "small_scale"]
    [parse_annotation(jam, path, annotation_id, level, metadata)
        for level in levels]

def get_duration(jam):
    """Obtains the duration from the jams annotations.

    Parameters
    ----------
    jam : object
        JAMS object with at least one annotation.

    Returns
    -------
    dur : float
        Duration of the file, by taking the max last boundary of all
        annotations (in seconds).
    """
    durs = []
    for annotation in jam.annotations:
        durs.append(annotation.data["time"].iloc[-1].total_seconds() +
                    annotation.data["duration"].iloc[-1].total_seconds())
    return np.max(durs)

def create_JAMS(in_dir, metadata, out_file):
    """Creates a JAMS file given the path to a SALAMI track.

    Parameters
    ----------
    in_dir : str
        Path to the input directory
    metadata : str
        Metadata read from the CSV file
    out_file : str
        Output JAMS file
    """
    path = os.path.join(in_dir, "annotations", metadata[0], )

    # Sanity check
    if not os.path.exists(path):
        logging.warning("Path not found %s", path)
        return

    # New JAMS and annotation
    jam = jams.JAMS()

    # Create Annotations if they exist
    # Maximum 3 annotations per file
    for annotation_id in range(1, 4):
        if os.path.isfile(os.path.join(path,
                                       "textfile" + str(annotation_id) +
                                       ".txt")):
            create_annotations(jam, path, annotation_id, metadata)

    # Get the duration from the annotations
    dur = get_duration(jam)

    # Global file metadata
    fill_global_metadata(jam, metadata, dur)

    # Save JAMS
    jam.save(out_file)

def process_one(metadata, in_dir, out_dir):
    """Processes one track given its metadata."""
    if metadata[0] == "SONG_ID":
        return

    # Create a JAMS file for this track
    logging.info("Parsing file %s..." % metadata[0])
    create_JAMS(in_dir, metadata,
                os.path.join(out_dir,
                    os.path.basename(metadata[0]) + ".jams"))

def process(in_dir, out_dir, n_jobs):
    """Converts the original SALAMI files into the JAMS format, and saves
    them in the out_dir folder."""

    # Check if output folder and create it if needed:
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)

    # Open CSV with metadata and parse
    with open(os.path.join(in_dir, "metadata", "metadata.csv")) as fh:
        csv_reader = csv.reader(fh)
        Parallel(n_jobs=n_jobs)(delayed(process_one)(
            metadata, in_dir, out_dir) for metadata in list(csv_reader))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description=
        "Converts the SALAMI dataset to the JAMS format",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("in_dir",
                        action="store",
                        help="SALAMI main folder")
    parser.add_argument("out_dir",
                        action="store",
                        help="Output JAMS folder")
    parser.add_argument("-j",
                        dest="n_jobs",
                        action="store",
                        type=int,
                        default=2,
                        help="Number of CPUs to run in parallel.")
    args = parser.parse_args()
    start_time = time.time()

    # Setup the logger
    logging.basicConfig(format='%(asctime)s: %(message)s', level=logging.INFO)

    # Run the parser
    process(args.in_dir, args.out_dir, args.n_jobs)

    # Done!
    logging.info("Done! Took %.2f seconds." % (time.time() - start_time))
<commit_msg>Some PEP8 stuff, updated version number<commit_after>#!/usr/bin/env python
"""
This script parses the SALAMI v2.0 dataset into the JAMS format.

Usage example:
    ./salami_parser.py salami-data-public/ OutputSalamiJAMS
"""

import argparse
import csv
from joblib import Parallel, delayed
import logging
import numpy as np
import os
import time

import jams

__author__ = "Oriol Nieto"
__copyright__ = "Copyright 2016, Music and Audio Research Lab (MARL)"
__license__ = "MIT"
__version__ = "1.2"
__email__ = "oriol@nyu.edu"

# Mapping of the labels to the actual labels allowed by the SALAMI guidelines
labels_map = {
    "a'/a''": "no_function",
    "b/c'": "no_function",
    "bagpipes": "instrumental",
    "banjo": "instrumental",
    "code": "coda",
    "dialog": "spoken",
    "female": "voice_female",
    "first": "main theme",
    "guitar": "instrumental",
    "hammond": "instrumental",
    "main_theme": "main theme",
    "male": "voice_male",
    "muted": "fade-out",
    "organ": "instrumental",
    "out": "outro",
    "piano": "instrumental",
    "recap": "recapitulation",
    "secondary_theme": "secondary theme",
    "spoken_voice": "spoken",
    "stage_speaking": "spoken",
    "steel": "instrumental",
    "tag": "head",
    "trumpet": "instrumental",
    "third": "third theme",
    "violin": "instrumental",
    "vocalizations": "voice",
    "vocals": "voice",
    "w/dialog": "spoken"
}

def fix_label(label):
    """Fixes the given label to comply the SALAMI guidelines and JAMS
    namespace."""
    label = label.lower()
    return labels_map.get(label, label)

def parse_annotation(jam, path, annotation_id, level, metadata):
    """Parses one annotation for the given level

    Parameters
    ----------
    jam: object
        The top-level JAMS object.
    path: str
        path to the track in the SALAMI dataset.
    annotation_id: int
        Whether to use the first or the second annotation.
    level: str
        Level of annotation.
    metadata: list
        List containing the information of the CSV file for the current track.
    """
    level_dict = {
        "function": "_functions",
        "large_scale": "_uppercase",
        "small_scale": "_lowercase"
    }
    namespace_dict = {
        "function": "segment_salami_function",
        "large_scale": "segment_salami_upper",
        "small_scale": "segment_salami_lower"
    }

    # File to open
    file_path = os.path.join(path, "parsed",
                             "textfile" + str(annotation_id) +
                             level_dict[level] + ".txt")

    # Open file
    try:
        f = open(file_path, "r")
    except IOError:
        logging.warning("Annotation missing in %s", file_path)
        return

    # Annotation Metadata
    curator = jams.Curator(name="Jordan Smith",
                           email="jblsmith@gmail.com")
    annotator = {
        "name": metadata[annotation_id + 1],
        "submission_date": metadata[annotation_id + 15]
    }
    ann_meta = jams.AnnotationMetadata(curator=curator,
                                       version="2.0",
                                       corpus="SALAMI",
                                       annotator=annotator,
                                       data_source=metadata[1],
                                       annotation_tools="Sonic Visualizer")

    # Create Annotation
    annot = jams.Annotation(namespace=namespace_dict[level],
                            annotation_metadata=ann_meta)

    # Actual Data
    lines = f.readlines()
    for i, line in enumerate(lines[:-1]):
        start_time, label = line.strip("\n").split("\t")
        end_time = lines[i + 1].split("\t")[0]
        start_time = float(start_time)
        end_time = float(end_time)
        dur = end_time - start_time
        if start_time - end_time == 0:
            continue

        if level == "function":
            label = fix_label(label)

        annot.data.add_observation(time=start_time, duration=dur, value=label)
    f.close()

    # Add annotation to the jams
    jam.annotations.append(annot)

def fill_global_metadata(jam, metadata, dur):
    """Fills the global metada into the JAMS jam."""
    meta = jams.FileMetadata(title=metadata[7],
                             artist=metadata[8],
                             duration=dur,
                             jams_version=jams.version.version)
    jam.file_metadata = meta

def create_annotations(jam, path, annotation_id, metadata):
    """Fills the JAMS annot annotation given a path to the original
    SALAMI annotations. The variable "annotator" let's you choose which
    SALAMI annotation to use.

    Parameters
    ----------
    jam: object
        The top-level JAMS object.
    path: str
        Patht to the file containing the references.
    annotation_id: int
        1 or 2 depending on which annotation to use
    metadata: list
        List containing the information of the CSV file for the current track.
    """
    # Parse all level annotations
    levels = ["function", "large_scale", "small_scale"]
    [parse_annotation(jam, path, annotation_id, level, metadata)
        for level in levels]

def get_duration(jam):
    """Obtains the duration from the jams annotations.

    Parameters
    ----------
    jam : object
        JAMS object with at least one annotation.

    Returns
    -------
    dur : float
        Duration of the file, by taking the max last boundary of all
        annotations (in seconds).
    """
    durs = []
    for annotation in jam.annotations:
        durs.append(annotation.data["time"].iloc[-1].total_seconds() +
                    annotation.data["duration"].iloc[-1].total_seconds())
    return np.max(durs)

def create_JAMS(in_dir, metadata, out_file):
    """Creates a JAMS file given the path to a SALAMI track.

    Parameters
    ----------
    in_dir : str
        Path to the input directory
    metadata : str
        Metadata read from the CSV file
    out_file : str
        Output JAMS file
    """
    path = os.path.join(in_dir, "annotations", metadata[0], )

    # Sanity check
    if not os.path.exists(path):
        logging.warning("Path not found %s", path)
        return

    # New JAMS and annotation
    jam = jams.JAMS()

    # Create Annotations if they exist
    # Maximum 3 annotations per file
    for annotation_id in range(1, 4):
        if os.path.isfile(os.path.join(path,
                                       "textfile" + str(annotation_id) +
                                       ".txt")):
            create_annotations(jam, path, annotation_id, metadata)

    # Get the duration from the annotations
    dur = get_duration(jam)

    # Global file metadata
    fill_global_metadata(jam, metadata, dur)

    # Save JAMS
    jam.save(out_file)

def process_one(metadata, in_dir, out_dir):
    """Processes one track given its metadata."""
    if metadata[0] == "SONG_ID":
        return

    # Create a JAMS file for this track
    logging.info("Parsing file %s..." % metadata[0])
    create_JAMS(in_dir, metadata,
                os.path.join(out_dir,
                             os.path.basename(metadata[0]) + ".jams"))

def process(in_dir, out_dir, n_jobs):
    """Converts the original SALAMI files into the JAMS format, and saves
    them in the out_dir folder."""

    # Check if output folder and create it if needed:
    if not os.path.exists(out_dir):
        os.makedirs(out_dir)

    # Open CSV with metadata and parse
    with open(os.path.join(in_dir, "metadata", "metadata.csv")) as fh:
        csv_reader = csv.reader(fh)
        Parallel(n_jobs=n_jobs)(delayed(process_one)(
            metadata, in_dir, out_dir) for metadata in list(csv_reader))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Converts the SALAMI dataset to the JAMS format",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("in_dir",
                        action="store",
                        help="SALAMI main folder")
    parser.add_argument("out_dir",
                        action="store",
                        help="Output JAMS folder")
    parser.add_argument("-j",
                        dest="n_jobs",
                        action="store",
                        type=int,
                        default=2,
                        help="Number of CPUs to run in parallel.")
    args = parser.parse_args()
    start_time = time.time()

    # Setup the logger
    logging.basicConfig(format='%(asctime)s: %(message)s', level=logging.INFO)

    # Run the parser
    process(args.in_dir, args.out_dir, args.n_jobs)

    # Done!
    logging.info("Done! Took %.2f seconds." % (time.time() - start_time))
<|endoftext|>


<commit_before>def minimum_number_of_coins(amount, wallet):
    solutions = [(0, wallet)]
    for intermediate_amount in range(1, amount + 1):
        solutions.append(find_best_solution(intermediate_amount, wallet, solutions))
    return solutions[-1][0]

def find_best_solution(amount, coin_values, solutions):
    return min(
        find_solutions(amount, coin_values, solutions),
        key=lambda solution: solution[0],
        default=("IMPOSSIBLE", {}))

def find_solutions(amount, coin_values, solutions):
    for coin_value in coin_values:
        if coin_value <= amount:
            coins_given, wallet = solutions[amount - coin_value]
            if wallet.get(coin_value, 0) > 0:
                yield coins_given + 1, remove_coin(wallet, coin_value)

def remove_coin(wallet, coin_value):
    return update(wallet, {coin_value: wallet[coin_value] - 1})

def update(a_dict, another_dict):
    new_dict = a_dict.copy()
    new_dict.update(another_dict)
    return new_dict

if __name__ == "__main__":
    amount = int(input())
    ncoin_types = int(input())
    wallet = {coin_value: coins_left for coins_left, coin_value in (map(int, input().split()) for _ in range(ncoin_types))}
    print(minimum_number_of_coins(amount, wallet))
<commit_msg>style: passes pylint<commit_after># pylint: disable=locally-disabled,missing-docstring

def minimum_number_of_coins(amount, wallet):
    solutions = [(0, wallet)]
    for intermediate_amount in range(1, amount + 1):
        solutions.append(find_best_solution(intermediate_amount, wallet, solutions))
    return solutions[-1][0]

def find_best_solution(amount, coin_values, solutions):
    return min(
        find_solutions(amount, coin_values, solutions),
        key=lambda solution: solution[0],
        default=("IMPOSSIBLE", {}))

def find_solutions(amount, coin_values, solutions):
    for coin_value in coin_values:
        if coin_value <= amount:
            coins_given, wallet = solutions[amount - coin_value]
            if wallet.get(coin_value, 0) > 0:
                yield coins_given + 1, remove_coin(wallet, coin_value)

def remove_coin(wallet, coin_value):
    return update(wallet, {coin_value: wallet[coin_value] - 1})

def update(a_dict, another_dict):
    new_dict = a_dict.copy()
    new_dict.update(another_dict)
    return new_dict

def main():
    amount = int(input())
    ncoin_types = int(input())
    wallet = {
        int(coin_value): int(coins_left) for coins_left, coin_value in
        (input().split() for _ in range(ncoin_types))}
    print(minimum_number_of_coins(amount, wallet))

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


<commit_before>from uuid import uuid4
import json
import logging
import marshal
import pickle

from bndl.util import serialize, threads
from bndl.util.conf import Float
from bndl.util.funcs import identity
from cytoolz.functoolz import compose

min_block_size = Float(4)  # MB
max_block_size = Float(16)  # MB

logger = logging.getLogger(__name__)

MISSING = 'bndl.compute.broadcast.MISSING'

# TODO implement alternative encodings
# as with caching

# TODO move such 'caches' outside of modules
# much better to keep them local to Nodes
# if deleted, GC will kick in

download_coordinator = threads.Coordinator()

def broadcast(ctx, value, serialization='auto', deserialization=None):
    '''
    Broadcast data to workers
    :param value: The value to broadcast
    :param serialization: auto, pickle, marshal, json, binary or text
        The format to serialize the broadcast value into
    :param deserialization: None or function(fileobj)
    '''
    if serialization is not None:
        if deserialization is not None:
            raise ValueError("Can't specify both serialization and deserialization")
        elif serialization == 'auto':
            marshalled, data = serialize.dumps(value)
            deserialization = marshal.loads if marshalled else pickle.loads
        elif serialization == 'pickle':
            data = pickle.dumps(value)
            deserialization = pickle.loads
        elif serialization == 'marshal':
            data = marshal.dumps(value)
            deserialization = marshal.loads
        elif serialization == 'json':
            data = json.dumps(value).encode()
            deserialization = compose(json.loads, bytes.decode)
        elif serialization == 'binary':
            data = value
            deserialization = identity
        elif serialization == 'text':
            data = value.encode()
            deserialization = bytes.decode
        else:
            raise ValueError('Unsupported serialization %s' % serialization)
    elif not deserialization:
        raise ValueError('Must specify either serialization or deserialization')
    else:
        data = value

    key = str(uuid4())
    min_block_size = int(ctx.conf.get('bndl.compute.broadcast.min_block_size') * 1024 * 1024)
    max_block_size = int(ctx.conf.get('bndl.compute.broadcast.max_block_size') * 1024 * 1024)
    if min_block_size == max_block_size:
        block_spec = ctx.node.serve_data(key, data, max_block_size)
    else:
        block_spec = ctx.node.serve_data(key, data, (ctx.worker_count * 2, min_block_size, max_block_size))
    return BroadcastValue(ctx, ctx.node.name, block_spec, deserialization)

def broadcast_pickled(ctx, pickled_value):
    '''
    Broadcast data which is already pickled.
    '''
    return broadcast(ctx, pickled_value, serialization=None, deserialization=pickle.loads)

class BroadcastValue(object):
    def __init__(self, ctx, seeder, block_spec, deserialize):
        self.ctx = ctx
        self.seeder = seeder
        self.block_spec = block_spec
        self.deserialize = deserialize

    @property
    def value(self):
        return download_coordinator.coordinate(self._get, self.block_spec.name)

    def _get(self):
        node = self.ctx.node
        blocks = node.get_blocks(self.block_spec, node.peers.filter(node_type='worker'))

        val = self.deserialize(b''.join(blocks))

        if node.name != self.block_spec.seeder:
            node.remove_blocks(self.block_spec.name, from_peers=False)

        return val

    def unpersist(self):
        node = self.ctx.node
        name = self.block_spec.name
        assert node.name == self.block_spec.seeder
        node.unpersist_broadcast_values(node, name)
        for peer in node.peers.filter():
            peer.unpersist_broadcast_values(name)
            # response isn't waited on

    def __del__(self):
        if self.ctx.node and self.ctx.node.name == self.block_spec.seeder:
            self.unpersist()
<commit_msg>allow to block on unpersist of broadcast<commit_after>from uuid import uuid4
import concurrent.futures
import json
import logging
import marshal
import pickle

from bndl.util import serialize, threads
from bndl.util.conf import Float
from bndl.util.funcs import identity
from cytoolz.functoolz import compose

min_block_size = Float(4)  # MB
max_block_size = Float(16)  # MB

logger = logging.getLogger(__name__)

MISSING = 'bndl.compute.broadcast.MISSING'

# TODO implement alternative encodings
# as with caching

# TODO move such 'caches' outside of modules
# much better to keep them local to Nodes
# if deleted, GC will kick in

download_coordinator = threads.Coordinator()

def broadcast(ctx, value, serialization='auto', deserialization=None):
    '''
    Broadcast data to workers
    :param value: The value to broadcast
    :param serialization: auto, pickle, marshal, json, binary or text
        The format to serialize the broadcast value into
    :param deserialization: None or function(fileobj)
    '''
    if serialization is not None:
        if deserialization is not None:
            raise ValueError("Can't specify both serialization and deserialization")
        elif serialization == 'auto':
            marshalled, data = serialize.dumps(value)
            deserialization = marshal.loads if marshalled else pickle.loads
        elif serialization == 'pickle':
            data = pickle.dumps(value)
            deserialization = pickle.loads
        elif serialization == 'marshal':
            data = marshal.dumps(value)
            deserialization = marshal.loads
        elif serialization == 'json':
            data = json.dumps(value).encode()
            deserialization = compose(json.loads, bytes.decode)
        elif serialization == 'binary':
            data = value
            deserialization = identity
        elif serialization == 'text':
            data = value.encode()
            deserialization = bytes.decode
        else:
            raise ValueError('Unsupported serialization %s' % serialization)
    elif not deserialization:
        raise ValueError('Must specify either serialization or deserialization')
    else:
        data = value

    key = str(uuid4())
    min_block_size = int(ctx.conf.get('bndl.compute.broadcast.min_block_size') * 1024 * 1024)
    max_block_size = int(ctx.conf.get('bndl.compute.broadcast.max_block_size') * 1024 * 1024)
    if min_block_size == max_block_size:
        block_spec = ctx.node.serve_data(key, data, max_block_size)
    else:
        block_spec = ctx.node.serve_data(key, data, (ctx.worker_count * 2, min_block_size, max_block_size))
    return BroadcastValue(ctx, ctx.node.name, block_spec, deserialization)

def broadcast_pickled(ctx, pickled_value):
    '''
    Broadcast data which is already pickled.
    '''
    return broadcast(ctx, pickled_value, serialization=None, deserialization=pickle.loads)

class BroadcastValue(object):
    def __init__(self, ctx, seeder, block_spec, deserialize):
        self.ctx = ctx
        self.seeder = seeder
        self.block_spec = block_spec
        self.deserialize = deserialize

    @property
    def value(self):
        return download_coordinator.coordinate(self._get, self.block_spec.name)

    def _get(self):
        node = self.ctx.node
        blocks = node.get_blocks(self.block_spec, node.peers.filter(node_type='worker'))

        val = self.deserialize(b''.join(blocks))

        if node.name != self.block_spec.seeder:
            node.remove_blocks(self.block_spec.name, from_peers=False)

        return val

    def unpersist(self, block=False, timeout=None):
        node = self.ctx.node
        name = self.block_spec.name
        assert node.name == self.block_spec.seeder
        node.unpersist_broadcast_values(node, name)
        requests = [peer.unpersist_broadcast_values
                   for peer in node.peers.filter()]
        if timeout:
            requests = [request.with_timeout(timeout) for request in requests]
        requests = [request(name) for request in requests]
        if block:
            for request in requests:
                try:
                    request.result()
                except concurrent.futures.TimeoutError:
                    pass
                except Exception:
                    logger.warning('error while unpersisting %s', name, exc_info=True)

    def __del__(self):
        if self.ctx.node and self.ctx.node.name == self.block_spec.seeder:
            self.unpersist()
<|endoftext|>


<commit_before>"""
Plugin that takes pcap files and outputs stats

Created on 1 November 2019
@author: Charlie Lewis
"""

import datetime
import json
import os
import shlex
import subprocess
import sys

import pika

def connect_rabbit(host='messenger', port=5672, queue='task_queue'):
    params = pika.ConnectionParameters(host=host, port=port)
    connection = pika.BlockingConnection(params)
    channel = connection.channel()
    channel.queue_declare(queue=queue, durable=True)
    return channel

def send_rabbit_msg(msg, channel, exchange='', routing_key='task_queue'):
    channel.basic_publish(exchange=exchange,
                          routing_key=routing_key,
                          body=json.dumps(msg),
                          properties=pika.BasicProperties(
                          delivery_mode=2,
                         ))
    print(" [X] %s UTC %r %r" % (str(datetime.datetime.utcnow()),
                                 str(msg['id']), str(msg['file_path'])))
    return

def get_version():
    version = ''
    with open('VERSION', 'r') as f:
        for line in f:
            version = line.strip()
    return version

def get_path():
    path = None
    try:
        path = sys.argv[1]
    except Exception as e:
        print("No path provided: {0}, quitting".format(str(e)))
    return path

def parse_capinfos(output):
    results = {'capinfos':{}}
    num_interfaces = 0
    interface_dict = {}
    interface = 0
    interface_name = ''
    for line in output.split('\n'):
        if line == '':
            continue
        if line.startswith('Number of interfaces in file:'):
            num_interfaces = int(line.split(':', 1)[1].strip())
            continue
        if interface < num_interfaces:
            if line.startswith('Interface '):
                interface_name = line.split()[1]
                interface_dict[interface_name] = {}
                continue
            else:
                if line.startswith('Number of packets '):
                    interface += 1
                name, value = line.split(' = ')
                interface_dict[interface_name][name.strip()] = value.strip()
                continue
        name, value = line.split(':', 1)
        results['capinfos'][name.strip()] = value.strip()
    results['capinfos']['interfaces'] = interface_dict
    print(results)
    return results

def run_capinfos(path):
    if os.path.getsize(path) == 0:
       print("pcap file empty, no stats")
       return

    output = ''
    try:
        output = subprocess.check_output(shlex.split(' '.join(['capinfos', path])))
        output = output.decode("utf-8")
    except Exception as e:
        print(str(e))

    results = parse_capinfos(output)
    return results

def parse_tshark(output):
    results = {'tshark':{}}
    in_block = False
    name = None
    for line in output.split('\n'):
        if line.startswith('==='):
            if in_block:
                in_block = False
                name = None
                continue
            else:
                in_block = True
                continue
        if in_block:
            if not name:
                name = ''.join(line.split(':')).strip()
                results['tshark'][name] = ''
                continue
            elif not line.startswith('Filter:') and line != '':
                results['tshark'][name] += line + '\n'

    for result in results['tshark'].keys():
        if 'Conversations' in result:
            # handle conversation parsing
            conversations = []
            for line in results['tshark'][result].split('\n'):
                if line == '' or line.startswith(' '):
                    # header or padding, dicard
                    continue
                else:
                    src, _, dst, frames_l, bytes_l, frames_r, bytes_r, frames_total, bytes_total, rel_start, duration = line.split()
                    conversations.append({'Source': src, 'Destination': dst, 'Frames to Source': frames_l, 'Bytes to Source': bytes_l, 'Frames to Destination': frames_r, 'Bytes to Destination': bytes_r, 'Total Frames': frames_total, 'Total Bytes': bytes_total, 'Relative Start': rel_start, 'Duration': duration})
            results['tshark'][result] = conversations
        elif 'Endpoints' in result:
            # handle endpoint parsing
            endpoints = []
            for line in results['tshark'][result].split('\n'):
                if line == '' or line.startswith(' '):
                    # header or padding, dicard
                    continue
                else:
                    # handle endpoint services with ports
                    if result.startswith('UDP') or result.startswith('TCP') or result.startswith('STCP'):
                        endpoint, port, packet_count, byte_count, tx_packets, tx_bytes, rx_packets, rx_bytes = line.split()
                        endpoints.append({'Endpoint': endpoint, 'Port': port, 'Packets': packet_count, 'Bytes': byte_count, 'Tx Packets': tx_packets, 'Tx Bytes': tx_bytes, 'Rx Packets': rx_packets, 'Rx Bytes': rx_bytes})
                    else:
                        endpoint, packet_count, byte_count, tx_packets, tx_bytes, rx_packets, rx_bytes = line.split()
                        endpoints.append({'Endpoint': endpoint, 'Packets': packet_count, 'Bytes': byte_count, 'Tx Packets': tx_packets, 'Tx Bytes': tx_bytes, 'Rx Packets': rx_packets, 'Rx Bytes': rx_bytes})
            results['tshark'][result] = endpoints
        else:
            # handle weird stuff
            for line in results['tshark'][result].split('\n'):
                if line == '' or line.startswith(' '):
                    # header or padding, dicard
                    continue
                else:
                    # handle icmp and icmpv6
                    if result.startswith('ICMP'):
                        if isinstance(results['tshark'][result], str):
                            results['tshark'][result] = {}
                        if line.startswith('Requests') or line.startswith('Minimum'):
                            # header
                            continue
                        else:
                            values = line.split()
                            if len(values) == 4:
                                requests, replies, lost, percent_loss = values
                                results['tshark'][result]['Requests'] = requests
                                results['tshark'][result]['Replies'] = replies
                                results['tshark'][result]['Lost'] = lost
                                results['tshark'][result]['% Loss'] = percent_loss
                            else:
                                minimum, maximum, mean, median, s_deviation, min_frame, max_frame = values
                                results['tshark'][result]['Minimum'] = minimum
                                results['tshark'][result]['Maximum'] = maximum
                                results['tshark'][result]['Mean'] = mean
                                results['tshark'][result]['Median'] = median
                                results['tshark'][result]['Standard Deviation'] = s_deviation
                                results['tshark'][result]['Minimum Frame'] = min_frame
                                results['tshark'][result]['Maximum Frame'] = max_frame
                    elif result.startswith('Protocol'):
                        # TODO
                        continue
                    # handle dns
                    elif result.startswith('DNS'):
                        # TODO
                        continue

    # TODO temporarily remove until parsed
    if 'DNS' in results['tshark']:
        del results['tshark']['DNS']

    # handle protocol hierarchy stats
    a = []
    if 'Protocol Hierarchy Statistics' in results['tshark']:
        a = results['tshark']['Protocol Hierarchy Statistics'].split('\n')
    h = []
    for line in a:
        if line != '':
            name, frame_count, byte_count = line.rsplit(' ', 2)
            name = name.rstrip()
            frame_count = frame_count.split(':')[1]
            byte_count = byte_count.split(':')[1]
            h.append([name, frame_count, byte_count])

    i = 1
    spaces = 0
    if h:
        h[0][0] = '"' + h[0][0].strip()
        while i < len(h):
            prev_spaces = spaces
            spaces = h[i][0].count('  ')
            h[i-1][0] = h[i-1][0].strip() + '":{"Frames": "' + h[i-1][1] + '", "Bytes": "' + h[i-1][2]
            if spaces > prev_spaces:
                h[i-1][0] += '","'
            elif spaces == prev_spaces:
                h[i-1][0] += '"},"'
            else:
                h[i-1][0] += '"}' + ('}'*(prev_spaces-spaces)) + ',"'
            i += 1
        h[i-1][0] = h[i-1][0].strip() + '":{"Frames": "' + h[i-1][1] + '", "Bytes": "' + h[i-1][2] + '"}' + ('}'*(prev_spaces-spaces))

    protocol_str = '{'
    for record in h:
        protocol_str += record[0]
    protocol_str += '}'
    results['tshark']['Protocol Hierarchy Statistics'] = json.loads(protocol_str)

    # add in condensed conversation fields
    results['tshark']['Condensed TCP Conversations'] = condense_conversations(results, 'TCP Conversations')
    results['tshark']['Condensed UDP Conversations'] = condense_conversations(results, 'UDP Conversations')

    print(results)
    return results

def condense_conversations(results, conv_type):
    prot_ip_map = {}
    if conv_type in results['tshark']:
        for conversation in results['tshark'][conv_type]:
            src_ip, src_port = conversation['Source'].rsplit(':', 1)
            dst_ip, dst_port = conversation['Destination'].rsplit(':', 1)
            if not src_ip in prot_ip_map:
                prot_ip_map[src_ip] = {'Destinations': [], 'Source Ports': [], 'Destination Ports': []}
            if not src_port in prot_ip_map[src_ip]['Source Ports']:
                prot_ip_map[src_ip]['Source Ports'].append(src_port)
            if not dst_port in prot_ip_map[src_ip]['Destination Ports']:
                prot_ip_map[src_ip]['Destination Ports'].append(dst_port)
            if not dst_ip in prot_ip_map[src_ip]['Destinations']:
                prot_ip_map[src_ip]['Destinations'].append(dst_ip)
    return prot_ip_map

def run_tshark(path):
    if os.path.getsize(path) == 0:
       print("pcap file empty, no stats")
       return

    results = {}
    output = ''
    try:
        conv_endpoint_types = ['bluetooth', 'eth', 'fc', 'fddi', 'ip', 'ipv6', 'ipx', 'jxta', 'ncp', 'rsvp', 'sctp', 'tcp', 'tr', 'usb', 'udp', 'wlan']
        options = '-n -q -z dns,tree -z io,phs -z icmp,srt -z icmpv6,srt'
        options += ' -z conv,'.join(conv_endpoint_types)
        options += ' -z endpoints,'.join(conv_endpoint_types)
        output = subprocess.check_output(shlex.split(' '.join(['tshark', '-r', path, options])))
        output = output.decode("utf-8")
    except Exception as e:
        print(str(e))

    results = parse_tshark(output)
    return results

if __name__ == '__main__':  # pragma: no cover
    path = get_path()
    uid = ''
    if 'id' in os.environ:
        uid = os.environ['id']
    if path:
        if 'rabbit' in os.environ and os.environ['rabbit'] == 'true':
            try:
                channel = connect_rabbit()
                capinfos_results = run_capinfos(path)
                body = {'id': uid, 'type': 'metadata', 'file_path': path, 'data': capinfos_results, 'results': {'tool': 'pcap_stats', 'version': get_version()}}
                send_rabbit_msg(body, channel)
                tshark_results = run_tshark(path)
                body = {'id': uid, 'type': 'metadata', 'file_path': path, 'data': tshark_results, 'results': {'tool': 'pcap_stats', 'version': get_version()}}
                send_rabbit_msg(body, channel)
                body = {'id': uid, 'type': 'metadata', 'file_path': path, 'data': '', 'results': {'tool': 'pcap_stats', 'version': get_version()}}
                send_rabbit_msg(body, channel)
            except Exception as e:
                print(str(e))
        else:
            capinfos_results = run_capinfos(path)
            tshark_results = run_tshark(path)
<commit_msg>keep name consistent<commit_after>"""
Plugin that takes pcap files and outputs stats

Created on 1 November 2019
@author: Charlie Lewis
"""

import datetime
import json
import os
import shlex
import subprocess
import sys

import pika

def connect_rabbit(host='messenger', port=5672, queue='task_queue'):
    params = pika.ConnectionParameters(host=host, port=port)
    connection = pika.BlockingConnection(params)
    channel = connection.channel()
    channel.queue_declare(queue=queue, durable=True)
    return channel

def send_rabbit_msg(msg, channel, exchange='', routing_key='task_queue'):
    channel.basic_publish(exchange=exchange,
                          routing_key=routing_key,
                          body=json.dumps(msg),
                          properties=pika.BasicProperties(
                          delivery_mode=2,
                         ))
    print(" [X] %s UTC %r %r" % (str(datetime.datetime.utcnow()),
                                 str(msg['id']), str(msg['file_path'])))
    return

def get_version():
    version = ''
    with open('VERSION', 'r') as f:
        for line in f:
            version = line.strip()
    return version

def get_path():
    path = None
    try:
        path = sys.argv[1]
    except Exception as e:
        print("No path provided: {0}, quitting".format(str(e)))
    return path

def parse_capinfos(output):
    results = {'capinfos':{}}
    num_interfaces = 0
    interface_dict = {}
    interface = 0
    interface_name = ''
    for line in output.split('\n'):
        if line == '':
            continue
        if line.startswith('Number of interfaces in file:'):
            num_interfaces = int(line.split(':', 1)[1].strip())
            continue
        if interface < num_interfaces:
            if line.startswith('Interface '):
                interface_name = line.split()[1]
                interface_dict[interface_name] = {}
                continue
            else:
                if line.startswith('Number of packets '):
                    interface += 1
                name, value = line.split(' = ')
                interface_dict[interface_name][name.strip()] = value.strip()
                continue
        name, value = line.split(':', 1)
        results['capinfos'][name.strip()] = value.strip()
    results['capinfos']['interfaces'] = interface_dict
    print(results)
    return results

def run_capinfos(path):
    if os.path.getsize(path) == 0:
       print("pcap file empty, no stats")
       return

    output = ''
    try:
        output = subprocess.check_output(shlex.split(' '.join(['capinfos', path])))
        output = output.decode("utf-8")
    except Exception as e:
        print(str(e))

    results = parse_capinfos(output)
    return results

def parse_tshark(output):
    results = {'tshark':{}}
    in_block = False
    name = None
    for line in output.split('\n'):
        if line.startswith('==='):
            if in_block:
                in_block = False
                name = None
                continue
            else:
                in_block = True
                continue
        if in_block:
            if not name:
                name = ''.join(line.split(':')).strip()
                results['tshark'][name] = ''
                continue
            elif not line.startswith('Filter:') and line != '':
                results['tshark'][name] += line + '\n'

    for result in results['tshark'].keys():
        if 'Conversations' in result:
            # handle conversation parsing
            conversations = []
            for line in results['tshark'][result].split('\n'):
                if line == '' or line.startswith(' '):
                    # header or padding, dicard
                    continue
                else:
                    src, _, dst, frames_l, bytes_l, frames_r, bytes_r, frames_total, bytes_total, rel_start, duration = line.split()
                    conversations.append({'Source': src, 'Destination': dst, 'Frames to Source': frames_l, 'Bytes to Source': bytes_l, 'Frames to Destination': frames_r, 'Bytes to Destination': bytes_r, 'Total Frames': frames_total, 'Total Bytes': bytes_total, 'Relative Start': rel_start, 'Duration': duration})
            results['tshark'][result] = conversations
        elif 'Endpoints' in result:
            # handle endpoint parsing
            endpoints = []
            for line in results['tshark'][result].split('\n'):
                if line == '' or line.startswith(' '):
                    # header or padding, dicard
                    continue
                else:
                    # handle endpoint services with ports
                    if result.startswith('UDP') or result.startswith('TCP') or result.startswith('STCP'):
                        endpoint, port, packet_count, byte_count, tx_packets, tx_bytes, rx_packets, rx_bytes = line.split()
                        endpoints.append({'Endpoint': endpoint, 'Port': port, 'Packets': packet_count, 'Bytes': byte_count, 'Tx Packets': tx_packets, 'Tx Bytes': tx_bytes, 'Rx Packets': rx_packets, 'Rx Bytes': rx_bytes})
                    else:
                        endpoint, packet_count, byte_count, tx_packets, tx_bytes, rx_packets, rx_bytes = line.split()
                        endpoints.append({'Endpoint': endpoint, 'Packets': packet_count, 'Bytes': byte_count, 'Tx Packets': tx_packets, 'Tx Bytes': tx_bytes, 'Rx Packets': rx_packets, 'Rx Bytes': rx_bytes})
            results['tshark'][result] = endpoints
        else:
            # handle weird stuff
            for line in results['tshark'][result].split('\n'):
                if line == '' or line.startswith(' '):
                    # header or padding, dicard
                    continue
                else:
                    # handle icmp and icmpv6
                    if result.startswith('ICMP'):
                        if isinstance(results['tshark'][result], str):
                            results['tshark'][result] = {}
                        if line.startswith('Requests') or line.startswith('Minimum'):
                            # header
                            continue
                        else:
                            values = line.split()
                            if len(values) == 4:
                                requests, replies, lost, percent_loss = values
                                results['tshark'][result]['Requests'] = requests
                                results['tshark'][result]['Replies'] = replies
                                results['tshark'][result]['Lost'] = lost
                                results['tshark'][result]['% Loss'] = percent_loss
                            else:
                                minimum, maximum, mean, median, s_deviation, min_frame, max_frame = values
                                results['tshark'][result]['Minimum'] = minimum
                                results['tshark'][result]['Maximum'] = maximum
                                results['tshark'][result]['Mean'] = mean
                                results['tshark'][result]['Median'] = median
                                results['tshark'][result]['Standard Deviation'] = s_deviation
                                results['tshark'][result]['Minimum Frame'] = min_frame
                                results['tshark'][result]['Maximum Frame'] = max_frame
                    elif result.startswith('Protocol'):
                        # TODO
                        continue
                    # handle dns
                    elif result.startswith('DNS'):
                        # TODO
                        continue

    # TODO temporarily remove until parsed
    if 'DNS' in results['tshark']:
        del results['tshark']['DNS']

    # handle protocol hierarchy stats
    a = []
    if 'Protocol Hierarchy Statistics' in results['tshark']:
        a = results['tshark']['Protocol Hierarchy Statistics'].split('\n')
    h = []
    for line in a:
        if line != '':
            name, frame_count, byte_count = line.rsplit(' ', 2)
            name = name.rstrip()
            frame_count = frame_count.split(':')[1]
            byte_count = byte_count.split(':')[1]
            h.append([name, frame_count, byte_count])

    i = 1
    spaces = 0
    if h:
        h[0][0] = '"' + h[0][0].strip()
        while i < len(h):
            prev_spaces = spaces
            spaces = h[i][0].count('  ')
            h[i-1][0] = h[i-1][0].strip() + '":{"Frames": "' + h[i-1][1] + '", "Bytes": "' + h[i-1][2]
            if spaces > prev_spaces:
                h[i-1][0] += '","'
            elif spaces == prev_spaces:
                h[i-1][0] += '"},"'
            else:
                h[i-1][0] += '"}' + ('}'*(prev_spaces-spaces)) + ',"'
            i += 1
        h[i-1][0] = h[i-1][0].strip() + '":{"Frames": "' + h[i-1][1] + '", "Bytes": "' + h[i-1][2] + '"}' + ('}'*(prev_spaces-spaces))

    protocol_str = '{'
    for record in h:
        protocol_str += record[0]
    protocol_str += '}'
    results['tshark']['Protocol Hierarchy Statistics'] = json.loads(protocol_str)

    # add in condensed conversation fields
    results['tshark']['Condensed TCP Conversations'] = condense_conversations(results, 'TCP Conversations')
    results['tshark']['Condensed UDP Conversations'] = condense_conversations(results, 'UDP Conversations')

    print(results)
    return results

def condense_conversations(results, conv_type):
    prot_ip_map = {}
    if conv_type in results['tshark']:
        for conversation in results['tshark'][conv_type]:
            src_ip, src_port = conversation['Source'].rsplit(':', 1)
            dst_ip, dst_port = conversation['Destination'].rsplit(':', 1)
            if not src_ip in prot_ip_map:
                prot_ip_map[src_ip] = {'Destinations': [], 'Source Ports': [], 'Destination Ports': []}
            if not src_port in prot_ip_map[src_ip]['Source Ports']:
                prot_ip_map[src_ip]['Source Ports'].append(src_port)
            if not dst_port in prot_ip_map[src_ip]['Destination Ports']:
                prot_ip_map[src_ip]['Destination Ports'].append(dst_port)
            if not dst_ip in prot_ip_map[src_ip]['Destinations']:
                prot_ip_map[src_ip]['Destinations'].append(dst_ip)
    return prot_ip_map

def run_tshark(path):
    if os.path.getsize(path) == 0:
       print("pcap file empty, no stats")
       return

    results = {}
    output = ''
    try:
        conv_endpoint_types = ['bluetooth', 'eth', 'fc', 'fddi', 'ip', 'ipv6', 'ipx', 'jxta', 'ncp', 'rsvp', 'sctp', 'tcp', 'tr', 'usb', 'udp', 'wlan']
        options = '-n -q -z dns,tree -z io,phs -z icmp,srt -z icmpv6,srt'
        options += ' -z conv,'.join(conv_endpoint_types)
        options += ' -z endpoints,'.join(conv_endpoint_types)
        output = subprocess.check_output(shlex.split(' '.join(['tshark', '-r', path, options])))
        output = output.decode("utf-8")
    except Exception as e:
        print(str(e))

    results = parse_tshark(output)
    return results

if __name__ == '__main__':  # pragma: no cover
    path = get_path()
    uid = ''
    if 'id' in os.environ:
        uid = os.environ['id']
    if path:
        if 'rabbit' in os.environ and os.environ['rabbit'] == 'true':
            try:
                channel = connect_rabbit()
                capinfos_results = run_capinfos(path)
                body = {'id': uid, 'type': 'metadata', 'file_path': path, 'data': capinfos_results, 'results': {'tool': 'pcap-stats', 'version': get_version()}}
                send_rabbit_msg(body, channel)
                tshark_results = run_tshark(path)
                body = {'id': uid, 'type': 'metadata', 'file_path': path, 'data': tshark_results, 'results': {'tool': 'pcap-stats', 'version': get_version()}}
                send_rabbit_msg(body, channel)
                body = {'id': uid, 'type': 'metadata', 'file_path': path, 'data': '', 'results': {'tool': 'pcap-stats', 'version': get_version()}}
                send_rabbit_msg(body, channel)
            except Exception as e:
                print(str(e))
        else:
            capinfos_results = run_capinfos(path)
            tshark_results = run_tshark(path)
<|endoftext|>


<commit_before># Copyright 2016 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Transport adapter for httplib2."""

from __future__ import absolute_import

import logging

from google.auth import exceptions
from google.auth import transport
import httplib2

_LOGGER = logging.getLogger(__name__)
# Properties present in file-like streams / buffers.
_STREAM_PROPERTIES = ('read', 'seek', 'tell')

class _Response(transport.Response):
    """httplib2 transport response adapter.

    Args:
        response (httplib2.Response): The raw httplib2 response.
        data (bytes): The response body.
    """
    def __init__(self, response, data):
        self._response = response
        self._data = data

    @property
    def status(self):
        """int: The HTTP status code."""
        return self._response.status

    @property
    def headers(self):
        """Mapping[str, str]: The HTTP response headers."""
        return dict(self._response)

    @property
    def data(self):
        """bytes: The response body."""
        return self._data

class Request(transport.Request):
    """httplib2 request adapter.

    This class is used internally for making requests using various transports
    in a consistent way. If you use :class:`AuthorizedHttp` you do not need
    to construct or use this class directly.

    This class can be useful if you want to manually refresh a
    :class:`~google.auth.credentials.Credentials` instance::

        import google.auth.transport.httplib2
        import httplib2

        http = httplib2.Http()
        request = google.auth.transport.httplib2.Request(http)

        credentials.refresh(request)

    Args:
        http (httplib2.Http): The underlying http object to use to make
            requests.

    .. automethod:: __call__
    """
    def __init__(self, http):
        self.http = http

    def __call__(self, url, method='GET', body=None, headers=None,
                 timeout=None, **kwargs):
        """Make an HTTP request using httplib2.

        Args:
            url (str): The URI to be requested.
            method (str): The HTTP method to use for the request. Defaults
                to 'GET'.
            body (bytes): The payload / body in HTTP request.
            headers (Mapping[str, str]): Request headers.
            timeout (Optional[int]): The number of seconds to wait for a
                response from the server. This is ignored by httplib2 and will
                issue a warning.
            kwargs: Additional arguments passed throught to the underlying
                :meth:`httplib2.Http.request` method.

        Returns:
            google.auth.transport.Response: The HTTP response.

        Raises:
            google.auth.exceptions.TransportError: If any exception occurred.
        """
        if timeout is not None:
            _LOGGER.warning(
                'httplib2 transport does not support per-request timeout. '
                'Set the timeout when constructing the httplib2.Http instance.'
            )

        try:
            _LOGGER.debug('Making request: %s %s', method, url)
            response, data = self.http.request(
                url, method=method, body=body, headers=headers, **kwargs)
            return _Response(response, data)
        except httplib2.HttpLib2Error as exc:
            raise exceptions.TransportError(exc)

def _make_default_http():
    """Returns a default httplib2.Http instance."""
    return httplib2.Http()

class AuthorizedHttp(object):
    """A httplib2 HTTP class with credentials.

    This class is used to perform requests to API endpoints that require
    authorization::

        from google.auth.transport._httplib2 import AuthorizedHttp

        authed_http = AuthorizedHttp(credentials)

        response = authed_http.request(
            'https://www.googleapis.com/storage/v1/b')

    This class implements :meth:`request` in the same way as
    :class:`httplib2.Http` and can usually be used just like any other
    instance of :class:``httplib2.Http`.

    The underlying :meth:`request` implementation handles adding the
    credentials' headers to the request and refreshing credentials as needed.
    """
    def __init__(self, credentials, http=None,
                 refresh_status_codes=transport.DEFAULT_REFRESH_STATUS_CODES,
                 max_refresh_attempts=transport.DEFAULT_MAX_REFRESH_ATTEMPTS):
        """
        Args:
            credentials (google.auth.credentials.Credentials): The credentials
                to add to the request.
            http (httplib2.Http): The underlying HTTP object to
                use to make requests. If not specified, a
                :class:`httplib2.Http` instance will be constructed.
            refresh_status_codes (Sequence[int]): Which HTTP status codes
                indicate that credentials should be refreshed and the request
                should be retried.
            max_refresh_attempts (int): The maximum number of times to attempt
                to refresh the credentials and retry the request.
        """

        if http is None:
            http = _make_default_http()

        self.http = http
        self.credentials = credentials
        self._refresh_status_codes = refresh_status_codes
        self._max_refresh_attempts = max_refresh_attempts
        # Request instance used by internal methods (for example,
        # credentials.refresh).
        self._request = Request(self.http)

    def request(self, uri, method='GET', body=None, headers=None,
                **kwargs):
        """Implementation of httplib2's Http.request."""

        _credential_refresh_attempt = kwargs.pop(
            '_credential_refresh_attempt', 0)

        # Make a copy of the headers. They will be modified by the credentials
        # and we want to pass the original headers if we recurse.
        request_headers = headers.copy() if headers is not None else {}

        self.credentials.before_request(
            self._request, method, uri, request_headers)

        # Check if the body is a file-like stream, and if so, save the body
        # stream position so that it can be restored in case of refresh.
        body_stream_position = None
        if all(getattr(body, stream_prop, None) for stream_prop in
               _STREAM_PROPERTIES):
            body_stream_position = body.tell()

        # Make the request.
        response, content = self.http.request(
            uri, method, body=body, headers=request_headers, **kwargs)

        # If the response indicated that the credentials needed to be
        # refreshed, then refresh the credentials and re-attempt the
        # request.
        # A stored token may expire between the time it is retrieved and
        # the time the request is made, so we may need to try twice.
        if (response.status in self._refresh_status_codes
                and _credential_refresh_attempt < self._max_refresh_attempts):

            _LOGGER.info(
                'Refreshing credentials due to a %s response. Attempt %s/%s.',
                response.status, _credential_refresh_attempt + 1,
                self._max_refresh_attempts)

            self.credentials.refresh(self._request)

            # Restore the body's stream position if needed.
            if body_stream_position is not None:
                body.seek(body_stream_position)

            # Recurse. Pass in the original headers, not our modified set.
            return self.request(
                uri, method, body=body, headers=headers,
                _credential_refresh_attempt=_credential_refresh_attempt + 1,
                **kwargs)

        return response, content

    @property
    def connections(self):
        """Proxy to httplib2.Http.connections."""
        return self.http.connections

    @connections.setter
    def connections(self, value):
        """Proxy to httplib2.Http.connections."""
        self.http.connections = value
<commit_msg>Update docstring for Request after move. (#2)<commit_after># Copyright 2016 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Transport adapter for httplib2."""

from __future__ import absolute_import

import logging

from google.auth import exceptions
from google.auth import transport
import httplib2

_LOGGER = logging.getLogger(__name__)
# Properties present in file-like streams / buffers.
_STREAM_PROPERTIES = ('read', 'seek', 'tell')

class _Response(transport.Response):
    """httplib2 transport response adapter.

    Args:
        response (httplib2.Response): The raw httplib2 response.
        data (bytes): The response body.
    """
    def __init__(self, response, data):
        self._response = response
        self._data = data

    @property
    def status(self):
        """int: The HTTP status code."""
        return self._response.status

    @property
    def headers(self):
        """Mapping[str, str]: The HTTP response headers."""
        return dict(self._response)

    @property
    def data(self):
        """bytes: The response body."""
        return self._data

class Request(transport.Request):
    """httplib2 request adapter.

    This class is used internally for making requests using various transports
    in a consistent way. If you use :class:`AuthorizedHttp` you do not need
    to construct or use this class directly.

    This class can be useful if you want to manually refresh a
    :class:`~google.auth.credentials.Credentials` instance::

        import google_auth_httplib2
        import httplib2

        http = httplib2.Http()
        request = google_auth_httplib2.Request(http)

        credentials.refresh(request)

    Args:
        http (httplib2.Http): The underlying http object to use to make
            requests.

    .. automethod:: __call__
    """
    def __init__(self, http):
        self.http = http

    def __call__(self, url, method='GET', body=None, headers=None,
                 timeout=None, **kwargs):
        """Make an HTTP request using httplib2.

        Args:
            url (str): The URI to be requested.
            method (str): The HTTP method to use for the request. Defaults
                to 'GET'.
            body (bytes): The payload / body in HTTP request.
            headers (Mapping[str, str]): Request headers.
            timeout (Optional[int]): The number of seconds to wait for a
                response from the server. This is ignored by httplib2 and will
                issue a warning.
            kwargs: Additional arguments passed throught to the underlying
                :meth:`httplib2.Http.request` method.

        Returns:
            google.auth.transport.Response: The HTTP response.

        Raises:
            google.auth.exceptions.TransportError: If any exception occurred.
        """
        if timeout is not None:
            _LOGGER.warning(
                'httplib2 transport does not support per-request timeout. '
                'Set the timeout when constructing the httplib2.Http instance.'
            )

        try:
            _LOGGER.debug('Making request: %s %s', method, url)
            response, data = self.http.request(
                url, method=method, body=body, headers=headers, **kwargs)
            return _Response(response, data)
        except httplib2.HttpLib2Error as exc:
            raise exceptions.TransportError(exc)

def _make_default_http():
    """Returns a default httplib2.Http instance."""
    return httplib2.Http()

class AuthorizedHttp(object):
    """A httplib2 HTTP class with credentials.

    This class is used to perform requests to API endpoints that require
    authorization::

        from google.auth.transport._httplib2 import AuthorizedHttp

        authed_http = AuthorizedHttp(credentials)

        response = authed_http.request(
            'https://www.googleapis.com/storage/v1/b')

    This class implements :meth:`request` in the same way as
    :class:`httplib2.Http` and can usually be used just like any other
    instance of :class:``httplib2.Http`.

    The underlying :meth:`request` implementation handles adding the
    credentials' headers to the request and refreshing credentials as needed.
    """
    def __init__(self, credentials, http=None,
                 refresh_status_codes=transport.DEFAULT_REFRESH_STATUS_CODES,
                 max_refresh_attempts=transport.DEFAULT_MAX_REFRESH_ATTEMPTS):
        """
        Args:
            credentials (google.auth.credentials.Credentials): The credentials
                to add to the request.
            http (httplib2.Http): The underlying HTTP object to
                use to make requests. If not specified, a
                :class:`httplib2.Http` instance will be constructed.
            refresh_status_codes (Sequence[int]): Which HTTP status codes
                indicate that credentials should be refreshed and the request
                should be retried.
            max_refresh_attempts (int): The maximum number of times to attempt
                to refresh the credentials and retry the request.
        """

        if http is None:
            http = _make_default_http()

        self.http = http
        self.credentials = credentials
        self._refresh_status_codes = refresh_status_codes
        self._max_refresh_attempts = max_refresh_attempts
        # Request instance used by internal methods (for example,
        # credentials.refresh).
        self._request = Request(self.http)

    def request(self, uri, method='GET', body=None, headers=None,
                **kwargs):
        """Implementation of httplib2's Http.request."""

        _credential_refresh_attempt = kwargs.pop(
            '_credential_refresh_attempt', 0)

        # Make a copy of the headers. They will be modified by the credentials
        # and we want to pass the original headers if we recurse.
        request_headers = headers.copy() if headers is not None else {}

        self.credentials.before_request(
            self._request, method, uri, request_headers)

        # Check if the body is a file-like stream, and if so, save the body
        # stream position so that it can be restored in case of refresh.
        body_stream_position = None
        if all(getattr(body, stream_prop, None) for stream_prop in
               _STREAM_PROPERTIES):
            body_stream_position = body.tell()

        # Make the request.
        response, content = self.http.request(
            uri, method, body=body, headers=request_headers, **kwargs)

        # If the response indicated that the credentials needed to be
        # refreshed, then refresh the credentials and re-attempt the
        # request.
        # A stored token may expire between the time it is retrieved and
        # the time the request is made, so we may need to try twice.
        if (response.status in self._refresh_status_codes
                and _credential_refresh_attempt < self._max_refresh_attempts):

            _LOGGER.info(
                'Refreshing credentials due to a %s response. Attempt %s/%s.',
                response.status, _credential_refresh_attempt + 1,
                self._max_refresh_attempts)

            self.credentials.refresh(self._request)

            # Restore the body's stream position if needed.
            if body_stream_position is not None:
                body.seek(body_stream_position)

            # Recurse. Pass in the original headers, not our modified set.
            return self.request(
                uri, method, body=body, headers=headers,
                _credential_refresh_attempt=_credential_refresh_attempt + 1,
                **kwargs)

        return response, content

    @property
    def connections(self):
        """Proxy to httplib2.Http.connections."""
        return self.http.connections

    @connections.setter
    def connections(self, value):
        """Proxy to httplib2.Http.connections."""
        self.http.connections = value
<|endoftext|>


<commit_before>import re
import sublime
from sublime_plugin import TextCommand

from ..git_command import GitCommand
from ...common import util

RELEASE_REGEXP = re.compile(r"^([0-9A-Za-z-]*[A-Za-z-])?([0-9]+)\.([0-9]+)\.([0-9]+)(?:-([0-9]+))?$")

TAG_CREATE_PROMPT = "Enter tag:"
TAG_CREATE_MESSAGE = "Tag \"{}\" created."
TAG_CREATE_MESSAGE_PROMPT = "Enter message:"
TAG_PARSE_FAIL_MESSAGE = "The last tag cannot be parsed."

def smart_incremented_tag(tag, release_type):
    """
    Automatic increment of a given tag depending on the type of release.

    >>> smart_incremented_tag('v1.3.2', "prerelease") == 'v1.3.3-0'
    >>> smart_incremented_tag('v1.3.2', "prepatch") == 'v1.3.3-0'
    >>> smart_incremented_tag('v1.3.2', "preminor") == 'v1.4.0-0'
    >>> smart_incremented_tag('v1.3.2', "premajor") == 'v2.0.0-0'
    >>> smart_incremented_tag('v1.3.2', "patch") == 'v1.3.3'
    >>> smart_incremented_tag('v1.3.2', "minor") == 'v1.4.0'
    >>> smart_incremented_tag('v1.3.2', "major") == 'v2.0.0'
    >>> smart_incremented_tag('v1.3.2-1', "prerelease") == 'v1.3.2-2'
    >>> smart_incremented_tag('v1.3.2-1', "prepatch") == 'v1.3.3-0'
    >>> smart_incremented_tag('v1.3.2-1', "preminor") == 'v1.4.0-0'
    >>> smart_incremented_tag('v1.3.2-1', "premajor") == 'v2.0.0-0'
    >>> smart_incremented_tag('v1.3.2-1', "patch") == 'v1.3.2'
    >>> smart_incremented_tag('v1.3.2-1', "minor") == 'v1.4.0'
    >>> smart_incremented_tag('v1.3.2-1', "major") == 'v2.0.0'
    >>> smart_incremented_tag('v1.3.0-1', "patch") == 'v1.3.0'
    >>> smart_incremented_tag('v1.3.0-1', "minor") == 'v1.3.0'
    >>> smart_incremented_tag('v1.3.0-1', "major") == 'v2.0.0'
    >>> smart_incremented_tag('v1.0.0-1', "major") == 'v1.0.0'

    """

    m = RELEASE_REGEXP.match(tag)
    if m:
        prefix, major, minor, patch, prerelease = m.groups()
        prefix = "" if not prefix else prefix

        if release_type == "premajor" \
                or (not prerelease and release_type == "major") \
                or (prerelease and release_type == "major" and (minor != "0" or patch != "0")):
            major = str(int(major)+1)
            minor = patch = "0"
            prerelease = None
        elif release_type == "preminor" \
                or (not prerelease and release_type == "minor") \
                or (prerelease and release_type == "minor" and patch != "0"):
            minor = str(int(minor)+1)
            patch = "0"
            prerelease = None
        elif release_type == "prepatch" \
                or (not prerelease and release_type == "prerelease") \
                or (not prerelease and release_type == "patch"):
            patch = str(int(patch)+1)
            prerelease = None

        if "pre" in release_type[0:3]:
            prerelease = str(int(prerelease)+1) if prerelease else "0"
            return prefix + major + "." + minor + "." + patch + "-" + prerelease
        else:
            return prefix + major + "." + minor + "." + patch

    return None

class GsTagCreateCommand(TextCommand, GitCommand):

    """
    Through a series of panels, allow the user to add a tag and message.
    """

    def run(self, edit, tag_name=""):
        self.window = self.view.window()
        self.tag_name = tag_name
        sublime.set_timeout_async(self.run_async)

    def run_async(self):
        """
        Prompt the user for a tag name.
        """
        self.window.show_input_panel(TAG_CREATE_PROMPT, self.tag_name, self.on_entered_name, None, None)

    def on_entered_name(self, tag_name):
        """
        After the user has entered a tag name, validate the tag name and prompt for
        a tag message.
        """
        if not tag_name:
            return

        stdout = self.git(
            "check-ref-format",
            "--normalize",
            "refs/tags/" + tag_name,
            throw_on_stderr=False
            )

        if not stdout:
            return util.log.panel("\"{}\" is not a valid tag name.".format(tag_name))

        self.tag_name = stdout.strip()[10:]
        self.window.show_input_panel(
            TAG_CREATE_MESSAGE_PROMPT,
            sublime\
                .load_settings("GitSavvy.sublime-settings")\
                .get("default_tag_message")\
                .format(tag_name=tag_name),
            self.on_entered_message,
            None,
            None
            )

    def on_entered_message(self, message):
        """
        Create a tag with the specified tag name and message.
        """
        if not message:
            return

        self.git("tag", self.tag_name, "-F", "-", stdin=message)
        sublime.status_message(TAG_CREATE_MESSAGE.format(self.tag_name))
        util.view.refresh_gitsavvy(self.view)

class GsSmartTagCommand(TextCommand, GitCommand):
    """
    Displays a panel of possible smart tag options, based on the choice,
    tag the current commit with the corresponding tagname.
    """

    release_types = [
        "major",
        "minor",
        "patch",
        "premajor",
        "preminor",
        "prepatch",
        "prerelease"
    ]

    def run(self, edit):
        sublime.set_timeout_async(self.run_async)

    def run_async(self):
        self.view.window().show_quick_panel(self.release_types, self.on_tag)

    def on_tag(self, action):
        if action < 0:
            return

        release_type = self.release_types[action]

        tag_name = None
        last_tag_name = self.get_lastest_local_tag()
        if last_tag_name:
            tag_name = smart_incremented_tag(last_tag_name, release_type)

        if not tag_name:
            sublime.message_dialog(TAG_PARSE_FAIL_MESSAGE)
            return

        self.view.run_command("gs_tag_create", {"tag_name": tag_name})
<commit_msg>reorder smart tag order list<commit_after>import re
import sublime
from sublime_plugin import TextCommand

from ..git_command import GitCommand
from ...common import util

RELEASE_REGEXP = re.compile(r"^([0-9A-Za-z-]*[A-Za-z-])?([0-9]+)\.([0-9]+)\.([0-9]+)(?:-([0-9]+))?$")

TAG_CREATE_PROMPT = "Enter tag:"
TAG_CREATE_MESSAGE = "Tag \"{}\" created."
TAG_CREATE_MESSAGE_PROMPT = "Enter message:"
TAG_PARSE_FAIL_MESSAGE = "The last tag cannot be parsed."

def smart_incremented_tag(tag, release_type):
    """
    Automatic increment of a given tag depending on the type of release.

    >>> smart_incremented_tag('v1.3.2', "prerelease") == 'v1.3.3-0'
    >>> smart_incremented_tag('v1.3.2', "prepatch") == 'v1.3.3-0'
    >>> smart_incremented_tag('v1.3.2', "preminor") == 'v1.4.0-0'
    >>> smart_incremented_tag('v1.3.2', "premajor") == 'v2.0.0-0'
    >>> smart_incremented_tag('v1.3.2', "patch") == 'v1.3.3'
    >>> smart_incremented_tag('v1.3.2', "minor") == 'v1.4.0'
    >>> smart_incremented_tag('v1.3.2', "major") == 'v2.0.0'
    >>> smart_incremented_tag('v1.3.2-1', "prerelease") == 'v1.3.2-2'
    >>> smart_incremented_tag('v1.3.2-1', "prepatch") == 'v1.3.3-0'
    >>> smart_incremented_tag('v1.3.2-1', "preminor") == 'v1.4.0-0'
    >>> smart_incremented_tag('v1.3.2-1', "premajor") == 'v2.0.0-0'
    >>> smart_incremented_tag('v1.3.2-1', "patch") == 'v1.3.2'
    >>> smart_incremented_tag('v1.3.2-1', "minor") == 'v1.4.0'
    >>> smart_incremented_tag('v1.3.2-1', "major") == 'v2.0.0'
    >>> smart_incremented_tag('v1.3.0-1', "patch") == 'v1.3.0'
    >>> smart_incremented_tag('v1.3.0-1', "minor") == 'v1.3.0'
    >>> smart_incremented_tag('v1.3.0-1', "major") == 'v2.0.0'
    >>> smart_incremented_tag('v1.0.0-1', "major") == 'v1.0.0'

    """

    m = RELEASE_REGEXP.match(tag)
    if m:
        prefix, major, minor, patch, prerelease = m.groups()
        prefix = "" if not prefix else prefix

        if release_type == "premajor" \
                or (not prerelease and release_type == "major") \
                or (prerelease and release_type == "major" and (minor != "0" or patch != "0")):
            major = str(int(major)+1)
            minor = patch = "0"
            prerelease = None
        elif release_type == "preminor" \
                or (not prerelease and release_type == "minor") \
                or (prerelease and release_type == "minor" and patch != "0"):
            minor = str(int(minor)+1)
            patch = "0"
            prerelease = None
        elif release_type == "prepatch" \
                or (not prerelease and release_type == "prerelease") \
                or (not prerelease and release_type == "patch"):
            patch = str(int(patch)+1)
            prerelease = None

        if "pre" in release_type[0:3]:
            prerelease = str(int(prerelease)+1) if prerelease else "0"
            return prefix + major + "." + minor + "." + patch + "-" + prerelease
        else:
            return prefix + major + "." + minor + "." + patch

    return None

class GsTagCreateCommand(TextCommand, GitCommand):

    """
    Through a series of panels, allow the user to add a tag and message.
    """

    def run(self, edit, tag_name=""):
        self.window = self.view.window()
        self.tag_name = tag_name
        sublime.set_timeout_async(self.run_async)

    def run_async(self):
        """
        Prompt the user for a tag name.
        """
        self.window.show_input_panel(TAG_CREATE_PROMPT, self.tag_name, self.on_entered_name, None, None)

    def on_entered_name(self, tag_name):
        """
        After the user has entered a tag name, validate the tag name and prompt for
        a tag message.
        """
        if not tag_name:
            return

        stdout = self.git(
            "check-ref-format",
            "--normalize",
            "refs/tags/" + tag_name,
            throw_on_stderr=False
            )

        if not stdout:
            return util.log.panel("\"{}\" is not a valid tag name.".format(tag_name))

        self.tag_name = stdout.strip()[10:]
        self.window.show_input_panel(
            TAG_CREATE_MESSAGE_PROMPT,
            sublime\
                .load_settings("GitSavvy.sublime-settings")\
                .get("default_tag_message")\
                .format(tag_name=tag_name),
            self.on_entered_message,
            None,
            None
            )

    def on_entered_message(self, message):
        """
        Create a tag with the specified tag name and message.
        """
        if not message:
            return

        self.git("tag", self.tag_name, "-F", "-", stdin=message)
        sublime.status_message(TAG_CREATE_MESSAGE.format(self.tag_name))
        util.view.refresh_gitsavvy(self.view)

class GsSmartTagCommand(TextCommand, GitCommand):
    """
    Displays a panel of possible smart tag options, based on the choice,
    tag the current commit with the corresponding tagname.
    """

    release_types = [
        "patch",
        "minor",
        "major",
        "prerelease",
        "prepatch",
        "preminor",
        "premajor"
    ]

    def run(self, edit):
        sublime.set_timeout_async(self.run_async)

    def run_async(self):
        self.view.window().show_quick_panel(self.release_types, self.on_tag)

    def on_tag(self, action):
        if action < 0:
            return

        release_type = self.release_types[action]

        tag_name = None
        last_tag_name = self.get_lastest_local_tag()
        if last_tag_name:
            tag_name = smart_incremented_tag(last_tag_name, release_type)

        if not tag_name:
            sublime.message_dialog(TAG_PARSE_FAIL_MESSAGE)
            return

        self.view.run_command("gs_tag_create", {"tag_name": tag_name})
<|endoftext|>


<commit_before>from celery import chain, group, chord
from celery.task import task
from api.wrappers.github import parse_github_url
from core.models import Job, Spec, Package, JobSpec
from api.wrappers import pypi, github, travis

@task
def process_job(job_pk):
    """ Process single batch of packages

        For each package spec we execute following tasks:
            * PyPI query
            * GitHub query
            * Travis query
    """
    job = Job.objects.get(pk=job_pk)
    job.do_start()

    return chord(process_spec.s(job_spec.pk)
                 for job_spec in JobSpec.objects.filter(job=job))(notify_completed_job.si(job_pk))

@task
def process_spec(job_spec_pk):
    """ Process jobspec"""
    job_spec = JobSpec.objects.get(pk=job_spec_pk)
    job_spec.do_start()

    package, version = job_spec.spec.package, job_spec.spec.version
    finished_specs = Spec.objects.filter(package=package, version=version, status='finished').exists()
    if finished_specs:  # not asking services if already parsed
        return notify_completed_spec.delay(job_spec_pk)
    else:
        return chain(query_pypi.s(job_spec.spec.pk),
                     query_github.s(job_spec.spec.package.pk),
                     notify_completed_spec.si(job_spec_pk)).delay()

@task
def notify_completed_job(job_pk):
    """ Job has finished, now we need to record the result"""
    job = Job.objects.get(pk=job_pk)
    job.do_finish()

@task
def notify_completed_spec(job_spec_pk):
    """ Job has finished, now we need to record the result
    """
    job_spec = JobSpec.objects.get(pk=job_spec_pk)
    job_spec.do_finish()
    spec = job_spec.spec
    spec.status = 'completed'
    spec.save(update_fields=['status'])

@task
def query_pypi(spec_pk):
    """ Query one spec of package on PyPI

        Sample result of PyPIWrapper:

        >>> pypi.PyPIWrapper().get_short_info('django', None)
        {'last_release_date': datetime.datetime(2013, 9, 15, 6, 30, 37),
         'name': 'Django',
         'python3_supported_versions': ['3', '3.2', '3.3'],
         'url': 'http://pypi.python.org/pypi/Django',
         'version': '1.5.4'}
    """
    spec = Spec.objects.get(pk=spec_pk)

    # query PyPI
    pypi_api = pypi.PyPIWrapper()
    pkg_data = pypi_api.get_short_info(spec.name, spec.version)
    #latest_pkg_data = pypi_api.get_short_info(spec.name)

    spec.release_date = pkg_data['last_release_date']
    spec.python_versions = pkg_data['python3_supported_versions']
    spec.save()

    # Canonical package name case for the win
    if pkg_data['name'] != spec.name:
        package = spec.package
        package.name = pkg_data['name']
        package.save()

    return pkg_data

@task
def search_github(package_name, url):
    """ Where are my sources, bro?

        Sometimes, package name is not the same as github repo name,
        but that should not be a problem, since Github search API
        is smart enough.

        >>> github.GithubSearchWrapper().get_most_popular_repo('moscowdjango')
        ('moscowdjango', 'futurecolors')
    """
    if url:
        repo_info = parse_github_url(url)
        if repo_info:
            # query directly to github account
            return repo_info['repo_name'], repo_info['owner']

    repo_name, owner = github.GithubSearchWrapper().get_most_popular_repo(package_name)
    return repo_name, owner

@task
def get_short_info(repo_tuple):
    """ What's with my sources, bro?

        >>> github.GithubWrapper().get_short_info('futurecolors', 'moscowdjango')
        {'updated_at': datetime.datetime(2013, 9, 15, 6, 30, 37),
         'url': 'https://api.github.com/repos/futurecolors/moscowdjango'}
    """
    repo_name, owner = repo_tuple  # FIXME
    return github.GithubWrapper().get_short_info(owner, repo_name)

@task
def get_forks(repo_tuple):
    repo_name, owner = repo_tuple
    return github.GithubWrapper().get_py3_fork_info(owner, repo_name)

@task
def get_issues(repo_tuple):
    repo_name, owner = repo_tuple
    return github.GithubWrapper().get_py3_issues_info(owner, repo_name)

@task
def get_pr(repo_tuple):
    repo_name, owner = repo_tuple
    return github.GithubWrapper().get_py3_pull_requests(owner, repo_name)

@task
def get_build_status(repo_tuple):
    """ How are my tests going? """
    repo_name, owner = repo_tuple
    travis_data = travis.TravisCI().get_build_status(owner, repo_name)

    return travis_data

@task
def notify_github_completed(r, package_pk):
    """ Update github package info in database"""
    package = Package.objects.get(pk=package_pk)
    package.repo_last_commit_date = r[0].get('updated_at')
    package.repo_url = r[0].get('url')

    # For now we'll take the first one available
    package.pr_url = r[1][0].get('url', '') if r[1] else ''
    package.pr_status = r[1][0].get('state') if r[1] else 'unknown'

    package.issue_url = r[2][0].get('url', '') if r[2] else ''
    package.issue_status = r[2][0].get('state') if r[2] else 'unknown'

    package.fork_url = r[3][0].pop(0, {}).get('url', '') if r[3] else ''

    # Travis
    package.ci_url = r[3].get('url', '') if r[4] else ''
    package.ci_status = r[3].get('last_build_status', 'unknown') if r[4] else 'unknown'
    package.save()
    return r

@task
def query_github(results, package_pk):
    """ Get all relevant info form Github"""
    package = Package.objects.get(pk=package_pk)

    # GH API queries running in parallel for speedup
    gh_queries = group(get_short_info.s(),
                       get_pr.s(),
                       get_issues.s(),
                       get_forks.s(),
                       get_build_status.s())

    # guessing github url and querying info
    return chain(search_github.s(package.name, results.get('url')),
                 gh_queries,
                 notify_github_completed.s(package_pk)).delay()
<commit_msg>Do not finish job before spec<commit_after>from celery import chain, group, chord
from celery.task import task
from api.wrappers.github import parse_github_url
from core.models import Job, Spec, Package, JobSpec
from api.wrappers import pypi, github, travis

@task
def process_job(job_pk):
    """ Process single batch of packages

        For each package spec we execute following tasks:
            * PyPI query
            * GitHub query
            * Travis query
    """
    job = Job.objects.get(pk=job_pk)
    job.do_start()

    return chord(process_spec.s(job_spec.pk)
                 for job_spec in JobSpec.objects.filter(job=job))(notify_completed_job.si(job_pk))

@task
def process_spec(job_spec_pk):
    """ Process jobspec"""
    job_spec = JobSpec.objects.get(pk=job_spec_pk)
    job_spec.do_start()

    package, version = job_spec.spec.package, job_spec.spec.version
    finished_specs = Spec.objects.filter(package=package, version=version, status='finished').exists()
    if finished_specs:  # not asking services if already parsed
        return notify_completed_spec.delay(job_spec_pk)
    else:
        return chain(query_pypi.s(job_spec.spec.pk),
                     query_github.s(job_spec.spec.package.pk),
                     notify_completed_spec.si(job_spec_pk)).delay()

@task
def notify_completed_job(job_pk):
    """ Job has finished, now we need to record the result"""
    job = Job.objects.get(pk=job_pk)
    job.do_finish()

@task
def notify_completed_spec(job_spec_pk):
    """ Job has finished, now we need to record the result
    """
    job_spec = JobSpec.objects.get(pk=job_spec_pk)
    spec = job_spec.spec
    spec.status = 'completed'
    spec.save(update_fields=['status'])
    job_spec.do_finish()

@task
def query_pypi(spec_pk):
    """ Query one spec of package on PyPI

        Sample result of PyPIWrapper:

        >>> pypi.PyPIWrapper().get_short_info('django', None)
        {'last_release_date': datetime.datetime(2013, 9, 15, 6, 30, 37),
         'name': 'Django',
         'python3_supported_versions': ['3', '3.2', '3.3'],
         'url': 'http://pypi.python.org/pypi/Django',
         'version': '1.5.4'}
    """
    spec = Spec.objects.get(pk=spec_pk)

    # query PyPI
    pypi_api = pypi.PyPIWrapper()
    pkg_data = pypi_api.get_short_info(spec.name, spec.version)
    #latest_pkg_data = pypi_api.get_short_info(spec.name)

    spec.release_date = pkg_data['last_release_date']
    spec.python_versions = pkg_data['python3_supported_versions']
    spec.save()

    # Canonical package name case for the win
    if pkg_data['name'] != spec.name:
        package = spec.package
        package.name = pkg_data['name']
        package.save()

    return pkg_data

@task
def search_github(package_name, url):
    """ Where are my sources, bro?

        Sometimes, package name is not the same as github repo name,
        but that should not be a problem, since Github search API
        is smart enough.

        >>> github.GithubSearchWrapper().get_most_popular_repo('moscowdjango')
        ('moscowdjango', 'futurecolors')
    """
    if url:
        repo_info = parse_github_url(url)
        if repo_info:
            # query directly to github account
            return repo_info['repo_name'], repo_info['owner']

    repo_name, owner = github.GithubSearchWrapper().get_most_popular_repo(package_name)
    return repo_name, owner

@task
def get_short_info(repo_tuple):
    """ What's with my sources, bro?

        >>> github.GithubWrapper().get_short_info('futurecolors', 'moscowdjango')
        {'updated_at': datetime.datetime(2013, 9, 15, 6, 30, 37),
         'url': 'https://api.github.com/repos/futurecolors/moscowdjango'}
    """
    repo_name, owner = repo_tuple  # FIXME
    return github.GithubWrapper().get_short_info(owner, repo_name)

@task
def get_forks(repo_tuple):
    repo_name, owner = repo_tuple
    return github.GithubWrapper().get_py3_fork_info(owner, repo_name)

@task
def get_issues(repo_tuple):
    repo_name, owner = repo_tuple
    return github.GithubWrapper().get_py3_issues_info(owner, repo_name)

@task
def get_pr(repo_tuple):
    repo_name, owner = repo_tuple
    return github.GithubWrapper().get_py3_pull_requests(owner, repo_name)

@task
def get_build_status(repo_tuple):
    """ How are my tests going? """
    repo_name, owner = repo_tuple
    travis_data = travis.TravisCI().get_build_status(owner, repo_name)

    return travis_data

@task
def notify_github_completed(r, package_pk):
    """ Update github package info in database"""
    package = Package.objects.get(pk=package_pk)
    package.repo_last_commit_date = r[0].get('updated_at')
    package.repo_url = r[0].get('url')

    # For now we'll take the first one available
    package.pr_url = r[1][0].get('url', '') if r[1] else ''
    package.pr_status = r[1][0].get('state') if r[1] else 'unknown'

    package.issue_url = r[2][0].get('url', '') if r[2] else ''
    package.issue_status = r[2][0].get('state') if r[2] else 'unknown'

    package.fork_url = r[3][0].pop(0, {}).get('url', '') if r[3] else ''

    # Travis
    package.ci_url = r[3].get('url', '') if r[4] else ''
    package.ci_status = r[3].get('last_build_status', 'unknown') if r[4] else 'unknown'
    package.save()
    return r

@task
def query_github(results, package_pk):
    """ Get all relevant info form Github"""
    package = Package.objects.get(pk=package_pk)

    # GH API queries running in parallel for speedup
    gh_queries = group(get_short_info.s(),
                       get_pr.s(),
                       get_issues.s(),
                       get_forks.s(),
                       get_build_status.s())

    # guessing github url and querying info
    return chain(search_github.s(package.name, results.get('url')),
                 gh_queries,
                 notify_github_completed.s(package_pk)).delay()
<|endoftext|>