<commit_before>"""
Plot calculated levels of a given sequence across channels or groups.
"""

from __future__ import division

# Built-ins
from collections import OrderedDict
import itertools
import logging
import os
import re

# Core data analysis libraries
from matplotlib import pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from scipy.stats import ttest_ind

import pyproteome as pyp

LOGGER = logging.getLogger("pyproteome.plot")

def plot(
    data,
    title=None,
    folder_name=None,
    ax=None,
    log_2=True,
):
    """
    Plot the levels of a sequence across multiple channels.

    Parameters
    ----------
    data : :class:`pyproteome.data_sets.DataSet`
    title : str, optional
    folder_name : str, optional
    figsize : tuple of (int, int), optional

    Returns
    -------
    figs : list of :class:`matplotlib.figure.Figure`
    """
    folder_name = pyp.utils.make_folder(
        data=data,
        folder_name=folder_name,
        sub="Peptides",
    )

    channel_names = [
        channel_name
        for group in data.groups.values()
        for channel_name in group
        if channel_name in data.channels
    ]
    channel_names = list(OrderedDict.fromkeys(channel_names))
    channels = [
        data.channels[channel_name]
        for channel_name in channel_names
    ]

    figures = []

    for _, row in data.psms.iterrows():
        seq = str(row["Sequence"])

        values = row[channels]
        mask = ~pd.isnull(row[channels])

        names = pd.Series(channel_names, index=values.index)[mask]
        values = values[mask]
        cmp_groups = data.cmp_groups or [[
            grp
            for grp, channels in data.groups.items()
            if any([i in data.channels for i in channels])
        ]]
        values = values / (
            values[[
                data.channels[i]
                for grp in cmp_groups
                for i in data.groups[grp[0]]
                if i in data.channels and data.channels[i] in values.index
            ]].median()
            if len(cmp_groups) > 0 else
            values[0]
        )

        if ax is None:
            fig, ax_i = plt.subplots(
                figsize=(len(values) / 2, 6 / 2),
            )
        else:
            ax_i = ax

        df = pd.DataFrame(
            [
                (
                    name,
                    val,
                    [
                        group_name
                        for group_name, group in data.groups.items()
                        if name in group
                    ][0],
                )
                for name, val in zip(names, values)
            ],
            columns=("name", "val", "group"),
        )

        if log_2:
            df["val"] = df["val"].apply(np.log2)

        sns.barplot(
            x="name",
            y="val",
            hue="group",
            data=df,
            ax=ax_i,
            dodge=False,
        )
        ax_i.set_xticklabels(
            ax_i.get_xticklabels(),
            rotation=45,
            horizontalalignment="right",
        )
        ax_i.set_xlabel("")
        ax_i.get_legend().set_title("")

        ax_i.axhline(np.log2(1), linestyle=":", color="k", alpha=.5)

        mod_str = row["Modifications"].__str__(prot_index=0)

        ax_i.set_title(
            title
            if title else
            "{} ({}{})".format(
                seq,
                " / ".join(row["Proteins"].genes)[:20],
                (" " + mod_str) if mod_str else "",
            ),
        )

        ylabel = "Intensity"
        # if data.cmp_groups:
        # else:
        #     ylabel = (
        #         "Cummulative Intensity" +
        #         (" (Normalized)" if data.intra_normalized else "")
        #     )

        ax_i.set_ylabel(
            ylabel,
        )

        ax_i.get_figure().savefig(
            os.path.join(
                folder_name,
                re.sub(
                    "[?/]",
                    "_",
                    "{} - {} - all.png".format(
                        "+".join(row["Proteins"].genes)[:100],
                        row["Sequence"],
                    ),
                ),
            ),
            bbox_inches="tight",
            dpi=pyp.DEFAULT_DPI,
            transparent=True,
        )

        figures.append((ax_i.get_figure(), ax_i))

    return figures

def plot_group(
    data,
    cmp_groups=None,
    title=None,
    folder_name=None,
    ax=None,
):
    """
    Plot the levels of a sequence across each group.

    Parameters
    ----------
    data : :class:`pyproteome.data_sets.DataSet`
    cmp_groups : list of tuple, optional
    title : str, optional
    folder_name : str, optional
    figsize : tuple of (int, int), optional

    Returns
    -------
    figs : list of :class:`matplotlib.figure.Figure`
    """
    folder_name = pyp.utils.make_folder(
        data=data,
        folder_name=folder_name,
        sub="Peptides",
    )

    if cmp_groups is None:
        cmp_groups = data.cmp_groups or [list(data.groups.keys())]

    figures = []

    for _, row in data.psms.iterrows():
        values = []

        for groups in cmp_groups:
            group_vals = pd.Series([
                row[[
                    data.channels[name]
                    for name in data.groups[group]
                    if name in data.channels
                ]]
                for group in groups
            ], index=groups, dtype=object)

            group_vals = pd.Series([
                group[~pd.isnull(group)]
                for group in group_vals
            ], index=group_vals.index, dtype=object)

            group_vals = group_vals[
                group_vals.apply(lambda x: x.shape[0] > 0)
            ]

            # Check normalization group is not null and at least one other
            # group of channels is not null
            if (
                group_vals.shape[0] < 1 or (
                    len(cmp_groups) > 1 and
                    groups[0] not in group_vals.index
                ) or all(
                    group not in group_vals.index
                    for group in groups[1:]
                )
            ):
                continue

            normalize = group_vals.iloc[0].median()

            group_vals = pd.Series([
                group / normalize
                for group in group_vals
            ], index=group_vals.index, dtype=object)

            values.append(group_vals)

        labels = [
            name
            for group in values
            for name in group.index
        ]

        if ax is None:
            _, plot_ax = plt.subplots(
                figsize=(len(labels) * .75, 4),
            )
        else:
            plot_ax = ax

        x = [
            ind
            for ind, l in enumerate(
                j
                for i in values
                for j in i.values
            )
            for k in l
        ]
        y = np.concatenate([
            np.log2(j.astype(float))
            for i in values
            for j in i.values
        ])
        df = pd.DataFrame(
            [
                (
                    np.log2(k),
                    label,
                    {
                        0: "#e19153",
                        1: "#60ae47",
                        2: "#ffae47",
                        3: "#e191ff",
                    }.get(max([ind for lst in cmp_groups for ind, val in enumerate(lst) if val == label]))
                )
                for i in values
                for label, j in i.iteritems()
                for k in j.values
            ],
            columns=("y", "label", "color"),
        )
        sns.boxplot(
            x="label",
            y="y",
            hue="color",
            data=df,
            ax=plot_ax,
            dodge=False,
            boxprops=dict(alpha=.3),
        )
        sns.swarmplot(
            x=x,
            y=y,
            color=".25",
            ax=plot_ax,
            # size=10,
        )
        plot_ax.axhline(
            0,
            linestyle="--",
            alpha=.25,
        )

        mod_str = row["Modifications"].__str__(prot_index=0)

        plot_ax.set_title(
            title
            if title else
            "{} ({}{})".format(
                row["Sequence"],
                " / ".join(row["Proteins"].genes)[:20],
                (" " + mod_str) if mod_str else "",
            ),
        )
        plot_ax.xaxis.grid(False)

        y_max = y.max()

        def stars(p):
            if p < 0.0001:
                return "****"
            elif (p < 0.001):
                return "***"
            elif (p < 0.01):
                return "**"
            elif (p < 0.05):
                return "*"
            else:
                return "-"

        v = [
            vals
            for group_vals in values
            for vals in group_vals
        ]

        for grp_set in cmp_groups:
            offset = y_max / 10

            for label_a, label_b in itertools.combinations(grp_set, 2):
                if label_a not in labels or label_b not in labels:
                    continue

                index_a = labels.index(label_a)
                index_b = labels.index(label_b)

                values_a = v[index_a]
                values_b = v[index_b]

                if values_a.shape[0] < 2 or values_b.shape[0] < 2:
                    continue

                pval = ttest_ind(
                    values_a.values,
                    values_b.values,
                ).pvalue

                if pval < 0.05:
                    plot_ax.annotate(
                        "",
                        xy=(
                            index_a,
                            y_max + offset,
                        ),
                        xytext=(
                            index_b,
                            y_max + offset,
                        ),
                        xycoords='data',
                        textcoords='data',
                        arrowprops=dict(
                            arrowstyle="-",
                            ec='#000000',
                        ),
                    )
                    plot_ax.text(
                        x=np.mean([index_a, index_b]),
                        y=y_max + offset + y_max / 40,
                        s=stars(pval),
                        horizontalalignment='center',
                        verticalalignment='center',
                    )
                    offset += y_max / 10

            plot_ax.set_ylim(
                bottom=plot_ax.get_ylim()[0],
                top=max([
                    plot_ax.get_ylim()[1],
                    y_max + offset + y_max / 10,
                ]),
            )

        plot_ax.set_xlabel("")
        plot_ax.set_ylabel(
            "{} Signal".format(
                "Relative" if cmp_groups else "Cumulative",
            ),
        )
        plot_ax.get_legend().set_visible(False)

        plot_ax.set_yticklabels(
            ["{:.2f}".format(i) for i in np.power(2, plot_ax.get_yticks())],
        )

        plot_ax.set_xticklabels(
            labels,
            rotation=45,
            horizontalalignment="right",
        )

        plot_ax.get_figure().savefig(
            os.path.join(
                folder_name,
                re.sub(
                    "[?/]",
                    "_",
                    "{} - {} - groups.png".format(
                        "+".join(row["Proteins"].genes)[:50],
                        row["Sequence"],
                    ),
                ),
            ),
            bbox_inches="tight",
            dpi=pyp.DEFAULT_DPI,
            transparent=True,
        )

        figures.append((plot_ax.get_figure(), plot_ax))

    return figures

