<commit_before>from ..preprocessing import filter_and_normalize, moments
from . import velocity, velocity_graph, velocity_embedding

def run_all(data, basis=None, mode='deterministic', min_counts=3, n_pcs=30, n_neighbors=30, copy=False):
    from time import time
    start = time()

    adata = data.copy() if copy else data
    filter_and_normalize(adata, min_counts=min_counts)
    print("Number of genes to be used:", adata.X.shape[1])
    moments(adata, n_neighbors=n_neighbors, n_pcs=n_pcs)

    velocity(adata, mode=mode)
    print("Number of genes to be used:", adata.var.velocity_genes.sum())

    velocity_graph(adata)
    if basis is not None: velocity_embedding(adata, basis=basis)

    print('Total time (seconds): ' + str(round(time() - start, 2)))

    return adata if copy else None

def convert_to_loom(adata, basis=None):
    from scipy.sparse import issparse
    import numpy as np
    import velocyto

    class VelocytoLoom(velocyto.VelocytoLoom):
        def __init__(self, adata, basis=None):
            self.S = adata.layers['spliced'].T
            self.U = adata.layers['unspliced'].T
            self.S = self.S.A if issparse(self.S) else self.S
            self.U = self.U.A if issparse(self.U) else self.U

            if 'ambiguous' in adata.layers.keys():
                self.A = adata.layers['ambiguous'].T
                if issparse(self.A): self.A = self.A.A

            self.initial_cell_size = self.S.sum(0)
            self.initial_Ucell_size = self.U.sum(0)

            self.ca = dict()
            self.ra = dict()
            for key in adata.obs.keys():
                self.ca[key] = np.array(adata.obs[key].values)
            for key in adata.var.keys():
                self.ra[key] = np.array(adata.var[key].values)
            if 'X_' + basis in adata.obsm.keys():
                self.ts = adata.obsm['X_' + basis]

            if 'clusters' in self.ca:
                self.set_clusters(self.ca['clusters'])
            elif 'louvain' in self.ca:
                self.set_clusters(self.ca['louvain'])

        def filter_and_normalize(self, min_counts=3, min_counts_u=3, n_top_genes=None):
            # counterpart to scv.pp.filter_and_normalize()
            self.score_detection_levels(min_expr_counts=min_counts, min_expr_counts_U=min_counts_u,
                                        min_cells_express=0, min_cells_express_U=0)
            self.filter_genes(by_detection_levels=True)

            self.normalize('both', size=True, log=False)
            if n_top_genes is not None:
                self.score_cv_vs_mean(n_top_genes)
                self.filter_genes(by_cv_vs_mean=True)

            self.normalize_by_total()
            print("Number of genes to be used:", self.S.shape[0])

        def impute(self, n_pcs=30, n_neighbors=30):
            # counterpart to scv.pp.moments()
            self.perform_PCA(n_components=n_pcs)
            self.knn_imputation(k=n_neighbors, n_pca_dims=n_pcs)
            self.normalize_median()

        def velocity(self, fit_offset=False):
            self.fit_gammas(fit_offset=fit_offset)
            self.filter_genes_good_fit()

            self.predict_U()
            self.calculate_velocity()
            self.calculate_shift()
            self.extrapolate_cell_at_t()
            print("Number of genes to be used:", self.S.shape[0])

        def velocity_graph(self, n_neighbors_graph=100):
            if not hasattr(self, 'ts'):
                raise ValueError('Compute embedding first.')
            else:
                # counterpart to scv.tl.velocity_graph()
                self.estimate_transition_prob(hidim="Sx_sz", embed="ts", transform="linear",
                                              n_neighbors=n_neighbors_graph, knn_random=True, sampled_fraction=1)

                # counterpart to scv.tl.velocity_embedding()
                self.calculate_embedding_shift(sigma_corr=0.1, expression_scaling=True)

        def run_all(self, min_counts=3, min_counts_u=3, n_pcs=30, n_neighbors=30, n_neighbors_graph=100,
                    n_top_genes=None, fit_offset=False):
            from time import time
            start = time()

            self.filter_and_normalize(min_counts=min_counts, min_counts_u=min_counts_u, n_top_genes=n_top_genes)
            print('Preprocessing: ' + str(round(time() - start, 2)))
            timestamp = time()

            self.impute(n_pcs=n_pcs, n_neighbors=n_neighbors)
            print('Imputation: ' + str(round(time() - timestamp, 2)))
            timestamp = time()

            self.velocity(fit_offset=fit_offset)
            print('Velocity Estimation: ' + str(round(time() - timestamp, 2)))
            timestamp = time()

            self.velocity_graph(n_neighbors_graph=n_neighbors_graph)
            print('Velocity Graph: ' + str(round(time() - timestamp, 2)))

            print('Total: ' + str(round(time() - start, 2)))

    return VelocytoLoom(adata, basis)
<commit_msg>for better comparison<commit_after>from ..preprocessing import filter_and_normalize, moments
from . import velocity, velocity_graph, velocity_embedding

def run_all(data, basis=None, mode='deterministic', min_counts=3, n_pcs=30, n_neighbors=30, copy=False):
    from time import time
    start = time()

    adata = data.copy() if copy else data
    filter_and_normalize(adata, min_counts=min_counts)
    print("Number of genes to be used:", adata.X.shape[1])
    moments(adata, n_neighbors=n_neighbors, n_pcs=n_pcs)

    velocity(adata, mode=mode)
    print("Number of genes to be used:", adata.var.velocity_genes.sum())

    velocity_graph(adata)
    if basis is not None: velocity_embedding(adata, basis=basis)

    print('Total time (seconds): ' + str(round(time() - start, 2)))

    return adata if copy else None

def convert_to_adata(vlm, basis=None):
    from collections import OrderedDict
    from .. import AnnData
    X = vlm.S_norm.T if hasattr(vlm, 'S_norm') else vlm.S_sz.T if hasattr(vlm, 'S_sz') else vlm.S.T  # .sparse().T.tocsr()

    layers = OrderedDict()
    layers['spliced'] = vlm.S_sz.T if hasattr(vlm, 'S_sz') else vlm.S.T
    layers['unspliced'] = vlm.U_sz.T if hasattr(vlm, 'U_sz') else vlm.U.T

    if hasattr(vlm, 'velocity'): layers['velocity'] = vlm.velocity.T
    if hasattr(vlm, 'Sx_sz'): layers['Ms'] = vlm.Sx_sz.T
    if hasattr(vlm, 'Ux_sz'): layers['Mu'] = vlm.Ux_sz.T

    obs = dict(vlm.ca)
    if 'CellID' in obs.keys(): obs['obs_names'] = obs.pop('CellID')

    var = dict(vlm.ra)
    if 'Gene' in var.keys(): var['var_names'] = var.pop('Gene')
    if hasattr(vlm, 'q'): var['velocity_offset'] = vlm.q
    if hasattr(vlm, 'gammas'): var['velocity_gamma'] = vlm.gammas
    if hasattr(vlm, 'R2'): var['velocity_r2'] = vlm.R2

    adata = AnnData(X, obs=obs, var=var, layers=layers)
    if basis is not None and hasattr(vlm, 'ts'): adata.obsm['X_' + basis] = vlm.ts

    return adata

def convert_to_loom(adata, basis=None):
    from scipy.sparse import issparse
    import numpy as np
    import velocyto

    class VelocytoLoom(velocyto.VelocytoLoom):
        def __init__(self, adata, basis=None):
            self.S = adata.layers['spliced'].T
            self.U = adata.layers['unspliced'].T
            self.S = self.S.A if issparse(self.S) else self.S
            self.U = self.U.A if issparse(self.U) else self.U

            if 'ambiguous' in adata.layers.keys():
                self.A = adata.layers['ambiguous'].T
                if issparse(self.A): self.A = self.A.A

            self.initial_cell_size = self.S.sum(0)
            self.initial_Ucell_size = self.U.sum(0)

            self.ca = dict()
            self.ra = dict()

            self.ca['CellID'] = np.array(adata.obs_names)
            self.ra['Gene'] = np.array(adata.var_names)

            for key in adata.obs.keys():
                self.ca[key] = np.array(adata.obs[key].values)
            for key in adata.var.keys():
                self.ra[key] = np.array(adata.var[key].values)
            if isinstance(basis, str) and 'X_' + basis in adata.obsm.keys():
                self.ts = adata.obsm['X_' + basis]

            if 'clusters' in self.ca:
                self.set_clusters(self.ca['clusters'])
            elif 'louvain' in self.ca:
                self.set_clusters(self.ca['louvain'])

        def filter_and_normalize(self, min_counts=3, min_counts_u=3, n_top_genes=None):
            # counterpart to scv.pp.filter_and_normalize()
            self.score_detection_levels(min_expr_counts=min_counts, min_expr_counts_U=min_counts_u,
                                        min_cells_express=0, min_cells_express_U=0)
            self.filter_genes(by_detection_levels=True)

            if n_top_genes is not None and n_top_genes < self.S.shape[0]:
                self.score_cv_vs_mean(n_top_genes)
                self.filter_genes(by_cv_vs_mean=True)

            self._normalize_S(relative_size=self.initial_cell_size, target_size=np.median(self.initial_cell_size))
            self._normalize_U(relative_size=self.initial_Ucell_size, target_size=np.median(self.initial_Ucell_size))

            #self.normalize_by_total()
            print("Number of genes to be used:", self.S.shape[0])

        def impute(self, n_pcs=30, n_neighbors=30):
            # counterpart to scv.pp.moments()
            self.perform_PCA(n_components=n_pcs)
            self.knn_imputation(k=n_neighbors, n_pca_dims=n_pcs)
            self.normalize_median()

        def velocity_estimation(self, limit_gamma=False, fit_offset=False):
            self.fit_gammas(limit_gamma=limit_gamma, fit_offset=fit_offset)
            self.filter_genes_good_fit()

            self.predict_U()
            self.calculate_velocity()
            self.calculate_shift()
            self.extrapolate_cell_at_t()
            print("Number of genes to be used:", self.S.shape[0])

        def velocity_graph(self, n_neighbors_graph=100, transform='linear', expression_scaling=False):
            if not hasattr(self, 'ts'):
                raise ValueError('Compute embedding first.')
            else:
                # counterpart to scv.tl.velocity_graph()
                self.estimate_transition_prob(hidim="Sx_sz", embed="ts", transform=transform,
                                              n_neighbors=n_neighbors_graph, knn_random=True, sampled_fraction=1)

                # counterpart to scv.tl.velocity_embedding()
                self.calculate_embedding_shift(sigma_corr=0.1, expression_scaling=expression_scaling)

        def run_all(self, min_counts=3, min_counts_u=3, n_pcs=30, n_neighbors=30, n_neighbors_graph=100,
                    n_top_genes=None, fit_offset=False, limit_gamma=False, transform='linear', expression_scaling=False,):
            from time import time
            start = time()

            self.filter_and_normalize(min_counts=min_counts, min_counts_u=min_counts_u, n_top_genes=n_top_genes)
            print('Preprocessing: ' + str(round(time() - start, 2)))
            timestamp = time()

            self.impute(n_pcs=n_pcs, n_neighbors=n_neighbors)
            print('Imputation: ' + str(round(time() - timestamp, 2)))
            timestamp = time()

            self.velocity(limit_gamma=limit_gamma, fit_offset=fit_offset)
            print('Velocity Estimation: ' + str(round(time() - timestamp, 2)))
            timestamp = time()

            self.velocity_graph(n_neighbors_graph=n_neighbors_graph, transform=transform, expression_scaling=expression_scaling)
            print('Velocity Graph: ' + str(round(time() - timestamp, 2)))

            print('Total: ' + str(round(time() - start, 2)))

    return VelocytoLoom(adata, basis)
