<commit_before># -*- coding: utf-8 -*-
"""
Set of encoding strategies for protein motifs sequences training with machine learning.
"""
import os
from itertools import product, izip_longest, islice, repeat

import pandas as pd
from Bio import SeqIO

__all__ = ['read_fasta', 'dataframe_sparse_encoder', 'dataframe_aaindex_encoder', 'window',
           'dataframe_k_mer_composition']

def read_fasta(fasta_file):
    """
    Parses a fasta file with BioPython to a dictionary

    :param fasta_file: input file in the fasta format
    :return dict: dictionary with [identifier, sequence] from the input file
    """
    with open(fasta_file) as open_input:
        seqs = SeqIO.to_dict(SeqIO.parse(open_input, 'fasta'),
                             key_function=lambda x: x.id.split()[0])
    return seqs

def read_positions(positions_file, columns=None):
    """
    Position for motifs within the protein sequences.

    :param positions_file: csv file with the positions
    :param array-like columns: columns with identifier and position, respectively
    :return dict :
    """
    if columns is None:
        columns = [0, 1]
    data = pd.read_csv(positions_file)
    try:
        data = data.loc[:, columns]
    except KeyError:  # named columns
        data = data.iloc[:, columns]

    return data.to_records()

def create_motifs_from_residues(sequence, motif_size=5, residues=['S', 'T'], extremity='X'):
    """
    Fragments the proteins in modifiable peptides.
    :param str sequence: protein sequence
    :param int motif_size: number of residues per motif
    :param srt extremity: character symbol for termination character
    :param array_like residues: residues' position within the protein sequences
    :rtype : Iterator of position collection.Iterable[tuple]

    .. example:
    >>> print(list(create_motifs_from_residues('VTIQHPWFKRTLGP'))) # doctest: +NORMALIZE_WHITESPACE
    [(2, 'XVTIQ'), (11, 'KRTLG')]
    """
    if len(extremity) > 1:
        extremity = extremity[0]
    if isinstance(residues, str):
        residues = [residues,]
    elif not hasattr(residues, "__contains__"):
        raise TypeError("residue paramenter should be a list-like object")

    if motif_size & 1:  # is odd
        start = motif_size // 2
        stop = (motif_size // 2) + 1
    else:
        start = motif_size // 2
        stop = start

    protein_len = len(sequence)
    for i, aa in enumerate(sequence):
        if aa in residues:
            if i - start < 0:
                motif = sequence[0: i].rjust(start, extremity) + sequence[i: i + stop]
            elif i + stop > protein_len:
                motif = sequence[i - start: i] + sequence[i: i + stop].ljust(stop, extremity)
            else:
                motif = sequence[i - start: i + stop]

            yield i + 1, motif

def dataframe_sparse_encoder(data, descriptor=None, col_names=None):
    """Sparsely encode the motifs with abstract descriptor. Similar to 'one-hot' encoding that
    is consistent across the motif columns.

    :param descriptor: dict-like descriptor to be applied to the columns
    :type descriptor: dict or None
    :param pd.DataFrame data: DataFrame with the data
    :param col_names: names for the new columns
    :type col_names: list[str] or None
    :return: table with encoded data
    :rtype : pd.DataFrame
    :raise ValueError: if the descriptor argument lack `get` method

    .. example:

        >>> import pandas as pd
        >>> df = pd.DataFrame(map(list, ['ABC', 'AAA', 'ACD']))
        >>> print(dataframe_sparse_encoder(df))  # doctest: +NORMALIZE_WHITESPACE
           0_A  1_A  2_A  0_B  1_B  2_B  0_C  1_C  2_C  0_D  1_D  2_D
        0    1    0    0    0    0    1    0    0    0    0    1    0
        1    1    0    0    0    1    0    0    0    1    0    0    0
        2    1    0    0    0    0    0    1    0    0    0    0    1

    .. note:: descriptor parameters can be any Python object with .get method, as pd.Series.
    """
    data = data.copy()
    if descriptor is None:
        descriptor = pd.Series(data.values.ravel()).unique()
        descriptor = pd.get_dummies(descriptor)
    elif hasattr(descriptor, 'get'):
        pass
    else:
        raise (ValueError("Descriptor should have be a dict or pandas.Series"))

    if col_names is None:
        col_names = ["{}_{}".format(col_name, letter) for letter, col_name in
                     product(descriptor.keys(), data.columns)]

    with_dummies = []
    for col in data:
        with_dummies.append(data[col].apply(descriptor.get))
    encoded = pd.concat(with_dummies, axis=1)
    encoded.columns = col_names
    return encoded

def dataframe_aaindex_encoder(data, contains=None, named_columns=False):
    """Encode the motifs with an a DataFrame of Motifs to correspondent AAindex vector
    Total number of descriptors 512.

    :param pd.DataFrame data: DataFrame with categorical data
    :param bool named_columns: Weather output dataframe has column_names
    :param str contains: Terms for reducing aaindex
    :return: table with encoded data
    :rtype: pd.DataFrame

    :Example:

    >>> import pandas as pd; df = pd.DataFrame(map(list, ['ABC', 'AAA', 'ACD']))
    >>> print(dataframe_aaindex_encoder(df).iloc[:, 0: 3])  # doctest: +NORMALIZE_WHITESPACE
        0     1     2
    0  4.35     B  4.65
    1  4.35  4.35  4.35
    2  4.35  4.65  4.76
    >>> print(dataframe_aaindex_encoder(df, contains='Hydrophobicity_index', \
    named_columns=True).iloc[:, 0:3])  # doctest: +NORMALIZE_WHITESPACE
        0_Hydrophobicity_index_(Argos_et_al.,_1982) \
    0                                         0.61
    1                                         0.61
    2                                         0.61
    <BLANKLINE>
        1_Hydrophobicity_index_(Argos_et_al.,_1982) \
    0                                           B
    1                                        0.61
    2                                        1.07
    <BLANKLINE>
        2_Hydrophobicity_index_(Argos_et_al.,_1982)
    0                                         1.07
    1                                         0.61
    2                                         0.46

    .. note:: If you use AAindex please cite one of their references: http://www.genome.jp/aaindex/
    .. warning:: It does not replace non-standard residues symbols
    .. todo:: other strategies to subset aa_index
    """
    data = data.copy()
    aa_index_path = os.path.join(os.path.dirname(__file__), 'aaindex.csv')
    aa_index = pd.read_csv(aa_index_path, index_col=(0, 1))
    aa_index = aa_index.T

    if contains:  # subset the aaindex based on the term
        names = aa_index.columns.get_level_values('name')
        mask = names.str.contains(contains, case=False)
        aa_index = aa_index.loc[:, mask]

    result = []
    for i in aa_index:
        to_replace = aa_index.loc[:, i]
        replaced = data.replace(to_replace)
        result.append(replaced)
    encoded = pd.concat(result, axis=1)

    if named_columns:
        pos = range(data.shape[1]) * aa_index.shape[1]
        names = [x for x in aa_index.columns.get_level_values('name')
                 for _ in range(data.shape[1])]
        encoded.columns = ['{}_{}'.format(a, b) for a, b in zip(pos, names)]

    return encoded

def window(seq, n=2):
    """
    Returns a sliding window (of width n) over data from the iterable
    s -> (s0,s1,...s[n-1]), (s1,s2,...,sn), ...
    :param seq: sequence to by cutted
    :type seq:  array-like or str
    :param n: motif size
    :return iterator:

    ..note:
    modified from https://docs.python.org/release/2.3.5/lib/itertools-example.html
    """
    it = iter(seq)
    result = tuple(islice(it, n))
    if len(result) == n:
        yield "".join(result)
    for elem in it:
        result = result[1:] + (elem,)
        yield "".join(result)

def dataframe_k_mer_composition(data, k=3, alphabet=None):
    """
    Return the k-mer composition for each sample (row) in the data.

    :param motif_df: motif in the format one residue per column
    :type motif_df: pandas.DataFrame
    :param k: length of the kmer
    :type k: int
    :param alphabet: all amino acids
    :type alphabet: str
    :return: matrix with counts of k-mer

    :Example:
    >>> import pandas as pd; df = pd.DataFrame(map(list, ['ABC', 'AAA', 'ACD']))
    >>> print(dataframe_k_mer_composition(df, k=1))  # doctest: +NORMALIZE_WHITESPACE
        A_comp    B_comp    C_comp    D_comp
    0  0.333333  0.333333  0.333333  0.000000
    1  1.000000  0.000000  0.000000  0.000000
    2  0.333333  0.000000  0.333333  0.333333
    >>> print(dataframe_k_mer_composition(df, k=2).iloc[:, :3])  # doctest: +NORMALIZE_WHITESPACE
        AA_comp   AB_comp   AC_comp
    0  0.000000  0.333333  0.000000
    1  0.666667  0.000000  0.000000
    2  0.000000  0.000000  0.333333
"""
    if k == 1:  # amino acid composition
        aa_comp = data.apply(pd.value_counts, axis=1).fillna(0)
        aa_comp /= data.shape[1]
        aa_comp.columns = aa_comp.columns + "_comp"

        return aa_comp

    if alphabet is None:
        alphabet = "".join(pd.unique(data.unstack().values))

    k_mers = data.apply(lambda x: list(window(x, n=k)), axis=1)
    k_mers = k_mers.apply(pd.value_counts)

    all_kmers = ["".join(aa) for aa in product(*repeat(list(alphabet), k))]
    k_mers = k_mers.loc[:, all_kmers].fillna(0)
    k_mers /= data.shape[1]
    k_mers.columns = k_mers.columns + "_comp"

    return k_mers

if __name__ == "__main__":
    import doctest

    doctest.testmod(verbose=True, optionflags=doctest.ELLIPSIS)
<commit_msg>add create_motifs and doct tests.<commit_after># -*- coding: utf-8 -*-
"""
Set of encoding strategies for protein motifs sequences training with machine learning.
"""
import os
from itertools import product, izip_longest, islice, repeat

import pandas as pd
from Bio import SeqIO

__all__ = ['read_fasta', 'dataframe_sparse_encoder', 'dataframe_aaindex_encoder', 'window',
           'dataframe_k_mer_composition']

def read_fasta(fasta_file):
    """
    Parses a fasta file with BioPython to a dictionary

    :param fasta_file: input file in the fasta format
    :return dict: dictionary with [identifier, sequence] from the input file
    """
    with open(fasta_file) as open_input:
        seqs = SeqIO.to_dict(SeqIO.parse(open_input, 'fasta'),
                             key_function=lambda x: x.id.split()[0])
    return seqs

def read_positions(positions_file, columns=None):
    """
    Position for motifs within the protein sequences.

    :param positions_file: csv file with the positions
    :param array-like columns: columns with identifier and position, respectively
    :return dict :
    """
    if columns is None:
        columns = [0, 1]
    data = pd.read_csv(positions_file)
    try:
        data = data.loc[:, columns]
    except KeyError:  # named columns
        data = data.iloc[:, columns]

    return data.to_records()

def create_motifs_from_residues(sequence, motif_size=5, residues=None, extremity='X'):
    """
    Fragments the proteins in motifs centred to certain residues.
    :param str sequence: protein sequence
    :param int motif_size: number of residues per motif
    :param srt extremity: character symbol for termination character
    :param array-like residues: residues' position within the protein sequences
    :rtype tuple[list, list]; list of positions and motifs, respectively

    .. example:
    >>> print(list(create_motifs_from_residues('VTIQHPWFKRTLGP'))) # doctest: +NORMALIZE_WHITESPACE
    [(2, 'XVTIQ'), (11, 'KRTLG')]
    """
    if residues is None:
        residues = ['S', 'T']
    if len(extremity) > 1:
        extremity = extremity[0]
    if isinstance(residues, str):
        residues = [residues, ]
    elif not hasattr(residues, "__contains__"):
        raise TypeError("residue paramenter should be a list-like object")

    if motif_size & 1:  # is odd
        start = motif_size // 2
        stop = (motif_size // 2) + 1
    else:
        start = motif_size // 2
        stop = start

    protein_len = len(sequence)
    positions = []
    motifs = []
    for i, aa in enumerate(sequence):
        if aa in residues:
            if i - start < 0:
                motif = sequence[0: i].rjust(start, extremity) + sequence[i: i + stop]
            elif i + stop > protein_len:
                motif = sequence[i - start: i] + sequence[i: i + stop].ljust(stop, extremity)
            else:
                motif = sequence[i - start: i + stop]
            positions.append(i + 1)
            motifs.append(motif)

    return positions, motifs

def built_motif_dataframe(sequences, **kwargs):
    """
    Transform a list of strings to a Motif DataFrame. Motif length should be consistent.
    Downstream arguments can be passed as keyword arguments. Each columns in a Motif DataFrame
    is a motif position and each row is an motif instance - position within the protein
    sequence.

    :param array-like sequences: list of protein sequences
    :return pandas.DataFrame: Motif DataFrame
    """
    positions, motifs, identifiers = [], [], []
    for identifier, sequence in sequences:
        positions_i, motifs_i = create_motifs_from_residues(sequence, **kwargs)
        identifiers += [identifier] * len(motifs_i)
        positions.extend(positions_i)
        motifs.extend(motifs_i)

    return pd.DataFrame([list(m) for m in motifs], index=[identifiers, positions])

def dataframe_sparse_encoder(data, descriptor=None, col_names=None):
    """Sparsely encode the motifs with abstract descriptor. Similar to 'one-hot' encoding that
    is consistent across the motif columns.

    :param descriptor: dict-like descriptor to be applied to the columns
    :type descriptor: dict or None
    :param pd.DataFrame data: DataFrame with the data
    :param col_names: names for the new columns
    :type col_names: list[str] or None
    :return: table with encoded data
    :rtype : pd.DataFrame
    :raise ValueError: if the descriptor argument lack `get` method

    .. example:

        >>> import pandas as pd
        >>> df = pd.DataFrame(map(list, ['ABC', 'AAA', 'ACD']))
        >>> print(dataframe_sparse_encoder(df))  # doctest: +NORMALIZE_WHITESPACE
           0_A  1_A  2_A  0_B  1_B  2_B  0_C  1_C  2_C  0_D  1_D  2_D
        0    1    0    0    0    0    1    0    0    0    0    1    0
        1    1    0    0    0    1    0    0    0    1    0    0    0
        2    1    0    0    0    0    0    1    0    0    0    0    1

    .. note:: descriptor parameters can be any Python object with .get method, as pd.Series.
    """
    data = data.copy()
    if descriptor is None:
        descriptor = pd.Series(data.values.ravel()).unique()
        descriptor = pd.get_dummies(descriptor)
    elif hasattr(descriptor, 'get'):
        pass
    else:
        raise (ValueError("Descriptor should have be a dict or pandas.Series"))

    if col_names is None:
        col_names = ["{}_{}".format(col_name, letter) for letter, col_name in
                     product(descriptor.keys(), data.columns)]

    with_dummies = []
    for col in data:
        with_dummies.append(data[col].apply(descriptor.get))
    encoded = pd.concat(with_dummies, axis=1)
    encoded.columns = col_names
    return encoded

def dataframe_aaindex_encoder(data, contains=None, named_columns=False):
    """Encode the motifs with an a DataFrame of Motifs to correspondent AAindex vector
    Total number of descriptors 512.

    :param pd.DataFrame data: DataFrame with categorical data
    :param bool named_columns: Weather output dataframe has column_names
    :param str contains: Terms for reducing aaindex
    :return: table with encoded data
    :rtype: pd.DataFrame

    :Example:

    >>> import pandas as pd; df = pd.DataFrame(map(list, ['ABC', 'AAA', 'ACD']))
    >>> print(dataframe_aaindex_encoder(df).iloc[:, 0: 3])  # doctest: +NORMALIZE_WHITESPACE
        0     1     2
    0  4.35     B  4.65
    1  4.35  4.35  4.35
    2  4.35  4.65  4.76
    >>> print(dataframe_aaindex_encoder(df, contains='Hydrophobicity_index', \
    named_columns=True).iloc[:, 0:3])  # doctest: +NORMALIZE_WHITESPACE
        0_Hydrophobicity_index_(Argos_et_al.,_1982) \
    0                                         0.61
    1                                         0.61
    2                                         0.61
    <BLANKLINE>
        1_Hydrophobicity_index_(Argos_et_al.,_1982) \
    0                                           B
    1                                        0.61
    2                                        1.07
    <BLANKLINE>
        2_Hydrophobicity_index_(Argos_et_al.,_1982)
    0                                         1.07
    1                                         0.61
    2                                         0.46

    .. note:: If you use AAindex please cite one of their references: http://www.genome.jp/aaindex/
    .. warning:: It does not replace non-standard residues symbols
    .. todo:: other strategies to subset aa_index
    """
    data = data.copy()
    aa_index_path = os.path.join(os.path.dirname(__file__), 'aaindex.csv')
    aa_index = pd.read_csv(aa_index_path, index_col=(0, 1))
    aa_index = aa_index.T

    if contains:  # subset the aaindex based on the term
        names = aa_index.columns.get_level_values('name')
        mask = names.str.contains(contains, case=False)
        aa_index = aa_index.loc[:, mask]

    result = []
    for i in aa_index:
        to_replace = aa_index.loc[:, i]
        replaced = data.replace(to_replace)
        result.append(replaced)
    encoded = pd.concat(result, axis=1)

    if named_columns:
        pos = range(data.shape[1]) * aa_index.shape[1]
        names = [x for x in aa_index.columns.get_level_values('name')
                 for _ in range(data.shape[1])]
        encoded.columns = ['{}_{}'.format(a, b) for a, b in zip(pos, names)]

    return encoded

def window(seq, n=2):
    """
    Returns a sliding window (of width n) over data from the iterable
    s -> (s0,s1,...s[n-1]), (s1,s2,...,sn), ...
    :param seq: sequence to by cutted
    :type seq:  array-like or str
    :param n: motif size
    :return iterator:

    ..note:
    modified from https://docs.python.org/release/2.3.5/lib/itertools-example.html
    """
    it = iter(seq)
    result = tuple(islice(it, n))
    if len(result) == n:
        yield "".join(result)
    for elem in it:
        result = result[1:] + (elem,)
        yield "".join(result)

def dataframe_k_mer_composition(data, k=3, alphabet=None):
    """
    Return the k-mer composition for each sample (row) in the data.

    :param motif_df: motif in the format one residue per column
    :type motif_df: pandas.DataFrame
    :param k: length of the kmer
    :type k: int
    :param alphabet: all amino acids
    :type alphabet: str
    :return: matrix with counts of k-mer

    :Example:
    >>> import pandas as pd; df = pd.DataFrame(map(list, ['ABC', 'AAA', 'ACD']))
    >>> print(dataframe_k_mer_composition(df, k=1))  # doctest: +NORMALIZE_WHITESPACE
        A_comp    B_comp    C_comp    D_comp
    0  0.333333  0.333333  0.333333  0.000000
    1  1.000000  0.000000  0.000000  0.000000
    2  0.333333  0.000000  0.333333  0.333333
    >>> print(dataframe_k_mer_composition(df, k=2).iloc[:, :3])  # doctest: +NORMALIZE_WHITESPACE
        AA_comp   AB_comp   AC_comp
    0  0.000000  0.333333  0.000000
    1  0.666667  0.000000  0.000000
    2  0.000000  0.000000  0.333333
"""
    if k == 1:  # amino acid composition
        aa_comp = data.apply(pd.value_counts, axis=1).fillna(0)
        aa_comp /= data.shape[1]
        aa_comp.columns = aa_comp.columns + "_comp"

        return aa_comp

    if alphabet is None:
        alphabet = "".join(pd.unique(data.unstack().values))

    k_mers = data.apply(lambda x: list(window(x, n=k)), axis=1)
    k_mers = k_mers.apply(pd.value_counts)

    all_kmers = ["".join(aa) for aa in product(*repeat(list(alphabet), k))]
    k_mers = k_mers.loc[:, all_kmers].fillna(0)
    k_mers /= data.shape[1]
    k_mers.columns = k_mers.columns + "_comp"

    return k_mers

if __name__ == "__main__":
    import doctest

    doctest.testmod(verbose=True, optionflags=doctest.ELLIPSIS)
<|endoftext|>


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

# Hive Appier Framework
# Copyright (c) 2008-2017 Hive Solutions Lda.
#
# This file is part of Hive Appier Framework.
#
# Hive Appier Framework is free software: you can redistribute it and/or modify
# it under the terms of the Apache License as published by the Apache
# Foundation, either version 2.0 of the License, or (at your option) any
# later version.
#
# Hive Appier Framework is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# Apache License for more details.
#
# You should have received a copy of the Apache License along with
# Hive Appier Framework. If not, see <http://www.apache.org/licenses/>.

__author__ = "João Magalhães <joamag@hive.pt>"
""" The author(s) of the module """

__version__ = "1.0.0"
""" The version of the module """

__revision__ = "$LastChangedRevision$"
""" The revision number of the module """

__date__ = "$LastChangedDate$"
""" The last change date of the module """

__copyright__ = "Copyright (c) 2008-2017 Hive Solutions Lda."
""" The copyright for the module """

__license__ = "Apache License, Version 2.0"
""" The license for the module """

import json
import uuid
import heapq

from . import amqp
from . import legacy
from . import exceptions

class Queue(object):

    def length(self):
        raise exceptions.NotImplementedError()

    def clear(self):
        raise exceptions.NotImplementedError()

    def push(self, value, priority = None, identify = False):
        raise exceptions.NotImplementedError()

    def pop(self, block = True, full = False):
        raise exceptions.NotImplementedError()

    def subscribe(self, callback):
        raise exceptions.NotImplementedError()

    def loop(self):
        raise exceptions.NotImplementedError()

    def unloop(self):
        raise exceptions.NotImplementedError()

    def build_value(
        self,
        value,
        priority = None,
        identify = False,
        reverse = False
    ):
        if identify: identifier = self.build_identifier()
        else: identifier = None
        if priority and reverse: priority *= -1
        return (priority, identifier, value), identifier

    def build_identifier(self):
        return str(uuid.uuid4())

class MemoryQueue(Queue):

    def __init__(self):
        Queue.__init__(self)
        self._queue = []

    def length(self):
        return len(self._queue)

    def clear(self):
        del self._queue[:]

    def push(self, value, priority = None, identify = False):
        value, identifier = self.build_value(
            value,
            priority = priority,
            identify = identify,
            reverse = True
        )
        heapq.heappush(self._queue, value)
        return identifier

    def pop(self, block = True, full = False):
        priority, identifier, value = heapq.heappop(self._queue)
        return (priority, identifier, value) if full else value

class MultiprocessQueue(Queue):

    def __init__(self):
        try: import queue
        except ImportError: import Queue as queue
        Queue.__init__(self)
        self._queue = queue.PriorityQueue()

    def length(self):
        return self._queue.qsize()

    def clear(self):
        try: import queue
        except ImportError: import Queue as queue
        self._queue = queue.PriorityQueue()

    def push(self, value, priority = None, identify = False):
        value, identifier = self.build_value(
            value,
            priority = priority,
            identify = identify,
            reverse = True
        )
        self._queue.put(value)
        return identifier

    def pop(self, block = True, full = False):
        priority, identifier, value = self._queue.get(block)
        return (priority, identifier, value) if full else value

class AMQPQueue(Queue):

    def __init__(
        self,
        url = None,
        name = "default",
        durable = True,
        max_priority = 256,
        encoder = "json",
        encoding = "utf-8",
        amqp = None
    ):
        self.url = url
        self.name = name
        self.durable = durable
        self.max_priority = max_priority
        self.encoder = encoder
        self.encoding = encoding
        self.amqp = amqp
        self._build()

    def clear(self):
        self.channel.queue_purge(queue = self.name)

    def push(self, value, priority = None, identify = False):
        value, identifier = self.build_value(
            value,
            priority = priority,
            identify = identify,
            reverse = False
        )
        body = self._dump(value)
        self.channel.basic_publish(
            exchange = "",
            routing_key = self.name,
            body = body,
            properties = amqp.properties(
                delivery_mode = 2,
                priority = value[0] or 0
            )
        )
        return identifier

    def pop(self, block = True, full = False):
        _method, _properties, body = self.channel.basic_get(
            queue = self.name, no_ack = True
        )
        priority, identifier, value = self._load(body)
        return (priority, identifier, value) if full else value

    def subscribe(self, callback, full = False):
        def handler(channel, method, properties, body):
            if legacy.is_bytes(body): body = body.decode(self.encoding)
            priority, identifier, value = json.loads(body)
            result = (priority, identifier, value) if full else value
            callback(result)

        self.channel.basic_consume(handler, queue = self.name, no_ack = True)

    def loop(self):
        self.channel.start_consuming()

    def unloop(self):
        self.channel.stop_consuming()

    def _dump(self, value):
        return self._dumper(value)

    def _dump_json(self, value):
        body = json.dumps(value)
        return legacy.bytes(body, encoding = self.encoding)

    def _load(self, body):
        return self._loader(body)

    def _load_json(self, body):
        if legacy.is_bytes(body): body = body.decode(self.encoding)
        return json.loads(body)

    def _build(self):
        if not self.amqp: self.amqp = amqp.AMQP(url = self.url)
        self.connection = self.amqp.get_connection()
        self.channel = self.connection.channel()
        self.channel.basic_qos(prefetch_count = 1, all_channels = True)
        self.queue = self.channel.queue_declare(
            queue = self.name,
            durable = self.durable,
            arguments = {
                "x-max-priority" : self.max_priority
            }
        )
        self._dumper = getattr(self, "_dump_" + self.encoder)
        self._loader = getattr(self, "_load_" + self.encoder)
<commit_msg>default picke engine changed<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

# Hive Appier Framework
# Copyright (c) 2008-2017 Hive Solutions Lda.
#
# This file is part of Hive Appier Framework.
#
# Hive Appier Framework is free software: you can redistribute it and/or modify
# it under the terms of the Apache License as published by the Apache
# Foundation, either version 2.0 of the License, or (at your option) any
# later version.
#
# Hive Appier Framework is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# Apache License for more details.
#
# You should have received a copy of the Apache License along with
# Hive Appier Framework. If not, see <http://www.apache.org/licenses/>.

__author__ = "João Magalhães <joamag@hive.pt>"
""" The author(s) of the module """

__version__ = "1.0.0"
""" The version of the module """

__revision__ = "$LastChangedRevision$"
""" The revision number of the module """

__date__ = "$LastChangedDate$"
""" The last change date of the module """

__copyright__ = "Copyright (c) 2008-2017 Hive Solutions Lda."
""" The copyright for the module """

__license__ = "Apache License, Version 2.0"
""" The license for the module """

import json
import uuid
import heapq

from . import amqp
from . import legacy
from . import exceptions

class Queue(object):

    def length(self):
        raise exceptions.NotImplementedError()

    def clear(self):
        raise exceptions.NotImplementedError()

    def push(self, value, priority = None, identify = False):
        raise exceptions.NotImplementedError()

    def pop(self, block = True, full = False):
        raise exceptions.NotImplementedError()

    def subscribe(self, callback):
        raise exceptions.NotImplementedError()

    def loop(self):
        raise exceptions.NotImplementedError()

    def unloop(self):
        raise exceptions.NotImplementedError()

    def build_value(
        self,
        value,
        priority = None,
        identify = False,
        reverse = False
    ):
        if identify: identifier = self.build_identifier()
        else: identifier = None
        if priority and reverse: priority *= -1
        return (priority, identifier, value), identifier

    def build_identifier(self):
        return str(uuid.uuid4())

class MemoryQueue(Queue):

    def __init__(self):
        Queue.__init__(self)
        self._queue = []

    def length(self):
        return len(self._queue)

    def clear(self):
        del self._queue[:]

    def push(self, value, priority = None, identify = False):
        value, identifier = self.build_value(
            value,
            priority = priority,
            identify = identify,
            reverse = True
        )
        heapq.heappush(self._queue, value)
        return identifier

    def pop(self, block = True, full = False):
        priority, identifier, value = heapq.heappop(self._queue)
        return (priority, identifier, value) if full else value

class MultiprocessQueue(Queue):

    def __init__(self):
        try: import queue
        except ImportError: import Queue as queue
        Queue.__init__(self)
        self._queue = queue.PriorityQueue()

    def length(self):
        return self._queue.qsize()

    def clear(self):
        try: import queue
        except ImportError: import Queue as queue
        self._queue = queue.PriorityQueue()

    def push(self, value, priority = None, identify = False):
        value, identifier = self.build_value(
            value,
            priority = priority,
            identify = identify,
            reverse = True
        )
        self._queue.put(value)
        return identifier

    def pop(self, block = True, full = False):
        priority, identifier, value = self._queue.get(block)
        return (priority, identifier, value) if full else value

class AMQPQueue(Queue):

    def __init__(
        self,
        url = None,
        name = "default",
        durable = True,
        max_priority = 256,
        encoder = "pickle",
        encoding = "utf-8",
        amqp = None
    ):
        self.url = url
        self.name = name
        self.durable = durable
        self.max_priority = max_priority
        self.encoder = encoder
        self.encoding = encoding
        self.amqp = amqp
        self._build()

    def clear(self):
        self.channel.queue_purge(queue = self.name)

    def push(self, value, priority = None, identify = False):
        value, identifier = self.build_value(
            value,
            priority = priority,
            identify = identify,
            reverse = False
        )
        body = self._dump(value)
        self.channel.basic_publish(
            exchange = "",
            routing_key = self.name,
            body = body,
            properties = amqp.properties(
                delivery_mode = 2,
                priority = value[0] or 0
            )
        )
        return identifier

    def pop(self, block = True, full = False):
        _method, _properties, body = self.channel.basic_get(
            queue = self.name, no_ack = True
        )
        priority, identifier, value = self._load(body)
        return (priority, identifier, value) if full else value

    def subscribe(self, callback, full = False):
        def handler(channel, method, properties, body):
            if legacy.is_bytes(body): body = body.decode(self.encoding)
            priority, identifier, value = json.loads(body)
            result = (priority, identifier, value) if full else value
            callback(result)

        self.channel.basic_consume(handler, queue = self.name, no_ack = True)

    def loop(self):
        self.channel.start_consuming()

    def unloop(self):
        self.channel.stop_consuming()

    def _dump(self, value):
        return self._dumper(value)

    def _dump_pickle(self, value):
        return legacy.cPickle.dumps(value)

    def _dump_json(self, value):
        body = json.dumps(value)
        return legacy.bytes(body, encoding = self.encoding)

    def _load(self, body):
        return self._loader(body)

    def _load_pickle(self, body):
        return legacy.cPickle.loads(body)

    def _load_json(self, body):
        if legacy.is_bytes(body): body = body.decode(self.encoding)
        return json.loads(body)

    def _build(self):
        if not self.amqp: self.amqp = amqp.AMQP(url = self.url)
        self.connection = self.amqp.get_connection()
        self.channel = self.connection.channel()
        self.channel.basic_qos(prefetch_count = 1, all_channels = True)
        self.queue = self.channel.queue_declare(
            queue = self.name,
            durable = self.durable,
            arguments = {
                "x-max-priority" : self.max_priority
            }
        )
        self._dumper = getattr(self, "_dump_" + self.encoder)
        self._loader = getattr(self, "_load_" + self.encoder)
<|endoftext|>


<commit_before>"""
Deis cloud provider implementation for Digital Ocean.
"""

from __future__ import unicode_literals

import json
import time

from dop import client
from api.models import Layer

REGION_LIST = {
    'san-francisco-1': 'San Francisco 1',
    'amsterdam-1': 'Amsterdam 1',
    'new-york-1': 'New York 1',
    'new-york-2': 'New York 2'
}

def seed_flavors():
    """Seed the database with default Flavors for digital ocean"""
    flavors = []
    for r in REGION_LIST.keys():
        flavors.append({
            'id': 'digitalocean-{}'.format(r),
            'provider': 'digitalocean',
            'params': json.dumps({
                'region': REGION_LIST[r],
                'size': '4GB',
                'image': 'deis-base'
            })
        })
    return flavors

def build_layer(layer):
    conn = _create_digitalocean_connection(layer['creds'])
    # create a new SSH key
    name = "deis-{formation}-{id}".format(**layer)
    # import a new keypair to the Digital Ocean Control Panel
    params = {'name': name, 'ssh_pub_key': layer['ssh_public_key']}
    conn.request('/ssh_keys/new/', method='GET', params=params)
    # Digital Ocean images only have the root user created by default
    l = Layer.objects.get(id=layer['id'], formation__id=layer['formation'])
    l.ssh_username = 'root'
    l.save()

def destroy_layer(layer):
    conn = _create_digitalocean_connection(layer['creds'])
    name = "deis-{formation}-{id}".format(**layer)
    # retrieve and delete the SSH key created
    for k in conn.all_ssh_keys():
        if k.name == name:
            conn.destroy_ssh_key(k.id)

def build_node(node):
    conn = _create_digitalocean_connection(node['creds'])
    # get digital ocean config for a new node
    kwargs = _get_droplet_kwargs(node, conn)
    # spawn the droplet
    droplet = conn.create_droplet(**kwargs)
    # wait until the droplet has become active
    while conn.show_droplet(droplet.id).status != 'active':
        time.sleep(1)
    provider_id = 'digitalocean'
    # the data in the old droplet is not there, so
    # we just create a new Droplet object
    fqdn = str(conn.show_droplet(droplet.id).ip_address)
    metadata = _get_droplet_metadata(conn.show_droplet(droplet.id))
    return provider_id, fqdn, metadata

def destroy_node(node):
    params, creds = node['params'], node['creds']
    conn = _create_digitalocean_connection(creds)
    name = node['id']
    for i in conn.images():
        if i.name == name:
            conn.destroy_droplet(i.id)
    return

def _get_droplet_kwargs(node, conn):
    params = node['params']
    return {
        'name': node['id'],
        'size_id': _get_id(conn.sizes(), params.get('size', '4GB')),
        'image_id': _get_id(conn.images(show_all=False), params.get('image', 'deis-base')),
        'region_id': _get_id(conn.regions(), params.get('region', 'San Francisco 1')),
        'ssh_key_ids': [str(_get_id(conn.all_ssh_keys(),
                        "deis-{formation}-{layer}".format(**node)))],
        'virtio': True
    }

def _get_droplet_metadata(droplet):
    return {
        'id': droplet.id,
        'name': droplet.name,
        'size_id': droplet.size_id,
        'image_id': droplet.image_id,
        'region_id': droplet.region_id,
        'event_id': droplet.event_id,
        'backups_active': droplet.backups_active,
        'status': droplet.status,
        'ip_address': droplet.ip_address
    }

def _create_digitalocean_connection(creds):
    if not creds:
        raise EnvironmentError('No credentials provided')
    return client.Client(creds['client_id'], creds['api_key'])

def _get_id(lst, name):
    for i in lst:
        if i.name == name:
            return i.id
    return None

def _get_name(lst, name):
    for i in lst:
        if i.name == name:
            return i.name
    return None
<commit_msg>DigitalOcean: Fix for #477<commit_after>"""
Deis cloud provider implementation for Digital Ocean.
"""

from __future__ import unicode_literals

import json
import time

from dop import client
from api.models import Layer

REGION_LIST = {
    'san-francisco-1': 'San Francisco 1',
    'amsterdam-1': 'Amsterdam 1',
    'new-york-1': 'New York 1',
    'new-york-2': 'New York 2'
}

def seed_flavors():
    """Seed the database with default Flavors for digital ocean"""
    flavors = []
    for r in REGION_LIST.keys():
        flavors.append({
            'id': 'digitalocean-{}'.format(r),
            'provider': 'digitalocean',
            'params': json.dumps({
                'region': REGION_LIST[r],
                'size': '4GB',
                'image': 'deis-base'
            })
        })
    return flavors

def build_layer(layer):
    conn = _create_digitalocean_connection(layer['creds'])
    # create a new SSH key
    name = "deis-{formation}-{id}".format(**layer)
    # import a new keypair to the Digital Ocean Control Panel
    params = {'name': name, 'ssh_pub_key': layer['ssh_public_key']}
    conn.request('/ssh_keys/new/', method='GET', params=params)
    # Digital Ocean images only have the root user created by default
    l = Layer.objects.get(id=layer['id'], formation__id=layer['formation'])
    l.ssh_username = 'root'
    l.save()

def destroy_layer(layer):
    conn = _create_digitalocean_connection(layer['creds'])
    name = "deis-{formation}-{id}".format(**layer)
    # retrieve and delete the SSH key created
    for k in conn.all_ssh_keys():
        if k.name == name:
            conn.destroy_ssh_key(k.id)

def build_node(node):
    conn = _create_digitalocean_connection(node['creds'])
    # get digital ocean config for a new node
    kwargs = _get_droplet_kwargs(node, conn)
    # spawn the droplet
    droplet = conn.create_droplet(**kwargs)
    # wait until the droplet has become active
    while conn.show_droplet(droplet.id).status != 'active':
        time.sleep(1)
    provider_id = 'digitalocean'
    # the data in the old droplet is not there, so
    # we just create a new Droplet object
    fqdn = str(conn.show_droplet(droplet.id).ip_address)
    metadata = _get_droplet_metadata(conn.show_droplet(droplet.id))
    return provider_id, fqdn, metadata

def destroy_node(node):
    params, creds = node['params'], node['creds']
    conn = _create_digitalocean_connection(creds)
    name = node['id']
    for i in conn.show_active_droplets():
        if i.name == name:
            conn.destroy_droplet(i.id)
    return

def _get_droplet_kwargs(node, conn):
    params = node['params']
    return {
        'name': node['id'],
        'size_id': _get_id(conn.sizes(), params.get('size', '4GB')),
        'image_id': _get_id(conn.images(show_all=False), params.get('image', 'deis-base')),
        'region_id': _get_id(conn.regions(), params.get('region', 'San Francisco 1')),
        'ssh_key_ids': [str(_get_id(conn.all_ssh_keys(),
                        "deis-{formation}-{layer}".format(**node)))],
        'virtio': True
    }

def _get_droplet_metadata(droplet):
    return {
        'id': droplet.id,
        'name': droplet.name,
        'size_id': droplet.size_id,
        'image_id': droplet.image_id,
        'region_id': droplet.region_id,
        'event_id': droplet.event_id,
        'backups_active': droplet.backups_active,
        'status': droplet.status,
        'ip_address': droplet.ip_address
    }

def _create_digitalocean_connection(creds):
    if not creds:
        raise EnvironmentError('No credentials provided')
    return client.Client(creds['client_id'], creds['api_key'])

def _get_id(lst, name):
    for i in lst:
        if i.name == name:
            return i.id
    return None

def _get_name(lst, name):
    for i in lst:
        if i.name == name:
            return i.name
    return None
<|endoftext|>


<commit_before>__version__ = '0.0.1'

class VasPy(object):
    def __init__(self, filename):
        "Base class to be inherited by all classes in VASPy."
        self.filename = filename

class CarfileValueError(Exception):
    "Exception raised for errors in the CONTCAR-like file."
    pass

class UnmatchedDataShape(Exception):
    "Exception raised for errors in unmatched data shape."
    pass
<commit_msg>modify version number<commit_after>__version__ = '0.0.2'

class VasPy(object):
    def __init__(self, filename):
        "Base class to be inherited by all classes in VASPy."
        self.filename = filename

class CarfileValueError(Exception):
    "Exception raised for errors in the CONTCAR-like file."
    pass

class UnmatchedDataShape(Exception):
    "Exception raised for errors in unmatched data shape."
    pass
<|endoftext|>


<commit_before>import requests
from django import forms
from django.core.exceptions import ValidationError
from django.db.models import Count, Q, Exists, OuterRef
from busstops.models import Operator, Service
from .models import Vehicle, VehicleType, VehicleFeature, Livery
from .fields import RegField

def get_livery_choices(operator):
    choices = {}
    liveries = Livery.objects.filter(vehicle__operator=operator).annotate(popularity=Count('vehicle'))
    for livery in liveries.order_by('-popularity').distinct():
        choices[livery.id] = livery
    for vehicle in operator.vehicle_set.distinct('colours'):
        if not vehicle.livery_id and vehicle.colours and vehicle.colours != 'Other':
            choices[vehicle.colours] = Livery(colours=vehicle.colours, name=f'Like {vehicle}')
    choices = [(key, livery.preview(name=True)) for key, livery in choices.items()]
    choices.append(('Other', 'Other'))
    return choices

class EditVehiclesForm(forms.Form):
    operator = forms.ModelChoiceField(queryset=None, label='Operator', empty_label='')
    vehicle_type = forms.ModelChoiceField(queryset=VehicleType.objects, label='Type', required=False, empty_label='')
    colours = forms.ChoiceField(label='Livery', widget=forms.RadioSelect, required=False)
    other_colour = forms.CharField(widget=forms.TextInput(attrs={"type": "color"}), required=False)
    branding = forms.CharField(label="Other branding", required=False, max_length=255)
    features = forms.ModelMultipleChoiceField(queryset=VehicleFeature.objects, label='Features',
                                              widget=forms.CheckboxSelectMultiple, required=False)
    depot = forms.ChoiceField(required=False)
    notes = forms.CharField(help_text="""Again, this should be blank in almost all cases. There’s no need to know a
vehicle’s previous owners""", required=False, max_length=255)
    withdrawn = forms.BooleanField(label='Permanently withdrawn', required=False)

    def clean_url(self):
        if self.cleaned_data['url']:
            try:
                response = requests.get(self.cleaned_data['url'], timeout=5)
                if response.ok:
                    return self.cleaned_data['url']
            except requests.RequestException:
                pass
            raise ValidationError('That URL doesn’t work for me. Maybe it’s too long, or Facebook')

    def clean_other_colour(self):
        if self.cleaned_data['other_colour']:
            if self.cleaned_data.get('colours') != 'Other':
                return
        return self.cleaned_data['other_colour']

    def has_really_changed(self):
        if not self.has_changed():
            return False
        if all(key == 'url' for key in self.changed_data):
            return False
        return True

    def __init__(self, *args, **kwargs):
        operator = kwargs.pop('operator', None)

        super().__init__(*args, **kwargs)

        if operator:
            self.fields['colours'].choices = get_livery_choices(operator)
            depots = operator.vehicle_set.distinct('data__Depot').values_list('data__Depot', flat=True)
            depots = [('', '')] + sorted([(depot, depot) for depot in depots if depot])
            if depots:
                self.fields['depot'].choices = depots
            else:
                del self.fields['depot']

        operators = None
        if operator and operator.parent:
            has_services = Exists(Service.objects.filter(current=True, operator=OuterRef('pk')).only('id'))
            has_vehicles = Exists(Vehicle.objects.filter(operator=OuterRef('pk')).only('id'))
            operators = Operator.objects.filter(has_services | has_vehicles | Q(pk=operator.pk), parent=operator.parent)
            self.fields['operator'].queryset = operators
        else:
            del(self.fields['operator'])

class EditVehicleForm(EditVehiclesForm):
    """With some extra fields, only applicable to editing a single vehicle
    """
    fleet_number = forms.CharField(required=False, max_length=14)
    reg = RegField(label='Registration', required=False, max_length=14)
    name = forms.CharField(label='Name', required=False, max_length=255)
    previous_reg = RegField(required=False, max_length=14)
    url = forms.URLField(label='URL', help_text='Link to a public (not Facebook) web page or photo showing changes',
                         required=False, max_length=255)
    field_order = ['operator', 'fleet_number', 'reg', 'vehicle_type', 'colours', 'other_colour', 'branding', 'name',
                   'previous_reg', 'features', 'depot', 'notes']

    def __init__(self, *args, **kwargs):
        vehicle = kwargs.pop('vehicle', None)

        super().__init__(*args, **kwargs)

        if str(vehicle.fleet_number) in vehicle.code:
            self.fields['fleet_number'].disabled = True
        if vehicle.reg and vehicle.reg in vehicle.code.replace('_', '').replace(' ', '').replace('-', ''):
            self.fields['reg'].disabled = True

        if not vehicle.notes:
            del self.fields['notes']
            if not vehicle.name:
                del self.fields['name']
            if not vehicle.branding:
                del self.fields['branding']
<commit_msg>fix hiding depot field when there are no depots<commit_after>import requests
from django import forms
from django.core.exceptions import ValidationError
from django.db.models import Count, Q, Exists, OuterRef
from busstops.models import Operator, Service
from .models import Vehicle, VehicleType, VehicleFeature, Livery
from .fields import RegField

def get_livery_choices(operator):
    choices = {}
    liveries = Livery.objects.filter(vehicle__operator=operator).annotate(popularity=Count('vehicle'))
    for livery in liveries.order_by('-popularity').distinct():
        choices[livery.id] = livery
    for vehicle in operator.vehicle_set.distinct('colours'):
        if not vehicle.livery_id and vehicle.colours and vehicle.colours != 'Other':
            choices[vehicle.colours] = Livery(colours=vehicle.colours, name=f'Like {vehicle}')
    choices = [(key, livery.preview(name=True)) for key, livery in choices.items()]
    choices.append(('Other', 'Other'))
    return choices

class EditVehiclesForm(forms.Form):
    operator = forms.ModelChoiceField(queryset=None, label='Operator', empty_label='')
    vehicle_type = forms.ModelChoiceField(queryset=VehicleType.objects, label='Type', required=False, empty_label='')
    colours = forms.ChoiceField(label='Livery', widget=forms.RadioSelect, required=False)
    other_colour = forms.CharField(widget=forms.TextInput(attrs={"type": "color"}), required=False)
    branding = forms.CharField(label="Other branding", required=False, max_length=255)
    features = forms.ModelMultipleChoiceField(queryset=VehicleFeature.objects, label='Features',
                                              widget=forms.CheckboxSelectMultiple, required=False)
    depot = forms.ChoiceField(required=False)
    notes = forms.CharField(help_text="""Again, this should be blank in almost all cases. There’s no need to know a
vehicle’s previous owners""", required=False, max_length=255)
    withdrawn = forms.BooleanField(label='Permanently withdrawn', required=False)

    def clean_url(self):
        if self.cleaned_data['url']:
            try:
                response = requests.get(self.cleaned_data['url'], timeout=5)
                if response.ok:
                    return self.cleaned_data['url']
            except requests.RequestException:
                pass
            raise ValidationError('That URL doesn’t work for me. Maybe it’s too long, or Facebook')

    def clean_other_colour(self):
        if self.cleaned_data['other_colour']:
            if self.cleaned_data.get('colours') != 'Other':
                return
        return self.cleaned_data['other_colour']

    def has_really_changed(self):
        if not self.has_changed():
            return False
        if all(key == 'url' for key in self.changed_data):
            return False
        return True

    def __init__(self, *args, **kwargs):
        operator = kwargs.pop('operator', None)

        super().__init__(*args, **kwargs)

        if operator:
            self.fields['colours'].choices = get_livery_choices(operator)
            depots = operator.vehicle_set.distinct('data__Depot').values_list('data__Depot', flat=True)
            depots = [(depot, depot) for depot in depots if depot]
            if depots:
                self.fields['depot'].choices = [('', '')] + sorted(depots)
            else:
                del self.fields['depot']

        operators = None
        if operator and operator.parent:
            has_services = Exists(Service.objects.filter(current=True, operator=OuterRef('pk')).only('id'))
            has_vehicles = Exists(Vehicle.objects.filter(operator=OuterRef('pk')).only('id'))
            operators = Operator.objects.filter(has_services | has_vehicles | Q(pk=operator.pk), parent=operator.parent)
            self.fields['operator'].queryset = operators
        else:
            del(self.fields['operator'])

class EditVehicleForm(EditVehiclesForm):
    """With some extra fields, only applicable to editing a single vehicle
    """
    fleet_number = forms.CharField(required=False, max_length=14)
    reg = RegField(label='Registration', required=False, max_length=14)
    name = forms.CharField(label='Name', required=False, max_length=255)
    previous_reg = RegField(required=False, max_length=14)
    url = forms.URLField(label='URL', help_text='Link to a public (not Facebook) web page or photo showing changes',
                         required=False, max_length=255)
    field_order = ['operator', 'fleet_number', 'reg', 'vehicle_type', 'colours', 'other_colour', 'branding', 'name',
                   'previous_reg', 'features', 'depot', 'notes']

    def __init__(self, *args, **kwargs):
        vehicle = kwargs.pop('vehicle', None)

        super().__init__(*args, **kwargs)

        if str(vehicle.fleet_number) in vehicle.code:
            self.fields['fleet_number'].disabled = True
        if vehicle.reg and vehicle.reg in vehicle.code.replace('_', '').replace(' ', '').replace('-', ''):
            self.fields['reg'].disabled = True

        if not vehicle.notes:
            del self.fields['notes']
            if not vehicle.name:
                del self.fields['name']
            if not vehicle.branding:
                del self.fields['branding']
<|endoftext|>


<commit_before>import pytest
from ftw.ftw import ruleset, util
import os

def get_rulesets(ruledir):
    """
    List of ruleset objects extracted from the yaml directory
    """
    if os.path.isdir(ruledir):
        yaml_files = util.get_files(ruledir, 'yaml')
    elif os.path.isfile(ruledir):
        yaml_files = [ruledir]
    extracted_files = util.extract_yaml(yaml_files)
    rulesets = []
    for extracted_yaml in extracted_files:
        rulesets.append(ruleset.Ruleset(extracted_yaml))
    return rulesets

def get_testdata(rulesets):
    """
    In order to do test-level parametrization (is this a word?), we have to
    bundle the test data from rulesets into tuples so py.test can understand
    how to run tests across the whole suite of rulesets
    """
    testdata = []
    for ruleset in rulesets:
        for test in ruleset.tests:
            testdata.append((ruleset, test))

    return testdata

def test_id(val):
    """
    Dynamically names tests, useful for when we are running dozens to hundreds
    of tests
    """
    if isinstance(val, (dict,ruleset.Test,)):
        # We must be carful here because errors are swallowed and defaults returned
        if 'name' in val.ruleset_meta.keys():
            return '%s -- %s' % (val.ruleset_meta['name'], val.test_title)
        else:
            return '%s -- %s' % ("Unnamed_Test", val.test_title)

@pytest.fixture
def destaddr(request):
    """
    Destination address override for tests
    """
    return request.config.getoption('--destaddr')

@pytest.fixture
def http_serv_obj():
    """
    Return an HTTP object listening on localhost port 80 for testing
    """
    return HTTPServer(('localhost', 80), SimpleHTTPRequestHandler)

def pytest_addoption(parser):
    """
    Adds command line options to py.test
    """
    parser.addoption('--ruledir', action='store', default=None,
        help='rule directory that holds YAML files for testing')
    parser.addoption('--destaddr', action='store', default=None,
        help='destination address to direct tests towards')
    parser.addoption('--rule', action='store', default=None,
        help='fully qualified path to one rule')

def pytest_generate_tests(metafunc):
    """
    Pre-test configurations, mostly used for parametrization
    """
    if metafunc.config.option.ruledir or metafunc.config.option.rule:
        if metafunc.config.option.ruledir:
            rulesets = get_rulesets(metafunc.config.option.ruledir)
        if metafunc.config.option.rule:
            rulesets = get_rulesets(metafunc.config.option.rule)
        if 'ruleset' in metafunc.fixturenames and 'test' in metafunc.fixturenames:
            metafunc.parametrize('ruleset,test', get_testdata(rulesets),
                ids=test_id)
<commit_msg>Adding support for ruledir_recurse<commit_after>import pytest
from ftw.ftw import ruleset, util
import os

def get_rulesets(ruledir, recurse):
    """
    List of ruleset objects extracted from the yaml directory
    """
    if os.path.isdir(ruledir) and recurse is True:
        yaml_files = []
        for root, dirs, files in os.walk(ruledir):
            for name in files:
                filename, file_extension = os.path.splitext(name)
                if file_extension == ".yaml":
                    yaml_files.append(os.path.join(root, name))
    if os.path.isdir(ruledir) and recurse is False:
        yaml_files = util.get_files(ruledir, 'yaml')
    elif os.path.isfile(ruledir):
        yaml_files = [ruledir]
    extracted_files = util.extract_yaml(yaml_files)
    rulesets = []
    for extracted_yaml in extracted_files:
        rulesets.append(ruleset.Ruleset(extracted_yaml))
    return rulesets

def get_testdata(rulesets):
    """
    In order to do test-level parametrization (is this a word?), we have to
    bundle the test data from rulesets into tuples so py.test can understand
    how to run tests across the whole suite of rulesets
    """
    testdata = []
    for ruleset in rulesets:
        for test in ruleset.tests:
            testdata.append((ruleset, test))

    return testdata

def test_id(val):
    """
    Dynamically names tests, useful for when we are running dozens to hundreds
    of tests
    """
    if isinstance(val, (dict,ruleset.Test,)):
        # We must be carful here because errors are swallowed and defaults returned
        if 'name' in val.ruleset_meta.keys():
            return '%s -- %s' % (val.ruleset_meta['name'], val.test_title)
        else:
            return '%s -- %s' % ("Unnamed_Test", val.test_title)

@pytest.fixture
def destaddr(request):
    """
    Destination address override for tests
    """
    return request.config.getoption('--destaddr')

@pytest.fixture
def http_serv_obj():
    """
    Return an HTTP object listening on localhost port 80 for testing
    """
    return HTTPServer(('localhost', 80), SimpleHTTPRequestHandler)

def pytest_addoption(parser):
    """
    Adds command line options to py.test
    """
    parser.addoption('--ruledir', action='store', default=None,
        help='rule directory that holds YAML files for testing')
    parser.addoption('--destaddr', action='store', default=None,
        help='destination address to direct tests towards')
    parser.addoption('--rule', action='store', default=None,
        help='fully qualified path to one rule')
    parser.addoption('--ruledir_recurse', action='store', default=None,
        help='walk the directory structure finding YAML files')

def pytest_generate_tests(metafunc):
    """
    Pre-test configurations, mostly used for parametrization
    """
    options = ['ruledir','ruledir_recurse','rule']
    args = metafunc.config.option.__dict__
    # Check if we have any arguments by creating a list of supplied args we want
    if [i for i in options if i in args and args[i] != None] :
        if metafunc.config.option.ruledir:
            rulesets = get_rulesets(metafunc.config.option.ruledir, False)
        if metafunc.config.option.ruledir_recurse:
            rulesets = get_rulesets(metafunc.config.option.ruledir_recurse, True)
        if metafunc.config.option.rule:
            rulesets = get_rulesets(metafunc.config.option.rule, False)
        if 'ruleset' in metafunc.fixturenames and 'test' in metafunc.fixturenames:
            metafunc.parametrize('ruleset,test', get_testdata(rulesets),
                ids=test_id)<|endoftext|>


<commit_before>"""Pytest fixtures for testing."""

import toml
import yaml
import json
import os
import shutil
import pytest

def patch_init(moulinette):
    """Configure moulinette to use the YunoHost namespace."""
    old_init = moulinette.core.Moulinette18n.__init__

    def monkey_path_i18n_init(self, package, default_locale="en"):
        old_init(self, package, default_locale)
        self.load_namespace("moulinette")

    moulinette.core.Moulinette18n.__init__ = monkey_path_i18n_init

def patch_translate(moulinette):
    """Configure translator to raise errors when there are missing keys."""
    old_translate = moulinette.core.Translator.translate

    def new_translate(self, key, *args, **kwargs):
        if key not in self._translations[self.default_locale].keys():
            message = "Unable to retrieve key %s for default locale!" % key
            raise KeyError(message)

        return old_translate(self, key, *args, **kwargs)

    moulinette.core.Translator.translate = new_translate

    def new_m18nn(self, key, *args, **kwargs):
        return self._global.translate(key, *args, **kwargs)

    moulinette.core.Moulinette18n.g = new_m18nn

def patch_logging(moulinette):
    """Configure logging to use the custom logger."""
    handlers = set(["tty", "api"])
    root_handlers = set(handlers)

    level = "INFO"
    tty_level = "INFO"

    return {
        "version": 1,
        "disable_existing_loggers": True,
        "formatters": {
            "tty-debug": {"format": "%(relativeCreated)-4d %(fmessage)s"},
            "precise": {
                "format": "%(asctime)-15s %(levelname)-8s %(name)s %(funcName)s - %(fmessage)s"  # noqa
            },
        },
        "filters": {"action": {"()": "moulinette.utils.log.ActionFilter",},},
        "handlers": {
            "api": {
                "level": level,
                "class": "moulinette.interfaces.api.APIQueueHandler",
            },
            "tty": {
                "level": tty_level,
                "class": "moulinette.interfaces.cli.TTYHandler",
                "formatter": "",
            },
        },
        "loggers": {
            "moulinette": {"level": level, "handlers": [], "propagate": True,},
            "moulinette.interface": {
                "level": level,
                "handlers": handlers,
                "propagate": False,
            },
        },
        "root": {"level": level, "handlers": root_handlers,},
    }

@pytest.fixture(scope="session", autouse=True)
def moulinette(tmp_path_factory):
    import moulinette

    # Can't call the namespace just 'test' because
    # that would lead to some "import test" not importing the right stuff
    namespace = "moulitest"
    tmp_cache = str(tmp_path_factory.mktemp("cache"))
    tmp_data = str(tmp_path_factory.mktemp("data"))
    tmp_lib = str(tmp_path_factory.mktemp("lib"))
    os.environ["MOULINETTE_CACHE_DIR"] = tmp_cache
    os.environ["MOULINETTE_DATA_DIR"] = tmp_data
    os.environ["MOULINETTE_LIB_DIR"] = tmp_lib
    shutil.copytree("./test/actionsmap", "%s/actionsmap" % tmp_data)
    shutil.copytree("./test/src", "%s/%s" % (tmp_lib, namespace))
    shutil.copytree("./test/locales", "%s/%s/locales" % (tmp_lib, namespace))

    patch_init(moulinette)
    patch_translate(moulinette)
    logging = patch_logging(moulinette)

    moulinette.init(logging_config=logging, _from_source=False)

    return moulinette

@pytest.fixture
def moulinette_webapi(moulinette):

    from webtest import TestApp
    from webtest.app import CookiePolicy

    # Dirty hack needed, otherwise cookies ain't reused between request .. not
    # sure why :|
    def return_true(self, cookie, request):
        return True

    CookiePolicy.return_ok_secure = return_true

    moulinette_webapi = moulinette.core.init_interface(
        "api",
        kwargs={"routes": {}, "use_websocket": False},
        actionsmap={"namespaces": ["moulitest"], "use_cache": True},
    )

    return TestApp(moulinette_webapi._app)

@pytest.fixture
def moulinette_cli(moulinette):
    # Dirty hack needed, otherwise cookies ain't reused between request .. not
    # sure why :|
    import argparse
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument('--debug',
                        action='store_true', default=False,
                        help="Log and print debug messages",
                        )
    moulinette_cli = moulinette.core.init_interface(
        "cli",
        actionsmap={"namespaces": ["moulitest"], "use_cache": False, "parser_kwargs": {'top_parser': parser}},
    )

    return moulinette_cli

@pytest.fixture
def test_file(tmp_path):
    test_text = "foo\nbar\n"
    test_file = tmp_path / "test.txt"
    test_file.write_bytes(test_text)
    return test_file

@pytest.fixture
def test_json(tmp_path):
    test_json = json.dumps({"foo": "bar"})
    test_file = tmp_path / "test.json"
    test_file.write_bytes(test_json)
    return test_file

@pytest.fixture
def test_yaml(tmp_path):
    test_yaml = yaml.dump({"foo": "bar"})
    test_file = tmp_path / "test.txt"
    test_file.write_bytes(test_yaml)
    return test_file

@pytest.fixture
def test_toml(tmp_path):
    test_toml = toml.dumps({"foo": "bar"})
    test_file = tmp_path / "test.txt"
    test_file.write_bytes(str(test_toml))
    return test_file

@pytest.fixture
def test_ldif(tmp_path):
    test_file = tmp_path / "test.txt"
    from ldif import LDIFWriter
    writer = LDIFWriter(open(str(test_file), 'wb'))

    writer.unparse('mail=alice@example.com', {
        'cn': ['Alice Alison'],
        'mail': ['alice@example.com'],
        'objectclass': ['top', 'person']
    })

    return test_file

@pytest.fixture
def user():
    return os.getlogin()

@pytest.fixture
def test_url():
    return "https://some.test.url/yolo.txt"
<commit_msg>fix moulinette cli _do_authenticate<commit_after>"""Pytest fixtures for testing."""

import toml
import yaml
import json
import os
import shutil
import pytest

def patch_init(moulinette):
    """Configure moulinette to use the YunoHost namespace."""
    old_init = moulinette.core.Moulinette18n.__init__

    def monkey_path_i18n_init(self, package, default_locale="en"):
        old_init(self, package, default_locale)
        self.load_namespace("moulinette")

    moulinette.core.Moulinette18n.__init__ = monkey_path_i18n_init

def patch_translate(moulinette):
    """Configure translator to raise errors when there are missing keys."""
    old_translate = moulinette.core.Translator.translate

    def new_translate(self, key, *args, **kwargs):
        if key not in self._translations[self.default_locale].keys():
            message = "Unable to retrieve key %s for default locale!" % key
            raise KeyError(message)

        return old_translate(self, key, *args, **kwargs)

    moulinette.core.Translator.translate = new_translate

    def new_m18nn(self, key, *args, **kwargs):
        return self._global.translate(key, *args, **kwargs)

    moulinette.core.Moulinette18n.g = new_m18nn

def patch_logging(moulinette):
    """Configure logging to use the custom logger."""
    handlers = set(["tty", "api"])
    root_handlers = set(handlers)

    level = "INFO"
    tty_level = "INFO"

    return {
        "version": 1,
        "disable_existing_loggers": True,
        "formatters": {
            "tty-debug": {"format": "%(relativeCreated)-4d %(fmessage)s"},
            "precise": {
                "format": "%(asctime)-15s %(levelname)-8s %(name)s %(funcName)s - %(fmessage)s"  # noqa
            },
        },
        "filters": {"action": {"()": "moulinette.utils.log.ActionFilter",},},
        "handlers": {
            "api": {
                "level": level,
                "class": "moulinette.interfaces.api.APIQueueHandler",
            },
            "tty": {
                "level": tty_level,
                "class": "moulinette.interfaces.cli.TTYHandler",
                "formatter": "",
            },
        },
        "loggers": {
            "moulinette": {"level": level, "handlers": [], "propagate": True,},
            "moulinette.interface": {
                "level": level,
                "handlers": handlers,
                "propagate": False,
            },
        },
        "root": {"level": level, "handlers": root_handlers,},
    }

@pytest.fixture(scope="session", autouse=True)
def moulinette(tmp_path_factory):
    import moulinette

    # Can't call the namespace just 'test' because
    # that would lead to some "import test" not importing the right stuff
    namespace = "moulitest"
    tmp_cache = str(tmp_path_factory.mktemp("cache"))
    tmp_data = str(tmp_path_factory.mktemp("data"))
    tmp_lib = str(tmp_path_factory.mktemp("lib"))
    os.environ["MOULINETTE_CACHE_DIR"] = tmp_cache
    os.environ["MOULINETTE_DATA_DIR"] = tmp_data
    os.environ["MOULINETTE_LIB_DIR"] = tmp_lib
    shutil.copytree("./test/actionsmap", "%s/actionsmap" % tmp_data)
    shutil.copytree("./test/src", "%s/%s" % (tmp_lib, namespace))
    shutil.copytree("./test/locales", "%s/%s/locales" % (tmp_lib, namespace))

    patch_init(moulinette)
    patch_translate(moulinette)
    logging = patch_logging(moulinette)

    moulinette.init(logging_config=logging, _from_source=False)

    return moulinette

@pytest.fixture
def moulinette_webapi(moulinette):

    from webtest import TestApp
    from webtest.app import CookiePolicy

    # Dirty hack needed, otherwise cookies ain't reused between request .. not
    # sure why :|
    def return_true(self, cookie, request):
        return True

    CookiePolicy.return_ok_secure = return_true

    moulinette_webapi = moulinette.core.init_interface(
        "api",
        kwargs={"routes": {}, "use_websocket": False},
        actionsmap={"namespaces": ["moulitest"], "use_cache": True},
    )

    return TestApp(moulinette_webapi._app)

@pytest.fixture
def moulinette_cli(moulinette, mocker):
    # Dirty hack needed, otherwise cookies ain't reused between request .. not
    # sure why :|
    import argparse
    parser = argparse.ArgumentParser(add_help=False)
    parser.add_argument('--debug',
                        action='store_true', default=False,
                        help="Log and print debug messages",
                        )
    mocker.patch("os.isatty", return_value=True)
    moulinette_cli = moulinette.core.init_interface(
        "cli",
        actionsmap={"namespaces": ["moulitest"], "use_cache": False, "parser_kwargs": {'top_parser': parser}},
    )
    mocker.stopall()

    return moulinette_cli

@pytest.fixture
def test_file(tmp_path):
    test_text = "foo\nbar\n"
    test_file = tmp_path / "test.txt"
    test_file.write_bytes(test_text)
    return test_file

@pytest.fixture
def test_json(tmp_path):
    test_json = json.dumps({"foo": "bar"})
    test_file = tmp_path / "test.json"
    test_file.write_bytes(test_json)
    return test_file

@pytest.fixture
def test_yaml(tmp_path):
    test_yaml = yaml.dump({"foo": "bar"})
    test_file = tmp_path / "test.txt"
    test_file.write_bytes(test_yaml)
    return test_file

@pytest.fixture
def test_toml(tmp_path):
    test_toml = toml.dumps({"foo": "bar"})
    test_file = tmp_path / "test.txt"
    test_file.write_bytes(str(test_toml))
    return test_file

@pytest.fixture
def test_ldif(tmp_path):
    test_file = tmp_path / "test.txt"
    from ldif import LDIFWriter
    writer = LDIFWriter(open(str(test_file), 'wb'))

    writer.unparse('mail=alice@example.com', {
        'cn': ['Alice Alison'],
        'mail': ['alice@example.com'],
        'objectclass': ['top', 'person']
    })

    return test_file

@pytest.fixture
def user():
    return os.getlogin()

@pytest.fixture
def test_url():
    return "https://some.test.url/yolo.txt"
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from ftplib import FTP_TLS, all_errors
import requests
import logging
import base64
import uuid
import io

# lib to know file extension (only for images)
import imghdr

_logger = logging.getLogger(__name__)
_debug_level = 1 # 1 Low or 2 High

#   _____ _____ _____    _____ _____ __    _____ _____ _____
#  |   __|_   _|  _  |  |  |  |   __|  |  |  _  |   __| __  |
#  |   __| | | |   __|  |     |   __|  |__|   __|   __|    -|
#  |__|    |_| |__|     |__|__|_____|_____|__|  |_____|__|__|
#

def _ftp_connect(settings):
	"""
	FTP TLS Connection from settings
	"""
	try:
		ftps = FTP_TLS(settings.get('server'), settings.get('user'), settings.get('password'))

		# Set debug (if debug mode is ON)
		if request.debug:
			ftps.set_debuglevel(_debug_level)
		# Secure connection
		ftps.prot_p()
		return ftps

	except Exception as e:
		_logger.error('Connection Error: %s', e)
	return False

def _img_extension(bytestream):
	"""
	Get base64 bytestream extension
	:param bytestream: List of bytes
	:return string: File extension
	"""
	return imghdr.what(None, h=bytestream).replace('jpeg', 'jpg')

def _redundant_check(filename):
	"""
	Redundant file check
	:param filename: Image complete name
	:return bool: File exists
	"""
	try:
		base_url = request.env['b2b.settings'].get_param('base_url')
		response = requests.get(base_url + filename, stream=True)
		if response.status_code == 200:
			return True
	except Exception as e:
		_logger.error('Redundant Check Error: %s', e)
	return False

def save_file(filename, settings=None):
	"""
	Save local file on FTP server
	:param filename: String, file name
	:return bool: File saved
	"""
	if filename:
		try:
			ftps = _ftp_connect(settings)
			if ftps:
				file = open(filename, 'rb')
				ftps.storbinary('STOR ' + filename, file)
				file.close() # Free file memory
				ftps.quit()
				_logger.info('File [%s] saved!', filename)
				return True
		except (all_errors, Exception) as e:
			_logger.error('File Error: %s', e)
	return False

def save_base64(base64_str, settings=None):
	"""
	Save decoded base64 on FTP server
	:param filename_without_ext: String, file name
	:param base64_str: String, base64 string
	:return str|bool: File name or False
	"""
	bytestream = base64.b64decode(base64_str)
	filename = '%s.%s' % (uuid.uuid4(), _img_extension(bytestream))

	if filename and bytestream:
		try:
			ftps = _ftp_connect(settings)
			if ftps:
				file = io.BytesIO(bytestream)
				ftps.storbinary('STOR ' + filename, file)
				file.close() # Free file memory
				ftps.quit()
				_logger.info('File stream [%s] saved!', filename)
				return filename
		except (all_errors, Exception) as e:
			_logger.error('Bytes Error: %s', e)
	return False

def delete_file(filename, settings=None):
	"""
	Delete file from FTP server
	:param filename: String, file name
	:return bool: File deleted
	"""
	if filename:
		try:
			ftps = _ftp_connect(settings)
			if ftps:
				ftps.delete(filename)
				ftps.quit()
				_logger.info('File [%s] deleted!', filename)
				return True
		except all_errors as e:
			_logger.error('Delete Error: %s', e)
	return False<commit_msg>Update ftp.py<commit_after># -*- coding: utf-8 -*-
from ftplib import FTP_TLS, all_errors
import requests
import logging
import base64
import uuid
import io

# lib to know file extension (only for images)
import imghdr

_logger = logging.getLogger(__name__)
_debug_level = 1 # 1 Low or 2 High

#   _____ _____ _____    _____ _____ __    _____ _____ _____
#  |   __|_   _|  _  |  |  |  |   __|  |  |  _  |   __| __  |
#  |   __| | | |   __|  |     |   __|  |__|   __|   __|    -|
#  |__|    |_| |__|     |__|__|_____|_____|__|  |_____|__|__|
#

def _ftp_connect(settings):
	"""
	FTP TLS Connection from settings
	"""
	try:
		ftps = FTP_TLS(settings.get('server'), settings.get('user'), settings.get('password'))

		# Set debug (if debug mode is ON)
		if request.debug:
			ftps.set_debuglevel(_debug_level)
		# Secure connection
		ftps.prot_p()
		return ftps

	except Exception as e:
		_logger.error('Connection Error: %s', e)
	return False

def _img_extension(bytestream):
	"""
	Get base64 bytestream extension
	:param bytestream: List of bytes
	:return string: File extension
	"""
	return imghdr.what(None, h=bytestream).replace('jpeg', 'jpg')

def _redundant_check(filename):
	"""
	Redundant file check
	:param filename: Image complete name
	:return bool: File exists
	"""
	try:
		base_url = request.env['b2b.settings'].get_param('base_url')
		response = requests.get(base_url + filename, stream=True)
		if response.status_code == 200:
			return True
	except Exception as e:
		_logger.error('Redundant Check Error: %s', e)
	return False

def save_file(filename, settings):
	"""
	Save local file on FTP server
	:param filename: String, file name
	:return bool: File saved
	"""
	if filename:
		try:
			ftps = _ftp_connect(settings)
			if ftps:
				file = open(filename, 'rb')
				ftps.storbinary('STOR ' + filename, file)
				file.close() # Free file memory
				ftps.quit()
				_logger.info('File [%s] saved!', filename)
				return True
		except (all_errors, Exception) as e:
			_logger.error('File Error: %s', e)
	return False

def save_base64(base64_str, settings):
	"""
	Save decoded base64 on FTP server
	:param filename_without_ext: String, file name
	:param base64_str: String, base64 string
	:return str|bool: File name or False
	"""
	bytestream = base64.b64decode(base64_str)
	filename = '%s.%s' % (uuid.uuid4(), _img_extension(bytestream))

	if filename and bytestream:
		try:
			ftps = _ftp_connect(settings)
			if ftps:
				file = io.BytesIO(bytestream)
				ftps.storbinary('STOR ' + filename, file)
				file.close() # Free file memory
				ftps.quit()
				_logger.info('File stream [%s] saved!', filename)
				return filename
		except (all_errors, Exception) as e:
			_logger.error('Bytes Error: %s', e)
	return False

def delete_file(filename, settings):
	"""
	Delete file from FTP server
	:param filename: String, file name
	:return bool: File deleted
	"""
	if filename:
		try:
			ftps = _ftp_connect(settings)
			if ftps:
				ftps.delete(filename)
				ftps.quit()
				_logger.info('File [%s] deleted!', filename)
				return True
		except all_errors as e:
			_logger.error('Delete Error: %s', e)
	return False<|endoftext|>


<commit_before>import fcntl
import os
import platform
import re
import shutil
import signal
import socket
import stat
import subprocess
import sys
import tempfile
import time
from multiprocessing import Process

import pytest

from unit.check.go import check_go
from unit.check.node import check_node
from unit.check.tls import check_openssl
from unit.check.isolation import check_isolation
from unit.option import option
from unit.utils import public_dir
from unit.utils import waitforfiles

def pytest_addoption(parser):
    parser.addoption(
        "--detailed",
        default=False,
        action="store_true",
        help="Detailed output for tests",
    )
    parser.addoption(
        "--print-log",
        default=False,
        action="store_true",
        help="Print unit.log to stdout in case of errors",
    )
    parser.addoption(
        "--save-log",
        default=False,
        action="store_true",
        help="Save unit.log after the test execution",
    )
    parser.addoption(
        "--unsafe",
        default=False,
        action="store_true",
        help="Run unsafe tests",
    )

unit_instance = {}
_processes = []

def pytest_configure(config):
    option.config = config.option

    option.detailed = config.option.detailed
    option.print_log = config.option.print_log
    option.save_log = config.option.save_log
    option.unsafe = config.option.unsafe

    option.generated_tests = {}
    option.current_dir = os.path.abspath(
        os.path.join(os.path.dirname(__file__), os.pardir)
    )
    option.test_dir = option.current_dir + '/test'
    option.architecture = platform.architecture()[0]
    option.system = platform.system()

    option.cache_dir = tempfile.mkdtemp(prefix='unit-test-cache-')
    public_dir(option.cache_dir)

    # set stdout to non-blocking

    if option.detailed or option.print_log:
        fcntl.fcntl(sys.stdout.fileno(), fcntl.F_SETFL, 0)

def pytest_generate_tests(metafunc):
    cls = metafunc.cls
    if (not hasattr(cls, 'application_type')
            or cls.application_type == None
            or cls.application_type == 'external'):
        return

    type = cls.application_type

    def generate_tests(versions):
        metafunc.fixturenames.append('tmp_ct')
        metafunc.parametrize('tmp_ct', versions)

        for version in versions:
            option.generated_tests[
                metafunc.function.__name__ + '[{}]'.format(version)
            ] = (type + ' ' + version)

    # take available module from option and generate tests for each version

    for module, prereq_version in cls.prerequisites['modules'].items():
        if module in option.available['modules']:
            available_versions = option.available['modules'][module]

            if prereq_version == 'all':
                generate_tests(available_versions)

            elif prereq_version == 'any':
                option.generated_tests[metafunc.function.__name__] = (
                    type + ' ' + available_versions[0]
                )
            elif callable(prereq_version):
                generate_tests(
                    list(filter(prereq_version, available_versions))
                )

            else:
                raise ValueError(
                    """
Unexpected prerequisite version "%s" for module "%s" in %s.
'all', 'any' or callable expected."""
                    % (str(prereq_version), module, str(cls))
                )

def pytest_sessionstart(session):
    option.available = {'modules': {}, 'features': {}}

    unit = unit_run()
    option.temp_dir = unit['temp_dir']

    # read unit.log

    for i in range(50):
        with open(unit['temp_dir'] + '/unit.log', 'r') as f:
            log = f.read()
            m = re.search('controller started', log)

            if m is None:
                time.sleep(0.1)
            else:
                break

    if m is None:
        _print_log(log)
        exit("Unit is writing log too long")

    # discover available modules from unit.log

    for module in re.findall(r'module: ([a-zA-Z]+) (.*) ".*"$', log, re.M):
        versions = option.available['modules'].setdefault(module[0], [])
        if module[1] not in versions:
            versions.append(module[1])

    # discover modules from check

    option.available['modules']['openssl'] = check_openssl(unit['unitd'])
    option.available['modules']['go'] = check_go(
        option.current_dir, unit['temp_dir'], option.test_dir
    )
    option.available['modules']['node'] = check_node(option.current_dir)

    # remove None values

    option.available['modules'] = {
        k: v for k, v in option.available['modules'].items() if v is not None
    }

    check_isolation()

    unit_stop()

    _check_alerts()

    shutil.rmtree(unit_instance['temp_dir'])

@pytest.hookimpl(tryfirst=True, hookwrapper=True)
def pytest_runtest_makereport(item, call):
    # execute all other hooks to obtain the report object
    outcome = yield
    rep = outcome.get_result()

    # set a report attribute for each phase of a call, which can
    # be "setup", "call", "teardown"

    setattr(item, "rep_" + rep.when, rep)

@pytest.fixture(scope='class', autouse=True)
def check_prerequisites(request):
    cls = request.cls
    missed = []

    # check modules

    if 'modules' in cls.prerequisites:
        available_modules = list(option.available['modules'].keys())

        for module in cls.prerequisites['modules']:
            if module in available_modules:
                continue

            missed.append(module)

    if missed:
        pytest.skip('Unit has no ' + ', '.join(missed) + ' module(s)')

    # check features

    if 'features' in cls.prerequisites:
        available_features = list(option.available['features'].keys())

        for feature in cls.prerequisites['features']:
            if feature in available_features:
                continue

            missed.append(feature)

    if missed:
        pytest.skip(', '.join(missed) + ' feature(s) not supported')

@pytest.fixture(autouse=True)
def run(request):
    unit = unit_run()
    option.temp_dir = unit['temp_dir']

    option.skip_alerts = [
        r'read signalfd\(4\) failed',
        r'sendmsg.+failed',
        r'recvmsg.+failed',
    ]
    option.skip_sanitizer = False

    yield

    # stop unit

    error = unit_stop()

    if error:
        _print_log()

    assert error is None, 'stop unit'

    # stop all processes

    error = stop_processes()

    if error:
        _print_log()

    assert error is None, 'stop unit'

    # check unit.log for alerts

    _check_alerts()

    # print unit.log in case of error

    if hasattr(request.node, 'rep_call') and request.node.rep_call.failed:
        _print_log()

    # remove unit.log

    if not option.save_log:
        shutil.rmtree(unit['temp_dir'])

def unit_run():
    global unit_instance
    build_dir = option.current_dir + '/build'
    unitd = build_dir + '/unitd'

    if not os.path.isfile(unitd):
        exit('Could not find unit')

    temp_dir = tempfile.mkdtemp(prefix='unit-test-')
    public_dir(temp_dir)

    if oct(stat.S_IMODE(os.stat(build_dir).st_mode)) != '0o777':
        public_dir(build_dir)

    os.mkdir(temp_dir + '/state')

    with open(temp_dir + '/unit.log', 'w') as log:
        unit_instance['process'] = subprocess.Popen(
            [
                unitd,
                '--no-daemon',
                '--modules',
                build_dir,
                '--state',
                temp_dir + '/state',
                '--pid',
                temp_dir + '/unit.pid',
                '--log',
                temp_dir + '/unit.log',
                '--control',
                'unix:' + temp_dir + '/control.unit.sock',
                '--tmp',
                temp_dir,
            ],
            stderr=log,
        )

    if not waitforfiles(temp_dir + '/control.unit.sock'):
        _print_log()
        exit('Could not start unit')

    unit_instance['temp_dir'] = temp_dir
    unit_instance['log'] = temp_dir + '/unit.log'
    unit_instance['control_sock'] = temp_dir + '/control.unit.sock'
    unit_instance['unitd'] = unitd

    return unit_instance

def unit_stop():
    p = unit_instance['process']

    if p.poll() is not None:
        return

    p.send_signal(signal.SIGQUIT)

    try:
        retcode = p.wait(15)
        if retcode:
            return 'Child process terminated with code ' + str(retcode)

    except KeyboardInterrupt:
        p.kill()
        raise

    except:
        p.kill()
        return 'Could not terminate unit'

def _check_alerts(path=None):
    if path is None:
        path = unit_instance['log']

    with open(path, 'r', encoding='utf-8', errors='ignore') as f:
        log = f.read()

    found = False

    alerts = re.findall(r'.+\[alert\].+', log)

    if alerts:
        print('\nAll alerts/sanitizer errors found in log:')
        [print(alert) for alert in alerts]
        found = True

    if option.skip_alerts:
        for skip in option.skip_alerts:
            alerts = [al for al in alerts if re.search(skip, al) is None]

    if alerts:
        _print_log(log)
        assert not alerts, 'alert(s)'

    if not option.skip_sanitizer:
        sanitizer_errors = re.findall('.+Sanitizer.+', log)

        if sanitizer_errors:
            _print_log(log)
            assert not sanitizer_errors, 'sanitizer error(s)'

    if found:
        print('skipped.')

def _print_log(data=None):
    path = unit_instance['log']

    print('Path to unit.log:\n' + path + '\n')

    if option.print_log:
        os.set_blocking(sys.stdout.fileno(), True)
        sys.stdout.flush()

        if data is None:
            with open(path, 'r', encoding='utf-8', errors='ignore') as f:
                shutil.copyfileobj(f, sys.stdout)
        else:
            sys.stdout.write(data)

def run_process(target, *args):
    global _processes

    process = Process(target=target, args=args)
    process.start()

    _processes.append(process)

def stop_processes():
    if not _processes:
        return

    fail = False
    for process in _processes:
        if process.is_alive():
            process.terminate()
            process.join(timeout=15)

            if process.is_alive():
                fail = True

    if fail:
        return 'Fail to stop process(es)'

@pytest.fixture()
def skip_alert():
    def _skip(*alerts):
        option.skip_alerts.extend(alerts)

    return _skip

@pytest.fixture
def temp_dir(request):
    return unit_instance['temp_dir']

@pytest.fixture
def is_unsafe(request):
    return request.config.getoption("--unsafe")

@pytest.fixture
def is_su(request):
    return os.geteuid() == 0

def pytest_sessionfinish(session):
    unit_stop()
    shutil.rmtree(option.cache_dir)
<commit_msg>Tests: ability to run unitd with specified "--user" option.<commit_after>import fcntl
import os
import platform
import re
import shutil
import signal
import socket
import stat
import subprocess
import sys
import tempfile
import time
from multiprocessing import Process

import pytest

from unit.check.go import check_go
from unit.check.node import check_node
from unit.check.tls import check_openssl
from unit.check.isolation import check_isolation
from unit.option import option
from unit.utils import public_dir
from unit.utils import waitforfiles

def pytest_addoption(parser):
    parser.addoption(
        "--detailed",
        default=False,
        action="store_true",
        help="Detailed output for tests",
    )
    parser.addoption(
        "--print-log",
        default=False,
        action="store_true",
        help="Print unit.log to stdout in case of errors",
    )
    parser.addoption(
        "--save-log",
        default=False,
        action="store_true",
        help="Save unit.log after the test execution",
    )
    parser.addoption(
        "--unsafe",
        default=False,
        action="store_true",
        help="Run unsafe tests",
    )
    parser.addoption(
        "--user",
        type=str,
        help="Default user for non-privileged processes of unitd",
    )

unit_instance = {}
_processes = []

def pytest_configure(config):
    option.config = config.option

    option.detailed = config.option.detailed
    option.print_log = config.option.print_log
    option.save_log = config.option.save_log
    option.unsafe = config.option.unsafe
    option.user = config.option.user

    option.generated_tests = {}
    option.current_dir = os.path.abspath(
        os.path.join(os.path.dirname(__file__), os.pardir)
    )
    option.test_dir = option.current_dir + '/test'
    option.architecture = platform.architecture()[0]
    option.system = platform.system()

    option.cache_dir = tempfile.mkdtemp(prefix='unit-test-cache-')
    public_dir(option.cache_dir)

    # set stdout to non-blocking

    if option.detailed or option.print_log:
        fcntl.fcntl(sys.stdout.fileno(), fcntl.F_SETFL, 0)

def pytest_generate_tests(metafunc):
    cls = metafunc.cls
    if (not hasattr(cls, 'application_type')
            or cls.application_type == None
            or cls.application_type == 'external'):
        return

    type = cls.application_type

    def generate_tests(versions):
        metafunc.fixturenames.append('tmp_ct')
        metafunc.parametrize('tmp_ct', versions)

        for version in versions:
            option.generated_tests[
                metafunc.function.__name__ + '[{}]'.format(version)
            ] = (type + ' ' + version)

    # take available module from option and generate tests for each version

    for module, prereq_version in cls.prerequisites['modules'].items():
        if module in option.available['modules']:
            available_versions = option.available['modules'][module]

            if prereq_version == 'all':
                generate_tests(available_versions)

            elif prereq_version == 'any':
                option.generated_tests[metafunc.function.__name__] = (
                    type + ' ' + available_versions[0]
                )
            elif callable(prereq_version):
                generate_tests(
                    list(filter(prereq_version, available_versions))
                )

            else:
                raise ValueError(
                    """
Unexpected prerequisite version "%s" for module "%s" in %s.
'all', 'any' or callable expected."""
                    % (str(prereq_version), module, str(cls))
                )

def pytest_sessionstart(session):
    option.available = {'modules': {}, 'features': {}}

    unit = unit_run()
    option.temp_dir = unit['temp_dir']

    # read unit.log

    for i in range(50):
        with open(unit['temp_dir'] + '/unit.log', 'r') as f:
            log = f.read()
            m = re.search('controller started', log)

            if m is None:
                time.sleep(0.1)
            else:
                break

    if m is None:
        _print_log(log)
        exit("Unit is writing log too long")

    # discover available modules from unit.log

    for module in re.findall(r'module: ([a-zA-Z]+) (.*) ".*"$', log, re.M):
        versions = option.available['modules'].setdefault(module[0], [])
        if module[1] not in versions:
            versions.append(module[1])

    # discover modules from check

    option.available['modules']['openssl'] = check_openssl(unit['unitd'])
    option.available['modules']['go'] = check_go(
        option.current_dir, unit['temp_dir'], option.test_dir
    )
    option.available['modules']['node'] = check_node(option.current_dir)

    # remove None values

    option.available['modules'] = {
        k: v for k, v in option.available['modules'].items() if v is not None
    }

    check_isolation()

    unit_stop()

    _check_alerts()

    shutil.rmtree(unit_instance['temp_dir'])

@pytest.hookimpl(tryfirst=True, hookwrapper=True)
def pytest_runtest_makereport(item, call):
    # execute all other hooks to obtain the report object
    outcome = yield
    rep = outcome.get_result()

    # set a report attribute for each phase of a call, which can
    # be "setup", "call", "teardown"

    setattr(item, "rep_" + rep.when, rep)

@pytest.fixture(scope='class', autouse=True)
def check_prerequisites(request):
    cls = request.cls
    missed = []

    # check modules

    if 'modules' in cls.prerequisites:
        available_modules = list(option.available['modules'].keys())

        for module in cls.prerequisites['modules']:
            if module in available_modules:
                continue

            missed.append(module)

    if missed:
        pytest.skip('Unit has no ' + ', '.join(missed) + ' module(s)')

    # check features

    if 'features' in cls.prerequisites:
        available_features = list(option.available['features'].keys())

        for feature in cls.prerequisites['features']:
            if feature in available_features:
                continue

            missed.append(feature)

    if missed:
        pytest.skip(', '.join(missed) + ' feature(s) not supported')

@pytest.fixture(autouse=True)
def run(request):
    unit = unit_run()
    option.temp_dir = unit['temp_dir']

    option.skip_alerts = [
        r'read signalfd\(4\) failed',
        r'sendmsg.+failed',
        r'recvmsg.+failed',
    ]
    option.skip_sanitizer = False

    yield

    # stop unit

    error = unit_stop()

    if error:
        _print_log()

    assert error is None, 'stop unit'

    # stop all processes

    error = stop_processes()

    if error:
        _print_log()

    assert error is None, 'stop unit'

    # check unit.log for alerts

    _check_alerts()

    # print unit.log in case of error

    if hasattr(request.node, 'rep_call') and request.node.rep_call.failed:
        _print_log()

    # remove unit.log

    if not option.save_log:
        shutil.rmtree(unit['temp_dir'])

def unit_run():
    global unit_instance
    build_dir = option.current_dir + '/build'
    unitd = build_dir + '/unitd'

    if not os.path.isfile(unitd):
        exit('Could not find unit')

    temp_dir = tempfile.mkdtemp(prefix='unit-test-')
    public_dir(temp_dir)

    if oct(stat.S_IMODE(os.stat(build_dir).st_mode)) != '0o777':
        public_dir(build_dir)

    os.mkdir(temp_dir + '/state')

    unitd_args = [
        unitd,
        '--no-daemon',
        '--modules',
        build_dir,
        '--state',
        temp_dir + '/state',
        '--pid',
        temp_dir + '/unit.pid',
        '--log',
        temp_dir + '/unit.log',
        '--control',
        'unix:' + temp_dir + '/control.unit.sock',
        '--tmp',
        temp_dir,
    ]

    if option.user:
        unitd_args.extend(['--user', option.user])

    with open(temp_dir + '/unit.log', 'w') as log:
        unit_instance['process'] = subprocess.Popen(unitd_args, stderr=log)

    if not waitforfiles(temp_dir + '/control.unit.sock'):
        _print_log()
        exit('Could not start unit')

    unit_instance['temp_dir'] = temp_dir
    unit_instance['log'] = temp_dir + '/unit.log'
    unit_instance['control_sock'] = temp_dir + '/control.unit.sock'
    unit_instance['unitd'] = unitd

    return unit_instance

def unit_stop():
    p = unit_instance['process']

    if p.poll() is not None:
        return

    p.send_signal(signal.SIGQUIT)

    try:
        retcode = p.wait(15)
        if retcode:
            return 'Child process terminated with code ' + str(retcode)

    except KeyboardInterrupt:
        p.kill()
        raise

    except:
        p.kill()
        return 'Could not terminate unit'

def _check_alerts(path=None):
    if path is None:
        path = unit_instance['log']

    with open(path, 'r', encoding='utf-8', errors='ignore') as f:
        log = f.read()

    found = False

    alerts = re.findall(r'.+\[alert\].+', log)

    if alerts:
        print('\nAll alerts/sanitizer errors found in log:')
        [print(alert) for alert in alerts]
        found = True

    if option.skip_alerts:
        for skip in option.skip_alerts:
            alerts = [al for al in alerts if re.search(skip, al) is None]

    if alerts:
        _print_log(log)
        assert not alerts, 'alert(s)'

    if not option.skip_sanitizer:
        sanitizer_errors = re.findall('.+Sanitizer.+', log)

        if sanitizer_errors:
            _print_log(log)
            assert not sanitizer_errors, 'sanitizer error(s)'

    if found:
        print('skipped.')

def _print_log(data=None):
    path = unit_instance['log']

    print('Path to unit.log:\n' + path + '\n')

    if option.print_log:
        os.set_blocking(sys.stdout.fileno(), True)
        sys.stdout.flush()

        if data is None:
            with open(path, 'r', encoding='utf-8', errors='ignore') as f:
                shutil.copyfileobj(f, sys.stdout)
        else:
            sys.stdout.write(data)

def run_process(target, *args):
    global _processes

    process = Process(target=target, args=args)
    process.start()

    _processes.append(process)

def stop_processes():
    if not _processes:
        return

    fail = False
    for process in _processes:
        if process.is_alive():
            process.terminate()
            process.join(timeout=15)

            if process.is_alive():
                fail = True

    if fail:
        return 'Fail to stop process(es)'

@pytest.fixture()
def skip_alert():
    def _skip(*alerts):
        option.skip_alerts.extend(alerts)

    return _skip

@pytest.fixture
def temp_dir(request):
    return unit_instance['temp_dir']

@pytest.fixture
def is_unsafe(request):
    return request.config.getoption("--unsafe")

@pytest.fixture
def is_su(request):
    return os.geteuid() == 0

def pytest_sessionfinish(session):
    unit_stop()
    shutil.rmtree(option.cache_dir)
<|endoftext|>