<commit_before>import unittest

import numpy
try:
    import scipy.sparse  # NOQA
    scipy_available = True
except ImportError:
    scipy_available = False

import cupy
from cupy import testing
from cupyx.scipy import sparse

def _make(xp, sp, dtype):
    data = xp.array([[0, 1, 2], [3, 4, 5]], dtype)
    offsets = xp.array([0, -1], 'i')
    # 0, 0, 0, 0
    # 3, 1, 0, 0
    # 0, 4, 2, 0
    return sp.dia_matrix((data, offsets), shape=(3, 4))

def _make_complex(xp, sp, dtype):
    data = xp.array([[0, 1, 2], [3, 4, 5]], dtype)
    if dtype in [numpy.complex64, numpy.complex128]:
        data = data - 1j
    offsets = xp.array([0, -1], 'i')
    # 0, 0, 0, 0
    # 3 - 1j, 1 - 1j, 0, 0
    # 0, 4 - 1j, 2 - 1j, 0
    return sp.dia_matrix((data, offsets), shape=(3, 4))

def _make_empty(xp, sp, dtype):
    data = xp.array([[]], 'f')
    offsets = xp.array([0], 'i')
    return sp.dia_matrix((data, offsets), shape=(3, 4))

@testing.parameterize(*testing.product({
    'dtype': [numpy.float32, numpy.float64, numpy.complex64, numpy.complex128],
}))
class TestDiaMatrix(unittest.TestCase):

    def setUp(self):
        self.m = _make(cupy, sparse, self.dtype)

    def test_dtype(self):
        self.assertEqual(self.m.dtype, self.dtype)

    def test_data(self):
        self.assertEqual(self.m.data.dtype, self.dtype)
        testing.assert_array_equal(
            self.m.data, cupy.array([[0, 1, 2], [3, 4, 5]], self.dtype))

    def test_offsets(self):
        self.assertEqual(self.m.offsets.dtype, numpy.int32)
        testing.assert_array_equal(
            self.m.offsets, cupy.array([0, -1], self.dtype))

    def test_shape(self):
        self.assertEqual(self.m.shape, (3, 4))

    def test_ndim(self):
        self.assertEqual(self.m.ndim, 2)

    def test_nnz(self):
        self.assertEqual(self.m.nnz, 5)

    def test_conj(self):
        n = _make_complex(cupy, sparse, self.dtype)
        cupy.testing.assert_array_equal(n.conj().data, n.data.conj())

    def test_conjugate(self):
        n = _make_complex(cupy, sparse, self.dtype)
        cupy.testing.assert_array_equal(n.conjugate().data, n.data.conj())

    @unittest.skipUnless(scipy_available, 'requires scipy')
    def test_str(self):
        if numpy.dtype(self.dtype).kind == 'f':
            expect = '''  (1, 1)\t1.0
  (2, 2)\t2.0
  (1, 0)\t3.0
  (2, 1)\t4.0'''
        else:
            expect = '''  (1, 1)\t(1+0j)
  (2, 2)\t(2+0j)
  (1, 0)\t(3+0j)
  (2, 1)\t(4+0j)'''
        self.assertEqual(str(self.m), expect)

    def test_toarray(self):
        m = self.m.toarray()
        expect = [
            [0, 0, 0, 0],
            [3, 1, 0, 0],
            [0, 4, 2, 0]
        ]
        self.assertTrue(m.flags.c_contiguous)
        cupy.testing.assert_allclose(m, expect)

    def test_diagonal(self):
        testing.assert_array_equal(
            self.m.diagonal(-2), cupy.array([0], self.dtype))
        testing.assert_array_equal(
            self.m.diagonal(-1), cupy.array([3, 4], self.dtype))
        testing.assert_array_equal(
            self.m.diagonal(), cupy.array([0, 1, 2], self.dtype))
        testing.assert_array_equal(
            self.m.diagonal(1), cupy.array([0, 0, 0], self.dtype))
        testing.assert_array_equal(
            self.m.diagonal(2), cupy.array([0, 0], self.dtype))
        testing.assert_array_equal(
            self.m.diagonal(3), cupy.array([0], self.dtype))

@testing.parameterize(*testing.product({
    'dtype': [numpy.float32, numpy.float64, numpy.complex64, numpy.complex128],
}))
@unittest.skipUnless(scipy_available, 'requires scipy')
class TestDiaMatrixInit(unittest.TestCase):

    def setUp(self):
        self.shape = (3, 4)

    def data(self, xp):
        return xp.array([[1, 2, 3], [4, 5, 6]], self.dtype)

    def offsets(self, xp):
        return xp.array([0, -1], 'i')

    @testing.numpy_cupy_raises(sp_name='sp', accept_error=ValueError)
    def test_shape_none(self, xp, sp):
        sp.dia_matrix(
            (self.data(xp), self.offsets(xp)), shape=None)

    @testing.numpy_cupy_allclose(sp_name='sp', atol=1e-5)
    def test_intlike_shape(self, xp, sp):
        s = sp.dia_matrix((self.data(xp), self.offsets(xp)),
                          shape=(xp.array(self.shape[0]),
                                 xp.int32(self.shape[1])))
        assert isinstance(s.shape[0], int)
        assert isinstance(s.shape[1], int)
        return s

    @testing.numpy_cupy_raises(sp_name='sp', accept_error=ValueError)
    def test_large_rank_offset(self, xp, sp):
        sp.dia_matrix(
            (self.data(xp), self.offsets(xp)[None]), shape=self.shape)

    @testing.numpy_cupy_raises(sp_name='sp', accept_error=ValueError)
    def test_large_rank_data(self, xp, sp):
        sp.dia_matrix(
            (self.data(xp)[None], self.offsets(xp)), shape=self.shape)

    @testing.numpy_cupy_raises(sp_name='sp', accept_error=ValueError)
    def test_data_offsets_different_size(self, xp, sp):
        offsets = xp.array([0, -1, 1], 'i')
        sp.dia_matrix(
            (self.data(xp), offsets), shape=self.shape)

    @testing.numpy_cupy_raises(sp_name='sp', accept_error=ValueError)
    def test_duplicated_offsets(self, xp, sp):
        offsets = xp.array([1, 1], 'i')
        sp.dia_matrix(
            (self.data(xp), offsets), shape=self.shape)

    @testing.numpy_cupy_equal(sp_name='sp')
    def test_conj(self, xp, sp):
        n = _make_complex(xp, sp, self.dtype)
        cupy.testing.assert_array_equal(n.conj().data, n.data.conj())

@testing.parameterize(*testing.product({
    'make_method': ['_make', '_make_empty'],
    'dtype': [numpy.float32, numpy.float64, numpy.complex64, numpy.complex128],
}))
@unittest.skipUnless(scipy_available, 'requires scipy')
class TestDiaMatrixScipyComparison(unittest.TestCase):

    @property
    def make(self):
        return globals()[self.make_method]

    @testing.numpy_cupy_equal(sp_name='sp')
    def test_nnz_axis(self, xp, sp):
        m = self.make(xp, sp, self.dtype)
        return m.nnz

    @testing.numpy_cupy_raises(sp_name='sp', accept_error=NotImplementedError)
    def test_nnz_axis_not_none(self, xp, sp):
        m = self.make(xp, sp, self.dtype)
        m.getnnz(axis=0)

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_toarray(self, xp, sp):
        m = self.make(xp, sp, self.dtype)
        return m.toarray()

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_A(self, xp, sp):
        m = self.make(xp, sp, self.dtype)
        return m.A

    @testing.numpy_cupy_raises(sp_name='sp')
    def test_sum_tuple_axis(self, xp, sp):
        m = _make(xp, sp, self.dtype)
        m.sum(axis=(0, 1))

    @testing.numpy_cupy_raises(sp_name='sp')
    def test_sum_float_axis(self, xp, sp):
        m = _make(xp, sp, self.dtype)
        m.sum(axis=0.0)

    @testing.numpy_cupy_raises(sp_name='sp')
    def test_sum_too_large_axis(self, xp, sp):
        m = _make(xp, sp, self.dtype)
        m.sum(axis=3)

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_tocoo(self, xp, sp):
        m = self.make(xp, sp, self.dtype)
        return m.tocoo()

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_tocoo_copy(self, xp, sp):
        m = self.make(xp, sp, self.dtype)
        n = m.tocoo(copy=True)
        self.assertIsNot(m.data, n.data)
        return n

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_tocsc(self, xp, sp):
        m = _make(xp, sp, self.dtype)
        return m.tocsc()

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_tocsc_copy(self, xp, sp):
        m = self.make(xp, sp, self.dtype)
        n = m.tocsc(copy=True)
        self.assertIsNot(m.data, n.data)
        return n

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_tocsr(self, xp, sp):
        m = self.make(xp, sp, self.dtype)
        return m.tocsr()

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_tocsr_copy(self, xp, sp):
        m = self.make(xp, sp, self.dtype)
        n = m.tocsr(copy=True)
        self.assertIsNot(m.data, n.data)
        return n

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_transpose(self, xp, sp):
        m = self.make(xp, sp, self.dtype)
        return m.transpose()

    @testing.numpy_cupy_raises(sp_name='sp')
    def test_diagonal_error(self, xp, sp):
        m = _make(xp, sp, self.dtype)
        m.diagonal(k=10)  # out of range

@testing.parameterize(*testing.product({
    'dtype': [numpy.float32, numpy.float64],
    'ret_dtype': [None, numpy.float32, numpy.float64],
    'axis': [None, 0, 1, -1, -2],
}))
@unittest.skipUnless(scipy_available, 'requires scipy')
class TestDiaMatrixSum(unittest.TestCase):

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_sum(self, xp, sp):
        m = _make(xp, sp, self.dtype)
        return m.sum(axis=self.axis, dtype=self.ret_dtype)

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_sum_with_out(self, xp, sp):
        m = _make(xp, sp, self.dtype)
        if self.axis is None:
            shape = ()
        else:
            shape = list(m.shape)
            shape[self.axis] = 1
            shape = tuple(shape)
        out = xp.empty(shape, dtype=self.ret_dtype)
        if xp is numpy:
            # TODO(unno): numpy.matrix is used for scipy.sparse though
            # cupy.ndarray is used for cupyx.scipy.sparse.
            out = xp.asmatrix(out)
        return m.sum(axis=self.axis, dtype=self.ret_dtype, out=out)

class TestIsspmatrixDia(unittest.TestCase):

    def test_dia(self):
        x = sparse.dia_matrix(
            (cupy.array([], 'f'),
             cupy.array([0], 'i')),
            shape=(0, 0), dtype='f')
        self.assertTrue(sparse.isspmatrix_dia(x))

    def test_csr(self):
        x = sparse.csr_matrix(
            (cupy.array([], 'f'),
             cupy.array([], 'i'),
             cupy.array([0], 'i')),
            shape=(0, 0), dtype='f')
        self.assertFalse(sparse.isspmatrix_dia(x))
<commit_msg>diagonal tests using k argument require scipy>=1.0<commit_after>import unittest