<|endoftext|>


<commit_before><commit_msg>fixed minor bug in _modelSingleParameters<commit_after><|endoftext|>


<commit_before>from rest_framework.test import APIClient
from rest_framework import status
from django.core.urlresolvers import reverse
from django.test import TestCase
from .factories import UserFactory, ProductFactory, StampFactory, VoucherFactory
from shop.models import Order, OrderDetails, Stamp, Voucher

class ShopAPITestCase(TestCase):
    def setUp(self):
        self.user = UserFactory.create()
        self.client = APIClient()
        self.client.force_authenticate(self.user)

    def test_order_post(self):
        prod_1 = ProductFactory.create(name='Widget', collect_stamp=True)
        prod_2 = ProductFactory.create(name='Gizmo', collect_stamp=False)

        url = reverse('shop-api:shop_orders', )
        data = [{"product": prod_1.id, "quantity": 26},
            {"product": prod_2.id, "quantity": 30}]

        response = self.client.post(url, data, format='json')
        self.assertEqual(response.status_code, status.HTTP_201_CREATED)
        self.assertEqual(response.data['success'], True)
        self.assertEqual(Order.objects.count(), 1)
        self.assertEqual(OrderDetails.objects.count(), 2)
        self.assertEqual(Stamp.objects.count(), 26)
        self.assertEqual(Stamp.objects.filter(redeemed=False).count(), 6)
        self.assertEqual(Voucher.objects.count(), 2)

    def test_order_post_invalid_product(self):
        ProductFactory.create(name='Widget', collect_stamp=True)
        prod_2 = ProductFactory.create(name='Gizmo', collect_stamp=False)

        url = reverse('shop-api:shop_orders', )
        data = [{"product": 1234, "quantity": 26},
            {"product": prod_2.id, "quantity": 30}]

        response = self.client.post(url, data, format='json')
        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)

    def test_get_available_stamps(self):
        StampFactory.create(user=self.user, redeemed=False)
        StampFactory.create(user=self.user, redeemed=False)
        StampFactory.create(user=self.user, redeemed=False)
        StampFactory.create(user=self.user, redeemed=False)
        StampFactory.create(user=self.user, redeemed=True)
        StampFactory.create(user=self.user, redeemed=True)

        url = reverse('shop-api:shop_stamps', )
        response = self.client.get(url, format='json')
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertEqual(response.data['stamps'], 4)

    def test_add_stamp_to_user(self):
        url = reverse('shop-api:shop_stamps', )
        response = self.client.post(url, format='json')
        self.assertEqual(Stamp.objects.count(), 1)
        self.assertEqual(response.data['success'], True)
        self.assertTrue(response.data['stamp'] > 0)
        self.assertEqual(response.status_code, status.HTTP_201_CREATED)

    def test_get_available_vouchers(self):
        VoucherFactory.create(user=self.user, redeemed=False)
        VoucherFactory.create(user=self.user, redeemed=False)
        VoucherFactory.create(user=self.user, redeemed=False)
        VoucherFactory.create(user=self.user, redeemed=True)
        VoucherFactory.create(user=self.user, redeemed=True)

        url = reverse('shop-api:shop_vouchers', )
        response = self.client.get(url, format='json')
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertEqual(response.data['vouchers'], 3)

    def test_add_voucher_to_user(self):
        url = reverse('shop-api:shop_vouchers', )
        response = self.client.post(url, format='json')
        self.assertEqual(Voucher.objects.count(), 1)
        self.assertEqual(response.data['success'], True)
        self.assertTrue(response.data['voucher'] > 0)
        self.assertEqual(response.status_code, status.HTTP_201_CREATED)

    def test_mark_voucher_redeemed(self):
        voucher = VoucherFactory.create(user=self.user)
        url = reverse('shop-api:shop_vouchers', )
        data = {'voucher': voucher.id}
        response = self.client.put(url, data, format='json')
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        voucher_modified = Voucher.objects.get(id=voucher.id)
        self.assertEqual(voucher_modified.redeemed, True)
<commit_msg>Verify that the serializer raises the correct error in case of unknown product<commit_after>from rest_framework.test import APIClient
from rest_framework import status
from django.core.urlresolvers import reverse
from django.test import TestCase
from .factories import UserFactory, ProductFactory, StampFactory, VoucherFactory
from shop.models import Order, OrderDetails, Stamp, Voucher

class ShopAPITestCase(TestCase):
    def setUp(self):
        self.user = UserFactory.create()
        self.client = APIClient()
        self.client.force_authenticate(self.user)

    def test_order_post(self):
        prod_1 = ProductFactory.create(name='Widget', collect_stamp=True)
        prod_2 = ProductFactory.create(name='Gizmo', collect_stamp=False)

        url = reverse('shop-api:shop_orders', )
        data = [{"product": prod_1.id, "quantity": 26},
            {"product": prod_2.id, "quantity": 30}]

        response = self.client.post(url, data, format='json')
        self.assertEqual(response.status_code, status.HTTP_201_CREATED)
        self.assertEqual(response.data['success'], True)
        self.assertEqual(Order.objects.count(), 1)
        self.assertEqual(OrderDetails.objects.count(), 2)
        self.assertEqual(Stamp.objects.count(), 26)
        self.assertEqual(Stamp.objects.filter(redeemed=False).count(), 6)
        self.assertEqual(Voucher.objects.count(), 2)

    def test_order_post_invalid_product(self):
        ProductFactory.create(name='Widget', collect_stamp=True)
        prod_2 = ProductFactory.create(name='Gizmo', collect_stamp=False)

        url = reverse('shop-api:shop_orders', )
        data = [{"product": 1234, "quantity": 26},
            {"product": prod_2.id, "quantity": 30}]

        response = self.client.post(url, data, format='json')
        self.assertEqual(response.status_code, status.HTTP_400_BAD_REQUEST)
        self.assertTrue('Invalid product' in response.content)

    def test_get_available_stamps(self):
        StampFactory.create(user=self.user, redeemed=False)
        StampFactory.create(user=self.user, redeemed=False)
        StampFactory.create(user=self.user, redeemed=False)
        StampFactory.create(user=self.user, redeemed=False)
        StampFactory.create(user=self.user, redeemed=True)
        StampFactory.create(user=self.user, redeemed=True)

        url = reverse('shop-api:shop_stamps', )
        response = self.client.get(url, format='json')
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertEqual(response.data['stamps'], 4)

    def test_add_stamp_to_user(self):
        url = reverse('shop-api:shop_stamps', )
        response = self.client.post(url, format='json')
        self.assertEqual(Stamp.objects.count(), 1)
        self.assertEqual(response.data['success'], True)
        self.assertTrue(response.data['stamp'] > 0)
        self.assertEqual(response.status_code, status.HTTP_201_CREATED)

    def test_get_available_vouchers(self):
        VoucherFactory.create(user=self.user, redeemed=False)
        VoucherFactory.create(user=self.user, redeemed=False)
        VoucherFactory.create(user=self.user, redeemed=False)
        VoucherFactory.create(user=self.user, redeemed=True)
        VoucherFactory.create(user=self.user, redeemed=True)

        url = reverse('shop-api:shop_vouchers', )
        response = self.client.get(url, format='json')
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        self.assertEqual(response.data['vouchers'], 3)

    def test_add_voucher_to_user(self):
        url = reverse('shop-api:shop_vouchers', )
        response = self.client.post(url, format='json')
        self.assertEqual(Voucher.objects.count(), 1)
        self.assertEqual(response.data['success'], True)
        self.assertTrue(response.data['voucher'] > 0)
        self.assertEqual(response.status_code, status.HTTP_201_CREATED)

    def test_mark_voucher_redeemed(self):
        voucher = VoucherFactory.create(user=self.user)
        url = reverse('shop-api:shop_vouchers', )
        data = {'voucher': voucher.id}
        response = self.client.put(url, data, format='json')
        self.assertEqual(response.status_code, status.HTTP_200_OK)
        voucher_modified = Voucher.objects.get(id=voucher.id)
        self.assertEqual(voucher_modified.redeemed, True)
<|endoftext|>