def plot_all(
    data,
    individual=True,
    between=True,
    cmp_groups=None,
    folder_name=None,
):
    """
    Runs :func:`.plot` and :func:`.plot_group` for all peptides in a data set.

    Parameters
    ----------
    data : :class:`pyproteome.data_sets.DataSet`
    figsize : tuple of (int, int), optional
    cmp_groups : list of tuple, optional
    folder_name : str, optional

    Returns
    -------
    figs : list of :class:`matplotlib.figure.Figure`
    """
    folder_name = pyp.utils.make_folder(
        data=data,
        folder_name=folder_name,
        sub="Peptides",
    )

    figures = []

    figures += plot(
        data,
        folder_name=folder_name,
    )

    figures += plot_group(
        data,
        cmp_groups=cmp_groups,
        folder_name=folder_name,
    )

    return figures
<commit_msg>Added log2 argument to plot_group Fixed bugs with missing channels / groups Improved color layout of plot_group<commit_after>"""
Plot calculated levels of a given sequence across channels or groups.
"""

from __future__ import division

# Built-ins
from collections import OrderedDict
import itertools
import logging
import os
import re

# Core data analysis libraries
from matplotlib import pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from scipy.stats import ttest_ind

import pyproteome as pyp

LOGGER = logging.getLogger("pyproteome.plot")

def plot(
    data,
    title=None,
    folder_name=None,
    ax=None,
    log_2=True,
):
    """
    Plot the levels of a sequence across multiple channels.

    Parameters
    ----------
    data : :class:`pyproteome.data_sets.DataSet`
    title : str, optional
    folder_name : str, optional
    figsize : tuple of (int, int), optional

    Returns
    -------
    figs : list of :class:`matplotlib.figure.Figure`
    """
    folder_name = pyp.utils.make_folder(
        data=data,
        folder_name=folder_name,
        sub="Peptides",
    )

    channel_names = [
        channel_name
        for group in data.groups.values()
        for channel_name in group
        if channel_name in data.channels
    ]
    channel_names = list(OrderedDict.fromkeys(channel_names))
    channels = [
        data.channels[channel_name]
        for channel_name in channel_names
    ]

    figures = []

    for _, row in data.psms.iterrows():
        seq = str(row["Sequence"])

        values = row[channels]
        mask = ~pd.isnull(row[channels])

        names = pd.Series(channel_names, index=values.index)[mask]
        values = values[mask]
        cmp_groups = data.cmp_groups or [[
            grp
            for grp, channels in data.groups.items()
            if any([i in data.channels for i in channels])
        ]]
        values = values / (
            values[[
                data.channels[i]
                for grp in cmp_groups
                for i in data.groups[grp[0]]
                if i in data.channels and data.channels[i] in values.index
            ]].median()
            if len(cmp_groups) > 0 else
            values[0]
        )

        if ax is None:
            fig, ax_i = plt.subplots(
                figsize=(len(values) / 2, 6 / 2),
            )
        else:
            ax_i = ax

        df = pd.DataFrame(
            [
                (
                    name,
                    val,
                    [
                        group_name
                        for group_name, group in data.groups.items()
                        if name in group
                    ][0],
                )
                for name, val in zip(names, values)
            ],
            columns=("name", "val", "group"),
        )

        if log_2:
            df["val"] = df["val"].apply(np.log2)

        sns.barplot(
            x="name",
            y="val",
            hue="group",
            data=df,
            ax=ax_i,
            dodge=False,
        )
        ax_i.set_xticklabels(
            ax_i.get_xticklabels(),
            rotation=45,
            horizontalalignment="right",
        )
        ax_i.set_xlabel("")
        ax_i.get_legend().set_title("")

        ax_i.axhline(np.log2(1), linestyle=":", color="k", alpha=.5)

        mod_str = row["Modifications"].__str__(prot_index=0)

        ax_i.set_title(
            title
            if title else
            "{} ({}{})".format(
                seq,
                " / ".join(row["Proteins"].genes)[:20],
                (" " + mod_str) if mod_str else "",
            ),
        )

        ylabel = "Intensity"
        # if data.cmp_groups:
        # else:
        #     ylabel = (
        #         "Cummulative Intensity" +
        #         (" (Normalized)" if data.intra_normalized else "")
        #     )

        ax_i.set_ylabel(
            ylabel,
        )

        ax_i.get_figure().savefig(
            os.path.join(
                folder_name,
                re.sub(
                    "[?/]",
                    "_",
                    "{} - {} - all.png".format(
                        "+".join(row["Proteins"].genes)[:100],
                        row["Sequence"],
                    ),
                ),
            ),
            bbox_inches="tight",
            dpi=pyp.DEFAULT_DPI,
            transparent=True,
        )

        figures.append((ax_i.get_figure(), ax_i))

    return figures

