<commit_before>from django.http import HttpResponseRedirect, Http404
from django.conf import settings
from django.core.urlresolvers import resolve
from exceptions import IndexError
from models import URL
import operator, re

class AliasFallbackMiddleware(object):
    def process_response(self, request, response):
        if response.status_code != 404:
            return response

        try:
            if hasattr(settings, 'ALIASES_MAP_ARGS') and settings.ALIASES_MAP_ARGS is False:
                alias = URL.objects.get(location=request.path_info)
            else:
                # I'm not a fan of raw SQL queries, but this should work in any SQL server (I think). If you know how to do
                # this in django's ORM, let me know. If this doesn't work and/or you don't need this functionality, you can
                # always set settings.ALIASES_MAP_ARGS to False and avoid this.

                alias = URL.objects.raw('SELECT * FROM aliases_url WHERE "%s" REGEXP CONCAT("^", location) ORDER BY LENGTH(location) LIMIT 1', [request.path_info])

                try:
                    alias = alias[0]
                except IndexError:
                    alias = None

            if alias is not None:
                if alias.get_related_url():
                    match = resolve(alias.get_related_url())
                else:
                    match = None

                # When we have extra data in the request, we pass it as args to the view
                if len(alias.location) < len(request.path_info):
                    additional_args = re.sub('\/$', '', request.path_info)
                    additional_args = additional_args[len(alias.location):].split('/')

                    # This is a tuple, which is immutable. Convert it to a list.
                    match.args = list(match.args)

                    match.args.extend(additional_args)

                if match:
                    return match.func(request, *match.args, **match.kwargs)
                else:
                    return HttpResponseRedirect(alias.get_related_url())
            else:
                raise Http404('The specified URL mapping does not exist.')

        except URL.DoesNotExist:
            return response

        except:
            if settings.DEBUG:
                raise

            return response

<commit_msg>Modified the query to use LIKE instead of REGEXP. I've been trying to find a way to get rid of REGEXP here...<commit_after>from django.http import HttpResponseRedirect, Http404
from django.conf import settings
from django.core.urlresolvers import resolve
from exceptions import IndexError
from models import URL
import operator, re

class AliasFallbackMiddleware(object):
    def process_response(self, request, response):
        if response.status_code != 404:
            return response

        try:
            if hasattr(settings, 'ALIASES_MAP_ARGS') and settings.ALIASES_MAP_ARGS is False:
                alias = URL.objects.get(location=request.path_info)
            else:
                # I'm not a fan of raw SQL queries, but this should work in any SQL server (I think). If you know how to do
                # this in django's ORM, let me know. If this doesn't work and/or you don't need this functionality, you can
                # always set settings.ALIASES_MAP_ARGS to False and avoid this.

                alias = URL.objects.raw('SELECT * FROM aliases_url WHERE "%s" LIKE CONCAT(location, "%") ORDER BY LENGTH(location) LIMIT 1', [request.path_info])

                try:
                    alias = alias[0]
                except IndexError:
                    alias = None

            if alias is not None:
                if alias.get_related_url():
                    match = resolve(alias.get_related_url())
                else:
                    match = None

                # When we have extra data in the request, we pass it as args to the view
                if len(alias.location) < len(request.path_info):
                    additional_args = re.sub('\/$', '', request.path_info)
                    additional_args = additional_args[len(alias.location):].split('/')

                    # This is a tuple, which is immutable. Convert it to a list.
                    match.args = list(match.args)

                    match.args.extend(additional_args)

                if match:
                    return match.func(request, *match.args, **match.kwargs)
                else:
                    return HttpResponseRedirect(alias.get_related_url())
            else:
                raise Http404('The specified URL mapping does not exist.')

        except URL.DoesNotExist:
            return response

        except:
            if settings.DEBUG:
                raise

            return response

<|endoftext|>


<commit_before>import numpy as np
from matplotlib import pyplot as plt
from bayessb.report import reporter, Result, FloatListResult, ThumbnailResult
from bayessb import convergence
from StringIO import StringIO
from matplotlib.figure import Figure
from matplotlib.backends.backend_agg import FigureCanvasAgg

reporter_group_name = "Estimation"
num_samples = 100

@reporter('Number of chains')
def num_chains(mcmc_set):
    return Result(len(mcmc_set.chains), None)

"""
@reporter('Estimation parameters')
def estimation_parameters(mcmc_set):
    chain = mcmc_set.chains[0]
    opts = chain.options
    output = StringIO()
    output.write("use_hessian: %s, " % opts.use_hessian)
    output.write("hessian_period: %s, " % opts.hessian_period)
    output.write("hessian_scale: %s, " % opts.hessian_scale)
    output.write("norm_step_size: %s, " % opts.norm_step_size)
    output.write("sigma_adj_interval: %s, " % opts.sigma_adj_interval)
    output.write("anneal_length: %s, " % opts.anneal_length)
    output.write("T_init: %s, " % opts.T_init)
    output.write("accept_rate_target: %s, " % opts.accept_rate_target)
    output.write("sigma_max: %s, " % opts.sigma_max)
    output.write("sigma_min: %s, " % opts.sigma_min)
    output.write("sigma_step: %s, " % opts.sigma_step)
    output.write("seed: %s, " % opts.seed)
    #output.write("accept_window: %s, " % opts.accept_window)
    return Result(output.getvalue(), None)
"""

@reporter('Conv. Criterion')
def convergence_criterion(mcmc_set):
    """Returns the vector of Gelman-Rubin convergence criterion values and a
    link to an HTML file containing plots of the traces of the walk for each
    parameter fitted."""

    # Prepare html for page showing plots of parameter traces
    html_str = "<title>Parameter traces for %s</title>\n" \
               % mcmc_set.name
    html_str += "
Parameter traces for %s
\n" \
                % mcmc_set.name
    img_str_list = []

    # Make plots of parameter traces
    for i in range(mcmc_set.chains[0].num_estimate):
        param_name = mcmc_set.chains[0].options.estimate_params[i].name
        fig = Figure()
        ax = fig.gca()
        for chain in mcmc_set.chains:
            if chain.pruned:
                ax.plot(chain.thinned_accept_steps, chain.positions[:,i])
            else:
                ax.plot(chain.positions[:, i])
        ax.set_title("Parameter: %s" % param_name)
        plot_filename = '%s_trace_%s.png' % (mcmc_set.name, param_name)
        canvas = FigureCanvasAgg(fig)
        fig.set_canvas(canvas)
        fig.savefig(plot_filename)
        img_str_list.append(plot_filename)

    # Make the html file
    html_str += '\n'.join([
        '' %
        (i, i) for i in img_str_list])
    html_str += ""
    html_filename = '%s_convergence.html' % mcmc_set.name
    with open(html_filename, 'w') as f:
        f.write(html_str)

    return FloatListResult(convergence.convergence_criterion(mcmc_set),
                           html_filename)

@reporter('Maximum likelihood')
def maximum_likelihood(mcmc_set):
    # Get the maximum likelihood
    (max_likelihood, max_likelihood_position) = mcmc_set.maximum_likelihood()

    return show_fit_at_position(mcmc_set, max_likelihood,
                                max_likelihood_position, 'max_likelihood')

@reporter('Maximum posterior')
def maximum_posterior(mcmc_set):
    # Get the maximum posterior
    (max_posterior, max_posterior_position) = mcmc_set.maximum_posterior()

    return show_fit_at_position(mcmc_set, max_posterior,
                                max_posterior_position, 'max_posterior')

def show_fit_at_position(mcmc_set, fit_value, position, fit_name):
    """Create the result page showing the fit quality at the given position.

    Parameters
    ----------
    mcmc_set : MCMCSet object
        The set of MCMC chains
    fit_value : float
        The quality of fit at the given position.
    position : numpy.array
        Array of (log10-transformed) parameter values at the given position.
    fit_name : string
        A shorthand name for the fit position, e.g., "max_likelihood". Should
        conform to rules of Python variable naming (no spaces, doesn't
        start with a number, etc.).

    Returns
    -------
    A result object containing the fit value and the link to the accompanying
    HTML plot page, if any.
    """

    # If the MCMC object does not have a fit_plotting_function defined
    # (for example, if it is a base MCMC object), then don't create a
    # plot for visualization.
    if not hasattr(mcmc_set.chains[0], 'fit_plotting_function'):
        return Result(fit_value, None)

    # Prepare html for page showing plots at position
    html_str = "<title>Simulation of %s " \
               "with %s parameter values</title>\n" \
               % (mcmc_set.name, fit_name)
    html_str += "
Simulation of %s with %s " \
                "parameter values
\n" % (mcmc_set.name, fit_name)

    # Show the plot vs. the data at the position
    fig = mcmc_set.chains[0].fit_plotting_function(position=position)
    img_filename = '%s_%s_plot.png' % (mcmc_set.name, fit_name)
    fig.savefig(img_filename)
    html_str += '

' % img_filename

    # Show the plot of all observables at the position
    chain0 = mcmc_set.chains[0]
    tspan = chain0.options.tspan
    observables = chain0.options.model.observables
    x = chain0.simulate(position=position, observables=True)

    fig = Figure()
    ax = fig.gca()
    lines = []
    for o in observables:
        line = ax.plot(tspan, x[o.name])
        lines.append(line)
    ax.set_title("Observables at %s" % fit_name)
    fig.legend(lines, [o.name for o in observables], 'lower right')
    canvas = FigureCanvasAgg(fig)
    fig.set_canvas(canvas)

    img_filename = '%s_%s_species.png' % (mcmc_set.name, fit_name)
    fig.savefig(img_filename)
    html_str += '

' % img_filename

    # Print the parameter values for the position as a dict that can be
    # used to override the initial values
    html_str += '