<commit_before>from sponge import app
from sponge.config import config
from sponge.exceptions import InvalidUsage
from flask import jsonify, request, abort
import arrow
import datetime
from orm import db, DbRelease, DbPackage, DbResults
from sponge.util import list_to_string

@app.errorhandler(InvalidUsage)
def handle_invalid_usage(error):
    response = jsonify(error.to_dict())
    response.status_code = error.status_code
    return response

def _create_release(request):
    """
    Create a DbRelease object from a request
    """

    references = request.json.get('references', [])
    if references and type(references) is list:
        references = list_to_string(references)

    platforms = request.json['platforms']
    if platforms and type(platforms) is list:
        platforms = list_to_string(platforms)

    return DbRelease(
        # Required attributes
        platforms=platforms,
        user=request.json['user'],
        # Not required
        notes=request.json.get('notes'),
        team=request.json.get('team'),
        references=references,
    )

def _create_package(release_id, request):
    """
    Create a package object for a release
    """

    return DbPackage(
        release_id,
        request.json['name'],
        request.json['version'],
    )

def _fetch_release(release_id):
    """
    Fetch a release by ID
    """
    rq = db.session.query(DbRelease).filter(DbRelease.id == release_id)
    dbRelease = rq.first()
    if not dbRelease:
        raise InvalidUsage("Release does not exist")
    return dbRelease

def _fetch_package(release_id, package_id):
    """
    Fetch a package, and validate it is part of the release
    """

    # Fetch release first, as it is a good sanity check
    _fetch_release(release_id)
    pq = db.session.query(DbPackage).filter(DbPackage.id == package_id)
    dbPackage = pq.first()

    if not dbPackage:
        raise InvalidUsage("Package does not exist")
    if str(dbPackage.release_id) != release_id:
        raise InvalidUsage("This package does not belong to this release")

    return dbPackage

def _validate_request_json(request):
    try:
        request.json
    except Exception:
        # This is pretty ugly, but we want something more user friendly than "Bad Request"
        raise InvalidUsage("Could not parse JSON document")

    if not request.json:
        raise InvalidUsage('Missing application/json header', status_code=400)

def _validate_release_input(request):
    _validate_request_json(request)
    if 'platforms' not in request.json:
        raise InvalidUsage('JSON doc missing platforms field', status_code=400)
    return True

def _validate_package_input(request, id):
    _validate_request_json(request)

    if not 'name' in request.json or not 'version' in request.json:
        raise InvalidUsage("Missing name / version in request body.")
    app.logger.debug("Request validated, id {}".format(id))
    return True

def _validate_package_stop_input(request):
    _validate_request_json(request)
    if 'success' not in request.json:
        raise InvalidUsage("Missing success key in JSON doc")

@app.route('/ping', methods=['GET'])
def ping():
    return 'pong'

@app.route('/releases', methods=['POST'])
def post_releases():
    """
    Posting a release - the first step in a deployment
    """
    _validate_release_input(request)
    release = _create_release(request)

    app.logger.info(
        'Create release {}, references: {}, platforms: {}'.format(
        release.id, release.notes, release.references, release.platforms)
    )

    release.start()

    db.session.add(release)
    db.session.commit()

    return jsonify(id=release.id)

@app.route('/releases/<release_id>/packages', methods=['POST'])
def post_packages(release_id):
    """
    Add a package to a release
    """
    _validate_package_input(request, id)

    dbRelease = _fetch_release(release_id)
    dbPackage = _create_package(dbRelease.id, request)

    app.logger.info(
        'Create package {}, release {}, name {}, version {}'.format(
            dbPackage.id, dbRelease.id, request.json['name'],
            request.json['version']))

    db.session.add(dbPackage)
    db.session.commit()

    return jsonify(id=dbPackage.id)

@app.route('/releases/<release_id>/packages/<package_id>/results',
           methods=['POST'])
def post_results(release_id, package_id):
    dbResults = DbResults(package_id, str(request.json))
    app.logger.info("Post results, release {}, package {}".format(
        release_id, package_id))
    db.session.add(dbResults)
    db.session.commit()
    return '', 204

@app.route('/releases/<release_id>/stop', methods=['POST'])
def post_releases_stop(release_id):
    """
    Indicate that a release has finished
    """
    dbRelease = _fetch_release(release_id)
    # TODO check that all packages have been finished
    app.logger.info("Release stop, release {}".format(release_id))

    dbRelease.stop()
    return '', 204

@app.route('/releases/<release_id>/packages/<package_id>/start',
           methods=['POST'])
def post_packages_start(release_id, package_id):
    """
    Indicate that a package has started deploying
    """
    dbPackage = _fetch_package(release_id, package_id)
    app.logger.info("Package start, release {}, package {}".format(
        release_id, package_id))
    dbPackage.start()

    db.session.add(dbPackage)
    db.session.commit()
    return '', 204

@app.route('/releases/<release_id>/packages/<package_id>/stop',
           methods=['POST'])
def post_packages_stop(release_id, package_id):
    """
    Indicate that a package has finished deploying
    """
    _validate_request_json(request)
    success = request.json['success']

    dbPackage = _fetch_package(release_id, package_id)
    app.logger.info("Package stop, release {}, package {}, success {}".format(
        release_id, package_id, success))
    dbPackage.stop(success=success)

    db.session.add(dbPackage)
    db.session.commit()
    return '', 204

@app.route('/releases', methods=['GET'])
def get_releases():
    """
    Return a list of releases to the client, filters optional
    """
    query = db.session.query(DbRelease).order_by(DbRelease.stime.asc())

    if 'package_name' in request.args:
        query = query.join(DbPackage).filter(DbPackage.name == request.args['package_name'])
    if 'user' in request.args:
        query = query.filter(DbRelease.user == request.args['user'])
    if 'platform' in request.args:
        query = query.filter(
            DbRelease.platforms.like('%{}%'.format(request.args['platform']))
        )
    if 'stime_before' in request.args:
        t = arrow.get(request.args['stime_before'])
        query = query.filter(DbRelease.stime <= t)
    if 'stime_after' in request.args:
        t = arrow.get(request.args['stime_after'])
        query = query.filter(DbRelease.stime >= t)
    if 'ftime_before' in request.args:
        t = arrow.get(request.args['ftime_before'])
        query = query.filter(DbRelease.ftime <= t)
    if 'ftime_after' in request.args:
        t = arrow.get(request.args['ftime_after'])
        query = query.filter(DbRelease.ftime >= t)
    if 'duration_less' in request.args:
        td = datetime.timedelta(seconds=int(request.args['duration_less']))
        query = query.filter(DbRelease.duration < td)
    if 'duration_greater' in request.args:
        td = datetime.timedelta(seconds=int(request.args['duration_greater']))
        query = query.filter(DbRelease.duration > td)
    if 'team' in request.args:
        query = query.filter(DbRelease.team == request.args['team'])

    releases = query.all()
    app.logger.debug("Returning {} releases".format(len(releases)))
    output = []
    for r in releases:
        output.append(r.to_dict())

    return jsonify(releases=output), 200

<commit_msg>Commit release stop to database<commit_after>from sponge import app
from sponge.config import config
from sponge.exceptions import InvalidUsage
from flask import jsonify, request, abort
import arrow
import datetime
from orm import db, DbRelease, DbPackage, DbResults
from sponge.util import list_to_string

@app.errorhandler(InvalidUsage)
def handle_invalid_usage(error):
    response = jsonify(error.to_dict())
    response.status_code = error.status_code
    return response

def _create_release(request):
    """
    Create a DbRelease object from a request
    """

    references = request.json.get('references', [])
    if references and type(references) is list:
        references = list_to_string(references)

    platforms = request.json['platforms']
    if platforms and type(platforms) is list:
        platforms = list_to_string(platforms)

    return DbRelease(
        # Required attributes
        platforms=platforms,
        user=request.json['user'],
        # Not required
        notes=request.json.get('notes'),
        team=request.json.get('team'),
        references=references,
    )

def _create_package(release_id, request):
    """
    Create a package object for a release
    """

    return DbPackage(
        release_id,
        request.json['name'],
        request.json['version'],
    )

def _fetch_release(release_id):
    """
    Fetch a release by ID
    """
    rq = db.session.query(DbRelease).filter(DbRelease.id == release_id)
    dbRelease = rq.first()
    if not dbRelease:
        raise InvalidUsage("Release does not exist")
    return dbRelease

def _fetch_package(release_id, package_id):
    """
    Fetch a package, and validate it is part of the release
    """

    # Fetch release first, as it is a good sanity check
    _fetch_release(release_id)
    pq = db.session.query(DbPackage).filter(DbPackage.id == package_id)
    dbPackage = pq.first()

    if not dbPackage:
        raise InvalidUsage("Package does not exist")
    if str(dbPackage.release_id) != release_id:
        raise InvalidUsage("This package does not belong to this release")

    return dbPackage

def _validate_request_json(request):
    try:
        request.json
    except Exception:
        # This is pretty ugly, but we want something more user friendly than "Bad Request"
        raise InvalidUsage("Could not parse JSON document")

    if not request.json:
        raise InvalidUsage('Missing application/json header', status_code=400)

def _validate_release_input(request):
    _validate_request_json(request)
    if 'platforms' not in request.json:
        raise InvalidUsage('JSON doc missing platforms field', status_code=400)
    return True

def _validate_package_input(request, id):
    _validate_request_json(request)

    if not 'name' in request.json or not 'version' in request.json:
        raise InvalidUsage("Missing name / version in request body.")
    app.logger.debug("Request validated, id {}".format(id))
    return True

def _validate_package_stop_input(request):
    _validate_request_json(request)
    if 'success' not in request.json:
        raise InvalidUsage("Missing success key in JSON doc")

@app.route('/ping', methods=['GET'])
def ping():
    return 'pong'

@app.route('/releases', methods=['POST'])
def post_releases():
    """
    Posting a release - the first step in a deployment
    """
    _validate_release_input(request)
    release = _create_release(request)

    app.logger.info(
        'Create release {}, references: {}, platforms: {}'.format(
        release.id, release.notes, release.references, release.platforms)
    )

    release.start()

    db.session.add(release)
    db.session.commit()

    return jsonify(id=release.id)