def plot_group(
    data,
    cmp_groups=None,
    title=None,
    folder_name=None,
    ax=None,
    show_p=True,
    log_2=True,
):
    """
    Plot the levels of a sequence across each group.

    Parameters
    ----------
    data : :class:`pyproteome.data_sets.DataSet`
    cmp_groups : list of tuple, optional
    title : str, optional
    folder_name : str, optional
    figsize : tuple of (int, int), optional

    Returns
    -------
    figs : list of :class:`matplotlib.figure.Figure`
    """
    folder_name = pyp.utils.make_folder(
        data=data,
        folder_name=folder_name,
        sub="Peptides",
    )

    if cmp_groups is None:
        cmp_groups = data.cmp_groups or [list(data.groups.keys())]

    figures = []

    for _, row in data.psms.iterrows():
        values = []

        for groups in cmp_groups:
            groups = [i for i in groups if i in data.groups]
            group_vals = pd.Series([
                row[[
                    data.channels[name]
                    for name in data.groups[group]
                    if name in data.channels
                ]]
                for group in groups
            ], index=groups, dtype=object)

            group_vals = pd.Series([
                group[~pd.isnull(group)]
                for group in group_vals
            ], index=group_vals.index, dtype=object)

            group_vals = group_vals[
                group_vals.apply(lambda x: x.shape[0] > 0)
            ]

            # Check normalization group is not null and at least one other
            # group of channels is not null
            if (
                group_vals.shape[0] < 1 or (
                    len(cmp_groups) > 1 and
                    groups[0] not in group_vals.index
                ) or all(
                    group not in group_vals.index
                    for group in groups[1:]
                )
            ):
                continue

            normalize = group_vals.iloc[0].median()

            group_vals = pd.Series([
                group / normalize
                for group in group_vals
            ], index=group_vals.index, dtype=object)

            values.append(group_vals)

        labels = [
            name
            for group in values
            for name in group.index
        ]

        if ax is None:
            _, plot_ax = plt.subplots(
                figsize=(len(labels) * .75, 4),
            )
        else:
            plot_ax = ax

        x = [
            ind
            for ind, l in enumerate(
                j
                for i in values
                for j in i.values
            )
            for k in l
        ]
        y = np.concatenate([
            j.astype(float)
            for i in values
            for j in i.values
        ])

        if log_2:
            y = np.log2(y)

        def _get_color(label):
            lst = [
                i
                for i in cmp_groups
                for ind, val in enumerate(i)
                if val == label
            ][0]
            return sns.color_palette("hls", len(lst)).as_hex()[lst.index(label)]

        df = pd.DataFrame(
            [
                (
                    k,
                    np.log2(k),
                    label,
                    _get_color(label),
                )
                for i in values
                for label, j in i.iteritems()
                for k in j.values
            ],
            columns=("y", 'log2_y', "label", "color"),
        )
        sns.boxplot(
            x="label",
            y="log2_y" if log_2 else 'y',
            hue="color",
            data=df,
            ax=plot_ax,
            dodge=False,
            boxprops=dict(alpha=.3),
        )
        sns.swarmplot(
            x=x,
            y=y,
            color=".25",
            ax=plot_ax,
            # size=10,
        )
        plot_ax.axhline(
            0,
            linestyle="--",
            alpha=.25,
        )

        mod_str = row["Modifications"].__str__(prot_index=0)

        plot_ax.set_title(
            title
            if title else
            "{}{}({}{})".format(
                row["Sequence"],
                ' ' if len(plot_ax.get_xticklabels()) > 2 else '\n',
                " / ".join(row["Proteins"].genes)[:20],
                (" " + mod_str) if mod_str else "",
            ),
        )
        plot_ax.xaxis.grid(False)

        if show_p:
            y_max = y.max()

            def stars(p):
                if p < 0.0001:
                    return "****"
                elif (p < 0.001):
                    return "***"
                elif (p < 0.01):
                    return "**"
                elif (p < 0.05):
                    return "*"
                else:
                    return "-"

            v = [
                vals
                for group_vals in values
                for vals in group_vals
            ]

            for grp_set in cmp_groups:
                offset = y_max / 10

                for label_a, label_b in itertools.combinations(grp_set, 2):
                    if label_a not in labels or label_b not in labels:
                        continue

                    index_a = labels.index(label_a)
                    index_b = labels.index(label_b)

                    values_a = v[index_a]
                    values_b = v[index_b]

                    if values_a.shape[0] < 2 or values_b.shape[0] < 2:
                        continue

                    pval = ttest_ind(
                        values_a.values,
                        values_b.values,
                    ).pvalue

                    if pval < 0.05:
                        plot_ax.annotate(
                            "",
                            xy=(
                                index_a,
                                y_max + offset,
                            ),
                            xytext=(
                                index_b,
                                y_max + offset,
                            ),
                            xycoords='data',
                            textcoords='data',
                            arrowprops=dict(
                                arrowstyle="-",
                                ec='#000000',
                            ),
                        )
                        plot_ax.text(
                            x=np.mean([index_a, index_b]),
                            y=y_max + offset + y_max / 40,
                            s=stars(pval),
                            horizontalalignment='center',
                            verticalalignment='center',
                        )
                        offset += y_max / 10

                plot_ax.set_ylim(
                    bottom=plot_ax.get_ylim()[0],
                    top=max([
                        plot_ax.get_ylim()[1],
                        y_max + offset + y_max / 10,
                    ]),
                )

        plot_ax.set_xlabel("")
        plot_ax.set_ylabel(
            "{} Signal".format(
                "Relative" if cmp_groups else "Cumulative",
            ),
        )
        plot_ax.get_legend().set_visible(False)

        if log_2:
            plot_ax.set_yticklabels(
                [
                    "{:.2f}".format(i)
                    for i in np.power(2, plot_ax.get_yticks())
                ],
            )

        plot_ax.set_xticklabels(
            labels,
            rotation=45,
            horizontalalignment="right",
        )

        plot_ax.get_figure().savefig(
            os.path.join(
                folder_name,
                re.sub(
                    "[?/]",
                    "_",
                    "{} - {} - groups.png".format(
                        "+".join(row["Proteins"].genes)[:50],
                        row["Sequence"],
                    ),
                ),
            ),
            bbox_inches="tight",
            dpi=pyp.DEFAULT_DPI,
            transparent=True,
        )

        figures.append((plot_ax.get_figure(), plot_ax))

    return figures

def plot_all(
    data,
    individual=True,
    between=True,
    cmp_groups=None,
    folder_name=None,
):
    """
    Runs :func:`.plot` and :func:`.plot_group` for all peptides in a data set.

    Parameters
    ----------
    data : :class:`pyproteome.data_sets.DataSet`
    figsize : tuple of (int, int), optional
    cmp_groups : list of tuple, optional
    folder_name : str, optional

    Returns
    -------
    figs : list of :class:`matplotlib.figure.Figure`
    """
    folder_name = pyp.utils.make_folder(
        data=data,
        folder_name=folder_name,
        sub="Peptides",
    )

    figures = []

    figures += plot(
        data,
        folder_name=folder_name,
    )

    figures += plot_group(
        data,
        cmp_groups=cmp_groups,
        folder_name=folder_name,
    )

    return figures
<|endoftext|>


<commit_before># Copyright (c) 2016-2018, Neil Booth
#
# All rights reserved.
#
# See the file "LICENCE" for information about the copyright
# and warranty status of this software.

'''Mempool handling.'''

import asyncio
import itertools
import time
from collections import defaultdict

import attr

from electrumx.lib.hash import hash_to_hex_str, hex_str_to_hash
from electrumx.lib.util import class_logger, chunks
from electrumx.server.db import UTXO

@attr.s(slots=True)
class MemPoolTx(object):
    prevouts = attr.ib()
    in_pairs = attr.ib()
    out_pairs = attr.ib()
    fee = attr.ib()
    size = attr.ib()