import numpy
try:
    import scipy.sparse  # NOQA
    scipy_available = True
except ImportError:
    scipy_available = False

import cupy
from cupy import testing
from cupyx.scipy import sparse

def _make(xp, sp, dtype):
    data = xp.array([[0, 1, 2], [3, 4, 5]], dtype)
    offsets = xp.array([0, -1], 'i')
    # 0, 0, 0, 0
    # 3, 1, 0, 0
    # 0, 4, 2, 0
    return sp.dia_matrix((data, offsets), shape=(3, 4))

def _make_complex(xp, sp, dtype):
    data = xp.array([[0, 1, 2], [3, 4, 5]], dtype)
    if dtype in [numpy.complex64, numpy.complex128]:
        data = data - 1j
    offsets = xp.array([0, -1], 'i')
    # 0, 0, 0, 0
    # 3 - 1j, 1 - 1j, 0, 0
    # 0, 4 - 1j, 2 - 1j, 0
    return sp.dia_matrix((data, offsets), shape=(3, 4))

def _make_empty(xp, sp, dtype):
    data = xp.array([[]], 'f')
    offsets = xp.array([0], 'i')
    return sp.dia_matrix((data, offsets), shape=(3, 4))

@testing.parameterize(*testing.product({
    'dtype': [numpy.float32, numpy.float64, numpy.complex64, numpy.complex128],
}))
class TestDiaMatrix(unittest.TestCase):

    def setUp(self):
        self.m = _make(cupy, sparse, self.dtype)

    def test_dtype(self):
        self.assertEqual(self.m.dtype, self.dtype)

    def test_data(self):
        self.assertEqual(self.m.data.dtype, self.dtype)
        testing.assert_array_equal(
            self.m.data, cupy.array([[0, 1, 2], [3, 4, 5]], self.dtype))

    def test_offsets(self):
        self.assertEqual(self.m.offsets.dtype, numpy.int32)
        testing.assert_array_equal(
            self.m.offsets, cupy.array([0, -1], self.dtype))

    def test_shape(self):
        self.assertEqual(self.m.shape, (3, 4))

    def test_ndim(self):
        self.assertEqual(self.m.ndim, 2)

    def test_nnz(self):
        self.assertEqual(self.m.nnz, 5)

    def test_conj(self):
        n = _make_complex(cupy, sparse, self.dtype)
        cupy.testing.assert_array_equal(n.conj().data, n.data.conj())

    def test_conjugate(self):
        n = _make_complex(cupy, sparse, self.dtype)
        cupy.testing.assert_array_equal(n.conjugate().data, n.data.conj())

    @unittest.skipUnless(scipy_available, 'requires scipy')
    def test_str(self):
        if numpy.dtype(self.dtype).kind == 'f':
            expect = '''  (1, 1)\t1.0
  (2, 2)\t2.0
  (1, 0)\t3.0
  (2, 1)\t4.0'''
        else:
            expect = '''  (1, 1)\t(1+0j)
  (2, 2)\t(2+0j)
  (1, 0)\t(3+0j)
  (2, 1)\t(4+0j)'''
        self.assertEqual(str(self.m), expect)

    def test_toarray(self):
        m = self.m.toarray()
        expect = [
            [0, 0, 0, 0],
            [3, 1, 0, 0],
            [0, 4, 2, 0]
        ]
        self.assertTrue(m.flags.c_contiguous)
        cupy.testing.assert_allclose(m, expect)

    def test_diagonal(self):
        testing.assert_array_equal(
            self.m.diagonal(-2), cupy.array([0], self.dtype))
        testing.assert_array_equal(
            self.m.diagonal(-1), cupy.array([3, 4], self.dtype))
        testing.assert_array_equal(
            self.m.diagonal(), cupy.array([0, 1, 2], self.dtype))
        testing.assert_array_equal(
            self.m.diagonal(1), cupy.array([0, 0, 0], self.dtype))
        testing.assert_array_equal(
            self.m.diagonal(2), cupy.array([0, 0], self.dtype))
        testing.assert_array_equal(
            self.m.diagonal(3), cupy.array([0], self.dtype))

@testing.parameterize(*testing.product({
    'dtype': [numpy.float32, numpy.float64, numpy.complex64, numpy.complex128],
}))
@unittest.skipUnless(scipy_available, 'requires scipy')
class TestDiaMatrixInit(unittest.TestCase):

    def setUp(self):
        self.shape = (3, 4)

    def data(self, xp):
        return xp.array([[1, 2, 3], [4, 5, 6]], self.dtype)

    def offsets(self, xp):
        return xp.array([0, -1], 'i')

    @testing.numpy_cupy_raises(sp_name='sp', accept_error=ValueError)
    def test_shape_none(self, xp, sp):
        sp.dia_matrix(
            (self.data(xp), self.offsets(xp)), shape=None)

    @testing.numpy_cupy_allclose(sp_name='sp', atol=1e-5)
    def test_intlike_shape(self, xp, sp):
        s = sp.dia_matrix((self.data(xp), self.offsets(xp)),
                          shape=(xp.array(self.shape[0]),
                                 xp.int32(self.shape[1])))
        assert isinstance(s.shape[0], int)
        assert isinstance(s.shape[1], int)
        return s

    @testing.numpy_cupy_raises(sp_name='sp', accept_error=ValueError)
    def test_large_rank_offset(self, xp, sp):
        sp.dia_matrix(
            (self.data(xp), self.offsets(xp)[None]), shape=self.shape)

    @testing.numpy_cupy_raises(sp_name='sp', accept_error=ValueError)
    def test_large_rank_data(self, xp, sp):
        sp.dia_matrix(
            (self.data(xp)[None], self.offsets(xp)), shape=self.shape)

    @testing.numpy_cupy_raises(sp_name='sp', accept_error=ValueError)
    def test_data_offsets_different_size(self, xp, sp):
        offsets = xp.array([0, -1, 1], 'i')
        sp.dia_matrix(
            (self.data(xp), offsets), shape=self.shape)

    @testing.numpy_cupy_raises(sp_name='sp', accept_error=ValueError)
    def test_duplicated_offsets(self, xp, sp):
        offsets = xp.array([1, 1], 'i')
        sp.dia_matrix(
            (self.data(xp), offsets), shape=self.shape)

    @testing.numpy_cupy_equal(sp_name='sp')
    def test_conj(self, xp, sp):
        n = _make_complex(xp, sp, self.dtype)
        cupy.testing.assert_array_equal(n.conj().data, n.data.conj())

@testing.parameterize(*testing.product({
    'make_method': ['_make', '_make_empty'],
    'dtype': [numpy.float32, numpy.float64, numpy.complex64, numpy.complex128],
}))
@unittest.skipUnless(scipy_available, 'requires scipy')
class TestDiaMatrixScipyComparison(unittest.TestCase):

    @property
    def make(self):
        return globals()[self.make_method]

    @testing.numpy_cupy_equal(sp_name='sp')
    def test_nnz_axis(self, xp, sp):
        m = self.make(xp, sp, self.dtype)
        return m.nnz

    @testing.numpy_cupy_raises(sp_name='sp', accept_error=NotImplementedError)
    def test_nnz_axis_not_none(self, xp, sp):
        m = self.make(xp, sp, self.dtype)
        m.getnnz(axis=0)

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_toarray(self, xp, sp):
        m = self.make(xp, sp, self.dtype)
        return m.toarray()

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_A(self, xp, sp):
        m = self.make(xp, sp, self.dtype)
        return m.A

    @testing.numpy_cupy_raises(sp_name='sp')
    def test_sum_tuple_axis(self, xp, sp):
        m = _make(xp, sp, self.dtype)
        m.sum(axis=(0, 1))

    @testing.numpy_cupy_raises(sp_name='sp')
    def test_sum_float_axis(self, xp, sp):
        m = _make(xp, sp, self.dtype)
        m.sum(axis=0.0)

    @testing.numpy_cupy_raises(sp_name='sp')
    def test_sum_too_large_axis(self, xp, sp):
        m = _make(xp, sp, self.dtype)
        m.sum(axis=3)

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_tocoo(self, xp, sp):
        m = self.make(xp, sp, self.dtype)
        return m.tocoo()

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_tocoo_copy(self, xp, sp):
        m = self.make(xp, sp, self.dtype)
        n = m.tocoo(copy=True)
        self.assertIsNot(m.data, n.data)
        return n

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_tocsc(self, xp, sp):
        m = _make(xp, sp, self.dtype)
        return m.tocsc()

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_tocsc_copy(self, xp, sp):
        m = self.make(xp, sp, self.dtype)
        n = m.tocsc(copy=True)
        self.assertIsNot(m.data, n.data)
        return n

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_tocsr(self, xp, sp):
        m = self.make(xp, sp, self.dtype)
        return m.tocsr()

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_tocsr_copy(self, xp, sp):
        m = self.make(xp, sp, self.dtype)
        n = m.tocsr(copy=True)
        self.assertIsNot(m.data, n.data)
        return n

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_transpose(self, xp, sp):
        m = self.make(xp, sp, self.dtype)
        return m.transpose()

    @testing.numpy_cupy_raises(sp_name='sp')
    @testing.with_requires('scipy>=1.0')
    def test_diagonal_error(self, xp, sp):
        m = _make(xp, sp, self.dtype)
        m.diagonal(k=10)  # out of range

@testing.parameterize(*testing.product({
    'dtype': [numpy.float32, numpy.float64],
    'ret_dtype': [None, numpy.float32, numpy.float64],
    'axis': [None, 0, 1, -1, -2],
}))
@unittest.skipUnless(scipy_available, 'requires scipy')
class TestDiaMatrixSum(unittest.TestCase):

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_sum(self, xp, sp):
        m = _make(xp, sp, self.dtype)
        return m.sum(axis=self.axis, dtype=self.ret_dtype)

    @testing.numpy_cupy_allclose(sp_name='sp')
    def test_sum_with_out(self, xp, sp):
        m = _make(xp, sp, self.dtype)
        if self.axis is None:
            shape = ()
        else:
            shape = list(m.shape)
            shape[self.axis] = 1
            shape = tuple(shape)
        out = xp.empty(shape, dtype=self.ret_dtype)
        if xp is numpy:
            # TODO(unno): numpy.matrix is used for scipy.sparse though
            # cupy.ndarray is used for cupyx.scipy.sparse.
            out = xp.asmatrix(out)
        return m.sum(axis=self.axis, dtype=self.ret_dtype, out=out)

class TestIsspmatrixDia(unittest.TestCase):

    def test_dia(self):
        x = sparse.dia_matrix(
            (cupy.array([], 'f'),
             cupy.array([0], 'i')),
            shape=(0, 0), dtype='f')
        self.assertTrue(sparse.isspmatrix_dia(x))

    def test_csr(self):
        x = sparse.csr_matrix(
            (cupy.array([], 'f'),
             cupy.array([], 'i'),
             cupy.array([0], 'i')),
            shape=(0, 0), dtype='f')
        self.assertFalse(sparse.isspmatrix_dia(x))