@app.route('/releases/<release_id>/packages', methods=['POST'])
def post_packages(release_id):
    """
    Add a package to a release
    """
    _validate_package_input(request, id)

    dbRelease = _fetch_release(release_id)
    dbPackage = _create_package(dbRelease.id, request)

    app.logger.info(
        'Create package {}, release {}, name {}, version {}'.format(
            dbPackage.id, dbRelease.id, request.json['name'],
            request.json['version']))

    db.session.add(dbPackage)
    db.session.commit()

    return jsonify(id=dbPackage.id)

@app.route('/releases/<release_id>/packages/<package_id>/results',
           methods=['POST'])
def post_results(release_id, package_id):
    dbResults = DbResults(package_id, str(request.json))
    app.logger.info("Post results, release {}, package {}".format(
        release_id, package_id))
    db.session.add(dbResults)
    db.session.commit()
    return '', 204

@app.route('/releases/<release_id>/stop', methods=['POST'])
def post_releases_stop(release_id):
    """
    Indicate that a release has finished
    """
    dbRelease = _fetch_release(release_id)
    # TODO check that all packages have been finished
    app.logger.info("Release stop, release {}".format(release_id))
    dbRelease.stop()

    db.session.add(dbRelease)
    db.session.commit()
    return '', 204

@app.route('/releases/<release_id>/packages/<package_id>/start',
           methods=['POST'])
def post_packages_start(release_id, package_id):
    """
    Indicate that a package has started deploying
    """
    dbPackage = _fetch_package(release_id, package_id)
    app.logger.info("Package start, release {}, package {}".format(
        release_id, package_id))
    dbPackage.start()

    db.session.add(dbPackage)
    db.session.commit()
    return '', 204

@app.route('/releases/<release_id>/packages/<package_id>/stop',
           methods=['POST'])
def post_packages_stop(release_id, package_id):
    """
    Indicate that a package has finished deploying
    """
    _validate_request_json(request)
    success = request.json['success']

    dbPackage = _fetch_package(release_id, package_id)
    app.logger.info("Package stop, release {}, package {}, success {}".format(
        release_id, package_id, success))
    dbPackage.stop(success=success)

    db.session.add(dbPackage)
    db.session.commit()
    return '', 204

@app.route('/releases', methods=['GET'])
def get_releases():
    """
    Return a list of releases to the client, filters optional
    """
    query = db.session.query(DbRelease).order_by(DbRelease.stime.asc())

    if 'package_name' in request.args:
        query = query.join(DbPackage).filter(DbPackage.name == request.args['package_name'])
    if 'user' in request.args:
        query = query.filter(DbRelease.user == request.args['user'])
    if 'platform' in request.args:
        query = query.filter(
            DbRelease.platforms.like('%{}%'.format(request.args['platform']))
        )
    if 'stime_before' in request.args:
        t = arrow.get(request.args['stime_before'])
        query = query.filter(DbRelease.stime <= t)
    if 'stime_after' in request.args:
        t = arrow.get(request.args['stime_after'])
        query = query.filter(DbRelease.stime >= t)
    if 'ftime_before' in request.args:
        t = arrow.get(request.args['ftime_before'])
        query = query.filter(DbRelease.ftime <= t)
    if 'ftime_after' in request.args:
        t = arrow.get(request.args['ftime_after'])
        query = query.filter(DbRelease.ftime >= t)
    if 'duration_less' in request.args:
        td = datetime.timedelta(seconds=int(request.args['duration_less']))
        query = query.filter(DbRelease.duration < td)
    if 'duration_greater' in request.args:
        td = datetime.timedelta(seconds=int(request.args['duration_greater']))
        query = query.filter(DbRelease.duration > td)
    if 'team' in request.args:
        query = query.filter(DbRelease.team == request.args['team'])

    releases = query.all()
    app.logger.debug("Returning {} releases".format(len(releases)))
    output = []
    for r in releases:
        output.append(r.to_dict())

    return jsonify(releases=output), 200

<|endoftext|>


<commit_before>"""
Library module for sprinter. To handle a lot of the typical
features of the library.

"""
import zipfile
import gzip
import logging
import inspect
import imp
import io
import os
import re
import shutil
import subprocess
import tarfile
import tempfile
import urllib2
import urllib
from StringIO import StringIO
from base64 import b64encode

from getpass import getpass
from subprocess import PIPE, STDOUT

from sprinter.formulabase import FormulaBase
from sprinter.exceptions import (CommandMissingException,
                                 BadCredentialsException,
                                 ExtractException)

LOGGER = logging.getLogger('sprinter')

DOMAIN_REGEX = re.compile("^https?://(\w+\.)?\w+\.\w+\/?")
COMMAND_WHITELIST = ["cd"]
BYTE_CHUNKS = 50

def get_formula_class(formula, environment):
    """
    Get the formula name and return an instance The formula path is a
    path to the module. get_formula_class performs reflection to find
    the first class that extends formulabase, and that is the class
    that an instance of it gets returned.
    """
    try:
        r = __recursive_import(formula)
        member_dict = dict(inspect.getmembers(r))
        sprinter_class = FormulaBase
        for v in member_dict.values():
            if inspect.isclass(v) and issubclass(v, FormulaBase) and v != FormulaBase:
                if sprinter_class is FormulaBase:
                    sprinter_class = v
        if sprinter_class is None:
            raise Exception("No formula %s exists in classpath!" % formula)
        return sprinter_class(environment)
    except ImportError as e:
        raise e

def call(command, stdin=None, env=os.environ, cwd=None, shell=False, output_log_level=logging.INFO, logger=LOGGER):
    """ Better, smarter call logic """
    args = command if shell else whitespace_smart_split(command)
    kw = {}
    if not shell and not which(args[0]):
        raise CommandMissingException(args[0])
    if shell:
        kw['shell'] = True
        kw['executable'] = '/bin/bash'
    process = subprocess.Popen(args, stdin=PIPE, stdout=PIPE, stderr=STDOUT,
                               env=env, cwd=cwd, **kw)
    output = process.communicate(input=stdin)[0]
    logger.log(output_log_level, output)
    return (process.returncode, output)

def __process(arg):
    """
    Process args for a bash shell
    """
    # assumes it's wrapped in quotes, or is a flag
    if arg[0] in ["'", '"', '-']:
        return arg
    else:
        return __escape(arg)

def __escape(s):
    """
    Custom escape module for bash
    """
    buf = StringIO()
    for c in s:
        if c in ['(', ')', ':']:
            buf.write("\\" + c)
        else:
            buf.write(c)
    return buf.getvalue()

def whitespace_smart_split(command):
    """
    Split a command by whitespace, taking care to not split on whitespace within quotes.

    >>> __whitespace_smart_split("test this \\\"in here\\\" again")
    ['test', 'this', '"in here"', 'again']
    """
    return_array = []
    s = ""
    in_double_quotes = False
    escape = False
    for c in command:
        if c == '"':
            if in_double_quotes:
                if escape:
                    s += c
                    escape = False
                else:
                    s += c
                    in_double_quotes = False
            else:
                in_double_quotes = True
                s += c
        else:
            if in_double_quotes:
                if c == '\\':
                    escape = True
                    s += c
                else:
                    escape = False
                    s += c
            else:
                if c == ' ':
                    return_array.append(s)
                    s = ""
                else:
                    s += c
    if s != "":
        return_array.append(s)
    return return_array

def authenticated_get(username, password, url):
    """
    Perform an authorized query to the url, and return the result
    """
    try:
        request = urllib2.Request(url)
        base64string = b64encode((b"%s:%s" % (username, password)).decode("ascii"))
        request.add_header("Authorization", "Basic %s" % base64string)
        result = urllib2.urlopen(request)
    except urllib2.HTTPError, e:
        if e.code == 401:
            raise BadCredentialsException(
                "Unable to authenticate user %s to %s with password provided!"
                % (username, url))
    return result.read()

def prompt(prompt_string, default=None, secret=False, boolean=False):
    """
    Prompt user for a string, with a default value

    * secret converts to password prompt
    * boolean converts return value to boolean, checking for starting with a Y
    """
    prompt_string += (" (default %s): " % default if default else ": ")
    if secret:
        val = getpass(prompt_string)
    else:
        val = raw_input(prompt_string)
    val = (val if val else default)
    if boolean:
        val = val.lower().startswith('y')
    return val

def __recursive_import(module_name):
    """
    Recursively looks for and imports the names, returning the
    module desired

    >>> __recursive_import("sprinter.formulas.unpack") # doctest: +ELLIPSIS
    <module 'unpack' from '...'>

    currently module with relative imports don't work.
    """
    names = module_name.split(".")
    path = None
    module = None
    while len(names) > 0:
        if module:
            path = module.__path__
        name = names.pop(0)
        (module_file, pathname, description) = imp.find_module(name, path)
        module = imp.load_module(name, module_file, pathname, description)
    return module

def is_affirmative(phrase):
    """
    Determine if a phrase is in the affirmative
    * start with a t or y, case insensitive
    """
    return phrase.lower()[0] in ['t', 'y']

# From:
# http://stackoverflow.com/questions/377017/test-if-executable-exists-in-python
def is_exe(fpath):
    return os.path.isfile(fpath) and os.access(fpath, os.X_OK)

def which(program):
    if program in COMMAND_WHITELIST:
        return True
    fpath, fname = os.path.split(program)
    if fpath:
        if is_exe(program):
            return program
    else:
        for path in os.environ["PATH"].split(os.pathsep):
            path = path.strip('"')
            exe_file = os.path.join(path, program)
            if is_exe(exe_file):
                return exe_file

    return None