class MemPool(object):
    '''Representation of the daemon's mempool.

    Updated regularly in caught-up state.  Goal is to enable efficient
    response to the value() and transactions() calls.

    To that end we maintain the following maps:

       tx_hash -> MemPoolTx
       hashX   -> set of all tx hashes in which the hashX appears

    A pair is a (hashX, value) tuple.  tx hashes are binary not strings.
    '''

    def __init__(self, coin, tasks, daemon, notifications, lookup_utxos):
        self.logger = class_logger(__name__, self.__class__.__name__)
        self.coin = coin
        self.lookup_utxos = lookup_utxos
        self.tasks = tasks
        self.daemon = daemon
        self.notifications = notifications
        self.txs = {}
        self.hashXs = defaultdict(set)  # None can be a key
        self.cached_compact_histogram = []

    async def _log_stats(self):
        while True:
            self.logger.info(f'{len(self.txs):,d} txs '
                             f'touching {len(self.hashXs):,d} addresses')
            await asyncio.sleep(120)

    def _update_histogram(self):
        # Build a histogram by fee rate
        histogram = defaultdict(int)
        for tx in self.txs.values():
            histogram[tx.fee // tx.size] += tx.size

        # Now compact it.  For efficiency, get_fees returns a
        # compact histogram with variable bin size.  The compact
        # histogram is an array of (fee_rate, vsize) values.
        # vsize_n is the cumulative virtual size of mempool
        # transactions with a fee rate in the interval
        # [rate_(n-1), rate_n)], and rate_(n-1) > rate_n.
        # Intervals are chosen to create tranches containing at
        # least 100kb of transactions
        compact = []
        cum_size = 0
        r = 0   # ?
        bin_size = 100 * 1000
        for fee_rate, size in sorted(histogram.items(), reverse=True):
            cum_size += size
            if cum_size + r > bin_size:
                compact.append((fee_rate, cum_size))
                r += cum_size - bin_size
                cum_size = 0
                bin_size *= 1.1
        self.logger.info(f'compact fee histogram: {compact}')
        self.cached_compact_histogram = compact

    def _accept_transactions(self, tx_map, utxo_map, touched):
        '''Accept transactions in tx_map to the mempool if all their inputs
        can be found in the existing mempool or a utxo_map from the
        DB.

        Returns an (unprocessed tx_map, unspent utxo_map) pair.
        '''
        hashXs = self.hashXs
        txs = self.txs

        deferred = {}
        unspent = set(utxo_map)
        # Try to find all prevouts so we can accept the TX
        for hash, tx in tx_map.items():
            in_pairs = []
            try:
                for prevout in tx.prevouts:
                    utxo = utxo_map.get(prevout)
                    if not utxo:
                        prev_hash, prev_index = prevout
                        # Raises KeyError if prev_hash is not in txs
                        utxo = txs[prev_hash].out_pairs[prev_index]
                    in_pairs.append(utxo)
            except KeyError:
                deferred[hash] = tx
                continue

            # Spend the prevouts
            unspent.difference_update(tx.prevouts)

            # Save the in_pairs, compute the fee,  and accept the TX
            tx.in_pairs = in_pairs
            tx.fee = (sum(v for hashX, v in tx.in_pairs) -
                      sum(v for hashX, v in tx.out_pairs))
            txs[hash] = tx
            for hashX, value in itertools.chain(tx.in_pairs, tx.out_pairs):
                touched.add(hashX)
                hashXs[hashX].add(hash)

        return deferred, {prevout: utxo_map[prevout] for prevout in unspent}

    async def _refresh_hashes(self, single_pass):
        '''Return a (hash set, height) pair when we're sure which height they
        are for.'''
        refresh_event = asyncio.Event()
        loop = self.tasks.loop
        for loop_count in itertools.count():
            height = self.daemon.cached_height()
            hex_hashes = await self.daemon.mempool_hashes()
            if height != await self.daemon.height():
                continue
            loop.call_later(5, refresh_event.set)
            hashes = set(hex_str_to_hash(hh) for hh in hex_hashes)
            touched = await self._process_mempool(hashes)
            await self.notifications.on_mempool(touched, height)
            # Refresh the cached histogram periodically.  Thread it as it
            # can be expensive.
            if loop_count % 100 == 0:
                await self.tasks.run_in_thread(self._update_histogram)
            if single_pass:
                return
            await refresh_event.wait()
            refresh_event.clear()

    async def _process_mempool(self, all_hashes):
        # Re-sync with the new set of hashes
        txs = self.txs
        hashXs = self.hashXs
        touched = set()

        # First handle txs that have disappeared
        for tx_hash in set(txs).difference(all_hashes):
            tx = txs.pop(tx_hash)
            tx_hashXs = set(hashX for hashX, value in tx.in_pairs)
            tx_hashXs.update(hashX for hashX, value in tx.out_pairs)
            for hashX in tx_hashXs:
                hashXs[hashX].remove(tx_hash)
                if not hashXs[hashX]:
                    del hashXs[hashX]
            touched.update(tx_hashXs)

        # Process new transactions
        new_hashes = list(all_hashes.difference(txs))
        jobs = [self.tasks.create_task(self._fetch_and_accept
                                       (hashes, all_hashes, touched))
                for hashes in chunks(new_hashes, 2000)]
        if jobs:
            await asyncio.wait(jobs)
            tx_map = {}
            utxo_map = {}
            for job in jobs:
                deferred, unspent = job.result()
                tx_map.update(deferred)
                utxo_map.update(unspent)

            # Handle the stragglers
            if len(tx_map) >= 10:
                self.logger.info(f'{len(tx_map)} stragglers')
            prior_count = 0
            # FIXME: this is not particularly efficient
            while tx_map and len(tx_map) != prior_count:
                prior_count = len(tx_map)
                tx_map, utxo_map = self._accept_transactions(tx_map, utxo_map,
                                                             touched)
            if tx_map:
                self.logger.info(f'{len(tx_map)} txs dropped')

        return touched

    async def _fetch_and_accept(self, hashes, all_hashes, touched):
        '''Fetch a list of mempool transactions.'''
        hex_hashes = [hash_to_hex_str(hash) for hash in hashes]
        raw_txs = await self.daemon.getrawtransactions(hex_hashes)
        count = len([raw_tx for raw_tx in raw_txs if raw_tx])

        def deserialize_txs():
            # This function is pure
            script_hashX = self.coin.hashX_from_script
            deserializer = self.coin.DESERIALIZER

            txs = {}
            for hash, raw_tx in zip(hashes, raw_txs):
                # The daemon may have evicted the tx from its
                # mempool or it may have gotten in a block
                if not raw_tx:
                    continue
                tx, tx_size = deserializer(raw_tx).read_tx_and_vsize()

                # Convert the tx outputs into (hashX, value) pairs
                txout_pairs = [(script_hashX(txout.pk_script), txout.value)
                               for txout in tx.outputs]

                # Convert the tx inputs to (prev_hash, prev_idx) pairs
                txin_pairs = [(txin.prev_hash, txin.prev_idx)
                              for txin in tx.inputs]

                txs[hash] = MemPoolTx(txin_pairs, None, txout_pairs,
                                      0, tx_size)
            return txs

        # Thread this potentially slow operation so as not to block
        tx_map = await self.tasks.run_in_thread(deserialize_txs)

        # Determine all prevouts not in the mempool, and fetch the
        # UTXO information from the database.  Failed prevout lookups
        # return None - concurrent database updates happen
        prevouts = [prevout for tx in tx_map.values()
                    for prevout in tx.prevouts
                    if prevout[0] not in all_hashes]
        utxos = await self.lookup_utxos(prevouts)
        utxo_map = {prevout: utxo for prevout, utxo in zip(prevouts, utxos)}

        # Attempt to complete processing of txs
        return self._accept_transactions(tx_map, utxo_map, touched)

    # External interface
    async def start_and_wait_for_sync(self):
        '''Starts the mempool synchronizer.

        Waits for an initial synchronization before returning.
        '''
        self.logger.info('beginning processing of daemon mempool.  '
                         'This can take some time...')
        start = time.time()
        await self._refresh_hashes(True)
        elapsed = time.time() - start
        self.logger.info(f'synced in {elapsed:.2f}s')
        self.tasks.create_task(self._log_stats())
        self.tasks.create_task(self._refresh_hashes(False))

    async def balance_delta(self, hashX):
        '''Return the unconfirmed amount in the mempool for hashX.

        Can be positive or negative.
        '''
        value = 0
        # hashXs is a defaultdict
        if hashX in self.hashXs:
            for hash in self.hashXs[hashX]:
                tx = self.txs[hash]
                value -= sum(v for h168, v in tx.in_pairs if h168 == hashX)
                value += sum(v for h168, v in tx.out_pairs if h168 == hashX)
        return value

    async def compact_fee_histogram(self):
        '''Return a compact fee histogram of the current mempool.'''
        return self.cached_compact_histogram

    async def potential_spends(self, hashX):
        '''Return a set of (prev_hash, prev_idx) pairs from mempool
        transactions that touch hashX.

        None, some or all of these may be spends of the hashX.
        '''
        result = set()
        for tx_hash in self.hashXs.get(hashX, []):
            tx = self.txs[tx_hash]
            result.update(tx.prevouts)
        return result

    async def transaction_summaries(self, hashX):
        '''Return a list of (tx_hash, tx_fee, unconfirmed) tuples for
        mempool entries for the hashX.

        unconfirmed is True if any txin is unconfirmed.
        '''
        # hashXs is a defaultdict, so use get() to query
        result = []
        for tx_hash in self.hashXs.get(hashX, []):
            tx = self.txs[tx_hash]
            unconfirmed = any(prev_hash in self.txs
                              for prev_hash, prev_idx in tx.prevouts)
            result.append((tx_hash, tx.fee, unconfirmed))
        return result

    async def unordered_UTXOs(self, hashX):
        '''Return an unordered list of UTXO named tuples from mempool
        transactions that pay to hashX.

        This does not consider if any other mempool transactions spend
        the outputs.
        '''
        utxos = []
        # hashXs is a defaultdict, so use get() to query
        for tx_hash in self.hashXs.get(hashX, []):
            tx = self.txs.get(tx_hash)
            if not tx:
                continue
            for pos, (hX, value) in enumerate(tx.out_pairs):
                if hX == hashX:
                    utxos.append(UTXO(-1, pos, tx_hash, 0, value))
        return utxos
<commit_msg>Reduce memory footprint<commit_after># Copyright (c) 2016-2018, Neil Booth
#
# All rights reserved.
#
# See the file "LICENCE" for information about the copyright
# and warranty status of this software.

'''Mempool handling.'''

import asyncio
import itertools
import time
from collections import defaultdict

import attr

from electrumx.lib.hash import hash_to_hex_str, hex_str_to_hash
from electrumx.lib.util import class_logger, chunks
from electrumx.server.db import UTXO

@attr.s(slots=True)
class MemPoolTx(object):
    prevouts = attr.ib()
    in_pairs = attr.ib()
    out_pairs = attr.ib()
    fee = attr.ib()
    size = attr.ib()

class MemPool(object):
    '''Representation of the daemon's mempool.

    Updated regularly in caught-up state.  Goal is to enable efficient
    response to the value() and transactions() calls.

    To that end we maintain the following maps:

       tx_hash -> MemPoolTx
       hashX   -> set of all tx hashes in which the hashX appears

    A pair is a (hashX, value) tuple.  tx hashes are binary not strings.
    '''

    def __init__(self, coin, tasks, daemon, notifications, lookup_utxos):
        self.logger = class_logger(__name__, self.__class__.__name__)
        self.coin = coin
        self.lookup_utxos = lookup_utxos
        self.tasks = tasks
        self.daemon = daemon
        self.notifications = notifications
        self.txs = {}
        self.hashXs = defaultdict(set)  # None can be a key
        self.cached_compact_histogram = []

    async def _log_stats(self):
        while True:
            self.logger.info(f'{len(self.txs):,d} txs '
                             f'touching {len(self.hashXs):,d} addresses')
            await asyncio.sleep(120)

    def _update_histogram(self):
        # Build a histogram by fee rate
        histogram = defaultdict(int)
        for tx in self.txs.values():
            histogram[tx.fee // tx.size] += tx.size

        # Now compact it.  For efficiency, get_fees returns a
        # compact histogram with variable bin size.  The compact
        # histogram is an array of (fee_rate, vsize) values.
        # vsize_n is the cumulative virtual size of mempool
        # transactions with a fee rate in the interval
        # [rate_(n-1), rate_n)], and rate_(n-1) > rate_n.
        # Intervals are chosen to create tranches containing at
        # least 100kb of transactions
        compact = []
        cum_size = 0
        r = 0   # ?
        bin_size = 100 * 1000
        for fee_rate, size in sorted(histogram.items(), reverse=True):
            cum_size += size
            if cum_size + r > bin_size:
                compact.append((fee_rate, cum_size))
                r += cum_size - bin_size
                cum_size = 0
                bin_size *= 1.1
        self.logger.info(f'compact fee histogram: {compact}')
        self.cached_compact_histogram = compact

    def _accept_transactions(self, tx_map, utxo_map, touched):
        '''Accept transactions in tx_map to the mempool if all their inputs
        can be found in the existing mempool or a utxo_map from the
        DB.

        Returns an (unprocessed tx_map, unspent utxo_map) pair.
        '''
        hashXs = self.hashXs
        txs = self.txs

        deferred = {}
        unspent = set(utxo_map)
        # Try to find all prevouts so we can accept the TX
        for hash, tx in tx_map.items():
            in_pairs = []
            try:
                for prevout in tx.prevouts:
                    utxo = utxo_map.get(prevout)
                    if not utxo:
                        prev_hash, prev_index = prevout
                        # Raises KeyError if prev_hash is not in txs
                        utxo = txs[prev_hash].out_pairs[prev_index]
                    in_pairs.append(utxo)
            except KeyError:
                deferred[hash] = tx
                continue

            # Spend the prevouts
            unspent.difference_update(tx.prevouts)

            # Save the in_pairs, compute the fee,  and accept the TX
            tx.in_pairs = tuple(in_pairs)
            tx.fee = (sum(v for hashX, v in tx.in_pairs) -
                      sum(v for hashX, v in tx.out_pairs))
            txs[hash] = tx
            for hashX, value in itertools.chain(tx.in_pairs, tx.out_pairs):
                touched.add(hashX)
                hashXs[hashX].add(hash)

        return deferred, {prevout: utxo_map[prevout] for prevout in unspent}

    async def _refresh_hashes(self, single_pass):
        '''Return a (hash set, height) pair when we're sure which height they
        are for.'''
        refresh_event = asyncio.Event()
        loop = self.tasks.loop
        for loop_count in itertools.count():
            height = self.daemon.cached_height()
            hex_hashes = await self.daemon.mempool_hashes()
            if height != await self.daemon.height():
                continue
            loop.call_later(5, refresh_event.set)
            hashes = set(hex_str_to_hash(hh) for hh in hex_hashes)
            touched = await self._process_mempool(hashes)
            await self.notifications.on_mempool(touched, height)
            # Refresh the cached histogram periodically.  Thread it as it
            # can be expensive.
            if loop_count % 100 == 0:
                await self.tasks.run_in_thread(self._update_histogram)
            if single_pass:
                return
            await refresh_event.wait()
            refresh_event.clear()

    async def _process_mempool(self, all_hashes):
        # Re-sync with the new set of hashes
        txs = self.txs
        hashXs = self.hashXs
        touched = set()

        # First handle txs that have disappeared
        for tx_hash in set(txs).difference(all_hashes):
            tx = txs.pop(tx_hash)
            tx_hashXs = set(hashX for hashX, value in tx.in_pairs)
            tx_hashXs.update(hashX for hashX, value in tx.out_pairs)
            for hashX in tx_hashXs:
                hashXs[hashX].remove(tx_hash)
                if not hashXs[hashX]:
                    del hashXs[hashX]
            touched.update(tx_hashXs)

        # Process new transactions
        new_hashes = list(all_hashes.difference(txs))
        jobs = [self.tasks.create_task(self._fetch_and_accept
                                       (hashes, all_hashes, touched))
                for hashes in chunks(new_hashes, 2000)]
        if jobs:
            await asyncio.wait(jobs)
            tx_map = {}
            utxo_map = {}
            for job in jobs:
                deferred, unspent = job.result()
                tx_map.update(deferred)
                utxo_map.update(unspent)

            # Handle the stragglers
            if len(tx_map) >= 10:
                self.logger.info(f'{len(tx_map)} stragglers')
            prior_count = 0
            # FIXME: this is not particularly efficient
            while tx_map and len(tx_map) != prior_count:
                prior_count = len(tx_map)
                tx_map, utxo_map = self._accept_transactions(tx_map, utxo_map,
                                                             touched)
            if tx_map:
                self.logger.info(f'{len(tx_map)} txs dropped')

        return touched

    async def _fetch_and_accept(self, hashes, all_hashes, touched):
        '''Fetch a list of mempool transactions.'''
        hex_hashes_iter = (hash_to_hex_str(hash) for hash in hashes)
        raw_txs = await self.daemon.getrawtransactions(hex_hashes_iter)

        def deserialize_txs():
            # This function is pure
            to_hashX = self.coin.hashX_from_script
            deserializer = self.coin.DESERIALIZER

            txs = {}
            for hash, raw_tx in zip(hashes, raw_txs):
                # The daemon may have evicted the tx from its
                # mempool or it may have gotten in a block
                if not raw_tx:
                    continue
                tx, tx_size = deserializer(raw_tx).read_tx_and_vsize()

                # Convert the tx outputs into (hashX, value) pairs
                txout_pairs = tuple((to_hashX(txout.pk_script), txout.value)
                                    for txout in tx.outputs)

                # Convert the tx inputs to (prev_hash, prev_idx) pairs
                txin_pairs = tuple((txin.prev_hash, txin.prev_idx)
                                   for txin in tx.inputs)

                txs[hash] = MemPoolTx(txin_pairs, None, txout_pairs,
                                      0, tx_size)
            return txs

        # Thread this potentially slow operation so as not to block
        tx_map = await self.tasks.run_in_thread(deserialize_txs)

        # Determine all prevouts not in the mempool, and fetch the
        # UTXO information from the database.  Failed prevout lookups
        # return None - concurrent database updates happen
        prevouts = tuple(prevout for tx in tx_map.values()
                         for prevout in tx.prevouts
                         if prevout[0] not in all_hashes)
        utxos = await self.lookup_utxos(prevouts)
        utxo_map = {prevout: utxo for prevout, utxo in zip(prevouts, utxos)}

        return self._accept_transactions(tx_map, utxo_map, touched)

    # External interface
    async def start_and_wait_for_sync(self):
        '''Starts the mempool synchronizer.

        Waits for an initial synchronization before returning.
        '''
        self.logger.info('beginning processing of daemon mempool.  '
                         'This can take some time...')
        start = time.time()
        await self._refresh_hashes(True)
        elapsed = time.time() - start
        self.logger.info(f'synced in {elapsed:.2f}s')
        self.tasks.create_task(self._log_stats())
        self.tasks.create_task(self._refresh_hashes(False))

    async def balance_delta(self, hashX):
        '''Return the unconfirmed amount in the mempool for hashX.

        Can be positive or negative.
        '''
        value = 0
        # hashXs is a defaultdict
        if hashX in self.hashXs:
            for hash in self.hashXs[hashX]:
                tx = self.txs[hash]
                value -= sum(v for h168, v in tx.in_pairs if h168 == hashX)
                value += sum(v for h168, v in tx.out_pairs if h168 == hashX)
        return value

    async def compact_fee_histogram(self):
        '''Return a compact fee histogram of the current mempool.'''
        return self.cached_compact_histogram

    async def potential_spends(self, hashX):
        '''Return a set of (prev_hash, prev_idx) pairs from mempool
        transactions that touch hashX.

        None, some or all of these may be spends of the hashX.
        '''
        result = set()
        for tx_hash in self.hashXs.get(hashX, ()):
            tx = self.txs[tx_hash]
            result.update(tx.prevouts)
        return result

    async def transaction_summaries(self, hashX):
        '''Return a list of (tx_hash, tx_fee, unconfirmed) tuples for
        mempool entries for the hashX.

        unconfirmed is True if any txin is unconfirmed.
        '''
        # hashXs is a defaultdict, so use get() to query
        result = []
        for tx_hash in self.hashXs.get(hashX, ()):
            tx = self.txs[tx_hash]
            unconfirmed = any(prev_hash in self.txs
                              for prev_hash, prev_idx in tx.prevouts)
            result.append((tx_hash, tx.fee, unconfirmed))
        return result

    async def unordered_UTXOs(self, hashX):
        '''Return an unordered list of UTXO named tuples from mempool
        transactions that pay to hashX.

        This does not consider if any other mempool transactions spend
        the outputs.
        '''
        utxos = []
        # hashXs is a defaultdict, so use get() to query
        for tx_hash in self.hashXs.get(hashX, ()):
            tx = self.txs.get(tx_hash)
            if not tx:
                continue
            for pos, (hX, value) in enumerate(tx.out_pairs):
                if hX == hashX:
                    utxos.append(UTXO(-1, pos, tx_hash, 0, value))
        return utxos
<|endoftext|>


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

from json import dumps as json_dumps
from .evaluator import Evaluator
from .parser import Parser

def parse(s, json=True, lex_options=None, yacc_options=None):
    """
    Parse a PuppetDBQuery-style query and transform it into a PuppetDB "AST"
    query.

    This function is intented to be the primary entry point for this package.
    It wraps up all the various components of this package into an easy to
    consume format. The output of this function is designed to be passed
    directly into :mod:`pypuppetdb`.

    For examples of the query syntax see `puppet-puppetdbquery
    <https://github.com/dalen/puppet-puppetdbquery>`_ and `node-puppetdbquery
    <https://github.com/dalen/node-puppetdbquery>`_ by `Erik Dalén
    <https://github.com/dalen>`_.

    :param str s: The query to parse and transform
    :param bool json: Whether to JSON-encode the PuppetDB AST result
    :param dict lex_options: Options passed to :func:`ply.lex.lex`
    :param dict yacc_options: Options passed to :func:`ply.yacc.yacc`
    """
    parser = Parser(lex_options=lex_options, yacc_options=yacc_options)
    evaluator = Evaluator()

    ast = parser.parse(s)
    raw = evaluator.evaluate(ast)

    if json and raw is not None:
        return json_dumps(raw)
    else:
        return raw
<commit_msg>pypuppetdbquery: add a module-level docstring<commit_after># -*- coding: utf-8 -*-
#
# This file is part of pypuppetdbquery.
# Copyright © 2016  Chris Boot <bootc@bootc.net>
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
This Python package contains a port of `Erik Dalén
<https://github.com/dalen>`__'s PuppetDB Query language to Python.

All the user-facing components of this package are intended to be found in this
module itself rather than any of the sub-modules.
"""

from json import dumps as json_dumps
from .evaluator import Evaluator
from .parser import Parser

def parse(s, json=True, lex_options=None, yacc_options=None):
    """
    Parse a PuppetDBQuery-style query and transform it into a PuppetDB "AST"
    query.

    This function is intented to be the primary entry point for this package.
    It wraps up all the various components of this package into an easy to
    consume format. The output of this function is designed to be passed
    directly into :mod:`pypuppetdb`.

    For examples of the query syntax see `puppet-puppetdbquery
    <https://github.com/dalen/puppet-puppetdbquery>`_ and `node-puppetdbquery
    <https://github.com/dalen/node-puppetdbquery>`_ by `Erik Dalén
    <https://github.com/dalen>`_.

    :param str s: The query to parse and transform
    :param bool json: Whether to JSON-encode the PuppetDB AST result
    :param dict lex_options: Options passed to :func:`ply.lex.lex`
    :param dict yacc_options: Options passed to :func:`ply.yacc.yacc`
    """
    parser = Parser(lex_options=lex_options, yacc_options=yacc_options)
    evaluator = Evaluator()

    ast = parser.parse(s)
    raw = evaluator.evaluate(ast)

    if json and raw is not None:
        return json_dumps(raw)
    else:
        return raw
<|endoftext|>


<commit_before>from .settings import reload_api_settings

__version__ = '0.13.1'

VERSION = __version__

def includeme(config):
    reload_api_settings(config.registry.settings)
<commit_msg>version bump<commit_after>from .settings import reload_api_settings

__version__ = '0.13.2'

VERSION = __version__

def includeme(config):
    reload_api_settings(config.registry.settings)
<|endoftext|>


<commit_before>import time
import socket
import zlib
import struct

import pysyncobj.pickle as pickle
import pysyncobj.win_inet_pton

from .poller import POLL_EVENT_TYPE
from .monotonic import monotonic as monotonicTime

class CONNECTION_STATE:
    DISCONNECTED = 0
    CONNECTING = 1
    CONNECTED = 2

def _getAddrType(addr):
    try:
        socket.inet_aton(addr)
        return socket.AF_INET
    except socket.error:
        pass
    try:
        socket.inet_pton(socket.AF_INET6, addr)
        return socket.AF_INET6
    except socket.error:
        pass
    raise Exception('unknown address type')

class TcpConnection(object):

    def __init__(self, poller, onMessageReceived = None, onConnected = None, onDisconnected = None,
                 socket=None, timeout=10.0, sendBufferSize = 2 ** 13, recvBufferSize = 2 ** 13):

        self.sendRandKey = None
        self.recvRandKey = None
        self.encryptor = None

        self.__socket = socket
        self.__readBuffer = bytes()
        self.__writeBuffer = bytes()
        self.__lastReadTime = monotonicTime()
        self.__timeout = timeout
        self.__poller = poller
        if socket is not None:
            self.__socket = socket
            self.__fileno = socket.fileno()
            self.__state = CONNECTION_STATE.CONNECTED
            self.__poller.subscribe(self.__fileno,
                                     self.__processConnection,
                                     POLL_EVENT_TYPE.READ | POLL_EVENT_TYPE.WRITE | POLL_EVENT_TYPE.ERROR)
        else:
            self.__state = CONNECTION_STATE.DISCONNECTED
            self.__fileno = None
            self.__socket = None

        self.__onMessageReceived = onMessageReceived
        self.__onConnected = onConnected
        self.__onDisconnected = onDisconnected
        self.__sendBufferSize = sendBufferSize
        self.__recvBufferSize = recvBufferSize

    def setOnConnectedCallback(self, onConnected):
        self.__onConnected = onConnected

    def setOnMessageReceivedCallback(self, onMessageReceived):
        self.__onMessageReceived = onMessageReceived

    def setOnDisconnectedCallback(self, onDisconnected):
        self.__onDisconnected = onDisconnected

    def connect(self, host, port):
        if host is None:
            return False
        self.__state = CONNECTION_STATE.DISCONNECTED
        self.__fileno = None
        self.__socket = socket.socket(_getAddrType(host), socket.SOCK_STREAM)
        self.__socket.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, self.__sendBufferSize)
        self.__socket.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, self.__recvBufferSize)
        self.__socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        self.__socket.setblocking(0)
        self.__readBuffer = bytes()
        self.__writeBuffer = bytes()
        self.__lastReadTime = monotonicTime()

        try:
            self.__socket.connect((host, port))
        except socket.error as e:
            if e.errno not in (socket.errno.EINPROGRESS, socket.errno.EWOULDBLOCK):
                return False
        self.__fileno = self.__socket.fileno()
        self.__state = CONNECTION_STATE.CONNECTING
        self.__poller.subscribe(self.__fileno,
                                 self.__processConnection,
                                 POLL_EVENT_TYPE.READ | POLL_EVENT_TYPE.WRITE | POLL_EVENT_TYPE.ERROR)
        return True

    def send(self, message):
        if self.sendRandKey:
            message = (self.sendRandKey, message)
        data = zlib.compress(pickle.dumps(message), 3)
        if self.encryptor:
            data = self.encryptor.encrypt(data)
        data = struct.pack('i', len(data)) + data
        self.__writeBuffer += data
        self.__trySendBuffer()

    def fileno(self):
        return self.__fileno

    def disconnect(self):
        print("Disconnect request")
        needCallDisconnect = False
        if self.__onDisconnected is not None and self.__state != CONNECTION_STATE.DISCONNECTED:
            needCallDisconnect = True
        self.sendRandKey = None
        self.recvRandKey = None
        if self.__socket is not None:
            self.__socket.close()
            self.__socket = None
        if self.__fileno is not None:
            self.__poller.unsubscribe(self.__fileno)
            self.__fileno = None
        self.__writeBuffer = bytes()
        self.__readBuffer = bytes()
        self.__state = CONNECTION_STATE.DISCONNECTED
        if needCallDisconnect:
            self.__onDisconnected()

    def getSendBufferSize(self):
        return len(self.__writeBuffer)

    def __processConnection(self, descr, eventType):
        poller = self.__poller
        if descr != self.__fileno:
            poller.unsubscribe(descr)
            return

        if eventType & POLL_EVENT_TYPE.ERROR:
            self.disconnect()
            return

        if monotonicTime() - self.__lastReadTime > self.__timeout:
            self.disconnect()
            return

        if eventType & POLL_EVENT_TYPE.READ or eventType & POLL_EVENT_TYPE.WRITE:
            if self.__socket.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR):
                self.disconnect()
                return

            if self.__state == CONNECTION_STATE.CONNECTING:
                if self.__onConnected is not None:
                    self.__onConnected()
                self.__state = CONNECTION_STATE.CONNECTED
                self.__lastReadTime = monotonicTime()
                return

        if eventType & POLL_EVENT_TYPE.WRITE:
            self.__trySendBuffer()
            if self.__state == CONNECTION_STATE.DISCONNECTED:
                return
            event = POLL_EVENT_TYPE.READ | POLL_EVENT_TYPE.ERROR
            if len(self.__writeBuffer) > 0:
                event |= POLL_EVENT_TYPE.WRITE
            poller.subscribe(descr, self.__processConnection, event)

        if eventType & POLL_EVENT_TYPE.READ:
            self.__tryReadBuffer()
            if self.__state == CONNECTION_STATE.DISCONNECTED:
                return

            while True:
                message = self.__processParseMessage()
                if message is None:
                    break
                if self.__onMessageReceived is not None:
                    self.__onMessageReceived(message)
                if self.__state == CONNECTION_STATE.DISCONNECTED:
                    return

    def __trySendBuffer(self):
        while self.__processSend():
            pass

    def __processSend(self):
        if not self.__writeBuffer:
            return False
        try:
            res = self.__socket.send(self.__writeBuffer)
            if res < 0:
                self.disconnect()
                return False
            if res == 0:
                return False
            self.__writeBuffer = self.__writeBuffer[res:]
            return True
        except socket.error as e:
            if e.errno not in (socket.errno.EAGAIN, socket.errno.EWOULDBLOCK):
                self.disconnect()
            return False

    def __tryReadBuffer(self):
        while self.__processRead():
            pass
        self.__lastReadTime = monotonicTime()

    def __processRead(self):
        try:
            incoming = self.__socket.recv(self.__recvBufferSize)
        except socket.error as e:
            if e.errno not in (socket.errno.EAGAIN, socket.errno.EWOULDBLOCK):
                self.disconnect()
            return False
        if self.__socket.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR):
            self.disconnect()
            return False
        if not incoming:
            self.disconnect()
            return False
        self.__readBuffer += incoming
        return True

    def __processParseMessage(self):
        if len(self.__readBuffer) < 4:
            return None
        l = struct.unpack('i', self.__readBuffer[:4])[0]
        if len(self.__readBuffer) - 4 < l:
            return None
        data = self.__readBuffer[4:4 + l]
        try:
            if self.encryptor:
                data = self.encryptor.decrypt(data)
            message = pickle.loads(zlib.decompress(data))
            if self.recvRandKey:
                randKey, message = message
                assert randKey == self.recvRandKey
        except:
            self.disconnect()
            return None
        self.__readBuffer = self.__readBuffer[4 + l:]
        return message

    @property
    def state(self):
        return self.__state
<commit_msg>Removed debug requests<commit_after>import time
import socket
import zlib
import struct

import pysyncobj.pickle as pickle
import pysyncobj.win_inet_pton

from .poller import POLL_EVENT_TYPE
from .monotonic import monotonic as monotonicTime

class CONNECTION_STATE:
    DISCONNECTED = 0
    CONNECTING = 1
    CONNECTED = 2

def _getAddrType(addr):
    try:
        socket.inet_aton(addr)
        return socket.AF_INET
    except socket.error:
        pass
    try:
        socket.inet_pton(socket.AF_INET6, addr)
        return socket.AF_INET6
    except socket.error:
        pass
    raise Exception('unknown address type')

class TcpConnection(object):

    def __init__(self, poller, onMessageReceived = None, onConnected = None, onDisconnected = None,
                 socket=None, timeout=10.0, sendBufferSize = 2 ** 13, recvBufferSize = 2 ** 13):

        self.sendRandKey = None
        self.recvRandKey = None
        self.encryptor = None

        self.__socket = socket
        self.__readBuffer = bytes()
        self.__writeBuffer = bytes()
        self.__lastReadTime = monotonicTime()
        self.__timeout = timeout
        self.__poller = poller
        if socket is not None:
            self.__socket = socket
            self.__fileno = socket.fileno()
            self.__state = CONNECTION_STATE.CONNECTED
            self.__poller.subscribe(self.__fileno,
                                     self.__processConnection,
                                     POLL_EVENT_TYPE.READ | POLL_EVENT_TYPE.WRITE | POLL_EVENT_TYPE.ERROR)
        else:
            self.__state = CONNECTION_STATE.DISCONNECTED
            self.__fileno = None
            self.__socket = None

        self.__onMessageReceived = onMessageReceived
        self.__onConnected = onConnected
        self.__onDisconnected = onDisconnected
        self.__sendBufferSize = sendBufferSize
        self.__recvBufferSize = recvBufferSize

    def setOnConnectedCallback(self, onConnected):
        self.__onConnected = onConnected

    def setOnMessageReceivedCallback(self, onMessageReceived):
        self.__onMessageReceived = onMessageReceived

    def setOnDisconnectedCallback(self, onDisconnected):
        self.__onDisconnected = onDisconnected

    def connect(self, host, port):
        if host is None:
            return False
        self.__state = CONNECTION_STATE.DISCONNECTED
        self.__fileno = None
        self.__socket = socket.socket(_getAddrType(host), socket.SOCK_STREAM)
        self.__socket.setsockopt(socket.SOL_SOCKET, socket.SO_SNDBUF, self.__sendBufferSize)
        self.__socket.setsockopt(socket.SOL_SOCKET, socket.SO_RCVBUF, self.__recvBufferSize)
        self.__socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        self.__socket.setblocking(0)
        self.__readBuffer = bytes()
        self.__writeBuffer = bytes()
        self.__lastReadTime = monotonicTime()

        try:
            self.__socket.connect((host, port))
        except socket.error as e:
            if e.errno not in (socket.errno.EINPROGRESS, socket.errno.EWOULDBLOCK):
                return False
        self.__fileno = self.__socket.fileno()
        self.__state = CONNECTION_STATE.CONNECTING
        self.__poller.subscribe(self.__fileno,
                                 self.__processConnection,
                                 POLL_EVENT_TYPE.READ | POLL_EVENT_TYPE.WRITE | POLL_EVENT_TYPE.ERROR)
        return True

    def send(self, message):
        if self.sendRandKey:
            message = (self.sendRandKey, message)
        data = zlib.compress(pickle.dumps(message), 3)
        if self.encryptor:
            data = self.encryptor.encrypt(data)
        data = struct.pack('i', len(data)) + data
        self.__writeBuffer += data
        self.__trySendBuffer()

    def fileno(self):
        return self.__fileno

    def disconnect(self):
        needCallDisconnect = False
        if self.__onDisconnected is not None and self.__state != CONNECTION_STATE.DISCONNECTED:
            needCallDisconnect = True
        self.sendRandKey = None
        self.recvRandKey = None
        if self.__socket is not None:
            self.__socket.close()
            self.__socket = None
        if self.__fileno is not None:
            self.__poller.unsubscribe(self.__fileno)
            self.__fileno = None
        self.__writeBuffer = bytes()
        self.__readBuffer = bytes()
        self.__state = CONNECTION_STATE.DISCONNECTED
        if needCallDisconnect:
            self.__onDisconnected()

    def getSendBufferSize(self):
        return len(self.__writeBuffer)

    def __processConnection(self, descr, eventType):
        poller = self.__poller
        if descr != self.__fileno:
            poller.unsubscribe(descr)
            return

        if eventType & POLL_EVENT_TYPE.ERROR:
            self.disconnect()
            return

        if monotonicTime() - self.__lastReadTime > self.__timeout:
            self.disconnect()
            return

        if eventType & POLL_EVENT_TYPE.READ or eventType & POLL_EVENT_TYPE.WRITE:
            if self.__socket.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR):
                self.disconnect()
                return

            if self.__state == CONNECTION_STATE.CONNECTING:
                if self.__onConnected is not None:
                    self.__onConnected()
                self.__state = CONNECTION_STATE.CONNECTED
                self.__lastReadTime = monotonicTime()
                return

        if eventType & POLL_EVENT_TYPE.WRITE:
            self.__trySendBuffer()
            if self.__state == CONNECTION_STATE.DISCONNECTED:
                return
            event = POLL_EVENT_TYPE.READ | POLL_EVENT_TYPE.ERROR
            if len(self.__writeBuffer) > 0:
                event |= POLL_EVENT_TYPE.WRITE
            poller.subscribe(descr, self.__processConnection, event)

        if eventType & POLL_EVENT_TYPE.READ:
            self.__tryReadBuffer()
            if self.__state == CONNECTION_STATE.DISCONNECTED:
                return

            while True:
                message = self.__processParseMessage()
                if message is None:
                    break
                if self.__onMessageReceived is not None:
                    self.__onMessageReceived(message)
                if self.__state == CONNECTION_STATE.DISCONNECTED:
                    return

    def __trySendBuffer(self):
        while self.__processSend():
            pass

    def __processSend(self):
        if not self.__writeBuffer:
            return False
        try:
            res = self.__socket.send(self.__writeBuffer)
            if res < 0:
                self.disconnect()
                return False
            if res == 0:
                return False
            self.__writeBuffer = self.__writeBuffer[res:]
            return True
        except socket.error as e:
            if e.errno not in (socket.errno.EAGAIN, socket.errno.EWOULDBLOCK):
                self.disconnect()
            return False

    def __tryReadBuffer(self):
        while self.__processRead():
            pass
        self.__lastReadTime = monotonicTime()

    def __processRead(self):
        try:
            incoming = self.__socket.recv(self.__recvBufferSize)
        except socket.error as e:
            if e.errno not in (socket.errno.EAGAIN, socket.errno.EWOULDBLOCK):
                self.disconnect()
            return False
        if self.__socket.getsockopt(socket.SOL_SOCKET, socket.SO_ERROR):
            self.disconnect()
            return False
        if not incoming:
            self.disconnect()
            return False
        self.__readBuffer += incoming
        return True

    def __processParseMessage(self):
        if len(self.__readBuffer) < 4:
            return None
        l = struct.unpack('i', self.__readBuffer[:4])[0]
        if len(self.__readBuffer) - 4 < l:
            return None
        data = self.__readBuffer[4:4 + l]
        try:
            if self.encryptor:
                data = self.encryptor.decrypt(data)
            message = pickle.loads(zlib.decompress(data))
            if self.recvRandKey:
                randKey, message = message
                assert randKey == self.recvRandKey
        except:
            self.disconnect()
            return None
        self.__readBuffer = self.__readBuffer[4 + l:]
        return message

    @property
    def state(self):
        return self.__state
<|endoftext|>


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

class SpinnakerPipelineManual(SpinnakerPipeline):
    """Manual JSON configured Spinnaker Pipelines."""

    def create_pipeline(self):
        """Use JSON files to create Pipelines."""
        self.log.info('Uploading manual Pipelines: %s')
        lookup = FileLookup(git_short=self.generated.gitlab()['main'], runway_dir=self.runway_dir)

        for json_file in self.settings['pipeline']['pipeline_files']:
            json_text = lookup.get(filename=json_file)
            self.post_pipeline(json_text)

        return True
<commit_msg>fix: Use filename for Pipeline name<commit_after>#   Foremast - Pipeline Tooling
#
#   Copyright 2016 Gogo, LLC
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.
"""Create manual Pipeline for Spinnaker."""
from ..utils.lookups import FileLookup
from .create_pipeline import SpinnakerPipeline

class SpinnakerPipelineManual(SpinnakerPipeline):
    """Manual JSON configured Spinnaker Pipelines."""

    def create_pipeline(self):
        """Use JSON files to create Pipelines."""
        self.log.info('Uploading manual Pipelines: %s')
        lookup = FileLookup(git_short=self.generated.gitlab()['main'], runway_dir=self.runway_dir)

        for json_file in self.settings['pipeline']['pipeline_files']:
            json_dict = lookup.json(filename=json_file)
            json_dict['name'] = json_file
            self.post_pipeline(json_dict)

        return True
<|endoftext|>


<commit_before>'''
Created on 18.10.2013

@author: marscher
'''
from emma2.util.log import getLogger
log = getLogger()

from emma2.util.pystallone import stallone_available,\
    stallone_array_to_ndarray, ndarray_to_stallone_array

if stallone_available:
    from emma2.util.pystallone import API, JavaError
else:
    log.warn('stallone not available. No functionality of TPT class!')

class TPT():
    """
        This class wraps around a stallone ITPTFlux class.
    """
    def __init__(self, T, A, B):
        """
        initializes this with given matrix
        Parameters
        ----------
        T : Transition matrix
            ndarray
        A : set of states
        # TODO: shape?
            ndarray( dtype = int, shape=?)
        B : set of states
            ndarray( dtype = int)
        # TODO: is this valid in sphinx?
        Throws
        ------
        RuntimeError, if stallone is not available
        """
        import numpy as np
        try:
            T = T.astype(np.float64)
            A = np.asarray(A).astype(np.int32)
            B = np.asarray(B).astype(np.int32)
            A = ndarray_to_stallone_array(A)
            B = ndarray_to_stallone_array(B)
            T = ndarray_to_stallone_array(T)
            log.debug('create TPT instance and calculate fluxes.')
            self.ITPT = API.msmNew.createTPT(T, A, B)
            log.debug('finished TPT calculation.')
        except JavaError as je:
            exception = je.getJavaException()
            msg = str(exception) + '\n' + \
                str(exception.getStackTrace()).replace(',', '\n')
            raise RuntimeError(msg)
        except Exception as e:
            log.error("tpt: unknown error occured: %s" %e)

    def getBackwardCommittor(self):
        """
        Returns
        -------
        Backward Committor : ndarray
        """
        return stallone_array_to_ndarray(self.ITPT.getBackwardCommittor())

    def getFlux(self):
        """
        Returns
        -------
        Flux : ndarray
        """
        return stallone_array_to_ndarray(self.ITPT.getFlux())

    def getForwardCommittor(self):
        """
        Returns
        -------
        Forward Committor : ndarray

        """
        return stallone_array_to_ndarray(self.ITPT.getForwardCommittor())

    def getNetFlux(self):
        """
        Returns
        -------
        Net flux : ndarray

        """
        return stallone_array_to_ndarray(self.ITPT.getNetFlux())

    def getRate(self):
        """
        Returns
        -------
        Rate : float
        """
        return self.ITPT.getRate()

    def getStationaryDistribution(self):
        """
        Returns
        -------
        Stationary distribution : ndarray

        """
        return stallone_array_to_ndarray(self.ITPT.getStationaryDistribution())

    def getTotalFlux(self):
        """
        Returns
        -------
        Total flux : float

        """
        return self.ITPT.getTotalFlux()

    def setRateMatrix(self, K):
        """
        Parameters
        ----------
        K (Rate matrix) : ndarray

        """
        self.K = ndarray_to_stallone_array(K)

    def setStationaryDistribution(self, pi):
        """
        Parameters
        ----------
        pi : stationary distribution

        """
        self.pi = ndarray_to_stallone_array(pi)

    def setTransitionMatrix(self, T):
        """
        Parameters
        ----------
        T : ndarray
        """
        self.T = ndarray_to_stallone_array(T)
<commit_msg>minor changes<commit_after>'''
Created on 18.10.2013

@author: marscher
'''
__all__ = ['TPT']

from emma2.util.log import getLogger
log = getLogger()

from emma2.util.pystallone import stallone_available,\
    stallone_array_to_ndarray, ndarray_to_stallone_array

if stallone_available:
    from emma2.util.pystallone import API, JavaError
else:
    log.warn('stallone not available. No functionality of TPT class!')

class TPT():
    """
        This class wraps around a stallone ITPTFlux class.
    """
    def __init__(self, T, A, B):
        """
        initializes this with given matrix
        Parameters
        ----------
        T : Transition matrix
            ndarray
        A : set of states
            ndarray( dtype = int)
        B : set of states
            ndarray( dtype = int)
        Throws
        ------
        RuntimeError, if stallone is not available
        """
        import numpy as np
        try:
            T = T.astype(np.float64)
            A = np.asarray(A).astype(np.int32)
            B = np.asarray(B).astype(np.int32)
            A = ndarray_to_stallone_array(A)
            B = ndarray_to_stallone_array(B)
            T = ndarray_to_stallone_array(T)
            log.debug('create TPT instance and calculate fluxes.')
            self.ITPT = API.msmNew.createTPT(T, A, B)
            log.debug('finished TPT calculation.')
        except JavaError as je:
            exception = je.getJavaException()
            msg = str(exception) + '\n' + \
                str(exception.getStackTrace()).replace(',', '\n')
            raise RuntimeError(msg)
        except Exception as e:
            log.error("tpt: unknown error occurred: %s" %e)

    def getBackwardCommittor(self):
        """
        Returns
        -------
        Backward Committor : ndarray
        """
        return stallone_array_to_ndarray(self.ITPT.getBackwardCommittor())

    def getFlux(self):
        """
        Returns
        -------
        Flux : ndarray
        """
        return stallone_array_to_ndarray(self.ITPT.getFlux())

    def getForwardCommittor(self):
        """
        Returns
        -------
        Forward Committor : ndarray

        """
        return stallone_array_to_ndarray(self.ITPT.getForwardCommittor())

    def getNetFlux(self):
        """
        Returns
        -------
        Net flux : ndarray

        """
        return stallone_array_to_ndarray(self.ITPT.getNetFlux())

    def getRate(self):
        """
        Returns
        -------
        Rate : float
        """
        return self.ITPT.getRate()

    def getStationaryDistribution(self):
        """
        Returns
        -------
        Stationary distribution : ndarray

        """
        return stallone_array_to_ndarray(self.ITPT.getStationaryDistribution())

    def getTotalFlux(self):
        """
        Returns
        -------
        Total flux : float

        """
        return self.ITPT.getTotalFlux()

    def setRateMatrix(self, K):
        """
        Parameters
        ----------
        K (Rate matrix) : ndarray

        """
        self.K = ndarray_to_stallone_array(K)

    def setStationaryDistribution(self, pi):
        """
        Parameters
        ----------
        pi : stationary distribution

        """
        self.pi = ndarray_to_stallone_array(pi)

    def setTransitionMatrix(self, T):
        """
        Parameters
        ----------
        T : ndarray
        """
        self.T = ndarray_to_stallone_array(T)
<|endoftext|>