<|endoftext|>


<commit_before>__version__ = '0.3.0'
<commit_msg>re-add (testing)<commit_after>__version__ = '0.3.0'

from emit.router.core import Router, CeleryRouter, RQRouter
<|endoftext|>


<commit_before>"""encbup.

Usage:
  encbup.py [options] --encrypt <outfile> <path> ...
  encbup.py [options] --decrypt <infile> <directory>

Options:
  -h --help                        Show this screen.
  -s --salt=<salt>                 The salt used for the key.
  -p --passphrase=<passphrase>     The encryption passphrase used.
  -r --rounds=<rounds>             The number of rounds used for PBKDF2i [default: 20000].
  -v --verbose                     Say more things.
     --version                     Show version.
"""

import base64
import hmac
import hashlib
import json
import logging
import os
import sys

from Crypto import Random
from Crypto.Protocol.KDF import PBKDF2
from Crypto.Cipher import AES
from docopt import docopt
from getpass import getpass

PROTOCOL_VERSION = 1
BLOCK_SIZE = 64 * 1024

class File:
    def __init__(self, key, root, filename):
        self.block_size = BLOCK_SIZE
        self._root = root
        self._filename = filename
        self._key = key

    def read_blocks(self, block_size=None):
        if block_size is None:
            block_size = self.block_size
        infile = open(self.absolute_path)
        while True:
            block = infile.read(block_size)
            if block:
                yield block
            else:
                return

    def pad(self, data):
        """
        Pad `data` as per PKCS#7.
        """
        length = 16 - (len(data) % 16)
        return data + chr(length) * length

    def hmac(self):
        if getattr(self, "_hmac", None) is None:
            h = hmac.new(self._key.key, digestmod=hashlib.sha512)
            infile = open(self.absolute_path)
            for block in self.read_blocks():
                h.update(block)
            infile.close()
            self._hmac = h
        return self._hmac

    def encrypted(self):
        # Set the IV as the first 16 bytes of the file's HMAC, so the same plaintext
        # always encrypts to the same ciphertext for a given key.
        self._cipher = AES.new(self._key.key, AES.MODE_CBC, IV=self.iv)
        infile = open(self.absolute_path)

        yield self.iv
        for block in self.read_blocks():
            if len(block) < self.block_size:
                yield self._cipher.encrypt(self.pad(block))
                break
            else:
                yield self._cipher.encrypt(block)
        else:
            yield self.pad("")

        infile.close()
        return

    @property
    def iv(self):
        return self.hmac().digest()[:16]

    @property
    def absolute_path(self):
        return os.path.join(self._root, self._filename)

    @property
    def relative_path(self):
        return self._filename

    @property
    def encrypted_absolute_path(self):
        iv = self.hmac().digest()[-16:]
        cipher = AES.new(self._key.key, AES.MODE_CBC, IV=iv)
        return iv + cipher.encrypt(self.pad(self.absolute_path))

    @property
    def encrypted_relative_path(self):
        iv = self.hmac().digest()[-16:]
        cipher = AES.new(self._key.key, AES.MODE_CBC, IV=iv)
        return iv + cipher.encrypt(self.pad(self.relative_path))

    @property
    def encrypted_size(self):
        padding_length = 16 - (self.size % 16)
        return 16 + self.size + padding_length

    @property
    def size(self):
        return os.path.getsize(self.absolute_path)

    @property
    def as_dict(self):
        filename = self.encrypted_absolute_path
        h = hmac.new(self._key.key, digestmod=hashlib.sha512)
        h.update(filename)
        filename_hmac = h.hexdigest()
        return {
            "plaintext_digest": self.hmac().hexdigest(),
            "size": self.encrypted_size,
            "filename": filename.encode("hex"),
            "filename_hmac": filename_hmac,
            }

class Key:
    def __init__(self, passphrase, salt=None, rounds=20000):
        if salt is None:
            salt = self.generate_salt()
        self.salt = salt
        self.rounds = rounds
        self.key = PBKDF2(passphrase, salt, count=rounds)

    def generate_salt(self):
        return base64.b64encode(Random.new().read(8))

    @property
    def hexdigest(self):
        return hashlib.sha512(self.key).hexdigest()

    @property
    def as_dict(self):
        return {"filename": ".encbup.json", "size": len(self.contents)}

    @property
    def contents(self):
        return json.dumps({"digest": self.hexdigest, "salt": self.salt, "rounds": self.rounds})

class Reader:
    def __init__(self, filename, base_dir):
        if filename == "-":
            self.file = sys.stdin
        else:
            self.file = open(filename)
        self._base_dir = base_dir

    def unpad(self, data):
        """
        Unpad `data` as per PKCS#7.
        """
        length = ord(data[-1])
        return data[:-length]

    def process_stream(self, passphrase):
        version = self.file.readline()
        state = 0

        # Read protocol version.
        if int(version) > PROTOCOL_VERSION:
            sys.exit("Unsupported protocol version.")

        # Read keyfile.
        metadata = json.loads(self.file.readline())
        keyfile = json.loads(self.file.read(metadata["size"] + 1))
        key = Key(passphrase, keyfile["salt"], rounds=keyfile["rounds"])
        if key.hexdigest != keyfile["digest"]:
            sys.exit("Wrong passphrase.")

        while True:
            if state == 0:
                # Read metadata.
                metadata = self.file.readline()
                if not metadata:
                    return
                metadata = json.loads(metadata)

                # Decrypt the filename.
                encrypted_filename = metadata["filename"].decode("hex")
                if hmac.new(key.key, msg=encrypted_filename, digestmod=hashlib.sha512).hexdigest() != metadata["filename_hmac"]:
                    sys.exit("Invalid filename HMAC, either it is corrupt or someone has tampered with it.")

                cipher = AES.new(key.key, AES.MODE_CBC, IV=encrypted_filename[:16])
                filename = self.unpad(cipher.decrypt(encrypted_filename[16:]))

                # Strip leading slashes.
                # TODO: Protect against ../../ etc.
                if filename.startswith("/"):
                    filename = filename[1:]
                logging.debug("Decrypting {0}...".format(filename))

                pathname = os.path.join(self._base_dir, filename)
                # Create parent directories.
                if not os.path.exists(os.path.dirname(pathname)):
                    os.makedirs(os.path.dirname(pathname))
                outfile = open(pathname, "wb")

                # Read the IV and prepare the cipher.
                cipher = AES.new(key.key, AES.MODE_CBC, IV=self.file.read(16))
                size = metadata["size"] - 16
                state = 1
            else:
                block_size = min(BLOCK_SIZE, size)
                data = self.file.read(block_size)
                plaintext = cipher.decrypt(data)
                if block_size < BLOCK_SIZE:
                    plaintext = self.unpad(plaintext)
                h = hmac.new(key.key, digestmod=hashlib.sha512)
                h.update(plaintext)
                outfile.write(plaintext)
                size -= block_size
                if size == 0:
                    if h.hexdigest() != metadata["plaintext_digest"]:
                        sys.exit("Invalid file HMAC, either it is corrupt or someone has tampered with it.")
                    self.file.read(1)  # Skip the newline.
                    state = 0
                    outfile.close()

class Writer:
    """
    A class for writing encrypted data to a file (or stdout).
    """
    def __init__(self, filename):
        if filename == "-":
            self.file = sys.stdout
        else:
            self.file = open(filename, "wb")

    def walk(self, path):
        """
        Recurse down the given `path`, returning filenames relative to it, if it is a directory.
        If not, return path itself.
        """
        if not os.path.isdir(path):
            yield os.path.split(path)
        else:
            for dirpath, dirnames, filenames in os.walk(path):
                for filename in filenames:
                    relpath = os.path.relpath(os.path.join(dirpath, filename), path)
                    yield path, relpath

    def write(self, line="", terminate=True):
        """
        Write a newline-terminated line to a file.
        """
        self.file.write(line)
        if terminate:
            self.file.write("\n")

    def write_preamble(self, key):
        self.write(str(PROTOCOL_VERSION))
        self.write(json.dumps(key.as_dict))
        self.write(key.contents)

    def write_files(self, paths):
        for path in paths:
            for root, filename in self.walk(path):
                logging.debug("Encrypting {0}...".format(filename))
                infile = File(key, root, filename)
                self.write(json.dumps(infile.as_dict))
                for block in infile.encrypted():
                    self.write(block, terminate=False)
                # Write a terminating blank line.
                self.write()

if __name__ == '__main__':
    arguments = docopt(__doc__, version='encbup 0.1')

    if arguments["--verbose"]:
        logging.basicConfig(level=logging.DEBUG)

    if arguments["--passphrase"] is None:
        passphrase = getpass("Please enter a passphrase: ")
        passphrase_again = getpass("Please confirm the passphrase: ")
        if passphrase != passphrase_again:
            sys.exit("The passphrases don't match.")
    else:
        passphrase = arguments["--passphrase"]

    try:
        rounds = int(arguments["--rounds"])
    except ValueError:
        sys.exit("Please enter an integer for the value of the rounds.")

    if arguments["--encrypt"]:
        key = Key(passphrase, arguments["--salt"], rounds=rounds)
        writer = Writer(arguments["<outfile>"])
        writer.write_preamble(key)
        writer.write_files(arguments["<path>"])
    elif arguments["--decrypt"]:
        reader = Reader(arguments["<infile>"], arguments["<directory>"])
        reader.process_stream(passphrase)
<commit_msg>Fix oddly stupid error.<commit_after>"""encbup.

Usage:
  encbup.py [options] --encrypt <outfile> <path> ...
  encbup.py [options] --decrypt <infile> <directory>

Options:
  -h --help                        Show this screen.
  -s --salt=<salt>                 The salt used for the key.
  -p --passphrase=<passphrase>     The encryption passphrase used.
  -r --rounds=<rounds>             The number of rounds used for PBKDF2i [default: 20000].
  -v --verbose                     Say more things.
     --version                     Show version.
"""

import base64
import hmac
import hashlib
import json
import logging
import os
import sys

from Crypto import Random
from Crypto.Protocol.KDF import PBKDF2
from Crypto.Cipher import AES
from docopt import docopt
from getpass import getpass

PROTOCOL_VERSION = 1
BLOCK_SIZE = 64 * 1024