def extract_targz(url, target_dir, remove_common_prefix=False, overwrite=False):
    """ extract a targz and install to the target directory """
    try:
        gz = gzip.GzipFile(fileobj=StringIO(urllib.urlopen(url).read()))
        tf = tarfile.TarFile(fileobj=gz)
        if not os.path.exists(target_dir):
            os.makedirs(target_dir)
        common_prefix = os.path.commonprefix(tf.getnames())
        if not common_prefix.endswith('/'):
            common_prefix += "/"
        for tfile in tf.getmembers():
            if remove_common_prefix:
                tfile.name = tfile.name.replace(common_prefix, "", 1)
            if tfile.name != "":
                target_path = os.path.join(tfile, target_dir)
                if target_path != target_dir and os.path.exists(target_path):
                    if overwrite:
                        remove_path(target_path)
                    else:
                        return
                tf.extract(tfile, target_dir)
    except OSError:
        raise ExtractException()
    except IOError:
        raise ExtractException()

def extract_zip(url, target_dir, remove_common_prefix=False, overwrite=False):
    try:
        memory_file = io.BytesIO(urllib.urlopen(url).read())
        zip_file = zipfile.ZipFile(memory_file)
        common_prefix = os.path.commonprefix(zip_file.namelist())
        for zip_file_info in zip_file.infolist():
            target_path = zip_file_info.filename
            if remove_common_prefix:
                target_path = target_path.replace(common_prefix, "", 1)
            if target_path != "":
                target_path = os.path.join(target_dir, target_path)
                if target_path != target_dir and os.path.exists(target_path):
                    if overwrite:
                        remove_path(target_path)
                    else:
                        return
                zip_file.extract(zip_file_info, target_path)
    except OSError:
        raise ExtractException()
    except IOError:
        raise ExtractException()

def extract_dmg(url, target_dir, remove_common_prefix=False, overwrite=False):
    if remove_common_prefix:
        raise("Remove common prefix for zip not implemented yet!")
    tmpdir = tempfile.mkdtemp()
    try:
        temp_file = os.path.join(tmpdir, "temp.dmg")
        urllib.urlretrieve(url, temp_file)
        call("hdiutil attach %s -mountpoint /Volumes/a/" % temp_file)
        for f in os.listdir("/Volumes/a/"):
            if not f.startswith(".") and f != ' ':
                source_path = os.path.join("/Volumes/a", f)
                target_path = os.path.join(target_dir, f)
                if target_path != target_dir and os.path.exists(target_path):
                    if overwrite:
                        remove_path(target_path)
                    else:
                        return
                if os.path.isdir(source_path):
                    shutil.copytree(source_path, target_path)
                else:
                    shutil.copy(source_path, target_path)
    except OSError:
        raise ExtractException()
    except IOError:
        raise ExtractException()
    finally:
        call("hdiutil unmount /Volumes/a")
        shutil.rmtree(tmpdir)

def remove_path(target_path):
    """ Delete the target path """
    if os.path.isdir(target_path):
        shutil.rmtree(target_path)
    else:
        os.unlink(target_path)

def _determine_overwrite(prompt, overwrite, path):
    """ Determines if an overwrite is desired """
    if overwrite:
        return True
    elif prompt:
        return prompt("Path %s already exist! Overwrite?" % path, boolean=True)
    return True

if __name__ == '__main__':
    import doctest
    doctest.testmod()
<commit_msg>Removing bash, doesn't seem to fix bugs<commit_after>"""
Library module for sprinter. To handle a lot of the typical
features of the library.

"""
import zipfile
import gzip
import logging
import inspect
import imp
import io
import os
import re
import shutil
import subprocess
import tarfile
import tempfile
import urllib2
import urllib
from StringIO import StringIO
from base64 import b64encode

from getpass import getpass
from subprocess import PIPE, STDOUT

from sprinter.formulabase import FormulaBase
from sprinter.exceptions import (CommandMissingException,
                                 BadCredentialsException,
                                 ExtractException)

LOGGER = logging.getLogger('sprinter')

DOMAIN_REGEX = re.compile("^https?://(\w+\.)?\w+\.\w+\/?")
COMMAND_WHITELIST = ["cd"]
BYTE_CHUNKS = 50

def get_formula_class(formula, environment):
    """
    Get the formula name and return an instance The formula path is a
    path to the module. get_formula_class performs reflection to find
    the first class that extends formulabase, and that is the class
    that an instance of it gets returned.
    """
    try:
        r = __recursive_import(formula)
        member_dict = dict(inspect.getmembers(r))
        sprinter_class = FormulaBase
        for v in member_dict.values():
            if inspect.isclass(v) and issubclass(v, FormulaBase) and v != FormulaBase:
                if sprinter_class is FormulaBase:
                    sprinter_class = v
        if sprinter_class is None:
            raise Exception("No formula %s exists in classpath!" % formula)
        return sprinter_class(environment)
    except ImportError as e:
        raise e

def call(command, stdin=None, env=os.environ, cwd=None, shell=False, output_log_level=logging.INFO, logger=LOGGER):
    """ Better, smarter call logic """
    args = command if shell else whitespace_smart_split(command)
    kw = {}
    if not shell and not which(args[0]):
        raise CommandMissingException(args[0])
    if shell:
        kw['shell'] = True
    process = subprocess.Popen(args, stdin=PIPE, stdout=PIPE, stderr=STDOUT,
                               env=env, cwd=cwd, **kw)
    output = process.communicate(input=stdin)[0]
    logger.log(output_log_level, output)
    return (process.returncode, output)

def __process(arg):
    """
    Process args for a bash shell
    """
    # assumes it's wrapped in quotes, or is a flag
    if arg[0] in ["'", '"', '-']:
        return arg
    else:
        return __escape(arg)

def __escape(s):
    """
    Custom escape module for bash
    """
    buf = StringIO()
    for c in s:
        if c in ['(', ')', ':']:
            buf.write("\\" + c)
        else:
            buf.write(c)
    return buf.getvalue()

def whitespace_smart_split(command):
    """
    Split a command by whitespace, taking care to not split on whitespace within quotes.

    >>> __whitespace_smart_split("test this \\\"in here\\\" again")
    ['test', 'this', '"in here"', 'again']
    """
    return_array = []
    s = ""
    in_double_quotes = False
    escape = False
    for c in command:
        if c == '"':
            if in_double_quotes:
                if escape:
                    s += c
                    escape = False
                else:
                    s += c
                    in_double_quotes = False
            else:
                in_double_quotes = True
                s += c
        else:
            if in_double_quotes:
                if c == '\\':
                    escape = True
                    s += c
                else:
                    escape = False
                    s += c
            else:
                if c == ' ':
                    return_array.append(s)
                    s = ""
                else:
                    s += c
    if s != "":
        return_array.append(s)
    return return_array

def authenticated_get(username, password, url):
    """
    Perform an authorized query to the url, and return the result
    """
    try:
        request = urllib2.Request(url)
        base64string = b64encode((b"%s:%s" % (username, password)).decode("ascii"))
        request.add_header("Authorization", "Basic %s" % base64string)
        result = urllib2.urlopen(request)
    except urllib2.HTTPError, e:
        if e.code == 401:
            raise BadCredentialsException(
                "Unable to authenticate user %s to %s with password provided!"
                % (username, url))
    return result.read()

def prompt(prompt_string, default=None, secret=False, boolean=False):
    """
    Prompt user for a string, with a default value

    * secret converts to password prompt
    * boolean converts return value to boolean, checking for starting with a Y
    """
    prompt_string += (" (default %s): " % default if default else ": ")
    if secret:
        val = getpass(prompt_string)
    else:
        val = raw_input(prompt_string)
    val = (val if val else default)
    if boolean:
        val = val.lower().startswith('y')
    return val

def __recursive_import(module_name):
    """
    Recursively looks for and imports the names, returning the
    module desired

    >>> __recursive_import("sprinter.formulas.unpack") # doctest: +ELLIPSIS
    <module 'unpack' from '...'>

    currently module with relative imports don't work.
    """
    names = module_name.split(".")
    path = None
    module = None
    while len(names) > 0:
        if module:
            path = module.__path__
        name = names.pop(0)
        (module_file, pathname, description) = imp.find_module(name, path)
        module = imp.load_module(name, module_file, pathname, description)
    return module

def is_affirmative(phrase):
    """
    Determine if a phrase is in the affirmative
    * start with a t or y, case insensitive
    """
    return phrase.lower()[0] in ['t', 'y']

# From:
# http://stackoverflow.com/questions/377017/test-if-executable-exists-in-python
def is_exe(fpath):
    return os.path.isfile(fpath) and os.access(fpath, os.X_OK)

def which(program):
    if program in COMMAND_WHITELIST:
        return True
    fpath, fname = os.path.split(program)
    if fpath:
        if is_exe(program):
            return program
    else:
        for path in os.environ["PATH"].split(os.pathsep):
            path = path.strip('"')
            exe_file = os.path.join(path, program)
            if is_exe(exe_file):
                return exe_file

    return None

def extract_targz(url, target_dir, remove_common_prefix=False, overwrite=False):
    """ extract a targz and install to the target directory """
    try:
        gz = gzip.GzipFile(fileobj=StringIO(urllib.urlopen(url).read()))
        tf = tarfile.TarFile(fileobj=gz)
        if not os.path.exists(target_dir):
            os.makedirs(target_dir)
        common_prefix = os.path.commonprefix(tf.getnames())
        if not common_prefix.endswith('/'):
            common_prefix += "/"
        for tfile in tf.getmembers():
            if remove_common_prefix:
                tfile.name = tfile.name.replace(common_prefix, "", 1)
            if tfile.name != "":
                target_path = os.path.join(tfile, target_dir)
                if target_path != target_dir and os.path.exists(target_path):
                    if overwrite:
                        remove_path(target_path)
                    else:
                        return
                tf.extract(tfile, target_dir)
    except OSError:
        raise ExtractException()
    except IOError:
        raise ExtractException()

def extract_zip(url, target_dir, remove_common_prefix=False, overwrite=False):
    try:
        memory_file = io.BytesIO(urllib.urlopen(url).read())
        zip_file = zipfile.ZipFile(memory_file)
        common_prefix = os.path.commonprefix(zip_file.namelist())
        for zip_file_info in zip_file.infolist():
            target_path = zip_file_info.filename
            if remove_common_prefix:
                target_path = target_path.replace(common_prefix, "", 1)
            if target_path != "":
                target_path = os.path.join(target_dir, target_path)
                if target_path != target_dir and os.path.exists(target_path):
                    if overwrite:
                        remove_path(target_path)
                    else:
                        return
                zip_file.extract(zip_file_info, target_path)
    except OSError:
        raise ExtractException()
    except IOError:
        raise ExtractException()

def extract_dmg(url, target_dir, remove_common_prefix=False, overwrite=False):
    if remove_common_prefix:
        raise("Remove common prefix for zip not implemented yet!")
    tmpdir = tempfile.mkdtemp()
    try:
        temp_file = os.path.join(tmpdir, "temp.dmg")
        urllib.urlretrieve(url, temp_file)
        call("hdiutil attach %s -mountpoint /Volumes/a/" % temp_file)
        for f in os.listdir("/Volumes/a/"):
            if not f.startswith(".") and f != ' ':
                source_path = os.path.join("/Volumes/a", f)
                target_path = os.path.join(target_dir, f)
                if target_path != target_dir and os.path.exists(target_path):
                    if overwrite:
                        remove_path(target_path)
                    else:
                        return
                if os.path.isdir(source_path):
                    shutil.copytree(source_path, target_path)
                else:
                    shutil.copy(source_path, target_path)
    except OSError:
        raise ExtractException()
    except IOError:
        raise ExtractException()
    finally:
        call("hdiutil unmount /Volumes/a")
        shutil.rmtree(tmpdir)

def remove_path(target_path):
    """ Delete the target path """
    if os.path.isdir(target_path):
        shutil.rmtree(target_path)
    else:
        os.unlink(target_path)

def _determine_overwrite(prompt, overwrite, path):
    """ Determines if an overwrite is desired """
    if overwrite:
        return True
    elif prompt:
        return prompt("Path %s already exist! Overwrite?" % path, boolean=True)
    return True

if __name__ == '__main__':
    import doctest
    doctest.testmod()
<|endoftext|>


<commit_before>#!/usr/bin/python

import re
import sys

def main(argv):
    examples = {}
    requires = set()
    for filename in argv[1:]:
        lines = open(filename, 'rU').readlines()
        if len(lines) > 0 and lines[0].startswith('// NOCOMPILE'):
            continue
        requires.update(line for line in lines if line.startswith('goog.require'))
        examples[filename] = [line for line in lines if \
            not (line.startswith('goog.require') or \
                 line.startswith('goog.provide'))]
    for require in sorted(requires):
        sys.stdout.write(require)
    for filename in sorted(examples.keys()):
        sys.stdout.write('// ' + filename + '\n')
        sys.stdout.write('(function(){\n')
        for line in examples[filename]:
            sys.stdout.write(line)
        sys.stdout.write('})();\n')

if __name__ == '__main__':
    sys.exit(main(sys.argv))
<commit_msg>Remove unused import<commit_after>#!/usr/bin/python

import sys

def main(argv):
    examples = {}
    requires = set()
    for filename in argv[1:]:
        lines = open(filename, 'rU').readlines()
        if len(lines) > 0 and lines[0].startswith('// NOCOMPILE'):
            continue
        requires.update(line for line in lines if line.startswith('goog.require'))
        examples[filename] = [line for line in lines if \
            not (line.startswith('goog.require') or \
                 line.startswith('goog.provide'))]
    for require in sorted(requires):
        sys.stdout.write(require)
    for filename in sorted(examples.keys()):
        sys.stdout.write('// ' + filename + '\n')
        sys.stdout.write('(function(){\n')
        for line in examples[filename]:
            sys.stdout.write(line)
        sys.stdout.write('})();\n')

if __name__ == '__main__':
    sys.exit(main(sys.argv))
<|endoftext|>


<commit_before>"""Custom implementation of multiprocessing.Pool with custom pickler