%s_params = {\n' % fit_name
    for i, p in enumerate(chain0.options.estimate_params):
        html_str += "\t'%s': %.17g,\n" % \
                    (p.name, 10 ** position[i])
    html_str += '}
'

    html_str += ''

    # Create the html file
    html_filename = '%s_%s_plot.html' % (mcmc_set.name, fit_name)
    with open(html_filename, 'w') as f:
        f.write(html_str)

    return Result(fit_value, html_filename)

@reporter('Sample fits')
def sample_fits(mcmc_set):
    tspan = mcmc_set.chains[0].options.tspan
    fig = Figure()
    ax = fig.gca()
    plot_filename = '%s_sample_fits.png' % mcmc_set.name
    thumbnail_filename = '%s_sample_fits_th.png' % mcmc_set.name

    # Make sure we can call the method 'get_observable_timecourse'
    if not hasattr(mcmc_set.chains[0], 'get_observable_timecourse') or \
       not hasattr(mcmc_set.chains[0], 'plot_data'):
        return Result('None', None)

    # Plot the original data
    mcmc_set.chains[0].plot_data(ax)

    # Plot a sampling of trajectories from the original parameter set
    for i in range(num_samples):
        position = mcmc_set.get_sample_position()
        x = mcmc_set.chains[0].get_observable_timecourse(position=position)
        ax.plot(tspan, x, color='g', alpha=0.1)

    canvas = FigureCanvasAgg(fig)
    fig.set_canvas(canvas)
    fig.savefig(plot_filename)
    fig.savefig(thumbnail_filename, dpi=10)

    return ThumbnailResult(thumbnail_filename, plot_filename)

@reporter('Marginals')
def marginals(mcmc_set):
    """Returns the vector of Gelman-Rubin convergence criterion values and a
    link to an HTML file containing plots of the traces of the walk for each
    parameter fitted."""

    # Prepare html for page showing plots of parameter traces
    html_str = "<title>Marginal distributions for " \
               "%s</title>\n" % mcmc_set.name
    html_str += "
Marginal distributions for for %s
\n" \
                % mcmc_set.name
    img_str_list = []

    # Make plots of parameter traces
    for i in range(mcmc_set.chains[0].num_estimate):
        param_name = mcmc_set.chains[0].options.estimate_params[i].name
        fig = Figure()
        ax = fig.gca()
        chains_for_param = [chain.positions[:,i] for chain in mcmc_set.chains]
        ax.hist(chains_for_param, histtype='step')
        ax.set_title("Parameter: %s" % param_name)
        plot_filename = '%s_marginal_%s.png' % (mcmc_set.name, param_name)
        canvas = FigureCanvasAgg(fig)
        fig.set_canvas(canvas)
        fig.savefig(plot_filename)
        img_str_list.append(plot_filename)

    # Make the html file
    html_str += '\n'.join([
        '' %
        (i, i) for i in img_str_list])
    html_str += ""
    html_filename = '%s_marginals.html' % mcmc_set.name
    with open(html_filename, 'w') as f:
        f.write(html_str)

    return Result(None, html_filename)

<commit_msg>Fixed get_observable_timecourse<commit_after>import numpy as np
from matplotlib import pyplot as plt
from bayessb.report import reporter, Result, FloatListResult, ThumbnailResult
from bayessb import convergence
from StringIO import StringIO
from matplotlib.figure import Figure
from matplotlib.backends.backend_agg import FigureCanvasAgg

reporter_group_name = "Estimation"
num_samples = 100

@reporter('Number of chains')
def num_chains(mcmc_set):
    return Result(len(mcmc_set.chains), None)

"""
@reporter('Estimation parameters')
def estimation_parameters(mcmc_set):
    chain = mcmc_set.chains[0]
    opts = chain.options
    output = StringIO()
    output.write("use_hessian: %s, " % opts.use_hessian)
    output.write("hessian_period: %s, " % opts.hessian_period)
    output.write("hessian_scale: %s, " % opts.hessian_scale)
    output.write("norm_step_size: %s, " % opts.norm_step_size)
    output.write("sigma_adj_interval: %s, " % opts.sigma_adj_interval)
    output.write("anneal_length: %s, " % opts.anneal_length)
    output.write("T_init: %s, " % opts.T_init)
    output.write("accept_rate_target: %s, " % opts.accept_rate_target)
    output.write("sigma_max: %s, " % opts.sigma_max)
    output.write("sigma_min: %s, " % opts.sigma_min)
    output.write("sigma_step: %s, " % opts.sigma_step)
    output.write("seed: %s, " % opts.seed)
    #output.write("accept_window: %s, " % opts.accept_window)
    return Result(output.getvalue(), None)
"""

@reporter('Conv. Criterion')
def convergence_criterion(mcmc_set):
    """Returns the vector of Gelman-Rubin convergence criterion values and a
    link to an HTML file containing plots of the traces of the walk for each
    parameter fitted."""

    # Prepare html for page showing plots of parameter traces
    html_str = "<title>Parameter traces for %s</title>\n" \
               % mcmc_set.name
    html_str += "
Parameter traces for %s
\n" \
                % mcmc_set.name
    img_str_list = []

    # Make plots of parameter traces
    for i in range(mcmc_set.chains[0].num_estimate):
        param_name = mcmc_set.chains[0].options.estimate_params[i].name
        fig = Figure()
        ax = fig.gca()
        for chain in mcmc_set.chains:
            if chain.pruned:
                ax.plot(chain.thinned_accept_steps, chain.positions[:,i])
            else:
                ax.plot(chain.positions[:, i])
        ax.set_title("Parameter: %s" % param_name)
        plot_filename = '%s_trace_%s.png' % (mcmc_set.name, param_name)
        canvas = FigureCanvasAgg(fig)
        fig.set_canvas(canvas)
        fig.savefig(plot_filename)
        img_str_list.append(plot_filename)

    # Make the html file
    html_str += '\n'.join([
        '' %
        (i, i) for i in img_str_list])
    html_str += ""
    html_filename = '%s_convergence.html' % mcmc_set.name
    with open(html_filename, 'w') as f:
        f.write(html_str)

    return FloatListResult(convergence.convergence_criterion(mcmc_set),
                           html_filename)

@reporter('Maximum likelihood')
def maximum_likelihood(mcmc_set):
    # Get the maximum likelihood
    (max_likelihood, max_likelihood_position) = mcmc_set.maximum_likelihood()

    return show_fit_at_position(mcmc_set, max_likelihood,
                                max_likelihood_position, 'max_likelihood')

@reporter('Maximum posterior')
def maximum_posterior(mcmc_set):
    # Get the maximum posterior
    (max_posterior, max_posterior_position) = mcmc_set.maximum_posterior()

    return show_fit_at_position(mcmc_set, max_posterior,
                                max_posterior_position, 'max_posterior')

def show_fit_at_position(mcmc_set, fit_value, position, fit_name):
    """Create the result page showing the fit quality at the given position.

    Parameters
    ----------
    mcmc_set : MCMCSet object
        The set of MCMC chains
    fit_value : float
        The quality of fit at the given position.
    position : numpy.array
        Array of (log10-transformed) parameter values at the given position.
    fit_name : string
        A shorthand name for the fit position, e.g., "max_likelihood". Should
        conform to rules of Python variable naming (no spaces, doesn't
        start with a number, etc.).

    Returns
    -------
    A result object containing the fit value and the link to the accompanying
    HTML plot page, if any.
    """

    # If the MCMC object does not have a fit_plotting_function defined
    # (for example, if it is a base MCMC object), then don't create a
    # plot for visualization.
    if not hasattr(mcmc_set.chains[0], 'fit_plotting_function'):
        return Result(fit_value, None)

    # Prepare html for page showing plots at position
    html_str = "<title>Simulation of %s " \
               "with %s parameter values</title>\n" \
               % (mcmc_set.name, fit_name)
    html_str += "
Simulation of %s with %s " \
                "parameter values
\n" % (mcmc_set.name, fit_name)

    # Show the plot vs. the data at the position
    fig = mcmc_set.chains[0].fit_plotting_function(position=position)
    img_filename = '%s_%s_plot.png' % (mcmc_set.name, fit_name)
    fig.savefig(img_filename)
    html_str += '

' % img_filename

    # Show the plot of all observables at the position
    chain0 = mcmc_set.chains[0]
    tspan = chain0.options.tspan
    observables = chain0.options.model.observables
    x = chain0.simulate(position=position, observables=True)

    fig = Figure()
    ax = fig.gca()
    lines = []
    for o in observables:
        line = ax.plot(tspan, x[o.name])
        lines.append(line)
    ax.set_title("Observables at %s" % fit_name)
    fig.legend(lines, [o.name for o in observables], 'lower right')
    canvas = FigureCanvasAgg(fig)
    fig.set_canvas(canvas)

    img_filename = '%s_%s_species.png' % (mcmc_set.name, fit_name)
    fig.savefig(img_filename)
    html_str += '

' % img_filename

    # Print the parameter values for the position as a dict that can be
    # used to override the initial values
    html_str += '
%s_params = {\n' % fit_name
    for i, p in enumerate(chain0.options.estimate_params):
        html_str += "\t'%s': %.17g,\n" % \
                    (p.name, 10 ** position[i])
    html_str += '}
'

    html_str += ''

    # Create the html file
    html_filename = '%s_%s_plot.html' % (mcmc_set.name, fit_name)
    with open(html_filename, 'w') as f:
        f.write(html_str)

    return Result(fit_value, html_filename)

@reporter('Sample fits')
def sample_fits(mcmc_set):
    tspan = mcmc_set.chains[0].options.tspan
    fig = Figure()
    ax = fig.gca()
    plot_filename = '%s_sample_fits.png' % mcmc_set.name
    thumbnail_filename = '%s_sample_fits_th.png' % mcmc_set.name

    # Make sure we can call the method 'get_observable_timecourses'
    if not hasattr(mcmc_set.chains[0], 'get_observable_timecourses') or \
       not hasattr(mcmc_set.chains[0], 'plot_data'):
        return Result('None', None)

    # Plot the original data
    mcmc_set.chains[0].plot_data(ax)

    # Plot a sampling of trajectories from the original parameter set
    for i in range(num_samples):
        position = mcmc_set.get_sample_position()
        timecourses = mcmc_set.chains[0].get_observable_timecourses(position=position)
        for obs_name, timecourse in timecourses.iteritems():
            ax.plot(tspan, timecourse, color='g', alpha=0.1, label=obs_name)

    canvas = FigureCanvasAgg(fig)
    fig.set_canvas(canvas)
    fig.savefig(plot_filename)
    fig.savefig(thumbnail_filename, dpi=10)

    return ThumbnailResult(thumbnail_filename, plot_filename)

@reporter('Marginals')
def marginals(mcmc_set):
    """Returns the vector of Gelman-Rubin convergence criterion values and a
    link to an HTML file containing plots of the traces of the walk for each
    parameter fitted."""

    # Prepare html for page showing plots of parameter traces
    html_str = "<title>Marginal distributions for " \
               "%s</title>\n" % mcmc_set.name
    html_str += "
Marginal distributions for for %s
\n" \
                % mcmc_set.name
    img_str_list = []

    # Make plots of parameter traces
    for i in range(mcmc_set.chains[0].num_estimate):
        param_name = mcmc_set.chains[0].options.estimate_params[i].name
        fig = Figure()
        ax = fig.gca()
        chains_for_param = [chain.positions[:,i] for chain in mcmc_set.chains]
        ax.hist(chains_for_param, histtype='step')
        ax.set_title("Parameter: %s" % param_name)
        plot_filename = '%s_marginal_%s.png' % (mcmc_set.name, param_name)
        canvas = FigureCanvasAgg(fig)
        fig.set_canvas(canvas)
        fig.savefig(plot_filename)
        img_str_list.append(plot_filename)

    # Make the html file
    html_str += '\n'.join([
        '' %
        (i, i) for i in img_str_list])
    html_str += ""
    html_filename = '%s_marginals.html' % mcmc_set.name
    with open(html_filename, 'w') as f:
        f.write(html_str)

    return Result(None, html_filename)

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import os
import time
import unittest
import unittest.mock as mock
from datetime import datetime

import intelmq.lib.test as test
from intelmq.bots.outputs.elasticsearch.output import ElasticsearchOutputBot, get_event_date

if os.environ.get('INTELMQ_TEST_DATABASES'):
    import elasticsearch

INPUT1 = {"__type": "Event",
          "classification.type": "infected-system",
          "source.asn": 64496,
          "source.ip": "192.0.2.1",
          "feed.name": "Example Feed",
          "extra": '{"foo.bar": "test"}'
          }
OUTPUT1 = {'classification.type': 'infected-system',
           'extra.foo.bar': 'test',
           'feed.name': 'Example Feed',
           'source.asn': 64496,
           'source.ip': '192.0.2.1',
           }
OUTPUT1_REPLACEMENT_CHARS = {
    'classification_type': 'infected-system',
    'extra_foo_bar': 'test',
    'feed_name': 'Example Feed',
    'source_asn': 64496,
    'source_ip': '192.0.2.1',
    }
ES_SEARCH = {
    "query": {
        "constant_score": {
            "filter": {
                "term": {
                    "source.asn": 64496
                }
            }
        }
    }
}
ES_SEARCH_REPLACEMENT_CHARS = {
    "query": {
        "constant_score": {
            "filter": {
                "term": {
                    "source_asn": 64496
                }
            }
        }
    }
}

SAMPLE_TEMPLATE = {
    "mappings": {
        "events": {
            "properties": {
                "time.observation": {
                    "type": "date"
                },
                "time.source": {
                    "type": "date"
                },
                "classification.type": {
                    "type": "keyword"
                },
                "source.asn": {
                    "type": "integer"
                },
                "feed.name": {
                    "type": "text"
                },
                "source.ip": {
                    "type": "ip"
                }
            }
        }
    },
    "template": "intelmq-*",
    "index_patterns": [
        "intelmq-*"
    ]
}

TIMESTAMP_1 = "1869-12-02T00:00:00+00:00"
TIMESTAMP_2 = "2020-02-02T01:23:45+00:00"
TIMESTAMP_3 = "2018-09-09T01:23:45+00:00"
INPUT_TIME_SOURCE = {
    "__type": "Event",
    "classification.type": "infected-system",
    "source.asn": 64496,
    "source.ip": "192.0.2.1",
    "feed.name": "Example Feed",
    "time.source": TIMESTAMP_1,
    "extra": '{"foo.bar": "test"}'
}
INPUT_TIME_OBSERVATION = {
    "__type": "Event",
    "classification.type": "infected-system",
    "source.asn": 64496,
    "source.ip": "192.0.2.1",
    "feed.name": "Example Feed",
    "time.observation": TIMESTAMP_2,
    "extra": '{"foo.bar": "test"}'
}

ROTATE_OPTIONS = {
    'never': None,
    'daily': '%Y-%m-%d',
    'weekly': '%Y-%W',
    'monthly': '%Y-%m',
    'yearly': '%Y'
}

@test.skip_database()
class TestElasticsearchOutputBot(test.BotTestCase, unittest.TestCase):

    @classmethod
    def set_bot(cls):
        cls.bot_reference = ElasticsearchOutputBot
        cls.default_input_message = INPUT1
        cls.sysconfig = {"flatten_fields": "extra",
                         "rotate_index": "never"}
        if os.environ.get('INTELMQ_TEST_DATABASES'):
            cls.con = elasticsearch.Elasticsearch()

    def test_event(self):
        self.run_bot()
        time.sleep(1)  # ES needs some time between inserting and searching
        result = self.con.search(index='intelmq', body=ES_SEARCH)['hits']['hits'][0]
        self.con.delete(index='intelmq', doc_type='events', id=result['_id'])
        self.assertDictEqual(OUTPUT1, result['_source'])

    def test_raise_when_no_template(self):
        """
        Test that a bot raises a RuntimeError if 'rotate_index' is set, but a matching template doesn't exist in ES.
        """
        self.sysconfig = {"flatten_fields": "extra",
                          "elastic_index": "intelmq",
                          "elastic_doctype": "events",
                          "rotate_index": "daily"}
        self.assertRaises(RuntimeError, self.run_bot)

    def test_get_event_date(self):
        """
        Test whether get_event_date detects the time.source and time.observation fields in an event.
        """
        self.assertEqual(get_event_date(INPUT_TIME_SOURCE),
                         datetime.strptime(TIMESTAMP_1, '%Y-%m-%dT%H:%M:%S+00:00').date())
        self.assertEqual(get_event_date(INPUT_TIME_OBSERVATION),
                         datetime.strptime(TIMESTAMP_2, '%Y-%m-%dT%H:%M:%S+00:00').date())

    def test_replacement_characters(self):
        """
        Checks that dots in field names are replaced with the replacement character if one is set.
        """
        self.sysconfig = {"flatten_fields": "extra",
                          "elastic_index": "intelmq",
                          "elastic_doctype": "events",
                          "replacement_char": "_",
                          "rotate_index": "never"}
        self.run_bot()
        time.sleep(1)  # ES needs some time between inserting and searching
        result = self.con.search(index=self.sysconfig.get('elastic_index'),
                                 body=ES_SEARCH_REPLACEMENT_CHARS)['hits']['hits'][0]

        self.con.delete(index=self.sysconfig.get('elastic_index'),
                        doc_type=self.sysconfig.get('elastic_doctype'),
                        id=result['_id'])

        self.assertDictEqual(OUTPUT1_REPLACEMENT_CHARS, result['_source'])

    def test_index_detected_from_time_source(self):
        """
        Tests whether an input event with a time.source field is indexed according to its time.source date.
        """
        self.sysconfig = {"flatten_fields": "extra",
                          "elastic_index": "intelmq",
                          "elastic_doctype": "events",
                          "rotate_index": "daily"}
        expected_index_name = "{}-1869-12-02".format(self.sysconfig.get('elastic_index'))
        self.base_check_expected_index_created(INPUT_TIME_SOURCE, expected_index_name)

    def test_index_detected_from_time_observation(self):
        """
        Tests whether an input event with a time.observation field (and no time.source field) is indexed according to
        its time.observation date.
        """
        self.sysconfig = {"flatten_fields": "extra",
                          "elastic_index": "intelmq",
                          "elastic_doctype": "events",
                          "rotate_index": "daily"}
        expected_index_name = "{}-2020-02-02".format(self.sysconfig.get('elastic_index'))
        self.base_check_expected_index_created(INPUT_TIME_OBSERVATION, expected_index_name)

    def test_index_falls_back_to_default_date(self):
        """
        Tests whether get_index returns an expected default date
         if no time.source or time.observation is present, and
         that the OutputBot will use the current date if none is present.
        """

        self.sysconfig = {"flatten_fields": "extra",
                          "elastic_index": "intelmq",
                          "elastic_doctype": "events",
                          "rotate_index": "daily"}

        class FakeDateTime(datetime):
            """
            Passed to bot to force expected datetime value for test.
            """
            @classmethod
            def today(cls):
                return datetime.strptime('2018-09-09T01:23:45+00:00', '%Y-%m-%dT%H:%M:%S+00:00')

        expected_index_name = "{}-{}".format(self.sysconfig.get('elastic_index'), "2018-09-09")

        # Patch datetime with FakeDateTime, run the bot, and check the created index.
        with mock.patch('intelmq.bots.outputs.elasticsearch.output.datetime', new=FakeDateTime):
            self.base_check_expected_index_created(INPUT1, expected_index_name)

    def test_index_falls_back_to_default_string(self):
        """
        Tests whether get_index returns an expected default string
         if no time.source or time.observation is present, and
         that the OutputBot will use the current date if none is present.
        """

        self.sysconfig = {"flatten_fields": "extra",
                          "elastic_index": "intelmq",
                          "elastic_doctype": "events",
                          "rotate_index": "daily"}

        self.prepare_bot()
        index = self.bot.get_index(INPUT1, default_string='test-default')
        self.assertEqual(index, 'intelmq-test-default')  # Check that get_index honors the supplied default string
        self.run_bot()  # Run to clear output queue. Will still index according to current date.

    def base_check_expected_index_created(self, input_event, expected_index_name):
        self.input_message = input_event

        self.assertFalse(self.con.indices.exists(expected_index_name))  # Index should not already exist

        # Create the template mapping in Elasticsearch
        self.con.indices.put_template(name=self.sysconfig.get('elastic_index'), body=SAMPLE_TEMPLATE)

        self.run_bot()
        time.sleep(1)  # Let ES store the event. Can also force this with ES API

        self.assertTrue(self.con.indices.exists(expected_index_name))

        result = self.con.search(index=expected_index_name, body=ES_SEARCH)['hits']['hits'][0]
        result_index_name = result["_index"]

        # Clean up test event and check that the index name was set correctly
        self.con.delete(index=result_index_name, doc_type=self.sysconfig.get('elastic_doctype'), id=result['_id'])
        self.assertEqual(result_index_name, expected_index_name)

if __name__ == '__main__':  # pragma: no cover
    unittest.main()
<commit_msg>Remove doc types from elasticsearch output test<commit_after># -*- coding: utf-8 -*-
import os
import time
import unittest
import unittest.mock as mock
from datetime import datetime

import intelmq.lib.test as test
from intelmq.bots.outputs.elasticsearch.output import ElasticsearchOutputBot, get_event_date

if os.environ.get('INTELMQ_TEST_DATABASES'):
    import elasticsearch

INPUT1 = {"__type": "Event",
          "classification.type": "infected-system",
          "source.asn": 64496,
          "source.ip": "192.0.2.1",
          "feed.name": "Example Feed",
          "extra": '{"foo.bar": "test"}'
          }
OUTPUT1 = {'classification.type': 'infected-system',
           'extra.foo.bar': 'test',
           'feed.name': 'Example Feed',
           'source.asn': 64496,
           'source.ip': '192.0.2.1',
           }
OUTPUT1_REPLACEMENT_CHARS = {
    'classification_type': 'infected-system',
    'extra_foo_bar': 'test',
    'feed_name': 'Example Feed',
    'source_asn': 64496,
    'source_ip': '192.0.2.1',
    }
ES_SEARCH = {
    "query": {
        "constant_score": {
            "filter": {
                "term": {
                    "source.asn": 64496
                }
            }
        }
    }
}
ES_SEARCH_REPLACEMENT_CHARS = {
    "query": {
        "constant_score": {
            "filter": {
                "term": {
                    "source_asn": 64496
                }
            }
        }
    }
}

SAMPLE_TEMPLATE = {
    "mappings": {
        "properties": {
            "time.observation": {
                "type": "date"
            },
            "time.source": {
                "type": "date"
            },
            "classification.type": {
                "type": "keyword"
            },
            "source.asn": {
                "type": "integer"
            },
            "feed.name": {
                "type": "text"
            },
            "source.ip": {
                "type": "ip"
            }
        }
    },
    "template": "intelmq-*",
    "index_patterns": [
        "intelmq-*"
    ]
}

TIMESTAMP_1 = "1869-12-02T00:00:00+00:00"
TIMESTAMP_2 = "2020-02-02T01:23:45+00:00"
TIMESTAMP_3 = "2018-09-09T01:23:45+00:00"
INPUT_TIME_SOURCE = {
    "classification.type": "infected-system",
    "source.asn": 64496,
    "source.ip": "192.0.2.1",
    "feed.name": "Example Feed",
    "time.source": TIMESTAMP_1,
    "extra": '{"foo.bar": "test"}'
}
INPUT_TIME_OBSERVATION = {
    "classification.type": "infected-system",
    "source.asn": 64496,
    "source.ip": "192.0.2.1",
    "feed.name": "Example Feed",
    "time.observation": TIMESTAMP_2,
    "extra": '{"foo.bar": "test"}'
}

ROTATE_OPTIONS = {
    'never': None,
    'daily': '%Y-%m-%d',
    'weekly': '%Y-%W',
    'monthly': '%Y-%m',
    'yearly': '%Y'
}

@test.skip_database()
class TestElasticsearchOutputBot(test.BotTestCase, unittest.TestCase):

    @classmethod
    def set_bot(cls):
        cls.bot_reference = ElasticsearchOutputBot
        cls.default_input_message = INPUT1
        cls.sysconfig = {"flatten_fields": "extra",
                         "rotate_index": "never"}
        if os.environ.get('INTELMQ_TEST_DATABASES'):
            cls.con = elasticsearch.Elasticsearch()

    def test_event(self):
        self.run_bot()
        time.sleep(1)  # ES needs some time between inserting and searching
        result = self.con.search(index='intelmq', body=ES_SEARCH)['hits']['hits'][0]
        self.con.delete(index='intelmq',
                        # doc_type='events',
                        id=result['_id'])
        self.assertDictEqual(OUTPUT1, result['_source'])

    def test_raise_when_no_template(self):
        """
        Test that a bot raises a RuntimeError if 'rotate_index' is set, but a matching template doesn't exist in ES.
        """
        self.sysconfig = {"flatten_fields": "extra",
                          "elastic_index": "intelmq",
                          "elastic_doctype": "events",
                          "rotate_index": "daily"}
        self.assertRaises(RuntimeError, self.run_bot)

    def test_get_event_date(self):
        """
        Test whether get_event_date detects the time.source and time.observation fields in an event.
        """
        self.assertEqual(get_event_date(INPUT_TIME_SOURCE),
                         datetime.strptime(TIMESTAMP_1, '%Y-%m-%dT%H:%M:%S+00:00').date())
        self.assertEqual(get_event_date(INPUT_TIME_OBSERVATION),
                         datetime.strptime(TIMESTAMP_2, '%Y-%m-%dT%H:%M:%S+00:00').date())

    def test_replacement_characters(self):
        """
        Checks that dots in field names are replaced with the replacement character if one is set.
        """
        self.sysconfig = {"flatten_fields": "extra",
                          "elastic_index": "intelmq",
                          "elastic_doctype": "events",
                          "replacement_char": "_",
                          "rotate_index": "never"}
        self.run_bot()
        time.sleep(1)  # ES needs some time between inserting and searching
        result = self.con.search(index=self.sysconfig.get('elastic_index'),
                                 body=ES_SEARCH_REPLACEMENT_CHARS)['hits']['hits'][0]

        self.con.delete(index=self.sysconfig.get('elastic_index'),
                        # doc_type=self.sysconfig.get('elastic_doctype'),
                        id=result['_id'])

        self.assertDictEqual(OUTPUT1_REPLACEMENT_CHARS, result['_source'])

    def test_index_detected_from_time_source(self):
        """
        Tests whether an input event with a time.source field is indexed according to its time.source date.
        """
        self.sysconfig = {"flatten_fields": "extra",
                          "elastic_index": "intelmq",
                          "elastic_doctype": "events",
                          "rotate_index": "daily"}
        expected_index_name = "{}-1869-12-02".format(self.sysconfig.get('elastic_index'))
        self.base_check_expected_index_created(INPUT_TIME_SOURCE, expected_index_name)

    def test_index_detected_from_time_observation(self):
        """
        Tests whether an input event with a time.observation field (and no time.source field) is indexed according to
        its time.observation date.
        """
        self.sysconfig = {"flatten_fields": "extra",
                          "elastic_index": "intelmq",
                          "elastic_doctype": "events",
                          "rotate_index": "daily"}
        expected_index_name = "{}-2020-02-02".format(self.sysconfig.get('elastic_index'))
        self.base_check_expected_index_created(INPUT_TIME_OBSERVATION, expected_index_name)

    def test_index_falls_back_to_default_date(self):
        """
        Tests whether get_index returns an expected default date
         if no time.source or time.observation is present, and
         that the OutputBot will use the current date if none is present.
        """

        self.sysconfig = {"flatten_fields": "extra",
                          "elastic_index": "intelmq",
                          "elastic_doctype": "events",
                          "rotate_index": "daily"}

        class FakeDateTime(datetime):
            """
            Passed to bot to force expected datetime value for test.
            """
            @classmethod
            def today(cls):
                return datetime.strptime('2018-09-09T01:23:45+00:00', '%Y-%m-%dT%H:%M:%S+00:00')

        expected_index_name = "{}-{}".format(self.sysconfig.get('elastic_index'), "2018-09-09")

        # Patch datetime with FakeDateTime, run the bot, and check the created index.
        with mock.patch('intelmq.bots.outputs.elasticsearch.output.datetime', new=FakeDateTime):
            self.base_check_expected_index_created(INPUT1, expected_index_name)

    def test_index_falls_back_to_default_string(self):
        """
        Tests whether get_index returns an expected default string
         if no time.source or time.observation is present, and
         that the OutputBot will use the current date if none is present.
        """

        self.sysconfig = {"flatten_fields": "extra",
                          "elastic_index": "intelmq",
                          "elastic_doctype": "events",
                          "rotate_index": "daily"}

        self.prepare_bot()
        index = self.bot.get_index(INPUT1, default_string='test-default')
        self.assertEqual(index, 'intelmq-test-default')  # Check that get_index honors the supplied default string
        self.run_bot()  # Run to clear output queue. Will still index according to current date.

    def base_check_expected_index_created(self, input_event, expected_index_name):
        self.input_message = input_event

        self.assertFalse(self.con.indices.exists(expected_index_name))  # Index should not already exist

        # Create the template mapping in Elasticsearch
        self.con.indices.put_template(name=self.sysconfig.get('elastic_index'), body=SAMPLE_TEMPLATE)

        self.run_bot()
        time.sleep(1)  # Let ES store the event. Can also force this with ES API

        self.assertTrue(self.con.indices.exists(expected_index_name))

        result = self.con.search(index=expected_index_name, body=ES_SEARCH)['hits']['hits'][0]
        result_index_name = result["_index"]

        # Clean up test event and check that the index name was set correctly
        self.con.delete(index=result_index_name,
                        # doc_type=self.sysconfig.get('elastic_doctype'),
                        id=result['_id'])
        self.assertEqual(result_index_name, expected_index_name)

if __name__ == '__main__':  # pragma: no cover
    unittest.main()
<|endoftext|>


<commit_before>""" This module implements the idea of finding out emotions similarities
by using the experiments similar to what Hinton describes in his NRelu paper."""

import restrictedBoltzmannMachine as rbm

import numpy as np
import theano
from theano import tensor as T

from common import *
from similarity_utils import *

theanoFloat  = theano.config.floatX

class Trainer(object):

  def __init__(self, input1, input2, net):

    self.w = theano.shared(value=np.float32(0))
    self.b = theano.shared(value=np.float32(0))
    self.net = net

    self.oldMeanSquarew = theano.shared(value=np.float32(0))
    self.oldMeanSquareb = theano.shared(value=np.float32(0))
    self.oldMeanSquareWeights = theano.shared(value=np.zeros(self.net.weights.shape , dtype=theanoFloat))
    self.oldMeanSquareBias = theano.shared(value=np.zeros(self.net.biases[1].shape , dtype=theanoFloat))

    self.oldDw = theano.shared(value=np.float32(0))
    self.oldDb = theano.shared(value=np.float32(0))
    self.oldDWeights = theano.shared(value=np.zeros(self.net.weights.shape , dtype=theanoFloat))
    self.oldDBias = theano.shared(value=np.zeros(self.net.biases[1].shape , dtype=theanoFloat))

    hiddenBias = net.sharedBiases[1]
    # Do I need to add all biases? Probably only the hidden ones
    self.params = [self.w, self.b, self.net.sharedWeights, hiddenBias]
    self.oldDParams = [self.oldDw, self.oldDb, self.oldDWeights, self.oldDBias]
    self.oldMeanSquares =  [self.oldMeanSquarew, self.oldMeanSquareb, self.oldMeanSquareWeights, self.oldMeanSquareBias]

    _, weightForHidden = rbm.testWeights(self.net.sharedWeights,
          visibleDropout=self.net.visibleDropout, hiddenDropout=self.net.hiddenDropout)

    hiddenActivations1 = net.hiddenActivationFunction.deterministic(T.dot(input1, weightForHidden) + hiddenBias)
    hiddenActivations2 = net.hiddenActivationFunction.deterministic(T.dot(input2, weightForHidden) + hiddenBias)

    # Use the cosine distance between the two vectors of of activations
    cos = cosineDistance(hiddenActivations1, hiddenActivations2)

    self.cos = cos
    prob = 1.0 /( 1.0 + T.exp(self.w * cos + self.b))

    self.output = prob

class SimilarityNet(object):

  def __init__(self, learningRate, maxMomentum, rbmNrVis, rbmNrHid, rbmLearningRate,
                visibleActivationFunction, hiddenActivationFunction,
                rbmDropoutVis, rbmDropoutHid, rmsprop,trainingEpochsRBM,
                nesterovRbm,
                momentumFactorForLearningRateRBM,
                sparsityConstraint, sparsityRegularization, sparsityTraget):

    self.learningRate = np.float32(learningRate)
    self.rmsprop = rmsprop
    self.rbmNrVis = rbmNrVis
    self.maxMomentum = np.float32(maxMomentum)
    self.rbmNrHid = rbmNrHid
    self.rbmLearningRate = np.float32(rbmLearningRate)
    self.rbmDropoutHid = rbmDropoutHid
    self.rbmDropoutVis = rbmDropoutVis
    self.trainingEpochsRBM = trainingEpochsRBM
    self.visibleActivationFunction = visibleActivationFunction
    self.hiddenActivationFunction = hiddenActivationFunction
    self.nesterovRbm = nesterovRbm
    self.momentumFactorForLearningRateRBM = momentumFactorForLearningRateRBM

    self.sparsityConstraint = sparsityConstraint
    self.sparsityRegularization = sparsityRegularization
    self.sparsityTraget = sparsityTraget

    if rmsprop:
      print "training similarity net with rmsprop"

  def _trainRBM(self, data1, data2):
    data = np.vstack([data1, data2])

    net = rbm.RBM(self.rbmNrVis, self.rbmNrHid,
                    self.rbmLearningRate,
                    hiddenDropout=self.rbmDropoutHid,
                    visibleDropout=self.rbmDropoutVis,
                    visibleActivationFunction=self.visibleActivationFunction,
                    hiddenActivationFunction=self.hiddenActivationFunction,
                    rmsprop=True,
                    momentumFactorForLearningRate=self.momentumFactorForLearningRateRBM,
                    nesterov=self.nesterovRbm,
                    trainingEpochs=self.trainingEpochsRBM,
                    sparsityConstraint=self.sparsityConstraint,
                    sparsityCostFunction=squaredDiff,
                    sparsityRegularization=self.sparsityRegularization,
                    sparsityTraget=self.sparsityTraget)
    net.train(data)

    return net

  def train(self, data1, data2, similarities, miniBatchSize=20, epochs=200):
    nrMiniBatches = len(data1) / miniBatchSize
    miniBatchIndex = T.lscalar()
    momentum = T.fscalar()
    learningRate = T.fscalar()

    learningRateMiniBatch = np.float32(self.learningRate / miniBatchSize)
    print "learningRateMiniBatch in similarity net"
    print learningRateMiniBatch

    net = self._trainRBM(data1, data2)

    data1  = theano.shared(np.asarray(data1,dtype=theanoFloat))
    data2  = theano.shared(np.asarray(data2,dtype=theanoFloat))
    similarities = theano.shared(np.asarray(similarities,dtype=theanoFloat))

    # The mini-batch data is a matrix
    x = T.matrix('x', dtype=theanoFloat)
    y = T.matrix('y', dtype=theanoFloat)
    self.x = x
    self.y = y

    z = T.vector('z', dtype=theanoFloat)

    trainer = Trainer(x, y, net)
    self.trainer = trainer

    # error = T.sum(T.sqr(trainer.output-z))
    error = T.sum(T.nnet.binary_crossentropy(trainer.output, z))

    updates = self.buildUpdates(trainer, error, learningRate, momentum)

    # Now you have to define the theano function
    discriminativeTraining = theano.function(
      inputs=[miniBatchIndex, learningRate, momentum],
      outputs=[trainer.output, trainer.cos, error],
      updates=updates,
      givens={
            x: data1[miniBatchIndex * miniBatchSize:(miniBatchIndex + 1) * miniBatchSize],
            y: data2[miniBatchIndex * miniBatchSize:(miniBatchIndex + 1) * miniBatchSize],
            z: similarities[miniBatchIndex * miniBatchSize:(miniBatchIndex + 1) * miniBatchSize],
            })

    try:
      for epoch in xrange(epochs):
        print "epoch", epoch
        momentum = np.float32(min(np.float32(0.5) + epoch * np.float32(0.05),
                         np.float32(self.maxMomentum)))

        for miniBatch in xrange(nrMiniBatches):
          output, cos, error = discriminativeTraining(miniBatch, learningRateMiniBatch, momentum)

        print error
    except KeyboardInterrupt:
      print "you have decided to interrupt training"
      print "we continue testing"

    print trainer.w.get_value()
    print trainer.b.get_value()

  def test(self, testData1, testData2):
    # If it is too slow try adding mini batches
    testData1 = np.array(testData1, dtype=theanoFloat)
    testData2 = np.array(testData2, dtype=theanoFloat)

    # TODO : think of making data1 and data2 shared
    testFunction = theano.function(
      inputs=[],
      outputs=self.trainer.output,
      givens={self.x: testData1,
            self.y: testData2
            })

    return testFunction()

  def buildUpdates(self, trainer, error, learningRate, momentum):
    if self.rmsprop:
      return self.buildUpdatesRmsprop(trainer, error, learningRate, momentum)
    else:
      return self.buildUpdatesNoRmsprop(trainer, error, learningRate, momentum)

  def buildUpdatesNoRmsprop(self, trainer, error, learningRate, momentum):
    updates = []
    gradients = T.grad(error, trainer.params)

    for param, oldParamUpdate, gradient in zip(trainer.params, trainer.oldDParams, gradients):
      paramUpdate = momentum * oldParamUpdate - learningRate * gradient
      updates.append((param, param + paramUpdate))
      updates.append((oldParamUpdate, paramUpdate))

    return updates

  def buildUpdatesRmsprop(self, trainer, error, learningRate, momentum):
    updates = []
    gradients = T.grad(error, trainer.params)

    for param, oldParamUpdate, oldMeanSquare, gradient in zip(trainer.params, trainer.oldDParams, trainer.oldMeanSquares, gradients):
      meanSquare = 0.9 * oldMeanSquare + 0.1 * gradient ** 2
      paramUpdate = momentum * oldParamUpdate - learningRate * gradient / T.sqrt(meanSquare + 1e-08)
      updates.append((param, param + paramUpdate))
      updates.append((oldParamUpdate, paramUpdate))
      updates.append((oldMeanSquare, meanSquare))

    return updates

def cosineDistance(first, second):
  normFirst = T.sqrt(T.sum(T.sqr(first), axis=1))
  normSecond = T.sqrt(T.sum(T.sqr(second), axis=1))
  return 1.0 - T.sum(first * second, axis=1) / (normFirst * normSecond)

<commit_msg>print lr div by mini batch size in similarity<commit_after>""" This module implements the idea of finding out emotions similarities
by using the experiments similar to what Hinton describes in his NRelu paper."""

import restrictedBoltzmannMachine as rbm

import numpy as np
import theano
from theano import tensor as T

from common import *
from similarity_utils import *

theanoFloat  = theano.config.floatX

class Trainer(object):

  def __init__(self, input1, input2, net):

    self.w = theano.shared(value=np.float32(0))
    self.b = theano.shared(value=np.float32(0))
    self.net = net

    self.oldMeanSquarew = theano.shared(value=np.float32(0))
    self.oldMeanSquareb = theano.shared(value=np.float32(0))
    self.oldMeanSquareWeights = theano.shared(value=np.zeros(self.net.weights.shape , dtype=theanoFloat))
    self.oldMeanSquareBias = theano.shared(value=np.zeros(self.net.biases[1].shape , dtype=theanoFloat))

    self.oldDw = theano.shared(value=np.float32(0))
    self.oldDb = theano.shared(value=np.float32(0))
    self.oldDWeights = theano.shared(value=np.zeros(self.net.weights.shape , dtype=theanoFloat))
    self.oldDBias = theano.shared(value=np.zeros(self.net.biases[1].shape , dtype=theanoFloat))

    hiddenBias = net.sharedBiases[1]
    # Do I need to add all biases? Probably only the hidden ones
    self.params = [self.w, self.b, self.net.sharedWeights, hiddenBias]
    self.oldDParams = [self.oldDw, self.oldDb, self.oldDWeights, self.oldDBias]
    self.oldMeanSquares =  [self.oldMeanSquarew, self.oldMeanSquareb, self.oldMeanSquareWeights, self.oldMeanSquareBias]

    _, weightForHidden = rbm.testWeights(self.net.sharedWeights,
          visibleDropout=self.net.visibleDropout, hiddenDropout=self.net.hiddenDropout)

    hiddenActivations1 = net.hiddenActivationFunction.deterministic(T.dot(input1, weightForHidden) + hiddenBias)
    hiddenActivations2 = net.hiddenActivationFunction.deterministic(T.dot(input2, weightForHidden) + hiddenBias)

    # Use the cosine distance between the two vectors of of activations
    cos = cosineDistance(hiddenActivations1, hiddenActivations2)

    self.cos = cos
    prob = 1.0 /( 1.0 + T.exp(self.w * cos + self.b))

    self.output = prob

class SimilarityNet(object):

  def __init__(self, learningRate, maxMomentum, rbmNrVis, rbmNrHid, rbmLearningRate,
                visibleActivationFunction, hiddenActivationFunction,
                rbmDropoutVis, rbmDropoutHid, rmsprop,trainingEpochsRBM,
                nesterovRbm,
                momentumFactorForLearningRateRBM,
                sparsityConstraint, sparsityRegularization, sparsityTraget):

    self.learningRate = np.float32(learningRate)
    self.rmsprop = rmsprop
    self.rbmNrVis = rbmNrVis
    self.maxMomentum = np.float32(maxMomentum)
    self.rbmNrHid = rbmNrHid
    self.rbmLearningRate = np.float32(rbmLearningRate)
    self.rbmDropoutHid = rbmDropoutHid
    self.rbmDropoutVis = rbmDropoutVis
    self.trainingEpochsRBM = trainingEpochsRBM
    self.visibleActivationFunction = visibleActivationFunction
    self.hiddenActivationFunction = hiddenActivationFunction
    self.nesterovRbm = nesterovRbm
    self.momentumFactorForLearningRateRBM = momentumFactorForLearningRateRBM

    self.sparsityConstraint = sparsityConstraint
    self.sparsityRegularization = sparsityRegularization
    self.sparsityTraget = sparsityTraget

    if rmsprop:
      print "training similarity net with rmsprop"

  def _trainRBM(self, data1, data2):
    data = np.vstack([data1, data2])

    net = rbm.RBM(self.rbmNrVis, self.rbmNrHid,
                    self.rbmLearningRate,
                    hiddenDropout=self.rbmDropoutHid,
                    visibleDropout=self.rbmDropoutVis,
                    visibleActivationFunction=self.visibleActivationFunction,
                    hiddenActivationFunction=self.hiddenActivationFunction,
                    rmsprop=True,
                    momentumFactorForLearningRate=self.momentumFactorForLearningRateRBM,
                    nesterov=self.nesterovRbm,
                    trainingEpochs=self.trainingEpochsRBM,
                    sparsityConstraint=self.sparsityConstraint,
                    sparsityCostFunction=squaredDiff,
                    sparsityRegularization=self.sparsityRegularization,
                    sparsityTraget=self.sparsityTraget)
    net.train(data)

    return net

  def train(self, data1, data2, similarities, miniBatchSize=20, epochs=200):
    self.miniBatchSize = miniBatchSize
    nrMiniBatches = len(data1) / miniBatchSize
    miniBatchIndex = T.lscalar()
    momentum = T.fscalar()
    learningRate = T.fscalar()

    learningRateMiniBatch = np.float32(self.learningRate / miniBatchSize)
    print "learningRateMiniBatch in similarity net"
    print learningRateMiniBatch

    net = self._trainRBM(data1, data2)

    data1  = theano.shared(np.asarray(data1,dtype=theanoFloat))
    data2  = theano.shared(np.asarray(data2,dtype=theanoFloat))
    similarities = theano.shared(np.asarray(similarities,dtype=theanoFloat))

    # The mini-batch data is a matrix
    x = T.matrix('x', dtype=theanoFloat)
    y = T.matrix('y', dtype=theanoFloat)
    self.x = x
    self.y = y

    z = T.vector('z', dtype=theanoFloat)

    trainer = Trainer(x, y, net)
    self.trainer = trainer

    # error = T.sum(T.sqr(trainer.output-z))
    error = T.sum(T.nnet.binary_crossentropy(trainer.output, z))

    updates = self.buildUpdates(trainer, error, learningRate, momentum)

    # Now you have to define the theano function
    discriminativeTraining = theano.function(
      inputs=[miniBatchIndex, learningRate, momentum],
      outputs=[trainer.output, trainer.cos, error],
      updates=updates,
      givens={
            x: data1[miniBatchIndex * miniBatchSize:(miniBatchIndex + 1) * miniBatchSize],
            y: data2[miniBatchIndex * miniBatchSize:(miniBatchIndex + 1) * miniBatchSize],
            z: similarities[miniBatchIndex * miniBatchSize:(miniBatchIndex + 1) * miniBatchSize],
            })

    try:
      for epoch in xrange(epochs):
        print "epoch", epoch
        momentum = np.float32(min(np.float32(0.5) + epoch * np.float32(0.05),
                         np.float32(self.maxMomentum)))

        for miniBatch in xrange(nrMiniBatches):
          output, cos, error = discriminativeTraining(miniBatch, learningRateMiniBatch, momentum)

        print error /self.miniBatchSize

    except KeyboardInterrupt:
      print "you have decided to interrupt training"
      print "we continue testing"

    print trainer.w.get_value()
    print trainer.b.get_value()

  def test(self, testData1, testData2):
    # If it is too slow try adding mini batches
    testData1 = np.array(testData1, dtype=theanoFloat)
    testData2 = np.array(testData2, dtype=theanoFloat)

    # TODO : think of making data1 and data2 shared
    testFunction = theano.function(
      inputs=[],
      outputs=self.trainer.output,
      givens={self.x: testData1,
            self.y: testData2
            })

    return testFunction()

  def buildUpdates(self, trainer, error, learningRate, momentum):
    if self.rmsprop:
      return self.buildUpdatesRmsprop(trainer, error, learningRate, momentum)
    else:
      return self.buildUpdatesNoRmsprop(trainer, error, learningRate, momentum)

  def buildUpdatesNoRmsprop(self, trainer, error, learningRate, momentum):
    updates = []
    gradients = T.grad(error, trainer.params)

    for param, oldParamUpdate, gradient in zip(trainer.params, trainer.oldDParams, gradients):
      paramUpdate = momentum * oldParamUpdate - learningRate * gradient
      updates.append((param, param + paramUpdate))
      updates.append((oldParamUpdate, paramUpdate))

    return updates

  def buildUpdatesRmsprop(self, trainer, error, learningRate, momentum):
    updates = []
    gradients = T.grad(error, trainer.params)

    for param, oldParamUpdate, oldMeanSquare, gradient in zip(trainer.params, trainer.oldDParams, trainer.oldMeanSquares, gradients):
      meanSquare = 0.9 * oldMeanSquare + 0.1 * gradient ** 2
      paramUpdate = momentum * oldParamUpdate - learningRate * gradient / T.sqrt(meanSquare + 1e-08)
      updates.append((param, param + paramUpdate))
      updates.append((oldParamUpdate, paramUpdate))
      updates.append((oldMeanSquare, meanSquare))

    return updates

def cosineDistance(first, second):
  normFirst = T.sqrt(T.sum(T.sqr(first), axis=1))
  normSecond = T.sqrt(T.sum(T.sqr(second), axis=1))
  return 1.0 - T.sum(first * second, axis=1) / (normFirst * normSecond)

<|endoftext|>


<commit_before>#
# Copyright (c) 2008 rPath, Inc.
#

import os
import stat
import zlib
import base64
import shutil
import urllib2
import urlparse
import tempfile
import subprocess

def encode(s):
    return base64.encodestring(s).replace('\n', '')

def decode(s):
    return base64.decodestring(s)

def compress(s, level=9):
    return zlib.compress(s, level)

def decompress(s):
    return zlib.decompress(s)

def call(cmd, stdout=None, stderr=None, stdin=None):
    null = open(os.devnull, 'w')
    if not stdout:
        stdout = null
    if not stderr:
        stderr = null
    if not stdin:
        stdin = null
    return subprocess.call(cmd, stdout=stdout, stderr=stderr, stdin=stdin)

def _splitPath(path):
    dirs = []
    current = os.sep
    for level in path.split(os.sep):
        current = os.path.join(current, level)
        dirs.append(current)
    return dirs

def mkdirChain(path):
    for dir in _splitPath(path):
        if not os.path.exists(dir):
            os.mkdir(dir)
            continue
        if not os.path.isdir(dir):
            raise OSError, 'File exists'

def _copystat(src, dst):
    shutil.copystat(src, dst)
    s = os.stat(src)
    uid = s[stat.ST_UID]
    gid = s[stat.ST_GID]
    os.chown(dst, uid, gid)

def copytree(src, dst, symlink=False):
    mkdirChain(os.path.dirname(dst))

    if not os.path.isdir(src):
        shutil.copy2(src, dst)
        return

    shutil.copytree(src, dst, symlink)
    _copystat(src, dst)

    for srcdir, dirs, files in os.walk(src):
        dstdir = os.path.join(dst, srcdir[len(src):])
        for elements in dirs + files:
            srcE = os.path.join(srcdir, element)
            dstE = os.path.join(dstdir, element)
            _copystat(srcE, dstE)

def movetree(src, dst, symlink=False):
    copytree(src, dst, symlink=symlink)
    shutil.rmtree(src)

def createUnlinkedTmpFile(path=None):
    if path:
        fd, name = tempfile.mkstemp(dir=path)
    else:
        fd, name = tempfile.mkstemp()
    os.unlink(name)
    fh = os.fdopen(fd, 'w')
    return fh

def growFile(fh, size):
    """ Fill a file with zeros up to a specified size. """
    # Convert to mBytes -> kBytes
    size = size * 1024
    kByte = '\x00' * 1024
    for i in range(size):
        fh.write(kByte)
    fh.flush()

def urlgrab(url, filename=None):
    if not filename:
        file = os.path.basename(urlparse.urlsplit(url)[2])
    else:
        file = filename
    fh = open(file, 'w')
    fh.write(urllib2.urlopen(url).read())
    fh.close()
    return file
<commit_msg>fix typo<commit_after>#
# Copyright (c) 2008 rPath, Inc.
#

import os
import stat
import zlib
import base64
import shutil
import urllib2
import urlparse
import tempfile
import subprocess

def encode(s):
    return base64.encodestring(s).replace('\n', '')

def decode(s):
    return base64.decodestring(s)

def compress(s, level=9):
    return zlib.compress(s, level)

def decompress(s):
    return zlib.decompress(s)

def call(cmd, stdout=None, stderr=None, stdin=None):
    null = open(os.devnull, 'w')
    if not stdout:
        stdout = null
    if not stderr:
        stderr = null
    if not stdin:
        stdin = null
    return subprocess.call(cmd, stdout=stdout, stderr=stderr, stdin=stdin)

def _splitPath(path):
    dirs = []
    current = os.sep
    for level in path.split(os.sep):
        current = os.path.join(current, level)
        dirs.append(current)
    return dirs

def mkdirChain(path):
    for dir in _splitPath(path):
        if not os.path.exists(dir):
            os.mkdir(dir)
            continue
        if not os.path.isdir(dir):
            raise OSError, 'File exists'

def _copystat(src, dst):
    shutil.copystat(src, dst)
    s = os.stat(src)
    uid = s[stat.ST_UID]
    gid = s[stat.ST_GID]
    os.chown(dst, uid, gid)

def copytree(src, dst, symlink=False):
    mkdirChain(os.path.dirname(dst))

    if not os.path.isdir(src):
        shutil.copy2(src, dst)
        return

    shutil.copytree(src, dst, symlink)
    _copystat(src, dst)

    for srcdir, dirs, files in os.walk(src):
        dstdir = os.path.join(dst, srcdir[len(src):])
        for element in dirs + files:
            srcE = os.path.join(srcdir, element)
            dstE = os.path.join(dstdir, element)
            _copystat(srcE, dstE)

def movetree(src, dst, symlink=False):
    copytree(src, dst, symlink=symlink)
    shutil.rmtree(src)

def createUnlinkedTmpFile(path=None):
    if path:
        fd, name = tempfile.mkstemp(dir=path)
    else:
        fd, name = tempfile.mkstemp()
    os.unlink(name)
    fh = os.fdopen(fd, 'w')
    return fh

def growFile(fh, size):
    """ Fill a file with zeros up to a specified size. """
    # Convert to mBytes -> kBytes
    size = size * 1024
    kByte = '\x00' * 1024
    for i in range(size):
        fh.write(kByte)
    fh.flush()

def urlgrab(url, filename=None):
    if not filename:
        file = os.path.basename(urlparse.urlsplit(url)[2])
    else:
        file = filename
    fh = open(file, 'w')
    fh.write(urllib2.urlopen(url).read())
    fh.close()
    return file
<|endoftext|>


<commit_before># Copyright 2016 gRPC authors.
#
# 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 is for the gRPC build system. This isn't intended to be used outsite of
# the BUILD file for gRPC. It contains the mapping for the template system we
# use to generate other platform's build system files.
#
# Please consider that there should be a high bar for additions and changes to
# this file.
# Each rule listed must be re-written for Google's internal build system, and
# each change must be ported from one to the other.
#

load("//bazel:cc_grpc_library.bzl", "cc_grpc_library")
load("@upb//bazel:upb_proto_library.bzl", "upb_proto_library")
load("@build_bazel_rules_apple//apple:ios.bzl", "ios_unit_test")

# The set of pollers to test against if a test exercises polling
POLLERS = ["epollex", "epoll1", "poll"]

def if_not_windows(a):
    return select({
        "//:windows": [],
        "//:windows_msvc": [],
        "//conditions:default": a,
    })

def if_mac(a):
    return select({
        "//:mac_x86_64": a,
        "//conditions:default": [],
    })

def _get_external_deps(external_deps):
    ret = []
    for dep in external_deps:
        if dep == "address_sorting":
            ret += ["//third_party/address_sorting"]
        elif dep == "cares":
            ret += select({
                "//:grpc_no_ares": [],
                "//conditions:default": ["//external:cares"],
            })
        elif dep == "cronet_c_for_grpc":
            ret += ["//third_party/objective_c/Cronet:cronet_c_for_grpc"]
        elif dep.startswith("absl/"):
            ret += ["@com_google_absl//" + dep]
        else:
            ret += ["//external:" + dep]
    return ret

def grpc_cc_library(
        name,
        srcs = [],
        public_hdrs = [],
        hdrs = [],
        external_deps = [],
        deps = [],
        standalone = False,
        language = "C++",
        testonly = False,
        visibility = None,
        alwayslink = 0,
        data = [],
        use_cfstream = False,
        tags = []):
    copts = []
    if use_cfstream:
        copts = if_mac(["-DGRPC_CFSTREAM"])
    if language.upper() == "C":
        copts = copts + if_not_windows(["-std=c99"])
    linkopts = if_not_windows(["-pthread"])
    if use_cfstream:
        linkopts = linkopts + if_mac(["-framework CoreFoundation"])

    native.cc_library(
        name = name,
        srcs = srcs,
        defines = select({
                      "//:grpc_no_ares": ["GRPC_ARES=0"],
                      "//conditions:default": [],
                  }) +
                  select({
                      "//:remote_execution": ["GRPC_PORT_ISOLATED_RUNTIME=1"],
                      "//conditions:default": [],
                  }) +
                  select({
                      "//:grpc_allow_exceptions": ["GRPC_ALLOW_EXCEPTIONS=1"],
                      "//:grpc_disallow_exceptions": ["GRPC_ALLOW_EXCEPTIONS=0"],
                      "//conditions:default": [],
                  }),
        hdrs = hdrs + public_hdrs,
        deps = deps + _get_external_deps(external_deps),
        copts = copts,
        visibility = visibility,
        testonly = testonly,
        linkopts = linkopts,
        includes = [
            "include",
            "src/core/ext/upb-generated",  # Once upb code-gen issue is resolved, remove this.
        ],
        alwayslink = alwayslink,
        data = data,
        tags = tags,
    )

def grpc_proto_plugin(name, srcs = [], deps = []):
    native.cc_binary(
        name = name,
        srcs = srcs,
        deps = deps,
    )

def grpc_proto_library(
        name,
        srcs = [],
        deps = [],
        well_known_protos = False,
        has_services = True,
        use_external = False,
        generate_mocks = False):
    cc_grpc_library(
        name = name,
        srcs = srcs,
        deps = deps,
        well_known_protos = well_known_protos,
        proto_only = not has_services,
        use_external = use_external,
        generate_mocks = generate_mocks,
    )

def ios_cc_test(
        name,
        tags = [],
        **kwargs):
    ios_test_adapter = "//third_party/objective_c/google_toolbox_for_mac:GTM_GoogleTestRunner_GTM_USING_XCTEST"

    test_lib_ios = name + "_test_lib_ios"
    ios_tags = tags + ["manual", "ios_cc_test"]
    if not any([t for t in tags if t.startswith("no_test_ios")]):
        native.objc_library(
            name = test_lib_ios,
            srcs = kwargs.get("srcs"),
            deps = kwargs.get("deps"),
            copts = kwargs.get("copts"),
            tags = ios_tags,
            alwayslink = 1,
            testonly = 1,
        )
        ios_test_deps = [ios_test_adapter, ":" + test_lib_ios]
        ios_unit_test(
            name = name + "_on_ios",
            size = kwargs.get("size"),
            tags = ios_tags,
            minimum_os_version = "9.0",
            deps = ios_test_deps,
        )

def grpc_cc_test(name, srcs = [], deps = [], external_deps = [], args = [], data = [], uses_polling = True, language = "C++", size = "medium", timeout = None, tags = [], exec_compatible_with = [], exec_properties = {}, shard_count = None):
    copts = if_mac(["-DGRPC_CFSTREAM"])
    if language.upper() == "C":
        copts = copts + if_not_windows(["-std=c99"])

    # NOTE: these attributes won't be used for the poller-specific versions of a test
    # automatically, you need to set them explicitly (if applicable)
    args = {
        "srcs": srcs,
        "args": args,
        "data": data,
        "deps": deps + _get_external_deps(external_deps),
        "copts": copts,
        "linkopts": if_not_windows(["-pthread"]),
        "size": size,
        "timeout": timeout,
        "exec_compatible_with": exec_compatible_with,
        "exec_properties": exec_properties,
        "shard_count": shard_count,
    }
    if uses_polling:
        # the vanilla version of the test should run on platforms that only
        # support a single poller
        native.cc_test(
            name = name,
            testonly = True,
            tags = (tags + [
                "no_linux",  # linux supports multiple pollers
            ]),
            **args
        )

        # on linux we run the same test multiple times, once for each poller
        for poller in POLLERS:
            native.sh_test(
                name = name + "@poller=" + poller,
                data = [name] + data,
                srcs = [
                    "//test/core/util:run_with_poller_sh",
                ],
                size = size,
                timeout = timeout,
                args = [
                    poller,
                    "$(location %s)" % name,
                ] + args["args"],
                tags = (tags + ["no_windows", "no_mac"]),
                exec_compatible_with = exec_compatible_with,
                exec_properties = exec_properties,
                shard_count = shard_count,
            )
    else:
        # the test behavior doesn't depend on polling, just generate the test
        native.cc_test(name = name, tags = tags + ["no_uses_polling"], **args)
    ios_cc_test(
        name = name,
        tags = tags,
        **args
    )

def grpc_cc_binary(name, srcs = [], deps = [], external_deps = [], args = [], data = [], language = "C++", testonly = False, linkshared = False, linkopts = [], tags = []):
    copts = []
    if language.upper() == "C":
        copts = ["-std=c99"]
    native.cc_binary(
        name = name,
        srcs = srcs,
        args = args,
        data = data,
        testonly = testonly,
        linkshared = linkshared,
        deps = deps + _get_external_deps(external_deps),
        copts = copts,
        linkopts = if_not_windows(["-pthread"]) + linkopts,
        tags = tags,
    )

def grpc_generate_one_off_targets():
    # In open-source, grpc_objc* libraries depend directly on //:grpc
    native.alias(
        name = "grpc_objc",
        actual = "//:grpc",
    )

def grpc_generate_objc_one_off_targets():
    pass

def grpc_sh_test(name, srcs, args = [], data = []):
    native.sh_test(
        name = name,
        srcs = srcs,
        args = args,
        data = data,
    )

def grpc_sh_binary(name, srcs, data = []):
    native.sh_binary(
        name = name,
        srcs = srcs,
        data = data,
    )

def grpc_py_binary(
        name,
        srcs,
        data = [],
        deps = [],
        external_deps = [],
        testonly = False,
        python_version = "PY2",
        **kwargs):
    native.py_binary(
        name = name,
        srcs = srcs,
        testonly = testonly,
        data = data,
        deps = deps + _get_external_deps(external_deps),
        python_version = python_version,
        **kwargs
    )

def grpc_package(name, visibility = "private", features = []):
    if visibility == "tests":
        visibility = ["//test:__subpackages__"]
    elif visibility == "public":
        visibility = ["//visibility:public"]
    elif visibility == "private":
        visibility = []
    else:
        fail("Unknown visibility " + visibility)

    if len(visibility) != 0:
        native.package(
            default_visibility = visibility,
            features = features,
        )

def grpc_objc_library(
        name,
        srcs = [],
        hdrs = [],
        textual_hdrs = [],
        data = [],
        deps = [],
        defines = [],
        includes = [],
        visibility = ["//visibility:public"]):
    """The grpc version of objc_library, only used for the Objective-C library compilation

    Args:
        name: name of target
        hdrs: public headers
        srcs: all source files (.m)
        textual_hdrs: private headers
        data: any other bundle resources
        defines: preprocessors
        includes: added to search path, always [the path to objc directory]
        deps: dependencies
        visibility: visibility, default to public
    """

    native.objc_library(
        name = name,
        hdrs = hdrs,
        srcs = srcs,
        textual_hdrs = textual_hdrs,
        data = data,
        deps = deps,
        defines = defines,
        includes = includes,
        visibility = visibility,
    )

def grpc_upb_proto_library(name, deps):
    upb_proto_library(name = name, deps = deps)

def python_config_settings():
    native.config_setting(
        name = "python3",
        flag_values = {"@bazel_tools//tools/python:python_version": "PY3"},
    )
<commit_msg>support Flaky=True for grpc_cc_tests<commit_after># Copyright 2016 gRPC authors.
#
# 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 is for the gRPC build system. This isn't intended to be used outsite of
# the BUILD file for gRPC. It contains the mapping for the template system we
# use to generate other platform's build system files.
#
# Please consider that there should be a high bar for additions and changes to
# this file.
# Each rule listed must be re-written for Google's internal build system, and
# each change must be ported from one to the other.
#

load("//bazel:cc_grpc_library.bzl", "cc_grpc_library")
load("@upb//bazel:upb_proto_library.bzl", "upb_proto_library")
load("@build_bazel_rules_apple//apple:ios.bzl", "ios_unit_test")

# The set of pollers to test against if a test exercises polling
POLLERS = ["epollex", "epoll1", "poll"]

def if_not_windows(a):
    return select({
        "//:windows": [],
        "//:windows_msvc": [],
        "//conditions:default": a,
    })

def if_mac(a):
    return select({
        "//:mac_x86_64": a,
        "//conditions:default": [],
    })

def _get_external_deps(external_deps):
    ret = []
    for dep in external_deps:
        if dep == "address_sorting":
            ret += ["//third_party/address_sorting"]
        elif dep == "cares":
            ret += select({
                "//:grpc_no_ares": [],
                "//conditions:default": ["//external:cares"],
            })
        elif dep == "cronet_c_for_grpc":
            ret += ["//third_party/objective_c/Cronet:cronet_c_for_grpc"]
        elif dep.startswith("absl/"):
            ret += ["@com_google_absl//" + dep]
        else:
            ret += ["//external:" + dep]
    return ret

def grpc_cc_library(
        name,
        srcs = [],
        public_hdrs = [],
        hdrs = [],
        external_deps = [],
        deps = [],
        standalone = False,
        language = "C++",
        testonly = False,
        visibility = None,
        alwayslink = 0,
        data = [],
        use_cfstream = False,
        tags = []):
    copts = []
    if use_cfstream:
        copts = if_mac(["-DGRPC_CFSTREAM"])
    if language.upper() == "C":
        copts = copts + if_not_windows(["-std=c99"])
    linkopts = if_not_windows(["-pthread"])
    if use_cfstream:
        linkopts = linkopts + if_mac(["-framework CoreFoundation"])

    native.cc_library(
        name = name,
        srcs = srcs,
        defines = select({
                      "//:grpc_no_ares": ["GRPC_ARES=0"],
                      "//conditions:default": [],
                  }) +
                  select({
                      "//:remote_execution": ["GRPC_PORT_ISOLATED_RUNTIME=1"],
                      "//conditions:default": [],
                  }) +
                  select({
                      "//:grpc_allow_exceptions": ["GRPC_ALLOW_EXCEPTIONS=1"],
                      "//:grpc_disallow_exceptions": ["GRPC_ALLOW_EXCEPTIONS=0"],
                      "//conditions:default": [],
                  }),
        hdrs = hdrs + public_hdrs,
        deps = deps + _get_external_deps(external_deps),
        copts = copts,
        visibility = visibility,
        testonly = testonly,
        linkopts = linkopts,
        includes = [
            "include",
            "src/core/ext/upb-generated",  # Once upb code-gen issue is resolved, remove this.
        ],
        alwayslink = alwayslink,
        data = data,
        tags = tags,
    )

def grpc_proto_plugin(name, srcs = [], deps = []):
    native.cc_binary(
        name = name,
        srcs = srcs,
        deps = deps,
    )

def grpc_proto_library(
        name,
        srcs = [],
        deps = [],
        well_known_protos = False,
        has_services = True,
        use_external = False,
        generate_mocks = False):
    cc_grpc_library(
        name = name,
        srcs = srcs,
        deps = deps,
        well_known_protos = well_known_protos,
        proto_only = not has_services,
        use_external = use_external,
        generate_mocks = generate_mocks,
    )

def ios_cc_test(
        name,
        tags = [],
        **kwargs):
    ios_test_adapter = "//third_party/objective_c/google_toolbox_for_mac:GTM_GoogleTestRunner_GTM_USING_XCTEST"

    test_lib_ios = name + "_test_lib_ios"
    ios_tags = tags + ["manual", "ios_cc_test"]
    if not any([t for t in tags if t.startswith("no_test_ios")]):
        native.objc_library(
            name = test_lib_ios,
            srcs = kwargs.get("srcs"),
            deps = kwargs.get("deps"),
            copts = kwargs.get("copts"),
            tags = ios_tags,
            alwayslink = 1,
            testonly = 1,
        )
        ios_test_deps = [ios_test_adapter, ":" + test_lib_ios]
        ios_unit_test(
            name = name + "_on_ios",
            size = kwargs.get("size"),
            tags = ios_tags,
            minimum_os_version = "9.0",
            deps = ios_test_deps,
        )

def grpc_cc_test(name, srcs = [], deps = [], external_deps = [], args = [], data = [], uses_polling = True, language = "C++", size = "medium", timeout = None, tags = [], exec_compatible_with = [], exec_properties = {}, shard_count = None, flaky = None):
    copts = if_mac(["-DGRPC_CFSTREAM"])
    if language.upper() == "C":
        copts = copts + if_not_windows(["-std=c99"])

    # NOTE: these attributes won't be used for the poller-specific versions of a test
    # automatically, you need to set them explicitly (if applicable)
    args = {
        "srcs": srcs,
        "args": args,
        "data": data,
        "deps": deps + _get_external_deps(external_deps),
        "copts": copts,
        "linkopts": if_not_windows(["-pthread"]),
        "size": size,
        "timeout": timeout,
        "exec_compatible_with": exec_compatible_with,
        "exec_properties": exec_properties,
        "shard_count": shard_count,
        "flaky": flaky,
    }
    if uses_polling:
        # the vanilla version of the test should run on platforms that only
        # support a single poller
        native.cc_test(
            name = name,
            testonly = True,
            tags = (tags + [
                "no_linux",  # linux supports multiple pollers
            ]),
            **args
        )

        # on linux we run the same test multiple times, once for each poller
        for poller in POLLERS:
            native.sh_test(
                name = name + "@poller=" + poller,
                data = [name] + data,
                srcs = [
                    "//test/core/util:run_with_poller_sh",
                ],
                size = size,
                timeout = timeout,
                args = [
                    poller,
                    "$(location %s)" % name,
                ] + args["args"],
                tags = (tags + ["no_windows", "no_mac"]),
                exec_compatible_with = exec_compatible_with,
                exec_properties = exec_properties,
                shard_count = shard_count,
                flaky = flaky,
            )
    else:
        # the test behavior doesn't depend on polling, just generate the test
        native.cc_test(name = name, tags = tags + ["no_uses_polling"], **args)
    ios_cc_test(
        name = name,
        tags = tags,
        **args
    )

def grpc_cc_binary(name, srcs = [], deps = [], external_deps = [], args = [], data = [], language = "C++", testonly = False, linkshared = False, linkopts = [], tags = []):
    copts = []
    if language.upper() == "C":
        copts = ["-std=c99"]
    native.cc_binary(
        name = name,
        srcs = srcs,
        args = args,
        data = data,
        testonly = testonly,
        linkshared = linkshared,
        deps = deps + _get_external_deps(external_deps),
        copts = copts,
        linkopts = if_not_windows(["-pthread"]) + linkopts,
        tags = tags,
    )

def grpc_generate_one_off_targets():
    # In open-source, grpc_objc* libraries depend directly on //:grpc
    native.alias(
        name = "grpc_objc",
        actual = "//:grpc",
    )

def grpc_generate_objc_one_off_targets():
    pass

def grpc_sh_test(name, srcs, args = [], data = []):
    native.sh_test(
        name = name,
        srcs = srcs,
        args = args,
        data = data,
    )

def grpc_sh_binary(name, srcs, data = []):
    native.sh_binary(
        name = name,
        srcs = srcs,
        data = data,
    )

def grpc_py_binary(
        name,
        srcs,
        data = [],
        deps = [],
        external_deps = [],
        testonly = False,
        python_version = "PY2",
        **kwargs):
    native.py_binary(
        name = name,
        srcs = srcs,
        testonly = testonly,
        data = data,
        deps = deps + _get_external_deps(external_deps),
        python_version = python_version,
        **kwargs
    )

def grpc_package(name, visibility = "private", features = []):
    if visibility == "tests":
        visibility = ["//test:__subpackages__"]
    elif visibility == "public":
        visibility = ["//visibility:public"]
    elif visibility == "private":
        visibility = []
    else:
        fail("Unknown visibility " + visibility)

    if len(visibility) != 0:
        native.package(
            default_visibility = visibility,
            features = features,
        )

def grpc_objc_library(
        name,
        srcs = [],
        hdrs = [],
        textual_hdrs = [],
        data = [],
        deps = [],
        defines = [],
        includes = [],
        visibility = ["//visibility:public"]):
    """The grpc version of objc_library, only used for the Objective-C library compilation

    Args:
        name: name of target
        hdrs: public headers
        srcs: all source files (.m)
        textual_hdrs: private headers
        data: any other bundle resources
        defines: preprocessors
        includes: added to search path, always [the path to objc directory]
        deps: dependencies
        visibility: visibility, default to public
    """

    native.objc_library(
        name = name,
        hdrs = hdrs,
        srcs = srcs,
        textual_hdrs = textual_hdrs,
        data = data,
        deps = deps,
        defines = defines,
        includes = includes,
        visibility = visibility,
    )

def grpc_upb_proto_library(name, deps):
    upb_proto_library(name = name, deps = deps)

def python_config_settings():
    native.config_setting(
        name = "python3",
        flag_values = {"@bazel_tools//tools/python:python_version": "PY3"},
    )
<|endoftext|>