class File:
    def __init__(self, key, root, filename):
        self.block_size = BLOCK_SIZE
        self._root = root
        self._filename = filename
        self._key = key

    def read_blocks(self, block_size=None):
        if block_size is None:
            block_size = self.block_size
        infile = open(self.absolute_path)
        while True:
            block = infile.read(block_size)
            if block:
                yield block
            else:
                return

    def pad(self, data):
        """
        Pad `data` as per PKCS#7.
        """
        length = 16 - (len(data) % 16)
        return data + chr(length) * length

    def hmac(self):
        if getattr(self, "_hmac", None) is None:
            h = hmac.new(self._key.key, digestmod=hashlib.sha512)
            infile = open(self.absolute_path)
            for block in self.read_blocks():
                h.update(block)
            infile.close()
            self._hmac = h
        return self._hmac

    def encrypted(self):
        # Set the IV as the first 16 bytes of the file's HMAC, so the same plaintext
        # always encrypts to the same ciphertext for a given key.
        self._cipher = AES.new(self._key.key, AES.MODE_CBC, IV=self.iv)
        infile = open(self.absolute_path)

        yield self.iv
        for block in self.read_blocks():
            if len(block) < self.block_size:
                yield self._cipher.encrypt(self.pad(block))
                break
            else:
                yield self._cipher.encrypt(block)
        else:
            yield self.pad("")

        infile.close()
        return

    @property
    def iv(self):
        return self.hmac().digest()[:16]

    @property
    def absolute_path(self):
        return os.path.join(self._root, self._filename)

    @property
    def relative_path(self):
        return self._filename

    @property
    def encrypted_absolute_path(self):
        iv = self.hmac().digest()[-16:]
        cipher = AES.new(self._key.key, AES.MODE_CBC, IV=iv)
        return iv + cipher.encrypt(self.pad(self.absolute_path))

    @property
    def encrypted_relative_path(self):
        iv = self.hmac().digest()[-16:]
        cipher = AES.new(self._key.key, AES.MODE_CBC, IV=iv)
        return iv + cipher.encrypt(self.pad(self.relative_path))

    @property
    def encrypted_size(self):
        padding_length = 16 - (self.size % 16)
        return 16 + self.size + padding_length

    @property
    def size(self):
        return os.path.getsize(self.absolute_path)

    @property
    def as_dict(self):
        filename = self.encrypted_absolute_path
        h = hmac.new(self._key.key, digestmod=hashlib.sha512)
        h.update(filename)
        filename_hmac = h.hexdigest()
        return {
            "plaintext_digest": self.hmac().hexdigest(),
            "size": self.encrypted_size,
            "filename": filename.encode("hex"),
            "filename_hmac": filename_hmac,
            }

class Key:
    def __init__(self, passphrase, salt=None, rounds=20000):
        if salt is None:
            salt = self.generate_salt()
        self.salt = salt
        self.rounds = rounds
        self.key = PBKDF2(passphrase, salt, count=rounds)

    def generate_salt(self):
        return base64.b64encode(Random.new().read(8))

    @property
    def hexdigest(self):
        return hashlib.sha512(self.key).hexdigest()

    @property
    def as_dict(self):
        return {"filename": ".encbup.json", "size": len(self.contents)}

    @property
    def contents(self):
        return json.dumps({"digest": self.hexdigest, "salt": self.salt, "rounds": self.rounds})

class Reader:
    def __init__(self, filename, base_dir):
        if filename == "-":
            self.file = sys.stdin
        else:
            self.file = open(filename)
        self._base_dir = base_dir

    def unpad(self, data):
        """
        Unpad `data` as per PKCS#7.
        """
        length = ord(data[-1])
        return data[:-length]

    def process_stream(self, passphrase):
        version = self.file.readline()
        state = 0

        # Read protocol version.
        if int(version) > PROTOCOL_VERSION:
            sys.exit("Unsupported protocol version.")

        # Read keyfile.
        metadata = json.loads(self.file.readline())
        keyfile = json.loads(self.file.read(metadata["size"] + 1))
        key = Key(passphrase, keyfile["salt"], rounds=keyfile["rounds"])
        if key.hexdigest != keyfile["digest"]:
            sys.exit("Wrong passphrase.")

        while True:
            if state == 0:
                # Read metadata.
                metadata = self.file.readline()
                if not metadata:
                    return
                metadata = json.loads(metadata)

                # Decrypt the filename.
                encrypted_filename = metadata["filename"].decode("hex")
                if hmac.new(key.key, msg=encrypted_filename, digestmod=hashlib.sha512).hexdigest() != metadata["filename_hmac"]:
                    sys.exit("Invalid filename HMAC, either it is corrupt or someone has tampered with it.")

                cipher = AES.new(key.key, AES.MODE_CBC, IV=encrypted_filename[:16])
                filename = self.unpad(cipher.decrypt(encrypted_filename[16:]))

                # Strip leading slashes.
                # TODO: Protect against ../../ etc.
                if filename.startswith("/"):
                    filename = filename[1:]
                logging.debug("Decrypting {0}...".format(filename))

                pathname = os.path.join(self._base_dir, filename)
                # Create parent directories.
                if not os.path.exists(os.path.dirname(pathname)):
                    os.makedirs(os.path.dirname(pathname))
                outfile = open(pathname, "wb")

                # Read the IV and prepare the cipher.
                cipher = AES.new(key.key, AES.MODE_CBC, IV=self.file.read(16))
                size = metadata["size"] - 16
                state = 1
            else:
                block_size = min(BLOCK_SIZE, size)
                data = self.file.read(block_size)
                plaintext = cipher.decrypt(data)
                if block_size < BLOCK_SIZE:
                    plaintext = self.unpad(plaintext)
                h = hmac.new(key.key, digestmod=hashlib.sha512)
                h.update(plaintext)
                outfile.write(plaintext)
                size -= block_size
                if size == 0:
                    if h.hexdigest() != metadata["plaintext_digest"]:
                        sys.exit("Invalid file HMAC, either it is corrupt or someone has tampered with it.")
                    self.file.read(1)  # Skip the newline.
                    state = 0
                    outfile.close()

class Writer:
    """
    A class for writing encrypted data to a file (or stdout).
    """
    def __init__(self, filename):
        if filename == "-":
            self.file = sys.stdout
        else:
            self.file = open(filename, "wb")

    def walk(self, path):
        """
        Recurse down the given `path`, returning filenames relative to it, if it is a directory.
        If not, return path itself.
        """
        if not os.path.isdir(path):
            yield os.path.split(path)
        else:
            for dirpath, dirnames, filenames in os.walk(path):
                for filename in filenames:
                    relpath = os.path.relpath(os.path.join(dirpath, filename), path)
                    yield path, relpath

    def write(self, line="", terminate=True):
        """
        Write a newline-terminated line to a file.
        """
        self.file.write(line)
        if terminate:
            self.file.write("\n")

    def write_preamble(self, key):
        self.write(str(PROTOCOL_VERSION))
        self.write(json.dumps(key.as_dict))
        self.write(key.contents)

    def write_files(self, key, paths):
        for path in paths:
            for root, filename in self.walk(path):
                logging.debug("Encrypting {0}...".format(filename))
                infile = File(key, root, filename)
                self.write(json.dumps(infile.as_dict))
                for block in infile.encrypted():
                    self.write(block, terminate=False)
                # Write a terminating blank line.
                self.write()

def main():
    arguments = docopt(__doc__, version='encbup 0.1')

    if arguments["--verbose"]:
        logging.basicConfig(level=logging.DEBUG)

    if arguments["--passphrase"] is None:
        passphrase = getpass("Please enter a passphrase: ")
        passphrase_again = getpass("Please confirm the passphrase: ")
        if passphrase != passphrase_again:
            sys.exit("The passphrases don't match.")
    else:
        passphrase = arguments["--passphrase"]

    try:
        rounds = int(arguments["--rounds"])
    except ValueError:
        sys.exit("Please enter an integer for the value of the rounds.")

    if arguments["--encrypt"]:
        key = Key(passphrase, arguments["--salt"], rounds=rounds)
        writer = Writer(arguments["<outfile>"])
        writer.write_preamble(key)
        writer.write_files(key, arguments["<path>"])
    elif arguments["--decrypt"]:
        reader = Reader(arguments["<infile>"], arguments["<directory>"])
        reader.process_stream(passphrase)

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


<commit_before>#!/usr/bin/env python
"""Handles encoding flac files"""
import subprocess
import config

HANDLES=['flac','fla']
DEFTYPE='flac'

def encode(inF, outF, options, metadata):
    cli_options = getCliOptions(inF, outF, options, metadata)

def getCliOptions(inF, outF, options, metadata):
    "Builds up a list of the command line options to pass to flac"
    cli_options = [ ]
    cli_options.append("--output-name=%s"%outF) #Output file
    #Metadata
    if 'Title' in metadata:
        cli_options.append("--tag=title=%s"%metadata['Title'])
    if 'Artist' in metadata:
        cli_options.append("--tag=artist=%s"%metadata['Artist'])
    if 'Album' in metadata:
        cli_options.append("--tag=album=%s"%metadata['Album'])
    if 'Genre' in metadata:
        cli_options.append("--tag=genre=%s"%metadata['Genre'])
    if 'Track Number' in metadata:
        cli_options.append("--tag=tracknumber=%s"%metadata['Track Number'])
    if 'Date' in metadata:
        cli_options.append("--tag=date=%s"%metadata['Date'])
    #Add user's options
    cli_options += options
    #Add input file
    cli_options.append(inF)

    return cli_options
<commit_msg>Wrong method name<commit_after>#!/usr/bin/env python
"""Handles encoding flac files"""
import subprocess
import config

HANDLES=['flac','fla']
DEFTYPE='flac'

def encode(inF, outF, options, metadata):
    #Get all the options to give to flac
    cli_options = getCliOptions(inF, outF, options, metadata)
    #Do the actual encoding
    st = subprocess.call(['flac'] + cli_options)

    if st == 0:
        return True
    else:
        return False

def getCliOptions(inF, outF, options, metadata):
    "Builds up a list of the command line options to pass to flac"
    cli_options = [ ]
    cli_options.append("--output-name=%s"%outF) #Output file
    cli_options.append("--totally-silent") #No output
    cli_options.append("--force") #Overwrite existing files
    #Metadata
    if 'Title' in metadata:
        cli_options.append("--tag=title=%s"%metadata['Title'])
    if 'Artist' in metadata:
        cli_options.append("--tag=artist=%s"%metadata['Artist'])
    if 'Album' in metadata:
        cli_options.append("--tag=album=%s"%metadata['Album'])
    if 'Genre' in metadata:
        cli_options.append("--tag=genre=%s"%metadata['Genre'])
    if 'Track Number' in metadata:
        cli_options.append("--tag=tracknumber=%s"%metadata['Track Number'])
    if 'Date' in metadata:
        cli_options.append("--tag=date=%s"%metadata['Date'])
    #Add user's options
    cli_options += options
    #Add input file
    cli_options.append(inF)

    return cli_options
<|endoftext|>


<commit_before>from .gobject import GObject
from .character import Trainer
from . import meta
from .signals import sighandler