This module provides efficient ways of working with data stored in
shared memory with numpy.memmap arrays without inducing any memory
copy between the parent and child processes.

This module should not be imported if multiprocessing is not
available as it implements subclasses of multiprocessing Pool
that uses a custom alternative to SimpleQueue.

"""
# Author: Olivier Grisel <olivier.grisel@ensta.org>
# Copyright: 2012, Olivier Grisel
# License: BSD 3 clause

import os
import sys
import threading
import atexit
import tempfile
import shutil
from cPickle import loads
from cPickle import dumps
from pickle import Pickler
from pickle import HIGHEST_PROTOCOL
try:
    from io import BytesIO
except ImportError:
    # Python 2.5 compat
    from StringIO import StringIO as BytesIO
try:
    from multiprocessing.pool import Pool
    from multiprocessing import Pipe
    from multiprocessing.synchronize import Lock
    from multiprocessing.forking import assert_spawning
except ImportError:
    class Pool(object):
        """Dummy class for python 2.5 backward compat"""
        pass
try:
    import numpy as np
except ImportError:
    np = None

from .numpy_pickle import load
from .numpy_pickle import dump
from .hashing import hash

###############################################################################
# Support for efficient transient pickling of numpy data structures

def reduce_memmap(a):
    """Pickle the descriptors of a memmap instance to reopen on same file"""
    mode = a.mode
    if mode == 'w+':
        # Do not make the subprocess erase the data from the parent memmap
        # inadvertently
        mode = 'r+'
    if a.flags['F_CONTIGUOUS']:
        order ='F'
    elif a.flags['C_CONTIGUOUS']:
        order = 'C'
    else:
        # This should never happen as I don't believe that a memmaped
        # array can be other than F or C contiguous
        raise NotImplementedError(
            'Arrays that are neither C nor F contiguous cannot be '
            'memmaped to disk.'
            )
    return (np.memmap, (a.filename, a.dtype, mode, a.offset, a.shape, order))

class ArrayMemmapReducer(object):
    """Reducer callable to dump large arrays to memmap files.

    Parameters
    ----------
    max_nbytes: int
        Threshold to trigger memmaping of large arrays to files created
        a folder.
    temp_folder: str
        Path of a folder where files for backing memmaped arrays are created.
    mmap_mode: 'r', 'r+' or 'c'
        Mode for the created memmap datastructure. See the documentation of
        numpy.memmap for more details. Note: 'w+' is coerced to 'r+'
        automatically to avoid zeroing the data on unpickling.
    verbose: int, optional, 0 by default
        If verbose > 0, memmap creations are logged.
        If verbose > 1, both memmap creations, reuse and array pickling are
        logged.

    """

    def __init__(self, max_nbytes, temp_folder, mmap_mode, verbose=0):
        self._max_nbytes = max_nbytes
        self._temp_folder = temp_folder
        self._mmap_mode = mmap_mode
        self.verbose = int(verbose)

    def __call__(self, a):
        if a.nbytes > self._max_nbytes:
            # check that the folder exists (lazily create the pool temp folder
            # if required)
            try:
                os.makedirs(self._temp_folder)
            except OSError, e:
                if e.errno != 17:
                    # Errno 17 corresponds to 'Directory exists'
                    raise e

            # Find a unique, concurrent safe filename for writing the
            # content of this array only once.
            basename = "%d-%d-%d-%s.pkl" % (
                os.getpid(), id(threading.current_thread()), id(a), hash(a))
            filename = os.path.join(self._temp_folder, basename)

            # In case the same array with the same content is passed several
            # times to the pool subprocess children, serialize it only once
            if not os.path.exists(filename):
                if self.verbose > 0:
                    print "Memmaping (shape=%r, dtype=%s) to new file %s" % (
                        a.shape, a.dtype, filename)
                dump(a, filename)
            elif self.verbose > 1:
                print "Memmaping (shape=%s, dtype=%s) to old file %s" % (
                        a.shape, a.dtype, filename)

            # Let's use the memmap reducer
            return reduce_memmap(load(filename, mmap_mode=self._mmap_mode))
        else:
            # do not convert a into memmap, let pickler do its usual copy with
            # the default system pickler
            if self.verbose > 1:
                print "Pickling array (shape=%r, dtype=%s)." % (
                    a.shape, a.dtype)
            return (loads, (dumps(a, protocol=HIGHEST_PROTOCOL),))

###############################################################################
# Enable custom pickling in Pool queues

class CustomizablePickler(Pickler):
    """Pickler that accepts custom reducers.

    HIGHEST_PROTOCOL is selected by default as this pickler is used
    to pickle ephemeral datastructures for interprocess communication
    hence no backward compatibility is required.

    `reducers` is expected to be a sequence of `(type, callable)`
    pairs where `callable` is a function that give an instance of
    `type` will return a tuple `(constructor, tuple_of_objects)`
    to rebuild an instance out of the pickled `tuple_of_objects`
    as would return a `__reduce__` method. See the standard library
    documentation on pickling for more details.

    """

    def __init__(self, writer, reducers=(), protocol=HIGHEST_PROTOCOL):
        Pickler.__init__(self, writer, protocol=protocol)
        # Make the dispatch registry an instance level attribute instead of a
        # reference to the class dictionary
        self.dispatch = Pickler.dispatch.copy()
        for type, reduce_func in reducers:
            self.register(type, reduce_func)

    def register(self, type, reduce_func):
        def dispatcher(self, obj):
            reduced = reduce_func(obj)
            self.save_reduce(obj=obj, *reduced)
        self.dispatch[type] = dispatcher

class CustomizablePicklingQueue(object):
    """Locked Pipe implementation that uses a customizable pickler.

    This class is an alternative to the multiprocessing implementation
    of SimpleQueue in order to make it possible to pass custom
    pickling reducers, for instance to avoid memory copy when passing
    memmory mapped datastructures.

    `reducers` is expected to be a sequence of `(type, callable)`
    pairs where `callable` is a function that give an instance of
    `type` will return a tuple `(constructor, tuple_of_objects)`
    to rebuild an instance out of the pickled `tuple_of_objects`
    as would return a `__reduce__` method. See the standard library
    documentation on pickling for more details.

    """

    def __init__(self, reducers=()):
        self._reducers = reducers
        self._reader, self._writer = Pipe(duplex=False)
        self._rlock = Lock()
        if sys.platform == 'win32':
            self._wlock = None
        else:
            self._wlock = Lock()
        self._make_methods()

    def __getstate__(self):
        assert_spawning(self)
        return (self._reader, self._writer, self._rlock, self._wlock,
                self._reducers)

    def __setstate__(self, state):
        (self._reader, self._writer, self._rlock, self._wlock,
         self._reducers) = state
        self._make_methods()

    def empty(self):
        return not self._reader.poll()

    def _make_methods(self):
        self._recv = recv = self._reader.recv
        racquire, rrelease = self._rlock.acquire, self._rlock.release

        def get():
            racquire()
            try:
                return recv()
            finally:
                rrelease()

        self.get = get

        if self._reducers:
            def send(obj):
                buffer = BytesIO()
                CustomizablePickler(buffer, self._reducers).dump(obj)
                self._writer.send_bytes(buffer.getvalue())
            self._send = send
        else:
            self._send = send = self._writer.send
        if self._wlock is None:
            # writes to a message oriented win32 pipe are atomic
            self.put = send
        else:
            wlock_acquire, wlock_release = self._wlock.acquire, \
                                                self._wlock.release

            def put(obj):
                wlock_acquire()
                try:
                    return send(obj)
                finally:
                    wlock_release()

            self.put = put

class PicklingPool(Pool):
    """Pool implementation with customizable pickling reducers.

    This is useful to control how data is shipped between processes
    and makes it possible to use shared memory without useless
    copies induces by the default pickling methods of the original
    objects passed as arguments to dispatch.

    `forward_reducers` and `backward_reducers` are expected to be
    sequences of `(type, callable)` pairs where `callable` is a
    function that give an instance of `type` will return a tuple
    `(constructor, tuple_of_objects)` to rebuild an instance out
    of the pickled `tuple_of_objects` as would return a `__reduce__`
    method. See the standard library documentation on pickling for
    more details.

    """

    def __init__(self, processes=None, initializer=None, initargs=(),
                 forward_reducers=(), backward_reducers=()):
        self._forward_reducers = forward_reducers
        self._backward_reducers = backward_reducers
        super(PicklingPool, self).__init__(processes=processes,
                                           initializer=initializer,
                                           initargs=initargs)

    def _setup_queues(self):
        self._inqueue = CustomizablePicklingQueue(self._forward_reducers)
        self._outqueue = CustomizablePicklingQueue(self._backward_reducers)
        self._quick_put = self._inqueue._send
        self._quick_get = self._outqueue._recv

class MemmapingPool(PicklingPool):
    """Process pool that shares large arrays to avoid memory copy.

    This drop-in replacement for `multiprocessing.pool.Pool` makes
    it possible to work efficiently with shared memory in a numpy
    context.

    Existing instances of numpy.memmap are preserved: the child
    suprocesses will have access to the same shared memory in the
    original mode except for the 'w+' mode that is automatically
    transformed as 'r+' to avoid zeroing the original data upon
    instanciation.

    Furthermore large arrays from the parent process are automatically
    dumped to a temporary folder on the filesystem such as child
    processes to access their content via memmaping (file system
    backed shared memory).

    Note: it is important to call the terminate method to collect
    the temporary folder used by the pool.

    Parameters
    ----------
    processes: int, optional
        Number of worker processes running concurrently in the pool.
    initializer: callable, optional
        Callable executed on worker process creation.
    initargs: tuple, optional
        Arguments passed to the initializer callable.
    temp_folder: str, optional
        Folder to be used by the pool for memmaping large arrays
        for sharing memory with worker processes. If None, this
        will use the system temporary folder or can be overridden
        with TMP, TMPDIR or TEMP environment variables.
    max_nbytes int or None, optional, 1e6 by default
        Threshold on the size of arrays passed to the workers that
        triggers automated memmory mapping in temp_folder.
        Use None to disable memmaping of large arrays.
    forward_reducers: sequence of tuples (see bellow), optional
        Reducers used to pickle objects passed from master to worker
        processes.
    backward_reducers: sequence of tuples (see bellow), optional
        Reducers used to pickle return values from workers back to the
        master process.
    verbose: int, optional
        Make it possible to monitor how the communication of numpy arrays
        with the subprocess is handled (pickling or memmaping)

    `forward_reducers` and `backward_reducers` are expected to be
    sequences of `(type, callable)` pairs where `callable` is a
    function that give an instance of `type` will return a tuple
    `(constructor, tuple_of_objects)` to rebuild an instance out
    of the pickled `tuple_of_objects` as would return a `__reduce__`
    method. See the standard library documentation on pickling for
    more details.

    """

    def __init__(self, processes=None, initializer=None, initargs=(),
                 temp_folder=None, max_nbytes=1e6, mmap_mode='c',
                 forward_reducers=(), backward_reducers=(),
                 verbose=0):
        # TODO: use dicts instead of sequences of pairs
        f_reducers = []
        b_reducers = []

        # Prepare a subfolder name for the serialization of this particular
        # pool instance (do not create in advance to spare FS write access if
        # no array is to be dumped):
        if temp_folder is None:
            temp_folder = tempfile.gettempdir()
        temp_folder = os.path.abspath(os.path.expanduser(temp_folder))
        self._temp_folder = temp_folder = os.path.join(
            temp_folder, "joblib_memmaping_pool_%d_%d" % (
                os.getpid(), id(self)))

        # Register the garbage collector at program exit in case caller
        # forget's to call it earlier
        atexit.register(self._collect_tempfile)

        if np is not None:
            if max_nbytes is not None:
                reduce_ndarray = ArrayMemmapReducer(
                    max_nbytes, temp_folder, mmap_mode, verbose)
                # We only register the automatic array to memmap reducer in
                # the forward direction (from parent to child processes)
                f_reducers.append((np.ndarray, reduce_ndarray))
            f_reducers.append((np.memmap, reduce_memmap))

            # Communication from child process to the parent process always
            # pickles numpy.ndarray without dumping them as memmap to avoid
            # confusing the caller and make it tricky to collect the temporary
            # folder
            b_reducers.append((np.memmap, reduce_memmap))
        f_reducers.extend(forward_reducers)
        b_reducers.extend(backward_reducers)

        super(MemmapingPool, self).__init__(processes=processes,
                                            initializer=initializer,
                                            initargs=initargs,
                                            forward_reducers=f_reducers,
                                            backward_reducers=b_reducers)

    def _collect_tempfile(self):
        if os.path.exists(self._temp_folder):
            shutil.rmtree(self._temp_folder, ignore_errors=True)

    def terminate(self):
        super(MemmapingPool, self).terminate()
        self._collect_tempfile()
<commit_msg>MISC: reducers: sequence -> dict<commit_after>"""Custom implementation of multiprocessing.Pool with custom pickler