"""
IDÉES

Supprimer méthode action de la Battle
Dans la boucle événementielle du jeu, si un combat est en cours (game.player.battle), celui-ci aura une fonction step appelée à chaque tour, et qui s'occupera de gérer les déplacements du joueur.
Ainsi, c'est la Battle qui s'occupe de spawner des choix aux utilisateurs, et d'attendre une réponse du dresseur / du pokemon via un signal retour.
Ainsi, un joueur réel peut temporiser la réponse en fonction des événements (l'event action sera par exemple associé à un Choice qui fera utiliser une attaque), tandis qu'un PNJ répondra tout de suite

Déroulement standard de la méthod Battle.step:
- return si précédent dresseur n'a pas répondu
- Appel de trainer.battle_action(self, beast) sur le prochain dresseur
Cette méthode est ensuite différentes suivant les dresseurs (Trainer, Player).
Dans le premier cas, elle fait appel à l'IA du dresseur pour lancer un signal.
Dans le second, elle crée une boîte de dialogue Choice associée aux bons signaux.
Les signaux auront comme receveur des méthodes permettant d'utiliser une attaque ou un objet.

Dans un premier temps, ne pas implémenter les déplacements en combat ni les objets, simplement un combat tour par tour.
"""

@meta.apply
class FakeTrainer(Trainer):
    pass

@meta.apply
class Battle(GObject):
    "Battle between trainers"

    __attributes__ = ('trainers', 'beasts')

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        for trainer in self.trainers:
            trainer.battle = self
        self.turn = 0
        self.turn_action = False

    @classmethod
    def from_args(cls, *args):
        args = (
            (obj, obj.beast) if isinstance(obj, Trainer)
            else (FakeTrainer(beast=obj), obj)
            for obj in args
        )
        trainers, beasts = zip(*args)
        return cls(trainers=trainers, beasts=beasts)

    def attack(self, beast, att):
        print(beast.name, 'uses', att.name)
        i = self.beasts.index(beast)
        beast.attack(att, self.beasts[not i])
        self.turn_action = False

    def end(self):
        for trainer in self.trainers:
            trainer.battle = None

    def step(self, game):
        if any(beast.ko for beast in self.beasts if beast):
            self.end()
            return
        if not self.turn_action:
            self.turn_action = True
            self.trainers[self.turn].battle_step(
                self,
                self.beasts[self.turn],
            )
            self.turn = (self.turn + 1) % len(self.trainers)

    @sighandler
    def action(self, game, player):
        pass
<commit_msg>Set turn/waiting as battle attributes<commit_after>from .gobject import GObject
from .character import Trainer
from . import meta
from .signals import sighandler

"""
IDÉES

Supprimer méthode action de la Battle
Dans la boucle événementielle du jeu, si un combat est en cours (game.player.battle), celui-ci aura une fonction step appelée à chaque tour, et qui s'occupera de gérer les déplacements du joueur.
Ainsi, c'est la Battle qui s'occupe de spawner des choix aux utilisateurs, et d'attendre une réponse du dresseur / du pokemon via un signal retour.
Ainsi, un joueur réel peut temporiser la réponse en fonction des événements (l'event action sera par exemple associé à un Choice qui fera utiliser une attaque), tandis qu'un PNJ répondra tout de suite

Déroulement standard de la méthod Battle.step:
- return si précédent dresseur n'a pas répondu
- Appel de trainer.battle_action(self, beast) sur le prochain dresseur
Cette méthode est ensuite différentes suivant les dresseurs (Trainer, Player).
Dans le premier cas, elle fait appel à l'IA du dresseur pour lancer un signal.
Dans le second, elle crée une boîte de dialogue Choice associée aux bons signaux.
Les signaux auront comme receveur des méthodes permettant d'utiliser une attaque ou un objet.

Dans un premier temps, ne pas implémenter les déplacements en combat ni les objets, simplement un combat tour par tour.
"""

@meta.apply
class FakeTrainer(Trainer):
    pass

@meta.apply
class Battle(GObject):
    "Battle between trainers"

    __attributes__ = ('trainers', 'beasts', 'turn', 'waiting')

    def __init__(self, **kwargs):
        kwargs.setdefault('turn', 0)
        kwargs.setdefault('waiting', False)
        super().__init__(**kwargs)
        for trainer in self.trainers:
            trainer.battle = self

    @classmethod
    def from_args(cls, *args):
        args = (
            (obj, obj.beast) if isinstance(obj, Trainer)
            else (FakeTrainer(beast=obj), obj)
            for obj in args
        )
        trainers, beasts = zip(*args)
        return cls(trainers=trainers, beasts=beasts)

    def attack(self, beast, att):
        print(beast.name, 'uses', att.name)
        i = self.beasts.index(beast)
        beast.attack(att, self.beasts[not i])
        self.waiting = False

    def end(self):
        for trainer in self.trainers:
            trainer.battle = None

    def step(self, game):
        if any(beast.ko for beast in self.beasts if beast):
            self.end()
            return
        if not self.waiting:
            self.waiting = True
            self.trainers[self.turn].battle_step(
                self,
                self.beasts[self.turn],
            )
            self.turn = (self.turn + 1) % len(self.trainers)

    @sighandler
    def action(self, game, player):
        pass
<|endoftext|>


<commit_before>#import
import os
import time
import logging
from goatools import obo_parser
from goatools.go_enrichment import GOEnrichmentStudy
import wget
from ftplib import FTP
import gzip
import Bio.UniProt.GOA as GOA
from collections import defaultdict
import pickle
import sqlite3

logging.basicConfig(format='%(levelname)s: %(name)s - %(message)s', level=logging.INFO)
logger = logging.getLogger('TFTA-GOEnrich')

_resource_dir = os.path.dirname(os.path.realpath(__file__)) + '/../enrichment/data'

class GOEnrich:
    def __init__(self, species=None):
        self.go = self.read_go()
        if self.go:
            logger.info('GOEnrich loaded GO data.')

        if not species:
            self.gaf_funcs = self.read_gaf()
        else:
            gaf_uri = '/pub/databases/GO/goa/' + species.upper() + '/goa_' + \
                       species.lower() + '.gaf.gz'
            self.gaf_funcs = self.read_gaf2(gaf_uri=gaf_uri)
        if self.gaf_funcs:
            logger.info("GOEnrich loaded GAF file.")
        else:
            logger.error("GAF file doesn't exist or couldn't be loaded.")

        #background population
        self.pop = self.gaf_funcs.keys()

        self.assoc = self.get_assoc_gene_go()

        #load go_gene.db
        db_file = os.path.join(_resource_dir, 'go_gene.db')
        #check file size to determine if it need regenerate
        if os.path.isfile(db_file):
            statinfo = os.stat(db_file)
        else:
            statinfo = None

        if statinfo and statinfo.st_size > 13000000:
            self.godb = sqlite3.connect(db_file, check_same_thread=False)
            logger.info('GOEnrich loaded go_gene database.')
        else:
            num = self.generate_go2gene_db(db_file)
            if os.path.isfile(db_file):
                self.godb = sqlite3.connect(db_file, check_same_thread=False)
            else:
                logger.error('GOEnrich could not load go_gene database.')
                self.godb = None;

    def read_go(self, go_obo_url=None, data_folder=_resource_dir):
        if not go_obo_url:
            go_obo_url = 'http://purl.obolibrary.org/obo/go/go-basic.obo'

        #read GO data as dictionary
        go_obo = self.download_file_http(go_obo_url, data_folder, 'go-basic.obo')
        if go_obo:
            go = obo_parser.GODag(go_obo)
            return go
        else:
            logger.error("The GO data file doesn't exist. You can manually download it from: {}".format(go_obo_url))
            return None

    def read_gaf(self, gaf_uri=None, gaf_fn=None, data_folder=_resource_dir):
        """
        Download gaf file from http://current.geneontology.org/products/pages/downloads.html
        Annotations for 2019-07-01 release
        """
        if not gaf_uri:
            gaf_uri = "http://geneontology.org/gene-associations/goa_human.gaf.gz"
        if not gaf_fn:
            gaf_fn = gaf_uri.split('/')[-1]

        gaf_file = self.download_file_http(gaf_uri, data_folder, gaf_fn)

        #read the annotations into a dictionary
        fn = os.path.join(_resource_dir, 'gaf_funcs.pickle')
        if os.path.isfile(fn):
            with open(fn, 'rb') as pickle_in:
                gaf_funcs = pickle.load(pickle_in)
        else:
            # File is a gunzip file, so we need to open it in this way
            with gzip.open(gaf_file, 'rt') as gaf_fp:
                gaf_funcs = defaultdict(list)
                # Iterate on each function using Bio.UniProt.GOA library.
                for entry in GOA.gafiterator(gaf_fp):
                    gene_symbol = entry.pop('DB_Object_Symbol')
                    gaf_funcs[gene_symbol].append(entry)

            #save gaf_funcs to file
            with open(fn, 'wb') as pickle_out:
                pickle.dump(gaf_funcs, pickle_out)

        return gaf_funcs

    def get_population(self):
        return self.population

    def get_annotations_keyword(self, keyword):
        """
        Generate a list of annotated genes which have the keyword in their names.

        parameter
        ------------
        keyword: str

        output
        ------------
        annots: dict
        """
        annots = dict()
        for g, entry in self.gaf_funcs.items():
            if keyword.lower() in entry[0]['DB_Object_Name'].lower():
                annots[g] = entry
        return annots

    def get_annotations_genes(self, gene_list):
        """
        Return the annotations for each gene in the gene_list
        """
        annots = {x: self.gaf_funcs[x] for x in gene_list}
        return annots

    def get_go_keyword(self, keyword):
        """
        return GO terms whose names have keyword
        """
        #check if the corresponding file exists
        data_folder = _resource_dir + '/go_files'
        fn = os.path.join(data_folder, keyword + ".pickle")
        if os.path.isfile(fn):
            with open(fn, 'rb') as pickle_in:
                go_name = pickle.load(pickle_in)
        else:
            go_name = self._go_keyword(keyword, data_folder)
        return go_name

    def _go_keyword(self, keyword, data_folder):
        """
        Generate a dict, its key is GO_id, value is GO name which contains the keyword
        """
        #data_folder = _resource_dir + '/go_files'
        if not os.path.isfile(data_folder):
            # Emulate mkdir -p (no error if folder exists)
            try:
                os.mkdir(data_folder)
            except OSError as e:
                if(e.errno != 17):
                    logger.error('GOEnrich could not make the go_files folder.')
                    return None
        else:
            logger.error('Data path (' + data_folder + ') exists as a file. '
                         'Please rename, remove or change the desired location of the data path.')
            return None

        go_name = dict()
        keyword = keyword.lower()
        if self.go:
            go_name = {go_id:self.go[go_id].name for go_id in self.go if keyword in self.go[go_id].name.lower()}
            #save to file
            fn = os.path.join(data_folder, keyword + ".pickle")
            with open(fn, 'wb') as pickle_out:
                pickle.dump(go_name, pickle_out)
        return go_name

    def go_enrichment_analysis(self, study, pop=None, assoc=None, go=None, propagate_counts=True,
                               alpha=0.05, methods=None, p_bonferroni=0.01, limit=30):
        """
        GO enrichment analysis

        output
        ----------
        res: list of GOEnrichmentRecord
        """
        if not pop:
            pop = self.pop
        if not assoc:
            assoc = self.assoc
        if not go:
            go = self.go

        #methods = ["bonferroni", "sidak", "holm", "fdr"]
        #In order to reduce delay, only use bonferroni by default
        if not methods:
            methods = ['bonferroni']
        g = GOEnrichmentStudy(pop, assoc, go,
                         propagate_counts=propagate_counts,
                         alpha=alpha,
                         methods=methods)
        g_res = g.run_study(study)

        #only return results with pvalue < pvalue
        res = []
        num = 1
        for x in g_res:
            if num > limit:
                break
            if x.p_bonferroni <= p_bonferroni:
                res.append(x)
                num += 1
        return res

    def get_assoc_gene_go(self):
        assoc = defaultdict(set)
        for x in self.gaf_funcs:
            for entry in self.gaf_funcs[x]:
                assoc[x].add(entry['GO_ID'])
        return assoc

    def download_file_http(self, url, data_folder, file_name):
        # Check if we have the ./data directory already
        if not os.path.isfile(data_folder):
            # Emulate mkdir -p (no error if folder exists)
            try:
                os.mkdir(data_folder)
            except OSError as e:
                if(e.errno != 17):
                    logger.error('GOEnrich could not make the folder.')
                    return None
        else:
            logger.error('Data path (' + data_folder + ') exists as a file. '
                         'Please rename, remove or change the desired location of the data path.')
            return None

        # Check if the file exists already
        if not os.path.isfile(os.path.join(data_folder, file_name)):
            downloaded_file = wget.download(url, os.path.join(data_folder, file_name))
        else:
            downloaded_file = os.path.join(data_folder, file_name)
        return downloaded_file

    def download_file_ftp(self, ftp_site, gaf_uri, gaf_fn, data_folder=_resource_dir):
        # Check if the file exists already
        gaf = os.path.join(data_folder, gaf_fn)
        if not os.path.isfile(gaf):
            # Login to FTP server
            ebi_ftp = FTP(ftp_site)
            ebi_ftp.login() # Logs in anonymously

            # Download
            with open(gaf,'wb') as fp:
                ebi_ftp.retrbinary('RETR {}'.format(gaf_uri), fp.write)

            # Logout from FTP server
            ebi_ftp.quit()

        return gaf

    def read_gaf2(self, ftp_site=None, gaf_uri=None, gaf_fn=None):
        """
        The gaf from ebi is 6/8/16 version. We will use the one from geneontology website.
        """
        if not ftp_site:
            ftp_site = 'ftp.ebi.ac.uk'
        if not gaf_uri:
            gaf_uri = '/pub/databases/GO/goa/HUMAN/goa_human.gaf.gz'
        if not gaf_fn:
            gaf_fn = gaf_uri.split('/')[-1]

        gaf_file = self.download_file_ftp(ftp_site, gaf_uri, gaf_fn)

        #read the annotations into a dictionary
        # File is a gunzip file, so we need to open it in this way
        with gzip.open(gaf_file, 'rt') as gaf_fp:
            gaf_funcs = defaultdict(list)
            # Iterate on each function using Bio.UniProt.GOA library.
            for entry in GOA.gafiterator(gaf_fp):
                gene_symbol = entry.pop('DB_Object_Symbol')
                gaf_funcs[gene_symbol].append(entry)
        return gaf_funcs

    def generate_go2gene_db(self, go_file_name, data_folder=_resource_dir):
        """
        Generate a sqlite db file which contains GO id and its associated genes.
        """
        t0 = time.perf_counter()

        db_file = os.path.join(data_folder, go_file_name)
        conn = sqlite3.connect(db_file)
        c = conn.cursor()
        c.execute('''CREATE TABLE go2genes
             (id integer, goid text, genesymbol text)''')
        t = []
        num = 1
        for gene,entry in self.gaf_funcs.items():
            for en in entry:
                go_term = en['GO_ID']
                #c.execute("INSERT INTO go2genes VALUES (?,?,?)", (num, go_terms, gene))
                t.append((num, go_term, gene))
                num += 1
        c.executemany('INSERT INTO go2genes VALUES (?,?,?)', t)
        conn.commit()
        time.sleep(0.1)
        c.close()
        conn.close()

        t1 = time.perf_counter()
        logger.info('Created go_gene.db. go2genes table has {} items.'.format(num))
        logger.info('Used {} seconds to generate go-gene db file.'.format(t1-t0))
        return num

    def generate_go2gene_file(self, data_folder=_resource_dir):
        num = 1
        fn = os.path.join(data_folder, 'go2genes_gaf.txt')
        with open(fn, 'w') as fw:
            for gene,entries in self.gaf_funcs.items():
                for entry in entries:
                    fw.write(str(num) + '\t' + entry['GO_ID'] + '\t' + gene + '\t')
                    name_space = self.go[entry['GO_ID']].namespace
                    #or
                    #name_space = entry['Aspect']
                    fw.write(name_space + '\n')
                    num += 1
        print('write {} rows to file.'.format(num))
        return num

<commit_msg>Comment out loading go-gene.db for now<commit_after>#import
import os
import time
import logging
from goatools import obo_parser
from goatools.go_enrichment import GOEnrichmentStudy
import wget
from ftplib import FTP
import gzip
import Bio.UniProt.GOA as GOA
from collections import defaultdict
import pickle
import sqlite3

logging.basicConfig(format='%(levelname)s: %(name)s - %(message)s', level=logging.INFO)
logger = logging.getLogger('TFTA-GOEnrich')

_resource_dir = os.path.dirname(os.path.realpath(__file__)) + '/../enrichment/data'

class GOEnrich:
    def __init__(self, species=None):
        self.go = self.read_go()
        if self.go:
            logger.info('GOEnrich loaded GO data.')

        if not species:
            self.gaf_funcs = self.read_gaf()
        else:
            gaf_uri = '/pub/databases/GO/goa/' + species.upper() + '/goa_' + \
                       species.lower() + '.gaf.gz'
            self.gaf_funcs = self.read_gaf2(gaf_uri=gaf_uri)
        if self.gaf_funcs:
            logger.info("GOEnrich loaded GAF file.")
        else:
            logger.error("GAF file doesn't exist or couldn't be loaded.")

        #background population
        self.pop = self.gaf_funcs.keys()

        self.assoc = self.get_assoc_gene_go()

        #load go_gene.db. Comment out for now since it's not used.
        #db_file = os.path.join(_resource_dir, 'go_gene.db')

        #check file size to determine if it need regenerate
        #if os.path.isfile(db_file):
        #    statinfo = os.stat(db_file)
        #else:
        #    statinfo = None

        #if statinfo and statinfo.st_size > 13000000:
        #    self.godb = sqlite3.connect(db_file, check_same_thread=False)
        #    logger.info('GOEnrich loaded go_gene database.')
        #else:
        #    num = self.generate_go2gene_db(db_file)
        #    if os.path.isfile(db_file):
        #        self.godb = sqlite3.connect(db_file, check_same_thread=False)
        #    else:
        #        logger.error('GOEnrich could not load go_gene database.')
        #        self.godb = None

    def read_go(self, go_obo_url=None, data_folder=_resource_dir):
        if not go_obo_url:
            go_obo_url = 'http://purl.obolibrary.org/obo/go/go-basic.obo'

        #read GO data as dictionary
        go_obo = self.download_file_http(go_obo_url, data_folder, 'go-basic.obo')
        if go_obo:
            go = obo_parser.GODag(go_obo)
            return go
        else:
            logger.error("The GO data file doesn't exist. You can manually download it from: {}".format(go_obo_url))
            return None

    def read_gaf(self, gaf_uri=None, gaf_fn=None, data_folder=_resource_dir):
        """
        Download gaf file from http://current.geneontology.org/products/pages/downloads.html
        Annotations for 2019-07-01 release
        """
        if not gaf_uri:
            gaf_uri = "http://geneontology.org/gene-associations/goa_human.gaf.gz"
        if not gaf_fn:
            gaf_fn = gaf_uri.split('/')[-1]

        gaf_file = self.download_file_http(gaf_uri, data_folder, gaf_fn)

        #read the annotations into a dictionary
        fn = os.path.join(_resource_dir, 'gaf_funcs.pickle')
        if os.path.isfile(fn):
            with open(fn, 'rb') as pickle_in:
                gaf_funcs = pickle.load(pickle_in)
        else:
            # File is a gunzip file, so we need to open it in this way
            with gzip.open(gaf_file, 'rt') as gaf_fp:
                gaf_funcs = defaultdict(list)
                # Iterate on each function using Bio.UniProt.GOA library.
                for entry in GOA.gafiterator(gaf_fp):
                    gene_symbol = entry.pop('DB_Object_Symbol')
                    gaf_funcs[gene_symbol].append(entry)

            #save gaf_funcs to file
            with open(fn, 'wb') as pickle_out:
                pickle.dump(gaf_funcs, pickle_out)

        return gaf_funcs

    def get_population(self):
        return self.population

    def get_annotations_keyword(self, keyword):
        """
        Generate a list of annotated genes which have the keyword in their names.

        parameter
        ------------
        keyword: str

        output
        ------------
        annots: dict
        """
        annots = dict()
        for g, entry in self.gaf_funcs.items():
            if keyword.lower() in entry[0]['DB_Object_Name'].lower():
                annots[g] = entry
        return annots

    def get_annotations_genes(self, gene_list):
        """
        Return the annotations for each gene in the gene_list
        """
        annots = {x: self.gaf_funcs[x] for x in gene_list}
        return annots

    def get_go_keyword(self, keyword):
        """
        return GO terms whose names have keyword
        """
        #check if the corresponding file exists
        data_folder = _resource_dir + '/go_files'
        fn = os.path.join(data_folder, keyword + ".pickle")
        if os.path.isfile(fn):
            with open(fn, 'rb') as pickle_in:
                go_name = pickle.load(pickle_in)
        else:
            go_name = self._go_keyword(keyword, data_folder)
        return go_name

    def _go_keyword(self, keyword, data_folder):
        """
        Generate a dict, its key is GO_id, value is GO name which contains the keyword
        """
        #data_folder = _resource_dir + '/go_files'
        if not os.path.isfile(data_folder):
            # Emulate mkdir -p (no error if folder exists)
            try:
                os.mkdir(data_folder)
            except OSError as e:
                if(e.errno != 17):
                    logger.error('GOEnrich could not make the go_files folder.')
                    return None
        else:
            logger.error('Data path (' + data_folder + ') exists as a file. '
                         'Please rename, remove or change the desired location of the data path.')
            return None

        go_name = dict()
        keyword = keyword.lower()
        if self.go:
            go_name = {go_id:self.go[go_id].name for go_id in self.go if keyword in self.go[go_id].name.lower()}
            #save to file
            fn = os.path.join(data_folder, keyword + ".pickle")
            with open(fn, 'wb') as pickle_out:
                pickle.dump(go_name, pickle_out)
        return go_name

    def go_enrichment_analysis(self, study, pop=None, assoc=None, go=None, propagate_counts=True,
                               alpha=0.05, methods=None, p_bonferroni=0.01, limit=30):
        """
        GO enrichment analysis

        output
        ----------
        res: list of GOEnrichmentRecord
        """
        if not pop:
            pop = self.pop
        if not assoc:
            assoc = self.assoc
        if not go:
            go = self.go

        #methods = ["bonferroni", "sidak", "holm", "fdr"]
        #In order to reduce delay, only use bonferroni by default
        if not methods:
            methods = ['bonferroni']
        g = GOEnrichmentStudy(pop, assoc, go,
                         propagate_counts=propagate_counts,
                         alpha=alpha,
                         methods=methods)
        g_res = g.run_study(study)

        #only return results with pvalue < pvalue
        res = []
        num = 1
        for x in g_res:
            if num > limit:
                break
            if x.p_bonferroni <= p_bonferroni:
                res.append(x)
                num += 1
            else:
                break
        return res

    def get_assoc_gene_go(self):
        assoc = defaultdict(set)
        for x in self.gaf_funcs:
            for entry in self.gaf_funcs[x]:
                assoc[x].add(entry['GO_ID'])
        return assoc

    def download_file_http(self, url, data_folder, file_name):
        # Check if we have the ./data directory already
        if not os.path.isfile(data_folder):
            # Emulate mkdir -p (no error if folder exists)
            try:
                os.mkdir(data_folder)
            except OSError as e:
                if(e.errno != 17):
                    logger.error('GOEnrich could not make the folder.')
                    return None
        else:
            logger.error('Data path (' + data_folder + ') exists as a file. '
                         'Please rename, remove or change the desired location of the data path.')
            return None

        # Check if the file exists already
        if not os.path.isfile(os.path.join(data_folder, file_name)):
            downloaded_file = wget.download(url, os.path.join(data_folder, file_name))
        else:
            downloaded_file = os.path.join(data_folder, file_name)
        return downloaded_file

    def download_file_ftp(self, ftp_site, gaf_uri, gaf_fn, data_folder=_resource_dir):
        # Check if the file exists already
        gaf = os.path.join(data_folder, gaf_fn)
        if not os.path.isfile(gaf):
            # Login to FTP server
            ebi_ftp = FTP(ftp_site)
            ebi_ftp.login() # Logs in anonymously

            # Download
            with open(gaf,'wb') as fp:
                ebi_ftp.retrbinary('RETR {}'.format(gaf_uri), fp.write)

            # Logout from FTP server
            ebi_ftp.quit()

        return gaf

    def read_gaf2(self, ftp_site=None, gaf_uri=None, gaf_fn=None):
        """
        The gaf from ebi is 6/8/16 version. We will use the one from geneontology website.
        """
        if not ftp_site:
            ftp_site = 'ftp.ebi.ac.uk'
        if not gaf_uri:
            gaf_uri = '/pub/databases/GO/goa/HUMAN/goa_human.gaf.gz'
        if not gaf_fn:
            gaf_fn = gaf_uri.split('/')[-1]

        gaf_file = self.download_file_ftp(ftp_site, gaf_uri, gaf_fn)

        #read the annotations into a dictionary
        # File is a gunzip file, so we need to open it in this way
        with gzip.open(gaf_file, 'rt') as gaf_fp:
            gaf_funcs = defaultdict(list)
            # Iterate on each function using Bio.UniProt.GOA library.
            for entry in GOA.gafiterator(gaf_fp):
                gene_symbol = entry.pop('DB_Object_Symbol')
                gaf_funcs[gene_symbol].append(entry)
        return gaf_funcs

    def generate_go2gene_db(self, go_file_name, data_folder=_resource_dir):
        """
        Generate a sqlite db file which contains GO id and its associated genes.
        """
        t0 = time.perf_counter()

        db_file = os.path.join(data_folder, go_file_name)
        conn = sqlite3.connect(db_file)
        c = conn.cursor()
        c.execute('''CREATE TABLE go2genes
             (id integer, goid text, genesymbol text)''')
        t = []
        num = 1
        for gene,entry in self.gaf_funcs.items():
            for en in entry:
                go_term = en['GO_ID']
                #c.execute("INSERT INTO go2genes VALUES (?,?,?)", (num, go_terms, gene))
                t.append((num, go_term, gene))
                num += 1
        c.executemany('INSERT INTO go2genes VALUES (?,?,?)', t)
        conn.commit()
        time.sleep(0.1)
        c.close()
        conn.close()

        t1 = time.perf_counter()
        logger.info('Created go_gene.db. go2genes table has {} items.'.format(num))
        logger.info('Used {} seconds to generate go-gene db file.'.format(t1-t0))
        return num

    def generate_go2gene_file(self, data_folder=_resource_dir):
        num = 1
        fn = os.path.join(data_folder, 'go2genes_gaf.txt')
        with open(fn, 'w') as fw:
            for gene,entries in self.gaf_funcs.items():
                for entry in entries:
                    fw.write(str(num) + '\t' + entry['GO_ID'] + '\t' + gene + '\t')
                    name_space = self.go[entry['GO_ID']].namespace
                    #or
                    #name_space = entry['Aspect']
                    fw.write(name_space + '\n')
                    num += 1
        print('write {} rows to file.'.format(num))
        return num