This module provides efficient ways of working with data stored in
shared memory with numpy.memmap arrays without inducing any memory
copy between the parent and child processes.

This module should not be imported if multiprocessing is not
available as it implements subclasses of multiprocessing Pool
that uses a custom alternative to SimpleQueue.

"""
# Author: Olivier Grisel <olivier.grisel@ensta.org>
# Copyright: 2012, Olivier Grisel
# License: BSD 3 clause

import os
import sys
import threading
import atexit
import tempfile
import shutil
from cPickle import loads
from cPickle import dumps
from pickle import Pickler
from pickle import HIGHEST_PROTOCOL
try:
    from io import BytesIO
except ImportError:
    # Python 2.5 compat
    from StringIO import StringIO as BytesIO
try:
    from multiprocessing.pool import Pool
    from multiprocessing import Pipe
    from multiprocessing.synchronize import Lock
    from multiprocessing.forking import assert_spawning
except ImportError:
    class Pool(object):
        """Dummy class for python 2.5 backward compat"""
        pass
try:
    import numpy as np
except ImportError:
    np = None

from .numpy_pickle import load
from .numpy_pickle import dump
from .hashing import hash

###############################################################################
# Support for efficient transient pickling of numpy data structures

def reduce_memmap(a):
    """Pickle the descriptors of a memmap instance to reopen on same file"""
    mode = a.mode
    if mode == 'w+':
        # Do not make the subprocess erase the data from the parent memmap
        # inadvertently
        mode = 'r+'
    if a.flags['F_CONTIGUOUS']:
        order ='F'
    elif a.flags['C_CONTIGUOUS']:
        order = 'C'
    else:
        # This should never happen as I don't believe that a memmaped
        # array can be other than F or C contiguous
        raise NotImplementedError(
            'Arrays that are neither C nor F contiguous cannot be '
            'memmaped to disk.'
            )
    return (np.memmap, (a.filename, a.dtype, mode, a.offset, a.shape, order))

class ArrayMemmapReducer(object):
    """Reducer callable to dump large arrays to memmap files.

    Parameters
    ----------
    max_nbytes: int
        Threshold to trigger memmaping of large arrays to files created
        a folder.
    temp_folder: str
        Path of a folder where files for backing memmaped arrays are created.
    mmap_mode: 'r', 'r+' or 'c'
        Mode for the created memmap datastructure. See the documentation of
        numpy.memmap for more details. Note: 'w+' is coerced to 'r+'
        automatically to avoid zeroing the data on unpickling.
    verbose: int, optional, 0 by default
        If verbose > 0, memmap creations are logged.
        If verbose > 1, both memmap creations, reuse and array pickling are
        logged.

    """

    def __init__(self, max_nbytes, temp_folder, mmap_mode, verbose=0):
        self._max_nbytes = max_nbytes
        self._temp_folder = temp_folder
        self._mmap_mode = mmap_mode
        self.verbose = int(verbose)

    def __call__(self, a):
        if a.nbytes > self._max_nbytes:
            # check that the folder exists (lazily create the pool temp folder
            # if required)
            try:
                os.makedirs(self._temp_folder)
            except OSError, e:
                if e.errno != 17:
                    # Errno 17 corresponds to 'Directory exists'
                    raise e

            # Find a unique, concurrent safe filename for writing the
            # content of this array only once.
            basename = "%d-%d-%d-%s.pkl" % (
                os.getpid(), id(threading.current_thread()), id(a), hash(a))
            filename = os.path.join(self._temp_folder, basename)

            # In case the same array with the same content is passed several
            # times to the pool subprocess children, serialize it only once
            if not os.path.exists(filename):
                if self.verbose > 0:
                    print "Memmaping (shape=%r, dtype=%s) to new file %s" % (
                        a.shape, a.dtype, filename)
                dump(a, filename)
            elif self.verbose > 1:
                print "Memmaping (shape=%s, dtype=%s) to old file %s" % (
                        a.shape, a.dtype, filename)

            # Let's use the memmap reducer
            return reduce_memmap(load(filename, mmap_mode=self._mmap_mode))
        else:
            # do not convert a into memmap, let pickler do its usual copy with
            # the default system pickler
            if self.verbose > 1:
                print "Pickling array (shape=%r, dtype=%s)." % (
                    a.shape, a.dtype)
            return (loads, (dumps(a, protocol=HIGHEST_PROTOCOL),))

###############################################################################
# Enable custom pickling in Pool queues

class CustomizablePickler(Pickler):
    """Pickler that accepts custom reducers.

    HIGHEST_PROTOCOL is selected by default as this pickler is used
    to pickle ephemeral datastructures for interprocess communication
    hence no backward compatibility is required.

    `reducers` is expected expected to be a dictionary with key/values
    being `(type, callable)` pairs where `callable` is a function that
    give an instance of `type` will return a tuple `(constructor,
    tuple_of_objects)` to rebuild an instance out of the pickled
    `tuple_of_objects` as would return a `__reduce__` method. See the
    standard library documentation on pickling for more details.
    """

    def __init__(self, writer, reducers=None, protocol=HIGHEST_PROTOCOL):
        Pickler.__init__(self, writer, protocol=protocol)
        # Make the dispatch registry an instance level attribute instead of a
        # reference to the class dictionary
        self.dispatch = Pickler.dispatch.copy()
        for type, reduce_func in reducers.iteritems():
            self.register(type, reduce_func)

    def register(self, type, reduce_func):
        def dispatcher(self, obj):
            reduced = reduce_func(obj)
            self.save_reduce(obj=obj, *reduced)
        self.dispatch[type] = dispatcher

class CustomizablePicklingQueue(object):
    """Locked Pipe implementation that uses a customizable pickler.

    This class is an alternative to the multiprocessing implementation
    of SimpleQueue in order to make it possible to pass custom
    pickling reducers, for instance to avoid memory copy when passing
    memmory mapped datastructures.

    `reducers` is expected expected to be a dictionary with key/values
    being `(type, callable)` pairs where `callable` is a function that
    give an instance of `type` will return a tuple `(constructor,
    tuple_of_objects)` to rebuild an instance out of the pickled
    `tuple_of_objects` as would return a `__reduce__` method. See the
    standard library documentation on pickling for more details.
    """

    def __init__(self, reducers=None):
        self._reducers = reducers
        self._reader, self._writer = Pipe(duplex=False)
        self._rlock = Lock()
        if sys.platform == 'win32':
            self._wlock = None
        else:
            self._wlock = Lock()
        self._make_methods()

    def __getstate__(self):
        assert_spawning(self)
        return (self._reader, self._writer, self._rlock, self._wlock,
                self._reducers)

    def __setstate__(self, state):
        (self._reader, self._writer, self._rlock, self._wlock,
         self._reducers) = state
        self._make_methods()

    def empty(self):
        return not self._reader.poll()

    def _make_methods(self):
        self._recv = recv = self._reader.recv
        racquire, rrelease = self._rlock.acquire, self._rlock.release

        def get():
            racquire()
            try:
                return recv()
            finally:
                rrelease()

        self.get = get

        if self._reducers:
            def send(obj):
                buffer = BytesIO()
                CustomizablePickler(buffer, self._reducers).dump(obj)
                self._writer.send_bytes(buffer.getvalue())
            self._send = send
        else:
            self._send = send = self._writer.send
        if self._wlock is None:
            # writes to a message oriented win32 pipe are atomic
            self.put = send
        else:
            wlock_acquire, wlock_release = self._wlock.acquire, \
                                                self._wlock.release

            def put(obj):
                wlock_acquire()
                try:
                    return send(obj)
                finally:
                    wlock_release()

            self.put = put

class PicklingPool(Pool):
    """Pool implementation with customizable pickling reducers.

    This is useful to control how data is shipped between processes
    and makes it possible to use shared memory without useless
    copies induces by the default pickling methods of the original
    objects passed as arguments to dispatch.

    `forward_reducers` and `backward_reducers` are expected to be
    dictionaries with key/values being `(type, callable)` pairs where
    `callable` is a function that give an instance of `type` will return
    a tuple `(constructor, tuple_of_objects)` to rebuild an instance out
    of the pickled `tuple_of_objects` as would return a `__reduce__`
    method. See the standard library documentation on pickling for more
    details.

    """

    def __init__(self, processes=None, initializer=None, initargs=(),
                 forward_reducers=None, backward_reducers=None):
        if forward_reducers is None:
            forward_reducers = dict()
        if backward_reducers is None:
            backward_reducers = dict()
        self._forward_reducers = forward_reducers
        self._backward_reducers = backward_reducers
        super(PicklingPool, self).__init__(processes=processes,
                                           initializer=initializer,
                                           initargs=initargs)

    def _setup_queues(self):
        self._inqueue = CustomizablePicklingQueue(self._forward_reducers)
        self._outqueue = CustomizablePicklingQueue(self._backward_reducers)
        self._quick_put = self._inqueue._send
        self._quick_get = self._outqueue._recv

class MemmapingPool(PicklingPool):
    """Process pool that shares large arrays to avoid memory copy.

    This drop-in replacement for `multiprocessing.pool.Pool` makes
    it possible to work efficiently with shared memory in a numpy
    context.

    Existing instances of numpy.memmap are preserved: the child
    suprocesses will have access to the same shared memory in the
    original mode except for the 'w+' mode that is automatically
    transformed as 'r+' to avoid zeroing the original data upon
    instanciation.

    Furthermore large arrays from the parent process are automatically
    dumped to a temporary folder on the filesystem such as child
    processes to access their content via memmaping (file system
    backed shared memory).

    Note: it is important to call the terminate method to collect
    the temporary folder used by the pool.

    Parameters
    ----------
    processes: int, optional
        Number of worker processes running concurrently in the pool.
    initializer: callable, optional
        Callable executed on worker process creation.
    initargs: tuple, optional
        Arguments passed to the initializer callable.
    temp_folder: str, optional
        Folder to be used by the pool for memmaping large arrays
        for sharing memory with worker processes. If None, this
        will use the system temporary folder or can be overridden
        with TMP, TMPDIR or TEMP environment variables.
    max_nbytes int or None, optional, 1e6 by default
        Threshold on the size of arrays passed to the workers that
        triggers automated memmory mapping in temp_folder.
        Use None to disable memmaping of large arrays.
    forward_reducers: dictionary, optional
        Reducers used to pickle objects passed from master to worker
        processes: see below.
    backward_reducers: dictionary, optional
        Reducers used to pickle return values from workers back to the
        master process.
    verbose: int, optional
        Make it possible to monitor how the communication of numpy arrays
        with the subprocess is handled (pickling or memmaping)

    `forward_reducers` and `backward_reducers` are expected to be
    dictionaries with key/values being `(type, callable)` pairs where
    `callable` is a function that give an instance of `type` will return
    a tuple `(constructor, tuple_of_objects)` to rebuild an instance out
    of the pickled `tuple_of_objects` as would return a `__reduce__`
    method. See the standard library documentation on pickling for more
    details.

    """

    def __init__(self, processes=None, initializer=None, initargs=(),
                 temp_folder=None, max_nbytes=1e6, mmap_mode='c',
                 forward_reducers=None, backward_reducers=None,
                 verbose=0):
        if forward_reducers is None:
            forward_reducers = dict()
        if backward_reducers is None:
            backward_reducers = dict()

        # Prepare a subfolder name for the serialization of this particular
        # pool instance (do not create in advance to spare FS write access if
        # no array is to be dumped):
        if temp_folder is None:
            temp_folder = tempfile.gettempdir()
        temp_folder = os.path.abspath(os.path.expanduser(temp_folder))
        self._temp_folder = temp_folder = os.path.join(
            temp_folder, "joblib_memmaping_pool_%d_%d" % (
                os.getpid(), id(self)))

        # Register the garbage collector at program exit in case caller
        # forget's to call it earlier
        atexit.register(self._collect_tempfile)

        if np is not None:
            if max_nbytes is not None:
                reduce_ndarray = ArrayMemmapReducer(
                    max_nbytes, temp_folder, mmap_mode, verbose)
                # We only register the automatic array to memmap reducer in
                # the forward direction (from parent to child processes)
                forward_reducers[np.ndarray] = reduce_ndarray
            forward_reducers[np.memmap] = reduce_memmap

            # Communication from child process to the parent process always
            # pickles numpy.ndarray without dumping them as memmap to avoid
            # confusing the caller and make it tricky to collect the temporary
            # folder
            backward_reducers[np.memmap] = reduce_memmap

        super(MemmapingPool, self).__init__(processes=processes,
                                        initializer=initializer,
                                        initargs=initargs,
                                        forward_reducers=forward_reducers,
                                        backward_reducers=backward_reducers)

    def _collect_tempfile(self):
        if os.path.exists(self._temp_folder):
            shutil.rmtree(self._temp_folder, ignore_errors=True)

    def terminate(self):
        super(MemmapingPool, self).terminate()
        self._collect_tempfile()
<|endoftext|>