<|endoftext|>


<commit_before>from __future__ import unicode_literals
from __future__ import print_function
from __future__ import division
from __future__ import absolute_import
from future import standard_library
standard_library.install_aliases()
def make_rendered(env, *render_args, **render_kwargs):
    base_step = env.step
    base_close = env.close

    def step(action):
        ret = base_step(action)
        env.render(*render_args, **render_kwargs)
        return ret

    def close():
        env.render(*render_args, close=True, **render_kwargs)
        base_close()

    env.step = step
    env.close = close

def make_timestep_limited(env, timestep_limit):
    t = 1
    old__step = env._step
    old__reset = env._reset

    def _step(action):
        nonlocal t
        observation, reward, done, info = old__step(action)
        if t >= timestep_limit:
            done = True
        t += 1
        return observation, reward, done, info

    def _reset():
        nonlocal t
        t = 1
        return old__reset()

    env._step = _step
    env._reset = _reset

def make_action_filtered(env, action_filter):
    old_step = env.step

    def step(action):
        return old_step(action_filter(action))

    env.step = step

def make_reward_filtered(env, reward_filter):
    old__step = env._step

    def _step(action):
        observation, reward, done, info = old__step(action)
        reward = reward_filter(reward)
        return observation, reward, done, info

    env._step = _step
<commit_msg>Support python2<commit_after>from __future__ import unicode_literals
from __future__ import print_function
from __future__ import division
from __future__ import absolute_import
from future import standard_library
standard_library.install_aliases()
def make_rendered(env, *render_args, **render_kwargs):
    base_step = env.step
    base_close = env.close

    def step(action):
        ret = base_step(action)
        env.render(*render_args, **render_kwargs)
        return ret

    def close():
        env.render(*render_args, close=True, **render_kwargs)
        base_close()

    env.step = step
    env.close = close

def make_timestep_limited(env, timestep_limit):
    t = [1]
    old__step = env._step
    old__reset = env._reset

    def _step(action):
        observation, reward, done, info = old__step(action)
        if t[0] >= timestep_limit:
            done = True
        t[0] += 1
        return observation, reward, done, info

    def _reset():
        t[0] = 1
        return old__reset()

    env._step = _step
    env._reset = _reset

def make_action_filtered(env, action_filter):
    old_step = env.step

    def step(action):
        return old_step(action_filter(action))

    env.step = step

def make_reward_filtered(env, reward_filter):
    old__step = env._step

    def _step(action):
        observation, reward, done, info = old__step(action)
        reward = reward_filter(reward)
        return observation, reward, done, info

    env._step = _step
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# Copyright 2001 Google Inc. All Rights Reserved.
#
# 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.

"""Simple web server for browsing dependency graph data.

This script is inlined into the final executable and spawned by
it when needed.
"""

from __future__ import print_function

try:
    import http.server as httpserver
except ImportError:
    import BaseHTTPServer as httpserver
import subprocess
import sys
import webbrowser
from collections import namedtuple

Node = namedtuple('Node', ['inputs', 'rule', 'target', 'outputs'])

# Ideally we'd allow you to navigate to a build edge or a build node,
# with appropriate views for each.  But there's no way to *name* a build
# edge so we can only display nodes.
#
# For a given node, it has at most one input edge, which has n
# different inputs.  This becomes node.inputs.  (We leave out the
# outputs of the input edge due to what follows.)  The node can have
# multiple dependent output edges.  Rather than attempting to display
# those, they are summarized by taking the union of all their outputs.
#
# This means there's no single view that shows you all inputs and outputs
# of an edge.  But I think it's less confusing than alternatives.

def match_strip(line, prefix):
    if not line.startswith(prefix):
        return (False, line)
    return (True, line[len(prefix):])

def parse(text):
    lines = iter(text.split('\n'))

    target = None
    rule = None
    inputs = []
    outputs = []

    try:
        target = next(lines)[:-1]  # strip trailing colon

        line = next(lines)
        (match, rule) = match_strip(line, '  input: ')
        if match:
            (match, line) = match_strip(next(lines), '    ')
            while match:
                type = None
                (match, line) = match_strip(line, '| ')
                if match:
                    type = 'implicit'
                (match, line) = match_strip(line, '|| ')
                if match:
                    type = 'order-only'
                inputs.append((line, type))
                (match, line) = match_strip(next(lines), '    ')

        match, _ = match_strip(line, '  outputs:')
        if match:
            (match, line) = match_strip(next(lines), '    ')
            while match:
                outputs.append(line)
                (match, line) = match_strip(next(lines), '    ')
    except StopIteration:
        pass

    return Node(inputs, rule, target, outputs)

def create_page(body):
    return '''<!DOCTYPE html>

''' + body

def generate_html(node):
    document = ['
<tt>%s</tt>
' % node.target]

    if node.inputs:
        document.append('
target is built using rule <tt>%s</tt> of
' %
                        node.rule)
        if len(node.inputs) > 0:
            document.append('
')
            for input, type in sorted(node.inputs):
                extra = ''
                if type:
                    extra = ' (%s)' % type
                document.append('<tt>%s%s</tt>
' %
                                (input, input, extra))
            document.append('
')

    if node.outputs:
        document.append('
dependent edges build:
')
        document.append('
')
        for output in sorted(node.outputs):
            document.append('<tt>%s</tt>
' %
                            (output, output))
        document.append('
')

    return '\n'.join(document)

def ninja_dump(target):
    proc = subprocess.Popen([sys.argv[1], '-t', 'query', target],
                            stdout=subprocess.PIPE, universal_newlines=True)
    return (proc.communicate()[0], proc.returncode)

class RequestHandler(httpserver.BaseHTTPRequestHandler):
    def do_GET(self):
        assert self.path[0] == '/'
        target = self.path[1:]

        if target == '':
            self.send_response(302)
            self.send_header('Location', '?' + sys.argv[2])
            self.end_headers()
            return

        if not target.startswith('?'):
            self.send_response(404)
            self.end_headers()
            return
        target = target[1:]

        input, exit_code = ninja_dump(target)
        if exit_code == 0:
            page_body = generate_html(parse(input.strip()))
        else:
            # Relay ninja's error message.
            page_body = '
<tt>%s</tt>
' % input

        self.send_response(200)
        self.end_headers()
        self.wfile.write(create_page(page_body).encode('utf-8'))

    def log_message(self, format, *args):
        pass  # Swallow console spam.

port = 8000
httpd = httpserver.HTTPServer(('',port), RequestHandler)
try:
    print('Web server running on port %d, ctl-C to abort...' % port)
    webbrowser.open_new('http://localhost:%s' % port)
    httpd.serve_forever()
except KeyboardInterrupt:
    print()
    pass  # Swallow console spam.

<commit_msg>browse: Read ninja's error text from stderr.<commit_after>#!/usr/bin/env python
#
# Copyright 2001 Google Inc. All Rights Reserved.
#
# 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.

"""Simple web server for browsing dependency graph data.

This script is inlined into the final executable and spawned by
it when needed.
"""

from __future__ import print_function

try:
    import http.server as httpserver
except ImportError:
    import BaseHTTPServer as httpserver
import subprocess
import sys
import webbrowser
from collections import namedtuple

Node = namedtuple('Node', ['inputs', 'rule', 'target', 'outputs'])

# Ideally we'd allow you to navigate to a build edge or a build node,
# with appropriate views for each.  But there's no way to *name* a build
# edge so we can only display nodes.
#
# For a given node, it has at most one input edge, which has n
# different inputs.  This becomes node.inputs.  (We leave out the
# outputs of the input edge due to what follows.)  The node can have
# multiple dependent output edges.  Rather than attempting to display
# those, they are summarized by taking the union of all their outputs.
#
# This means there's no single view that shows you all inputs and outputs
# of an edge.  But I think it's less confusing than alternatives.

def match_strip(line, prefix):
    if not line.startswith(prefix):
        return (False, line)
    return (True, line[len(prefix):])

def parse(text):
    lines = iter(text.split('\n'))

    target = None
    rule = None
    inputs = []
    outputs = []

    try:
        target = next(lines)[:-1]  # strip trailing colon

        line = next(lines)
        (match, rule) = match_strip(line, '  input: ')
        if match:
            (match, line) = match_strip(next(lines), '    ')
            while match:
                type = None
                (match, line) = match_strip(line, '| ')
                if match:
                    type = 'implicit'
                (match, line) = match_strip(line, '|| ')
                if match:
                    type = 'order-only'
                inputs.append((line, type))
                (match, line) = match_strip(next(lines), '    ')

        match, _ = match_strip(line, '  outputs:')
        if match:
            (match, line) = match_strip(next(lines), '    ')
            while match:
                outputs.append(line)
                (match, line) = match_strip(next(lines), '    ')
    except StopIteration:
        pass

    return Node(inputs, rule, target, outputs)

def create_page(body):
    return '''<!DOCTYPE html>

''' + body

def generate_html(node):
    document = ['
<tt>%s</tt>
' % node.target]

    if node.inputs:
        document.append('
target is built using rule <tt>%s</tt> of
' %
                        node.rule)
        if len(node.inputs) > 0:
            document.append('
')
            for input, type in sorted(node.inputs):
                extra = ''
                if type:
                    extra = ' (%s)' % type
                document.append('<tt>%s%s</tt>
' %
                                (input, input, extra))
            document.append('
')

    if node.outputs:
        document.append('
dependent edges build:
')
        document.append('
')
        for output in sorted(node.outputs):
            document.append('<tt>%s</tt>
' %
                            (output, output))
        document.append('
')

    return '\n'.join(document)

def ninja_dump(target):
    proc = subprocess.Popen([sys.argv[1], '-t', 'query', target],
                            stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                            universal_newlines=True)
    return proc.communicate() + (proc.returncode,)

class RequestHandler(httpserver.BaseHTTPRequestHandler):
    def do_GET(self):
        assert self.path[0] == '/'
        target = self.path[1:]

        if target == '':
            self.send_response(302)
            self.send_header('Location', '?' + sys.argv[2])
            self.end_headers()
            return

        if not target.startswith('?'):
            self.send_response(404)
            self.end_headers()
            return
        target = target[1:]

        ninja_output, ninja_error, exit_code = ninja_dump(target)
        if exit_code == 0:
            page_body = generate_html(parse(ninja_output.strip()))
        else:
            # Relay ninja's error message.
            page_body = '
<tt>%s</tt>
' % ninja_error

        self.send_response(200)
        self.end_headers()
        self.wfile.write(create_page(page_body).encode('utf-8'))

    def log_message(self, format, *args):
        pass  # Swallow console spam.

port = 8000
httpd = httpserver.HTTPServer(('',port), RequestHandler)
try:
    print('Web server running on port %d, ctl-C to abort...' % port)
    webbrowser.open_new('http://localhost:%s' % port)
    httpd.serve_forever()
except KeyboardInterrupt:
    print()
    pass  # Swallow console spam.

<|endoftext|>


<commit_before>import os
import shutil
import eos.archive
import eos.cache
import eos.log
import eos.post
import eos.repo
import eos.util

def bootstrap_library(json_obj, name, library_dir, postprocessing_dir):
    eos.log("Bootstrapping library '" + name + "' to " + library_dir)

    # create directory for library
    if not os.path.exists(library_dir):
        os.mkdir(library_dir)

    # get library

    src = json_obj.get('source', None)
    if not src:
        eos.log_warning("library '" + name + "' is missing source description")
        return False

    src_type = src.get('type', None)
    src_url = src.get('url', None)

    if not src_type or not src_url:
        eos.log_warning("library '" + name + "' is missing type or URL description")
        return False

    if src_type not in ['archive', 'git', 'hg', 'svn']:
        eos.log_warning("unknown source type for library '" + name)
        return False

    if src_type == "archive":
        sha1_hash = src.get('sha1', None)
        user_agent = src.get('user-agent', None)

        download_filename = eos.util.download_file(src_url, eos.cache.get_archive_dir(), sha1_hash, user_agent)
        if download_filename == "":
            eos.log_error("downloading of file for '" + name + "' from " + src_url + " failed")
            return False

        if os.path.exists(library_dir):
            shutil.rmtree(library_dir)

        if not eos.archive.extract_file(download_filename, library_dir):
            eos.log_error("extraction of file for '" + download_filename + "' failed")
            return False
    else:
        branch = src.get('branch', None)
        if not branch:
            branch = src.get('branch-follow', None)
        revision = src.get('revision', None)

        if branch and revision:
            eos.log_error("cannot specify both branch (to follow) and revision for repository '" + name + "'")
            return False

        if not eos.repo.update_state(src_type, src_url, name, library_dir, branch, revision):
            eos.log_error("updating repository state for '" + name + " failed")
            return False

    # post-process library

    post = json_obj.get('postprocess', None)
    if not post:
        return True  # it's optional

    post_type = post.get('type', None)
    if not post_type:
        eos.log_error("postprocessing object for library '" + name + "' must have a 'type'")
        return False

    post_file = post.get('file', None)
    if not post_file:
        eos.log_error("postprocessing object for library '" + name + "' must have a 'file'")
        return False

    if post_type not in ['patch', 'script']:
        eos.log_error("unknown postprocessing type for library '" + name + "'")
        return False

    # If we have a postprocessing directory specified, make it an absolute path
    if postprocessing_dir:
        post_file = os.path.join(postprocessing_dir, post_file)

    if post_type == "patch":
        pnum = post.get('pnum', 2)
        # Try to apply patch
        if not eos.post.apply_patch(name, library_dir, post_file, pnum):
            eos.log_error("patch application of " + post_file + " failed for library '" + name + "'")
            return False
    elif post_type == "script":
        # Try to run script
        if not eos.post.run_script(name, post_file):
            eos.log_error("script execution of " + post_file + " failed for library '" + name + "'")
            return False

    return True

<commit_msg>In script file string, replace variable $LIBRARY_DIR with the library directory. This enables passing the library directory as an argument to the script.<commit_after>import os
import shutil
import eos.archive
import eos.cache
import eos.log
import eos.post
import eos.repo
import eos.util

def bootstrap_library(json_obj, name, library_dir, postprocessing_dir):
    eos.log("Bootstrapping library '" + name + "' to " + library_dir)

    # create directory for library
    if not os.path.exists(library_dir):
        os.mkdir(library_dir)

    # get library

    src = json_obj.get('source', None)
    if not src:
        eos.log_warning("library '" + name + "' is missing source description")
        return False

    src_type = src.get('type', None)
    src_url = src.get('url', None)

    if not src_type or not src_url:
        eos.log_warning("library '" + name + "' is missing type or URL description")
        return False

    if src_type not in ['archive', 'git', 'hg', 'svn']:
        eos.log_warning("unknown source type for library '" + name)
        return False

    if src_type == "archive":
        sha1_hash = src.get('sha1', None)
        user_agent = src.get('user-agent', None)

        download_filename = eos.util.download_file(src_url, eos.cache.get_archive_dir(), sha1_hash, user_agent)
        if download_filename == "":
            eos.log_error("downloading of file for '" + name + "' from " + src_url + " failed")
            return False

        if os.path.exists(library_dir):
            shutil.rmtree(library_dir)

        if not eos.archive.extract_file(download_filename, library_dir):
            eos.log_error("extraction of file for '" + download_filename + "' failed")
            return False
    else:
        branch = src.get('branch', None)
        if not branch:
            branch = src.get('branch-follow', None)
        revision = src.get('revision', None)

        if branch and revision:
            eos.log_error("cannot specify both branch (to follow) and revision for repository '" + name + "'")
            return False

        if not eos.repo.update_state(src_type, src_url, name, library_dir, branch, revision):
            eos.log_error("updating repository state for '" + name + " failed")
            return False

    # post-process library

    post = json_obj.get('postprocess', None)
    if not post:
        return True  # it's optional

    post_type = post.get('type', None)
    if not post_type:
        eos.log_error("postprocessing object for library '" + name + "' must have a 'type'")
        return False

    post_file = post.get('file', None)
    if not post_file:
        eos.log_error("postprocessing object for library '" + name + "' must have a 'file'")
        return False

    if post_type not in ['patch', 'script']:
        eos.log_error("unknown postprocessing type for library '" + name + "'")
        return False

    # If we have a postprocessing directory specified, make it an absolute path
    if postprocessing_dir:
        post_file = os.path.join(postprocessing_dir, post_file)

    if post_type == "patch":
        pnum = post.get('pnum', 2)
        # Try to apply patch
        if not eos.post.apply_patch(name, library_dir, post_file, pnum):
            eos.log_error("patch application of " + post_file + " failed for library '" + name + "'")
            return False
    elif post_type == "script":
        # Replace variable strings with contents
        post_file = post_file.replace("$LIBRARY_DIR", os.path.abspath(library_dir))
        # Try to run script
        if not eos.post.run_script(name, post_file):
            eos.log_error("script execution of " + post_file + " failed for library '" + name + "'")
            return False

    return True

<|endoftext|>