<commit_before>from __future__ import print_function

import functools
import numpy as np

# sklearn imports

have_isotonic = False
try:
    from sklearn.isotonic import IsotonicRegression
    have_isotonic = True
except ImportError:
    raise ValueError('unable to import isotonic regression from sklearn, SLOPE subgradient projection will not work')

# regreg imports

from regreg.atoms.slope import _basic_proximal_map
import regreg.api as rr

from ..constraints.affine import constraints

from .randomization import randomization
from .base import restricted_estimator
from .lasso import highdim
from .query import (query,
                    multiple_queries,
                    langevin_sampler,
                    affine_gaussian_sampler)

class slope(highdim):

    def __init__(self,
                 loglike,
                 slope_weights,
                 ridge_term,
                 randomizer_scale,
                 perturb=None):
        r"""
        Create a new post-selection object for the SLOPE problem

        Parameters
        ----------

        loglike : `regreg.smooth.glm.glm`
            A (negative) log-likelihood as implemented in `regreg`.

        slope_weights : np.ndarray
            SLOPE weights for L-1 penalty. If a float,
            it is broadcast to all features.

        ridge_term : float
            How big a ridge term to add?

        randomizer_scale : float
            Scale for IID components of randomization.

        perturb : np.ndarray
            Random perturbation subtracted as a linear
            term in the objective function.
        """

        self.loglike = loglike
        self.nfeature = p = self.loglike.shape[0]

        if np.asarray(slope_weights).shape == ():
            slope_weights = np.ones(loglike.shape) * slope_weights
        self.slope_weights = np.asarray(slope_weights)

        self.randomizer = randomization.isotropic_gaussian((p,), randomizer_scale)
        self.ridge_term = ridge_term
        self.penalty = rr.slope(slope_weights, lagrange=1.)
        self._initial_omega = perturb  # random perturbation

    def fit(self,
            solve_args={'tol': 1.e-12, 'min_its': 50},
            perturb=None):

        p = self.nfeature

        # take a new perturbation if supplied
        if perturb is not None:
            self._initial_omega = perturb
        if self._initial_omega is None:
            self._initial_omega = self.randomizer.sample()

        quad = rr.identity_quadratic(self.ridge_term, 0, -self._initial_omega, 0)
        problem = rr.simple_problem(self.loglike, self.penalty)
        self.initial_soln = problem.solve(quad, **solve_args)

        # now we have to work out SLOPE details, clusters, etc.

        active_signs = np.sign(self.initial_soln)
        active = self._active = active_signs != 0

        print("check active terms", active.sum())

        self._overall = overall = active> 0
        self._inactive = inactive = ~self._overall

        _active_signs = active_signs.copy()
        self.selection_variable = {'sign': _active_signs,
                                   'variables': self._overall}

        initial_subgrad = -(self.loglike.smooth_objective(self.initial_soln, 'grad') +
                            quad.objective(self.initial_soln, 'grad'))
        self.initial_subgrad = initial_subgrad

        indices = np.argsort(-np.fabs(self.initial_soln))
        sorted_soln = self.initial_soln[indices]
        initial_scalings = np.sort(np.unique(np.fabs(self.initial_soln[active])))[::-1]
        self.observed_opt_state = initial_scalings

        self._unpenalized = np.zeros(p, np.bool)

        _beta_unpenalized = restricted_estimator(self.loglike, self._overall, solve_args=solve_args)

        beta_bar = np.zeros(p)
        beta_bar[overall] = _beta_unpenalized
        self._beta_full = beta_bar

        self.num_opt_var = self.observed_opt_state.shape[0]

        X, y = self.loglike.data
        W = self._W = self.loglike.saturated_loss.hessian(X.dot(beta_bar))
        _hessian_active = np.dot(X.T, X[:, active] * W[:, None])
        _score_linear_term = -_hessian_active
        self.score_transform = (_score_linear_term, np.zeros(_score_linear_term.shape[0]))

        self.observed_score_state = _score_linear_term.dot(_beta_unpenalized)
        self.observed_score_state[inactive] += self.loglike.smooth_objective(beta_bar, 'grad')[inactive]

        cur_indx_array = []
        cur_indx_array.append(0)
        cur_indx = 0
        pointer = 0
        signs_cluster = []
        for j in range(p - 1):
            if np.abs(sorted_soln[j + 1]) != np.abs(sorted_soln[cur_indx]):
                cur_indx_array.append(j + 1)
                cur_indx = j + 1
                sign_vec = np.zeros(p)
                sign_vec[np.arange(j + 1 - cur_indx_array[pointer]) + cur_indx_array[pointer]] = \
                    np.sign(self.initial_soln[indices[np.arange(j + 1 - cur_indx_array[pointer]) + cur_indx_array[pointer]]])
                signs_cluster.append(sign_vec)
                pointer = pointer + 1
                if sorted_soln[j + 1] == 0:
                    break

        signs_cluster = np.asarray(signs_cluster).T
        X_clustered = X[:, indices].dot(signs_cluster)
        _opt_linear_term = X.T.dot(X_clustered)
        self.opt_transform = (_opt_linear_term, self.initial_subgrad)

        cov, prec = self.randomizer.cov_prec
        opt_linear, opt_offset = self.opt_transform

        print("check if correct", np.allclose(self.observed_score_state + opt_offset + opt_linear.dot(initial_scalings),
                                              self._initial_omega, rtol=1e-05, atol=1e-08))

        cond_precision = opt_linear.T.dot(opt_linear) * prec
        cond_cov = np.linalg.inv(cond_precision)
        logdens_linear = cond_cov.dot(opt_linear.T) * prec
        cond_mean = -logdens_linear.dot(self.observed_score_state + opt_offset)

        logdens_transform = (logdens_linear, opt_offset)

        def log_density(logdens_linear, offset, cond_prec, score, opt):
            if score.ndim == 1:
                mean_term = logdens_linear.dot(score.T + offset).T
            else:
                mean_term = logdens_linear.dot(score.T + offset[:, None]).T
            arg = opt + mean_term
            return - 0.5 * np.sum(arg * cond_prec.dot(arg.T).T, 1)

        log_density = functools.partial(log_density, logdens_linear, opt_offset, cond_precision)

        # now make the constraints

        A_scaling_0 = -np.identity(self.num_opt_var)
        A_scaling_1 = -np.identity(self.num_opt_var)[:(self.num_opt_var-1), :]
        for k in range(A_scaling_1.shape[0]):
           A_scaling_1[k,k+1]= 1
        A_scaling = np.vstack([A_scaling_0, A_scaling_1])
        b_scaling = np.zeros(2*self.num_opt_var-1)

        # A_scaling = -np.identity(self.num_opt_var)
        # b_scaling = np.zeros(self.num_opt_var)

        affine_con = constraints(A_scaling,
                                 b_scaling,
                                 mean=cond_mean,
                                 covariance=cond_cov)

        self.sampler = affine_gaussian_sampler(affine_con,
                                               self.observed_opt_state,
                                               self.observed_score_state,
                                               log_density,
                                               logdens_transform,
                                               selection_info=self.selection_variable)
        return active_signs

    def selective_MLE(self,
                      target="selected",
                      features=None,
                      parameter=None,
                      level=0.9,
                      compute_intervals=False,
                      dispersion=None,
                      solve_args={'tol': 1.e-12}):
        """
        Parameters
        ----------
        target : one of ['selected', 'full']
        features : np.bool
            Binary encoding of which features to use in final
            model and targets.
        parameter : np.array
            Hypothesized value for parameter -- defaults to 0.
        level : float
            Confidence level.
        ndraw : int (optional)
            Defaults to 1000.
        burnin : int (optional)
            Defaults to 1000.
        compute_intervals : bool
            Compute confidence intervals?
        dispersion : float (optional)
            Use a known value for dispersion, or Pearson's X^2?
        """

        if parameter is None:
            parameter = np.zeros(self.loglike.shape[0])

        if target == 'selected':
            observed_target, cov_target, cov_target_score, alternatives = self.selected_targets(features=features,
                                                                                dispersion=dispersion)
        elif target == 'full':
            X, y = self.loglike.data
            n, p = X.shape
            if n > p:
                observed_target, cov_target, cov_target_score, alternatives = self.full_targets(features=features,
                                                                                dispersion=dispersion)
            else:
                observed_target, cov_target, cov_target_score, alternatives = self.debiased_targets(features=features,
                                                                                dispersion=dispersion)

        # working out conditional law of opt variables given
        # target after decomposing score wrt target

        return self.sampler.selective_MLE(observed_target,
                                          cov_target,
                                          cov_target_score,
                                          self.observed_opt_state,
                                          solve_args=solve_args)

    # Targets of inference
    # and covariance with score representation
    # are same as highdim LASSO

    @staticmethod
    def gaussian(X,
                 Y,
                 slope_weights,
                 sigma=1.,
                 quadratic=None,
                 ridge_term=0.,
                 randomizer_scale=None):

        loglike = rr.glm.gaussian(X, Y, coef=1. / sigma ** 2, quadratic=quadratic)
        n, p = X.shape

        mean_diag = np.mean((X ** 2).sum(0))
        if ridge_term is None:
            ridge_term = np.std(Y) * np.sqrt(mean_diag) / np.sqrt(n - 1)

        if randomizer_scale is None:
            randomizer_scale = np.sqrt(mean_diag) * 0.5 * np.std(Y) * np.sqrt(n / (n - 1.))

        return slope(loglike, np.asarray(slope_weights) / sigma ** 2, ridge_term, randomizer_scale)

# Projection onto selected subgradients of SLOPE

def _projection_onto_selected_subgradients(prox_arg,
                                           weights,
                                           ordering,
                                           cluster_sizes,
                                           active_signs,
                                           last_value_zero=True):
    """
    Compute the projection of a point onto the set of
    subgradients of the SLOPE penalty with a given
    clustering of the solution and signs of the variables.
    This is a projection onto a lower dimensional set. The dimension
    of this set is p -- the dimensions of the `prox_arg` minus
    the number of unique values in `ordered_clustering` + 1 if the
    last value of the solution was zero (i.e. solution was sparse).

    Parameters
    ----------

    prox_arg : np.ndarray(p, np.float)
        Point to project

    weights : np.ndarray(p, np.float)
        Weights of the SLOPE penalty.

    ordering : np.ndarray(p, np.int)
        Order of original argument to SLOPE prox.
        First entry corresponds to largest argument of SLOPE prox.

    cluster_sizes : sequence
        Sizes of clusters, starting with
        largest in absolute value.

    active_signs : np.ndarray(p, np.int)
         Signs of non-zero coefficients.

    last_value_zero : bool
        Is the last solution value equal to 0?
    """

    result = np.zeros_like(prox_arg)

    ordered_clustering = []
    cur_idx = 0
    for cluster_size in cluster_sizes:
        ordered_clustering.append([ordering[j + cur_idx] for j in range(cluster_size)])
        cur_idx += cluster_size

    # Now, run appropriate SLOPE prox on each cluster
    cur_idx = 0
    for i, cluster in enumerate(ordered_clustering):
        prox_subarg = np.array([prox_arg[j] for j in cluster])

        # If the value of the soln to the prox was non-zero
        # then we solve a SLOPE of size 1 smaller than the cluster

        # If the cluster size is 1, the value is just
        # the corresponding signed weight

        if i < len(ordered_clustering) - 1 or not last_value_zero:
            if len(cluster) == 1:
                result[cluster[0]] = weights[cur_idx] * active_signs[cluster[0]]
            else:
                indices = [j + cur_idx for j in range(len(cluster))]
                cluster_weights = weights[indices]

                ir = IsotonicRegression()
                _ir_result = ir.fit_transform(np.arange(len(cluster)), cluster_weights[::-1])[::-1]
                result[indices] = -np.multiply(active_signs[indices], _ir_result/2.)

        else:
            indices = np.array([j + cur_idx for j in range(len(cluster))])
            cluster_weights = weights[indices]

            slope_prox = _basic_proximal_map(prox_subarg, cluster_weights)
            result[indices] = prox_subarg - slope_prox

        cur_idx += len(cluster)

    return result
<commit_msg>highdim LASSO selective MLE should be reusable<commit_after>from __future__ import print_function

import functools
import numpy as np

# sklearn imports

have_isotonic = False
try:
    from sklearn.isotonic import IsotonicRegression
    have_isotonic = True
except ImportError:
    raise ValueError('unable to import isotonic regression from sklearn, SLOPE subgradient projection will not work')

# regreg imports

from regreg.atoms.slope import _basic_proximal_map
import regreg.api as rr

from ..constraints.affine import constraints

from .randomization import randomization
from .base import restricted_estimator
from .lasso import highdim
from .query import (query,
                    multiple_queries,
                    langevin_sampler,
                    affine_gaussian_sampler)

class slope(highdim):

    def __init__(self,
                 loglike,
                 slope_weights,
                 ridge_term,
                 randomizer_scale,
                 perturb=None):
        r"""
        Create a new post-selection object for the SLOPE problem

        Parameters
        ----------

        loglike : `regreg.smooth.glm.glm`
            A (negative) log-likelihood as implemented in `regreg`.

        slope_weights : np.ndarray
            SLOPE weights for L-1 penalty. If a float,
            it is broadcast to all features.

        ridge_term : float
            How big a ridge term to add?

        randomizer_scale : float
            Scale for IID components of randomization.

        perturb : np.ndarray
            Random perturbation subtracted as a linear
            term in the objective function.
        """

        self.loglike = loglike
        self.nfeature = p = self.loglike.shape[0]

        if np.asarray(slope_weights).shape == ():
            slope_weights = np.ones(loglike.shape) * slope_weights
        self.slope_weights = np.asarray(slope_weights)

        self.randomizer = randomization.isotropic_gaussian((p,), randomizer_scale)
        self.ridge_term = ridge_term
        self.penalty = rr.slope(slope_weights, lagrange=1.)
        self._initial_omega = perturb  # random perturbation

    def fit(self,
            solve_args={'tol': 1.e-12, 'min_its': 50},
            perturb=None):

        p = self.nfeature

        # take a new perturbation if supplied
        if perturb is not None:
            self._initial_omega = perturb
        if self._initial_omega is None:
            self._initial_omega = self.randomizer.sample()

        quad = rr.identity_quadratic(self.ridge_term, 0, -self._initial_omega, 0)
        problem = rr.simple_problem(self.loglike, self.penalty)
        self.initial_soln = problem.solve(quad, **solve_args)

        # now we have to work out SLOPE details, clusters, etc.

        active_signs = np.sign(self.initial_soln)
        active = self._active = active_signs != 0

        print("check active terms", active.sum())

        self._overall = overall = active> 0
        self._inactive = inactive = ~self._overall

        _active_signs = active_signs.copy()
        self.selection_variable = {'sign': _active_signs,
                                   'variables': self._overall}

        initial_subgrad = -(self.loglike.smooth_objective(self.initial_soln, 'grad') +
                            quad.objective(self.initial_soln, 'grad'))
        self.initial_subgrad = initial_subgrad

        indices = np.argsort(-np.fabs(self.initial_soln))
        sorted_soln = self.initial_soln[indices]
        initial_scalings = np.sort(np.unique(np.fabs(self.initial_soln[active])))[::-1]
        self.observed_opt_state = initial_scalings

        self._unpenalized = np.zeros(p, np.bool)

        _beta_unpenalized = restricted_estimator(self.loglike, self._overall, solve_args=solve_args)

        beta_bar = np.zeros(p)
        beta_bar[overall] = _beta_unpenalized
        self._beta_full = beta_bar

        self.num_opt_var = self.observed_opt_state.shape[0]

        X, y = self.loglike.data
        W = self._W = self.loglike.saturated_loss.hessian(X.dot(beta_bar))
        _hessian_active = np.dot(X.T, X[:, active] * W[:, None])
        _score_linear_term = -_hessian_active
        self.score_transform = (_score_linear_term, np.zeros(_score_linear_term.shape[0]))

        self.observed_score_state = _score_linear_term.dot(_beta_unpenalized)
        self.observed_score_state[inactive] += self.loglike.smooth_objective(beta_bar, 'grad')[inactive]

        cur_indx_array = []
        cur_indx_array.append(0)
        cur_indx = 0
        pointer = 0
        signs_cluster = []
        for j in range(p - 1):
            if np.abs(sorted_soln[j + 1]) != np.abs(sorted_soln[cur_indx]):
                cur_indx_array.append(j + 1)
                cur_indx = j + 1
                sign_vec = np.zeros(p)
                sign_vec[np.arange(j + 1 - cur_indx_array[pointer]) + cur_indx_array[pointer]] = \
                    np.sign(self.initial_soln[indices[np.arange(j + 1 - cur_indx_array[pointer]) + cur_indx_array[pointer]]])
                signs_cluster.append(sign_vec)
                pointer = pointer + 1
                if sorted_soln[j + 1] == 0:
                    break

        signs_cluster = np.asarray(signs_cluster).T
        X_clustered = X[:, indices].dot(signs_cluster)
        _opt_linear_term = X.T.dot(X_clustered)
        self.opt_transform = (_opt_linear_term, self.initial_subgrad)

        cov, prec = self.randomizer.cov_prec
        opt_linear, opt_offset = self.opt_transform

        print("check if correct", np.allclose(self.observed_score_state + opt_offset + opt_linear.dot(initial_scalings),
                                              self._initial_omega, rtol=1e-05, atol=1e-08))

        cond_precision = opt_linear.T.dot(opt_linear) * prec
        cond_cov = np.linalg.inv(cond_precision)
        logdens_linear = cond_cov.dot(opt_linear.T) * prec
        cond_mean = -logdens_linear.dot(self.observed_score_state + opt_offset)

        logdens_transform = (logdens_linear, opt_offset)

        def log_density(logdens_linear, offset, cond_prec, score, opt):
            if score.ndim == 1:
                mean_term = logdens_linear.dot(score.T + offset).T
            else:
                mean_term = logdens_linear.dot(score.T + offset[:, None]).T
            arg = opt + mean_term
            return - 0.5 * np.sum(arg * cond_prec.dot(arg.T).T, 1)

        log_density = functools.partial(log_density, logdens_linear, opt_offset, cond_precision)

        # now make the constraints

        A_scaling_0 = -np.identity(self.num_opt_var)
        A_scaling_1 = -np.identity(self.num_opt_var)[:(self.num_opt_var-1), :]
        for k in range(A_scaling_1.shape[0]):
           A_scaling_1[k,k+1]= 1
        A_scaling = np.vstack([A_scaling_0, A_scaling_1])
        b_scaling = np.zeros(2*self.num_opt_var-1)

        # A_scaling = -np.identity(self.num_opt_var)
        # b_scaling = np.zeros(self.num_opt_var)

        affine_con = constraints(A_scaling,
                                 b_scaling,
                                 mean=cond_mean,
                                 covariance=cond_cov)

        self.sampler = affine_gaussian_sampler(affine_con,
                                               self.observed_opt_state,
                                               self.observed_score_state,
                                               log_density,
                                               logdens_transform,
                                               selection_info=self.selection_variable)
        return active_signs

    # Targets of inference
    # and covariance with score representation
    # are same as highdim LASSO

    @staticmethod
    def gaussian(X,
                 Y,
                 slope_weights,
                 sigma=1.,
                 quadratic=None,
                 ridge_term=0.,
                 randomizer_scale=None):

        loglike = rr.glm.gaussian(X, Y, coef=1. / sigma ** 2, quadratic=quadratic)
        n, p = X.shape

        mean_diag = np.mean((X ** 2).sum(0))
        if ridge_term is None:
            ridge_term = np.std(Y) * np.sqrt(mean_diag) / np.sqrt(n - 1)

        if randomizer_scale is None:
            randomizer_scale = np.sqrt(mean_diag) * 0.5 * np.std(Y) * np.sqrt(n / (n - 1.))

        return slope(loglike, np.asarray(slope_weights) / sigma ** 2, ridge_term, randomizer_scale)

# Projection onto selected subgradients of SLOPE

def _projection_onto_selected_subgradients(prox_arg,
                                           weights,
                                           ordering,
                                           cluster_sizes,
                                           active_signs,
                                           last_value_zero=True):
    """
    Compute the projection of a point onto the set of
    subgradients of the SLOPE penalty with a given
    clustering of the solution and signs of the variables.
    This is a projection onto a lower dimensional set. The dimension
    of this set is p -- the dimensions of the `prox_arg` minus
    the number of unique values in `ordered_clustering` + 1 if the
    last value of the solution was zero (i.e. solution was sparse).

    Parameters
    ----------

    prox_arg : np.ndarray(p, np.float)
        Point to project

    weights : np.ndarray(p, np.float)
        Weights of the SLOPE penalty.

    ordering : np.ndarray(p, np.int)
        Order of original argument to SLOPE prox.
        First entry corresponds to largest argument of SLOPE prox.

    cluster_sizes : sequence
        Sizes of clusters, starting with
        largest in absolute value.

    active_signs : np.ndarray(p, np.int)
         Signs of non-zero coefficients.

    last_value_zero : bool
        Is the last solution value equal to 0?
    """

    result = np.zeros_like(prox_arg)

    ordered_clustering = []
    cur_idx = 0
    for cluster_size in cluster_sizes:
        ordered_clustering.append([ordering[j + cur_idx] for j in range(cluster_size)])
        cur_idx += cluster_size

    # Now, run appropriate SLOPE prox on each cluster
    cur_idx = 0
    for i, cluster in enumerate(ordered_clustering):
        prox_subarg = np.array([prox_arg[j] for j in cluster])

        # If the value of the soln to the prox was non-zero
        # then we solve a SLOPE of size 1 smaller than the cluster

        # If the cluster size is 1, the value is just
        # the corresponding signed weight

        if i < len(ordered_clustering) - 1 or not last_value_zero:
            if len(cluster) == 1:
                result[cluster[0]] = weights[cur_idx] * active_signs[cluster[0]]
            else:
                indices = [j + cur_idx for j in range(len(cluster))]
                cluster_weights = weights[indices]

                ir = IsotonicRegression()
                _ir_result = ir.fit_transform(np.arange(len(cluster)), cluster_weights[::-1])[::-1]
                result[indices] = -np.multiply(active_signs[indices], _ir_result/2.)

        else:
            indices = np.array([j + cur_idx for j in range(len(cluster))])
            cluster_weights = weights[indices]

            slope_prox = _basic_proximal_map(prox_subarg, cluster_weights)
            result[indices] = prox_subarg - slope_prox

        cur_idx += len(cluster)

    return result
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#Source: http://flask.pocoo.org/snippets/45/

from flask.ext.restful import request, abort
from functools import wraps

def request_wants_json():
    best = request.accept_mimetypes \
        .best_match(['application/json', 'text/html'])
    return best == 'application/json' and \
        request.accept_mimetypes[best] > \
        request.accept_mimetypes['text/html']

def json_mimetype_required(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        if not request_wants_json():
            abort(400, message="You need to accept json to complete your request")
        return f(*args, **kwargs)
    return wrapper
<commit_msg>Change request_wants_json<commit_after># -*- coding: utf-8 -*-
#Source: http://flask.pocoo.org/snippets/45/

from flask.ext.restful import request, abort
from functools import wraps

def request_wants_json():
    best = request.accept_mimetypes \
        .best_match(['application/json', 'text/html'])
    return best == 'application/json' and \
        request.accept_mimetypes[best] >= \
        request.accept_mimetypes['text/html']

def json_mimetype_required(f):
    @wraps(f)
    def wrapper(*args, **kwargs):
        if not request_wants_json():
            abort(400, message="You need to accept json to complete your request")
        return f(*args, **kwargs)
    return wrapper
<|endoftext|>


<commit_before>from collections import namedtuple
import itertools
import os
import string
import sys
import re
import socket
import random

# don't use the standard instance - we don't want to be seeded
true_random = random.Random()

def port_used(address='127.0.0.1', port=6633):
  s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  try:
    s.bind((address, port))
    s.listen(1)
    s.close()
    return False
  except Exception, e:
    # TODO(cs): catch specific errors
    return True

# TODO(cs): this function don't appear to be invoked?
def find_port(port_spec):
  if isinstance(port_spec, int):
    def port_gen():
      yield port_spec
      raise Exception("Fixed port %d is busy. Consider specifying a range or a lambda " % port_spec)
  elif isinstance(port_spec, list):
    def port_gen():
      cands = list(port_spec)
      true_random.shuffle(cands)
      for c in cands:
        yield c
      raise Exception("Port list/range %s exhausted" % str(port_spec))
  elif isinstance(port_spec, types.FunctionType) or isinstance(port_spec, types.LambdaType):
    port_gen = port_spec

  gen = port_gen()
  for attempt in range(0,100):
    candidate = gen.next()
    if not port_used(port=candidate):
      return candidate
  raise Exception("Could not find a port in 100 tries")

# TODO(cs): this function don't appear to be invoked?
def find_ports(**kwargs):
  return { k : find_port(v) for k, v in kwargs.iteritems() }

class ControllerConfig(object):
  _port_gen = itertools.count(6633)
  _controller_count_gen = itertools.count(1)
  _controller_labels = set()

  def __init__(self, cmdline="", address="127.0.0.1", port=None,
               additional_ports={}, cwd=None, sync=None, controller_type=None,
               label=None, config_file=None, config_template=None,
               try_new_ports=True):
    '''
    Store metadata for the controller.
      - cmdline is an array of command line tokens.

        Note: if you need to pass in the address and port to controller's
        command line, use the aliases __address__ and __port__ to have the
        values interpolated automatically
      - address and port are the sockets switches will bind to
      - controller_type: help us figure out by specifying controller_type as a
        string that represents the controller. If it is not specified, this
        method will try to guess if it is either pox or floodlight.
    '''
    if cmdline == "":
      raise RuntimeError("Must specify boot parameters.")
    self.cmdline = cmdline

    self.address = address
    orig_port = port
    if (re.match("[0-9]{1,3}.[0-9]{1,3}.[0-9]{1,3}.[0-9]{1,3}", address) or
        address == "localhost"):
      # Normal TCP socket
      if not port:
        port = self._port_gen.next()
      if try_new_ports:
        while port_used(port=port):
          print "Port %d in use... trying next" % port
          port += true_random.randint(0,50)
      self.port = port
      self._server_info = (self.address, port)
    else:
      # Unix domain socket
      self.port = None
      self._server_info = address

    # TODO(sam): we should either call them all controller_type or all 'name'
    # we only accept strings
    self.name = ""
    if isinstance(controller_type,str):
      self.name = controller_type
    elif "pox" in self.cmdline:
      self.name = "pox"
    elif "floodlight" in self.cmdline:
      self.name = "floodlight"

    self.cwd = cwd
    if not cwd:
        sys.stderr.write("""
        =======================================================================
        WARN - no working directory defined for controller with command line
        %s
        The controller is run in the STS base directory. This may result
        in unintended consequences (i.e., POX not logging correctly).
        =======================================================================
        \n""" % (self.cmdline) )

    self.sync = sync
    if label:
      label = label
    else:
      label = "c"+str(self._controller_count_gen.next())
    if label in self._controller_labels:
      raise ValueError("Label %s already registered!" % label)
    self._controller_labels.add(label)
    self.label = label

    self.config_file = config_file
    self.config_template = config_template
    self.additional_ports = additional_ports

  @property
  def cid(self):
    ''' Return this controller's id '''
    return self.label

  @property
  def server_info(self):
    """ information about the _real_ socket that the controller is listening on"""
    return self._server_info

  def _expand_vars(self, s):
    return reduce(lambda s, (name, val): s.replace("__%s_port__" % name, str(val)), self.additional_ports.iteritems(), s) \
            .replace("__port__", str(self.port)) \
            .replace("__address__", str(self.address)) \
            .replace("__config__", str(os.path.abspath(self.config_file) if self.config_file else ""))

  @property
  def expanded_cmdline(self):
    return map(self._expand_vars, self.cmdline.split())

  def generate_config_file(self, target_dir):
    if self.config_file is None:
      self.config_file = os.path.join(target_dir, os.path.basename(self.config_template).replace(".template", ""))

    with open(self.config_template, "r") as in_file:
      with open(self.config_file, "w") as out_file:
        out_file.write(self._expand_vars(in_file.read()))

  def __repr__(self):
    attributes = ("cmdline", "address", "port", "cwd", "sync")

    pairs = ( (attr, getattr(self, attr)) for attr in attributes)
    quoted = ( "%s=%s" % (attr, repr(value)) for (attr, value) in pairs if value)

    return self.__class__.__name__  + "(" + ", ".join(quoted) + ")"
<commit_msg>Add missing copyright header<commit_after># Copyright 2011-2013 Colin Scott
# Copyright 2011-2013 Andreas Wundsam
#
# 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 collections import namedtuple
import itertools
import os
import string
import sys
import re
import socket
import random

# don't use the standard instance - we don't want to be seeded
true_random = random.Random()

def port_used(address='127.0.0.1', port=6633):
  s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
  try:
    s.bind((address, port))
    s.listen(1)
    s.close()
    return False
  except Exception, e:
    # TODO(cs): catch specific errors
    return True

# TODO(cs): this function don't appear to be invoked?
def find_port(port_spec):
  if isinstance(port_spec, int):
    def port_gen():
      yield port_spec
      raise Exception("Fixed port %d is busy. Consider specifying a range or a lambda " % port_spec)
  elif isinstance(port_spec, list):
    def port_gen():
      cands = list(port_spec)
      true_random.shuffle(cands)
      for c in cands:
        yield c
      raise Exception("Port list/range %s exhausted" % str(port_spec))
  elif isinstance(port_spec, types.FunctionType) or isinstance(port_spec, types.LambdaType):
    port_gen = port_spec

  gen = port_gen()
  for attempt in range(0,100):
    candidate = gen.next()
    if not port_used(port=candidate):
      return candidate
  raise Exception("Could not find a port in 100 tries")

# TODO(cs): this function don't appear to be invoked?
def find_ports(**kwargs):
  return { k : find_port(v) for k, v in kwargs.iteritems() }

class ControllerConfig(object):
  _port_gen = itertools.count(6633)
  _controller_count_gen = itertools.count(1)
  _controller_labels = set()

  def __init__(self, cmdline="", address="127.0.0.1", port=None,
               additional_ports={}, cwd=None, sync=None, controller_type=None,
               label=None, config_file=None, config_template=None,
               try_new_ports=True):
    '''
    Store metadata for the controller.
      - cmdline is an array of command line tokens.

        Note: if you need to pass in the address and port to controller's
        command line, use the aliases __address__ and __port__ to have the
        values interpolated automatically
      - address and port are the sockets switches will bind to
      - controller_type: help us figure out by specifying controller_type as a
        string that represents the controller. If it is not specified, this
        method will try to guess if it is either pox or floodlight.
    '''
    if cmdline == "":
      raise RuntimeError("Must specify boot parameters.")
    self.cmdline = cmdline

    self.address = address
    orig_port = port
    if (re.match("[0-9]{1,3}.[0-9]{1,3}.[0-9]{1,3}.[0-9]{1,3}", address) or
        address == "localhost"):
      # Normal TCP socket
      if not port:
        port = self._port_gen.next()
      if try_new_ports:
        while port_used(port=port):
          print "Port %d in use... trying next" % port
          port += true_random.randint(0,50)
      self.port = port
      self._server_info = (self.address, port)
    else:
      # Unix domain socket
      self.port = None
      self._server_info = address

    # TODO(sam): we should either call them all controller_type or all 'name'
    # we only accept strings
    self.name = ""
    if isinstance(controller_type,str):
      self.name = controller_type
    elif "pox" in self.cmdline:
      self.name = "pox"
    elif "floodlight" in self.cmdline:
      self.name = "floodlight"

    self.cwd = cwd
    if not cwd:
        sys.stderr.write("""
        =======================================================================
        WARN - no working directory defined for controller with command line
        %s
        The controller is run in the STS base directory. This may result
        in unintended consequences (i.e., POX not logging correctly).
        =======================================================================
        \n""" % (self.cmdline) )

    self.sync = sync
    if label:
      label = label
    else:
      label = "c"+str(self._controller_count_gen.next())
    if label in self._controller_labels:
      raise ValueError("Label %s already registered!" % label)
    self._controller_labels.add(label)
    self.label = label

    self.config_file = config_file
    self.config_template = config_template
    self.additional_ports = additional_ports

  @property
  def cid(self):
    ''' Return this controller's id '''
    return self.label

  @property
  def server_info(self):
    """ information about the _real_ socket that the controller is listening on"""
    return self._server_info

  def _expand_vars(self, s):
    return reduce(lambda s, (name, val): s.replace("__%s_port__" % name, str(val)), self.additional_ports.iteritems(), s) \
            .replace("__port__", str(self.port)) \
            .replace("__address__", str(self.address)) \
            .replace("__config__", str(os.path.abspath(self.config_file) if self.config_file else ""))

  @property
  def expanded_cmdline(self):
    return map(self._expand_vars, self.cmdline.split())

  def generate_config_file(self, target_dir):
    if self.config_file is None:
      self.config_file = os.path.join(target_dir, os.path.basename(self.config_template).replace(".template", ""))

    with open(self.config_template, "r") as in_file:
      with open(self.config_file, "w") as out_file:
        out_file.write(self._expand_vars(in_file.read()))

  def __repr__(self):
    attributes = ("cmdline", "address", "port", "cwd", "sync")

    pairs = ( (attr, getattr(self, attr)) for attr in attributes)
    quoted = ( "%s=%s" % (attr, repr(value)) for (attr, value) in pairs if value)

    return self.__class__.__name__  + "(" + ", ".join(quoted) + ")"
<|endoftext|>


<commit_before>import os
import shutil
import numpy as np
# import pickle
import cPickle as pickle
from AdaptivePELE.atomset import RMSDCalculator
from AdaptivePELE.testing import runMarkovChainModel as run
import AdaptivePELE.atomset.atomset as atomset

def cleanup(tmpFolder):
    """
        Remove folder if exists

        :param tmpFolder: Folder to remove
        :type tmpFolder: str
    """
    if os.path.exists(tmpFolder):
        shutil.rmtree(tmpFolder)

def makeFolder(outputDir):
    """
        Makes folder

        :param outputDir: Folder filename
        :type outputDir: str
    """
    if not os.path.exists(outputDir):
        os.makedirs(outputDir)

def getSnapshots(trajectoryFile, verbose=False):
    """
        Gets the snapshots

        :param trajectoryFile: Trajectory filename
        :type trajectoryFile: str
        :param verbose: Add verbose to snapshots
        :type verbose: bool

        :returns: str -- Snapshots with information
    """
    inputFile = open(trajectoryFile, "r")
    inputFileContent = inputFile.read()

    snapshots = inputFileContent.split("ENDMDL")[:-1]
    if not verbose:
        return snapshots

    remarkInfo = "REMARK 000 File created using PELE++\nREMARK source            : %s\nREMARK original model nr : %d\nREMARK First snapshot is 1, not 0 (as opposed to report)\n"
    snapshotsWithInfo = [remarkInfo % (trajectoryFile, i+1)+snapshot for i, snapshot in enumerate(snapshots)]
    return snapshotsWithInfo

def getTrajNum(trajFilename):
    """
        Gets the trajectory number

        :param trajFilename: Trajectory filename
        :type trajFilename: str

        :returns: int -- Trajectory number
    """
    return int(trajFilename.split("_")[-1][:-4])

def calculateContactMapEigen(contactMap):
    """
        Calculates eigenvectors and values of an extended contact map

        :param contactMap: Contact map
        :type contactMap: np.array

        :returns: (numpy.ndarray, numpy.ndarray) -- eigenvalues, eigenvectors
    """
    nLig, nCA = contactMap.shape
    extendedCM = np.zeros((nLig+nCA, nLig+nCA))
    extendedCM[nCA:, :nCA] = contactMap
    extendedCM[:nCA, nCA:] = contactMap.T
    assert (extendedCM == extendedCM.T).all(), "Extended ContactMap not symmetric"
    eiv, eic = np.linalg.eigh(extendedCM)
    return eiv, eic

def assertSymmetriesDict(symmetries, PDB):
    """
        Asserts the symmetry list in a PDB

        :param symmetries: List of symmetry groups
        :type symmetries: list of str
        :param PDB: PDB object to check symmetry list against
        :type PDB: :py:class:`.PDB`

        :raise AssertionError: If an atom is not found in the structure
    """
    for group in symmetries:
        for key in group:
            assert key in PDB.atoms, "Symmetry atom %s not found in initial structure" % key
    if symmetries:
        print "Symmetry dictionary correctly defined!"

def getRMSD(traj, nativePDB, resname, symmetries):
    """
        Computes the RMSD of a trajectory, given a native and symmetries

        :param traj: Trajecotry filename
        :type traj: str
        :param nativePDB:  Native PDB object
        :type native PDB: :py:class:`.PDB`
        :param resname: Resname to compute its RMSD
        :type resname: str
        :param symmetries: Symmetries dictionary list with independent symmetry groups
        :type symmetries: list of dict

        :return: np.array -- Array with the rmsd values of the trajectory
    """

    snapshots = getSnapshots(traj)
    rmsds = np.zeros(len(snapshots))
    RMSDCalc = RMSDCalculator.RMSDCalculator(symmetries)
    for i, snapshot in enumerate(snapshots):
        snapshotPDB = atomset.PDB()
        snapshotPDB.initialise(snapshot, resname=resname)

        rmsds[i] = RMSDCalc.computeRMSD(nativePDB, snapshotPDB)

    return rmsds

def readClusteringObject(clusteringObjectPath):
    """
        Reads and returns a clustering object

        :param clusteringObjectPath: Clustering object path
        :type clusteringObjectPath: str

        :raise EOFError: If the object is empty

        :returns: :py:class:`.Clustering` -- clusteringObject
    """
    with open(clusteringObjectPath, 'rb') as f:
        try:
            return pickle.load(f)
        except EOFError:
            raise EOFError("Empty clustering object!")

def ensure_connectivity_msm(msm):
    if msm.nstates_full == msm.nstates:
        return msm.stationary_distribution
    else:
        counts = msm.count_matrix_full
        counts += 1/float(counts.shape[0])
        trans = run.buildRevTransitionMatrix(counts)
        eiv, eic = run.getSortedEigen(trans)
        return run.getStationaryDistr(eic[:, 0])

def get_epoch_folders(path):
    allFolders = os.listdir(path)
    return [epoch for epoch in allFolders if epoch.isdigit()]

def write_PDB_clusters(pmf_xyzg, title="clusters.pdb"):
    templateLine = "HETATM%s H%sCLT L 502    %s%s%s  0.75%s                H\n"

    content = ""
    for i, line in enumerate(pmf_xyzg):
        number = str(i).rjust(5)
        number3 = str(i).ljust(3)
        x = ("%.3f" % line[0]).rjust(8)
        y = ("%.3f" % line[1]).rjust(8)
        z = ("%.3f" % line[2]).rjust(8)
        g = 0
        content += templateLine % (number, number3, x, y, z, g)

    f = open(title, 'w')
    f.write(content)
    f.close()
<commit_msg>Fix minor code smell and improve EOFError exception in utilities<commit_after>import os
import sys
import shutil
import numpy as np
# import pickle
import cPickle as pickle
from AdaptivePELE.atomset import RMSDCalculator
from AdaptivePELE.testing import runMarkovChainModel as run
import AdaptivePELE.atomset.atomset as atomset

def cleanup(tmpFolder):
    """
        Remove folder if exists

        :param tmpFolder: Folder to remove
        :type tmpFolder: str
    """
    if os.path.exists(tmpFolder):
        shutil.rmtree(tmpFolder)

def makeFolder(outputDir):
    """
        Makes folder

        :param outputDir: Folder filename
        :type outputDir: str
    """
    if not os.path.exists(outputDir):
        os.makedirs(outputDir)

def getSnapshots(trajectoryFile, verbose=False):
    """
        Gets the snapshots

        :param trajectoryFile: Trajectory filename
        :type trajectoryFile: str
        :param verbose: Add verbose to snapshots
        :type verbose: bool

        :returns: str -- Snapshots with information
    """
    inputFile = open(trajectoryFile, "r")
    inputFileContent = inputFile.read()

    snapshots = inputFileContent.split("ENDMDL")[:-1]
    if not verbose:
        return snapshots

    remarkInfo = "REMARK 000 File created using PELE++\nREMARK source            : %s\nREMARK original model nr : %d\nREMARK First snapshot is 1, not 0 (as opposed to report)\n"
    snapshotsWithInfo = [remarkInfo % (trajectoryFile, i+1)+snapshot for i, snapshot in enumerate(snapshots)]
    return snapshotsWithInfo

def getTrajNum(trajFilename):
    """
        Gets the trajectory number

        :param trajFilename: Trajectory filename
        :type trajFilename: str

        :returns: int -- Trajectory number
    """
    return int(trajFilename.split("_")[-1][:-4])

def calculateContactMapEigen(contactMap):
    """
        Calculates eigenvectors and values of an extended contact map

        :param contactMap: Contact map
        :type contactMap: np.array

        :returns: (numpy.ndarray, numpy.ndarray) -- eigenvalues, eigenvectors
    """
    nLig, nCA = contactMap.shape
    extendedCM = np.zeros((nLig+nCA, nLig+nCA))
    extendedCM[nCA:, :nCA] = contactMap
    extendedCM[:nCA, nCA:] = contactMap.T
    assert (extendedCM == extendedCM.T).all(), "Extended ContactMap not symmetric"
    eiv, eic = np.linalg.eigh(extendedCM)
    return eiv, eic

def assertSymmetriesDict(symmetries, PDB):
    """
        Asserts the symmetry list in a PDB

        :param symmetries: List of symmetry groups
        :type symmetries: list of str
        :param PDB: PDB object to check symmetry list against
        :type PDB: :py:class:`.PDB`

        :raise AssertionError: If an atom is not found in the structure
    """
    for group in symmetries:
        for key in group:
            assert key in PDB.atoms, "Symmetry atom %s not found in initial structure" % key
    if symmetries:
        print "Symmetry dictionary correctly defined!"

def getRMSD(traj, nativePDB, resname, symmetries):
    """
        Computes the RMSD of a trajectory, given a native and symmetries

        :param traj: Trajecotry filename
        :type traj: str
        :param nativePDB:  Native PDB object
        :type native PDB: :py:class:`.PDB`
        :param resname: Resname to compute its RMSD
        :type resname: str
        :param symmetries: Symmetries dictionary list with independent symmetry groups
        :type symmetries: list of dict

        :return: np.array -- Array with the rmsd values of the trajectory
    """

    snapshots = getSnapshots(traj)
    rmsds = np.zeros(len(snapshots))
    RMSDCalc = RMSDCalculator.RMSDCalculator(symmetries)
    for i, snapshot in enumerate(snapshots):
        snapshotPDB = atomset.PDB()
        snapshotPDB.initialise(snapshot, resname=resname)

        rmsds[i] = RMSDCalc.computeRMSD(nativePDB, snapshotPDB)

    return rmsds

def readClusteringObject(clusteringObjectPath):
    """
        Reads and returns a clustering object

        :param clusteringObjectPath: Clustering object path
        :type clusteringObjectPath: str

        :raise EOFError: If the object is empty

        :returns: :py:class:`.Clustering` -- clusteringObject
    """
    with open(clusteringObjectPath, 'rb') as f:
        try:
            return pickle.load(f)
        except EOFError:
            raise EOFError, EOFError("Empty clustering object!"), sys.exc_info()[2]

def ensure_connectivity_msm(msm):
    if msm.nstates_full == msm.nstates:
        return msm.stationary_distribution
    else:
        counts = msm.count_matrix_full
        counts += 1/float(counts.shape[0])
        trans = run.buildRevTransitionMatrix(counts)
        _, eic = run.getSortedEigen(trans)
        return run.getStationaryDistr(eic[:, 0])

def get_epoch_folders(path):
    allFolders = os.listdir(path)
    return [epoch for epoch in allFolders if epoch.isdigit()]

def write_PDB_clusters(pmf_xyzg, title="clusters.pdb"):
    templateLine = "HETATM%s H%sCLT L 502    %s%s%s  0.75%s                H\n"

    content = ""
    for i, line in enumerate(pmf_xyzg):
        number = str(i).rjust(5)
        number3 = str(i).ljust(3)
        x = ("%.3f" % line[0]).rjust(8)
        y = ("%.3f" % line[1]).rjust(8)
        z = ("%.3f" % line[2]).rjust(8)
        g = 0
        content += templateLine % (number, number3, x, y, z, g)

    with open(title, 'w') as f:
        f.write(content)
<|endoftext|>


<commit_before>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Jan 25 22:34:05 2020

@author: mostafamousavi
"""

from EQTransformer.core.trainer import trainer
import pytest
import glob
import os

def test_generator():
    trainer(input_hdf5='test_data/waveforms.hdf5',
            input_csv='test_data/metadata.csv',
            output_name='test_trainer',
            cnn_blocks=2,
            lstm_blocks=1,
            drop_rate=0.2,
            label_type='gaussian',
            add_event_r=0.6,
            shift_event_r=0.9,
            add_noise_r=0.5,
            mode='generator',
            batch_size=200,
            epochs=1,
            patience=2,
            gpuid=None,
            gpu_limit=None)

    dir_list = [ev for ev in os.listdir('.') if ev.split('_')[-1] == 'outputs']
    if 'test_trainer_outputs'  in dir_list:
        successful = True
    else:
        successful = False
    assert successful == True

def test_report():
    report = glob.glob("test_trainer_outputs/X_report.txt")
    assert len(report) == 1

def test_tests():
    tests = glob.glob("test_trainer_outputs/test.npy")
    assert len(tests) == 1

def test_hist():
    hist = glob.glob("test_trainer_outputs/history.npy")
    assert len(hist) == 1

def test_models():
    models = glob.glob("test_trainer_outputs/*.h5")
    assert len(models) == 2

def test_plots():
    learning_cr = glob.glob("test_trainer_outputs/*.png")
    assert len(learning_cr) == 2<commit_msg>Delete test_trainer.py<commit_after><|endoftext|>


<commit_before>import signal, os
from gi.repository import Gtk as gtk
from gi.repository import AppIndicator3 as appindicator
import alignak_data as ad

APPINDICATOR_ID = 'appalignak'
img = os.path.abspath('images/alignak.png')

def main():
    indicator = appindicator.Indicator.new(APPINDICATOR_ID, img, appindicator.IndicatorCategory.SYSTEM_SERVICES)
    indicator.set_status(appindicator.IndicatorStatus.ACTIVE)
    indicator.set_menu(build_menu())
    gtk.main()

def alignak_link_menu(menu, value):
    item = gtk.MenuItem(value)
    menu.append(item)

def build_menu():
    menu = gtk.Menu()

    auth = ad.alignak_backend_auth()
    data = ad.get_host_state(auth)
    for key, v in data.items():
        value = key + ' -> ' + v
        alignak_link_menu(menu, value)

    item_quit = gtk.MenuItem('Quit')
    item_quit.connect('activate', quit_app)
    menu.append(item_quit)

    menu.show_all()
    return menu

def quit_app(source):
    gtk.main_quit()

if __name__ == "__main__":
    signal.signal(signal.SIGINT, signal.SIG_DFL)
    main()

<commit_msg>Add url and improve code<commit_after>import signal, os, webbrowser, gi
gi.require_version('Gtk', '3.0')
from gi.repository import Gtk as gtk
gi.require_version('AppIndicator3', '0.1')
from gi.repository import AppIndicator3 as appindicator
gi.require_version('Notify', '0.7')
from gi.repository import Notify as notify
import alignak_data as ad

# Alignak Variables
APPINDICATOR_ID = 'appalignak'
img = os.path.abspath('images/alignak.png')
auth = ad.alignak_backend_auth()
url = 'http://94.76.229.155:91'
# Gtk Object
menu = gtk.Menu()
up_item = gtk.MenuItem('')
down_item = gtk.MenuItem('')
menu.append(up_item)
menu.append(down_item)

def main():
    app = set_indicator()
    gtk.main()

def build_menu():
    # Add menu for hosts
    build_hosts_menu()

    # Add menu for Refresh
    build_check_menu()

    # Add menu for Quit
    build_menu_quit()

    menu.show_all()
    return menu

def set_indicator():
    indicator = appindicator.Indicator.new(APPINDICATOR_ID, img, appindicator.IndicatorCategory.SYSTEM_SERVICES)
    indicator.set_status(appindicator.IndicatorStatus.ACTIVE)
    indicator.set_menu(build_menu())
    notify.init(APPINDICATOR_ID)
    return indicator

def build_hosts_menu():
    UP, DOWN = check()
    if UP > 0:
        str_UP = 'Hosts UP (' + str(UP) + ')'
        up_item.set_label(str_UP)
        # menu.append(up_item)
    if DOWN > 0:
        str_NOK = 'Hosts DOWN (' + str(DOWN) + ')'
        down_item.set_label(str_NOK)
        # menu.append(down_item)

def update_menu(source):
    notify.Notification.new("<b>Update Alignak Data..</b>", build_hosts_menu(), None).show()

def check():
    UP = 0
    DOWN = 0
    data = ad.get_host_state(auth)
    for key, v in data.items():
        if 'UP' in v:
            UP += 1
        if 'DOWN' in v:
            DOWN += 1
    return UP, DOWN

def build_check_menu():
    check_item = gtk.MenuItem('Check')
    check_item.connect('activate', update_menu)
    menu.append(check_item)

def build_menu_quit():
    item_quit = gtk.MenuItem('Quit')
    item_quit.connect('activate', quit_app)
    menu.append(item_quit)

def open_url(source):
    webbrowser.open(url + '/hosts')

up_item.connect("activate", open_url)
down_item.connect("activate", open_url)

def print_hello():
    print('hello world')

def quit_app(source):
    notify.uninit()
    gtk.main_quit()

if __name__ == "__main__":
    signal.signal(signal.SIGINT, signal.SIG_DFL)
    main()

<|endoftext|>


<commit_before># Send UDP broadcast packets

MYPORT = 50000

import sys, time
from socket import *

s = socket(AF_INET, SOCK_DGRAM)
s.bind('', 0)
s.allowbroadcast(1)

while 1:
	data = `time.time()` + '\n'
	s.sendto(data, ('<broadcast>', MYPORT))
	time.sleep(2)

<commit_msg>Allow broadcast using setsockopt()<commit_after># Send UDP broadcast packets

MYPORT = 50000

import sys, time
from socket import *
from SOCKET import *

s = socket(AF_INET, SOCK_DGRAM)
s.bind('', 0)
s.setsockopt(SOL_SOCKET, SO_BROADCAST, 1)

while 1:
	data = `time.time()` + '\n'
	s.sendto(data, ('<broadcast>', MYPORT))
	time.sleep(2)

<|endoftext|>


<commit_before>
import discord
from discord.ext import commands, menus, tasks

import asyncio
import datetime
from typing import Optional

from parsedatetime import Calendar, VERSION_CONTEXT_STYLE

from utilities import checks
from utilities.menu import Menu

def setup(bot):
	bot.add_cog(Reminders(bot))

class Reminders(commands.Cog):

	def __init__(self, bot):
		self.bot = bot

		self.calendar = Calendar(version = VERSION_CONTEXT_STYLE)
		# Add mo as valid abbreviation for month
		self.calendar.ptc.units["months"].append("mo")

		# Add menu as subcommand of reminder and reminder list commands
		self.reminder_command.add_command(self.menu)
		self.list_reminders.add_command(self.menu)

		self.menus = []

		self.current_timer = None
		self.new_reminder = asyncio.Event()
		self.restarting_timer = False
		self.timer.start().set_name("Reminders")

	def cog_unload(self):
		for menu in self.menus:
			menu.stop()
		self.timer.cancel()

	async def initialize_database(self):
		await self.bot.connect_to_database()
		await self.bot.db.execute("CREATE SCHEMA IF NOT EXISTS reminders")
		await self.bot.db.execute(
			"""
			CREATE TABLE IF NOT EXISTS reminders.reminders (
				id				INT PRIMARY KEY GENERATED ALWAYS AS IDENTITY,
				user_id			BIGINT,
				channel_id		BIGINT,
				message_id		BIGINT,
				created_time	TIMESTAMPTZ DEFAULT NOW(),
				remind_time		TIMESTAMPTZ,
				reminder		TEXT,
				reminded		BOOL DEFAULT FALSE,
				cancelled 		BOOL DEFAULT FALSE,
				failed			BOOL DEFAULT FALSE
			)
			"""
		)

	async def cog_check(self, ctx):
		return await checks.not_forbidden().predicate(ctx)

	# TODO: Allow setting timezone (+ for time command as well)

	@commands.group(name = "reminder", aliases = ["remind", "reminders", "timer", "timers"],
					invoke_without_command = True, case_insensitive = True)
	async def reminder_command(self, ctx, *, reminder: Optional[commands.clean_content]):
		'''
		See and set reminders
		Times are in UTC
		'''
		# Invoke list subcommand for reminders and timers invocations
		if ctx.invoked_with in ("reminders", "timers"):
			offset = 0
			count = 10
			try:
				inputs = reminder.split()
				offset = int(inputs[0])
				count = int(inputs[1])
			except (AttributeError, IndexError, ValueError):
				pass
			return await ctx.invoke(self.list_reminders, offset, count)
		elif not reminder:
			raise commands.BadArgument("Time not specified")
		# Clean reminder input
		for prefix in ("me about ", "me to ", "me "):
			if reminder.startswith(prefix):
				reminder = reminder[len(prefix):]
		reminder = reminder.replace("from now", "")
		# Parse reminder
		now = datetime.datetime.now(datetime.timezone.utc)
		if not (matches := self.calendar.nlp(reminder, sourceTime = now)):
			raise commands.BadArgument("Time not specified")
		parsed_datetime, context, start_pos, end_pos, matched_text = matches[0]
		if not context.hasTime:
			parsed_datetime = parsed_datetime.replace(hour = now.hour, minute = now.minute,
														second = now.second, microsecond = now.microsecond)
		parsed_datetime = parsed_datetime.replace(tzinfo = datetime.timezone.utc)
		if parsed_datetime < now:
			raise commands.BadArgument("Time is in the past")
		# Respond
		reminder = reminder[:start_pos] + reminder[end_pos + 1:]
		reminder = reminder.strip()
		response = await ctx.embed_reply(fields = (("Reminder", reminder or ctx.bot.ZWS),),
											footer_text = f"Set for {parsed_datetime.isoformat(timespec = 'seconds').replace('+00:00', 'Z')}",
											timestamp = parsed_datetime)
		# Insert into database
		created_time = ctx.message.created_at.replace(tzinfo = datetime.timezone.utc)
		await self.bot.db.execute(
			"""
			INSERT INTO reminders.reminders (user_id, channel_id, message_id, created_time, remind_time, reminder)
			VALUES ($1, $2, $3, $4, $5, $6)
			""",
			ctx.author.id, ctx.channel.id, response.id, created_time, parsed_datetime, reminder
		)
		# Update timer
		if self.current_timer and parsed_datetime < self.current_timer["remind_time"]:
			self.restarting_timer = True
			self.timer.restart()
			self.timer.get_task().set_name("Reminders")
		else:
			self.new_reminder.set()

	@reminder_command.command(aliases = ["delete", "remove"])
	async def cancel(self, ctx, reminder_id: int):
		'''Cancel a reminder'''
		cancelled = await ctx.bot.db.fetchrow(
			"""
			UPDATE reminders.reminders
			SET cancelled = TRUE
			WHERE id = $1 AND user_id = $2 AND reminded = FALSE AND failed = FALSE
			RETURNING *
			""",
			reminder_id, ctx.author.id
		)
		if not cancelled:
			return await ctx.embed_reply(":no_entry: Error: Unable to find and cancel reminder")
		if self.current_timer and self.current_timer["id"] == reminder_id:
			self.restarting_timer = True
			self.timer.restart()
			self.timer.get_task().set_name("Reminders")
		await ctx.embed_reply(fields = (("Cancelled Reminder", cancelled["reminder"] or ctx.bot.ZWS),),
								footer_text = f"Set for {cancelled['remind_time'].isoformat(timespec = 'seconds').replace('+00:00', 'Z')}",
								timestamp = cancelled["remind_time"])

	@reminder_command.group(name = "list", invoke_without_command = True, case_insensitive = True)
	async def list_reminders(self, ctx, offset: Optional[int] = 0, count: Optional[int] = 10):
		'''
		List reminders
		Max count is 10
		'''
		records = await ctx.bot.db.fetch(
			"""
			SELECT id, channel_id, message_id, remind_time, reminder
			FROM reminders.reminders
			WHERE user_id = $1 AND reminded = FALSE AND cancelled = FALSE AND failed = FALSE
			ORDER BY remind_time
			LIMIT $2 OFFSET $3
			""",
			ctx.author.id, min(count, 10), offset
		)
		fields = []
		for record in records:
			value = record["reminder"] or "Reminder"
			if channel := ctx.bot.get_channel(record["channel_id"]):
				try:
					message = await channel.fetch_message(record["message_id"])
					value = f"[{value}]({message.jump_url})"
				except discord.NotFound:
					pass
				value += f"\nIn {getattr(channel, 'mention', 'DMs')}"
			# TODO: Attempt to fetch channel?
			value += f"\nAt {record['remind_time'].isoformat(timespec = 'seconds').replace('+00:00', 'Z')}"
			fields.append((f"ID: {record['id']}", value))
		await ctx.embed_reply(title = "Reminders", fields = fields)

	@commands.command(aliases = ['m', "menus", 'r', "reaction", "reactions"])
	async def menu(self, ctx, per_page: Optional[int] = 10):
		'''
		Reminders menu
		Max per_page is 10
		'''
		records = await ctx.bot.db.fetch(
			"""
			SELECT id, channel_id, message_id, remind_time, reminder
			FROM reminders.reminders
			WHERE user_id = $1 AND reminded = FALSE AND cancelled = FALSE AND failed = FALSE
			ORDER BY remind_time
			""",
			ctx.author.id
		)
		menu = RemindersMenu(records, per_page = min(per_page, 10))
		self.menus.append(menu)
		await menu.start(ctx, wait = True)
		self.menus.remove(menu)

	# TODO: clear subcommand

	# R/PT0S
	@tasks.loop()
	async def timer(self):
		record = await self.bot.db.fetchrow(
			"""
			SELECT * FROM reminders.reminders
			WHERE reminded = FALSE AND cancelled = FALSE AND failed = FALSE
			ORDER BY remind_time
			LIMIT 1
			"""
		)
		if not record:
			self.new_reminder.clear()
			return await self.new_reminder.wait()
		self.current_timer = record
		await discord.utils.sleep_until(record["remind_time"])
		if not (channel := self.bot.get_channel(record["channel_id"])):
			# TODO: Attempt to fetch channel?
			return await self.bot.db.execute("UPDATE reminders.reminders SET failed = TRUE WHERE id = $1", record["id"])
		user = self.bot.get_user(record["user_id"]) or self.bot.fetch_user(record["user_id"])
		# TODO: Handle user not found?
		embed = discord.Embed(color = self.bot.bot_color)
		try:
			message = await channel.fetch_message(record["message_id"])
			embed.description = f"[{record['reminder'] or 'Reminder'}]({message.jump_url})"
		except discord.NotFound:
			embed.description = record["reminder"] or "Reminder"
		embed.set_footer(text = "Reminder set")
		embed.timestamp = record["created_time"]
		try:
			await channel.send(user.mention, embed = embed)
		except discord.Forbidden:
			# TODO: Attempt to send without embed
			# TODO: Fall back to DM
			await self.bot.db.execute("UPDATE reminders.reminders SET failed = TRUE WHERE id = $1", record["id"])
		else:
			await self.bot.db.execute("UPDATE reminders.reminders SET reminded = TRUE WHERE id = $1", record["id"])

	@timer.before_loop
	async def before_timer(self):
		await self.initialize_database()
		await self.bot.wait_until_ready()

	@timer.after_loop
	async def after_timer(self):
		if self.restarting_timer:
			self.restarting_timer = False
		else:
			print(f"{self.bot.console_message_prefix}Reminders task cancelled @ {datetime.datetime.now().isoformat()}")

class RemindersMenu(Menu, menus.MenuPages):

	def __init__(self, records, per_page):
		super().__init__(RemindersSource(records, per_page), timeout = None, clear_reactions_after = True, check_embeds = True)

	async def send_initial_message(self, ctx, channel):
		message = await super().send_initial_message(ctx, channel)
		await ctx.bot.attempt_delete_message(ctx.message)
		return message

class RemindersSource(menus.ListPageSource):

	def __init__(self, records, per_page):
		super().__init__(records, per_page = per_page)

	async def format_page(self, menu, records):
		embed = discord.Embed(title = "Reminders", color = menu.bot.bot_color)
		embed.set_author(name = menu.ctx.author.display_name, icon_url = menu.ctx.author.avatar_url)
		if self.per_page == 1:
			records = [records]
		for record in records:
			value = record["reminder"] or "Reminder"
			if channel := menu.bot.get_channel(record["channel_id"]):
				try:
					message = await channel.fetch_message(record["message_id"])
					value = f"[{value}]({message.jump_url})"
				except discord.NotFound:
					pass
				value += f"\nIn {getattr(channel, 'mention', 'DMs')}"
			# TODO: Attempt to fetch channel?
			value += f"\nAt {record['remind_time'].isoformat(timespec = 'seconds').replace('+00:00', 'Z')}"
			embed.add_field(name = f"ID: {record['id']}", value = value)
		offset = menu.current_page * self.per_page
		start = offset + 1
		end = min(offset + self.per_page, len(self.entries))
		if start < end:
			embed.set_footer(text = f"Reminders {start} - {end} of {len(self.entries)}")
		elif start == end:
			embed.set_footer(text = f"Reminder {start} of {len(self.entries)}")
		else:
			return embed.set_footer(text = f"In response to: {menu.ctx.message.clean_content}")
		return {"content": f"In response to: `{menu.ctx.message.clean_content}`", "embed": embed}

<commit_msg>[Discord] Improve reminder menu command method name<commit_after>
import discord
from discord.ext import commands, menus, tasks

import asyncio
import datetime
from typing import Optional

from parsedatetime import Calendar, VERSION_CONTEXT_STYLE

from utilities import checks
from utilities.menu import Menu

def setup(bot):
	bot.add_cog(Reminders(bot))

class Reminders(commands.Cog):

	def __init__(self, bot):
		self.bot = bot

		self.calendar = Calendar(version = VERSION_CONTEXT_STYLE)
		# Add mo as valid abbreviation for month
		self.calendar.ptc.units["months"].append("mo")

		# Add menu as subcommand of reminder and reminder list commands
		self.reminder_command.add_command(self.menu)
		self.list_reminders.add_command(self.menu)

		self.menus = []

		self.current_timer = None
		self.new_reminder = asyncio.Event()
		self.restarting_timer = False
		self.timer.start().set_name("Reminders")

	def cog_unload(self):
		for menu in self.menus:
			menu.stop()
		self.timer.cancel()

	async def initialize_database(self):
		await self.bot.connect_to_database()
		await self.bot.db.execute("CREATE SCHEMA IF NOT EXISTS reminders")
		await self.bot.db.execute(
			"""
			CREATE TABLE IF NOT EXISTS reminders.reminders (
				id				INT PRIMARY KEY GENERATED ALWAYS AS IDENTITY,
				user_id			BIGINT,
				channel_id		BIGINT,
				message_id		BIGINT,
				created_time	TIMESTAMPTZ DEFAULT NOW(),
				remind_time		TIMESTAMPTZ,
				reminder		TEXT,
				reminded		BOOL DEFAULT FALSE,
				cancelled 		BOOL DEFAULT FALSE,
				failed			BOOL DEFAULT FALSE
			)
			"""
		)

	async def cog_check(self, ctx):
		return await checks.not_forbidden().predicate(ctx)

	# TODO: Allow setting timezone (+ for time command as well)

	@commands.group(name = "reminder", aliases = ["remind", "reminders", "timer", "timers"],
					invoke_without_command = True, case_insensitive = True)
	async def reminder_command(self, ctx, *, reminder: Optional[commands.clean_content]):
		'''
		See and set reminders
		Times are in UTC
		'''
		# Invoke list subcommand for reminders and timers invocations
		if ctx.invoked_with in ("reminders", "timers"):
			offset = 0
			count = 10
			try:
				inputs = reminder.split()
				offset = int(inputs[0])
				count = int(inputs[1])
			except (AttributeError, IndexError, ValueError):
				pass
			return await ctx.invoke(self.list_reminders, offset, count)
		elif not reminder:
			raise commands.BadArgument("Time not specified")
		# Clean reminder input
		for prefix in ("me about ", "me to ", "me "):
			if reminder.startswith(prefix):
				reminder = reminder[len(prefix):]
		reminder = reminder.replace("from now", "")
		# Parse reminder
		now = datetime.datetime.now(datetime.timezone.utc)
		if not (matches := self.calendar.nlp(reminder, sourceTime = now)):
			raise commands.BadArgument("Time not specified")
		parsed_datetime, context, start_pos, end_pos, matched_text = matches[0]
		if not context.hasTime:
			parsed_datetime = parsed_datetime.replace(hour = now.hour, minute = now.minute,
														second = now.second, microsecond = now.microsecond)
		parsed_datetime = parsed_datetime.replace(tzinfo = datetime.timezone.utc)
		if parsed_datetime < now:
			raise commands.BadArgument("Time is in the past")
		# Respond
		reminder = reminder[:start_pos] + reminder[end_pos + 1:]
		reminder = reminder.strip()
		response = await ctx.embed_reply(fields = (("Reminder", reminder or ctx.bot.ZWS),),
											footer_text = f"Set for {parsed_datetime.isoformat(timespec = 'seconds').replace('+00:00', 'Z')}",
											timestamp = parsed_datetime)
		# Insert into database
		created_time = ctx.message.created_at.replace(tzinfo = datetime.timezone.utc)
		await self.bot.db.execute(
			"""
			INSERT INTO reminders.reminders (user_id, channel_id, message_id, created_time, remind_time, reminder)
			VALUES ($1, $2, $3, $4, $5, $6)
			""",
			ctx.author.id, ctx.channel.id, response.id, created_time, parsed_datetime, reminder
		)
		# Update timer
		if self.current_timer and parsed_datetime < self.current_timer["remind_time"]:
			self.restarting_timer = True
			self.timer.restart()
			self.timer.get_task().set_name("Reminders")
		else:
			self.new_reminder.set()

	@reminder_command.command(aliases = ["delete", "remove"])
	async def cancel(self, ctx, reminder_id: int):
		'''Cancel a reminder'''
		cancelled = await ctx.bot.db.fetchrow(
			"""
			UPDATE reminders.reminders
			SET cancelled = TRUE
			WHERE id = $1 AND user_id = $2 AND reminded = FALSE AND failed = FALSE
			RETURNING *
			""",
			reminder_id, ctx.author.id
		)
		if not cancelled:
			return await ctx.embed_reply(":no_entry: Error: Unable to find and cancel reminder")
		if self.current_timer and self.current_timer["id"] == reminder_id:
			self.restarting_timer = True
			self.timer.restart()
			self.timer.get_task().set_name("Reminders")
		await ctx.embed_reply(fields = (("Cancelled Reminder", cancelled["reminder"] or ctx.bot.ZWS),),
								footer_text = f"Set for {cancelled['remind_time'].isoformat(timespec = 'seconds').replace('+00:00', 'Z')}",
								timestamp = cancelled["remind_time"])

	@reminder_command.group(name = "list", invoke_without_command = True, case_insensitive = True)
	async def list_reminders(self, ctx, offset: Optional[int] = 0, count: Optional[int] = 10):
		'''
		List reminders
		Max count is 10
		'''
		records = await ctx.bot.db.fetch(
			"""
			SELECT id, channel_id, message_id, remind_time, reminder
			FROM reminders.reminders
			WHERE user_id = $1 AND reminded = FALSE AND cancelled = FALSE AND failed = FALSE
			ORDER BY remind_time
			LIMIT $2 OFFSET $3
			""",
			ctx.author.id, min(count, 10), offset
		)
		fields = []
		for record in records:
			value = record["reminder"] or "Reminder"
			if channel := ctx.bot.get_channel(record["channel_id"]):
				try:
					message = await channel.fetch_message(record["message_id"])
					value = f"[{value}]({message.jump_url})"
				except discord.NotFound:
					pass
				value += f"\nIn {getattr(channel, 'mention', 'DMs')}"
			# TODO: Attempt to fetch channel?
			value += f"\nAt {record['remind_time'].isoformat(timespec = 'seconds').replace('+00:00', 'Z')}"
			fields.append((f"ID: {record['id']}", value))
		await ctx.embed_reply(title = "Reminders", fields = fields)

	@commands.command(name = "menu", aliases = ['m', "menus", 'r', "reaction", "reactions"])
	async def menu_command(self, ctx, per_page: Optional[int] = 10):
		'''
		Reminders menu
		Max per_page is 10
		'''
		records = await ctx.bot.db.fetch(
			"""
			SELECT id, channel_id, message_id, remind_time, reminder
			FROM reminders.reminders
			WHERE user_id = $1 AND reminded = FALSE AND cancelled = FALSE AND failed = FALSE
			ORDER BY remind_time
			""",
			ctx.author.id
		)
		menu = RemindersMenu(records, per_page = min(per_page, 10))
		self.menus.append(menu)
		await menu.start(ctx, wait = True)
		self.menus.remove(menu)

	# TODO: clear subcommand

	# R/PT0S
	@tasks.loop()
	async def timer(self):
		record = await self.bot.db.fetchrow(
			"""
			SELECT * FROM reminders.reminders
			WHERE reminded = FALSE AND cancelled = FALSE AND failed = FALSE
			ORDER BY remind_time
			LIMIT 1
			"""
		)
		if not record:
			self.new_reminder.clear()
			return await self.new_reminder.wait()
		self.current_timer = record
		await discord.utils.sleep_until(record["remind_time"])
		if not (channel := self.bot.get_channel(record["channel_id"])):
			# TODO: Attempt to fetch channel?
			return await self.bot.db.execute("UPDATE reminders.reminders SET failed = TRUE WHERE id = $1", record["id"])
		user = self.bot.get_user(record["user_id"]) or self.bot.fetch_user(record["user_id"])
		# TODO: Handle user not found?
		embed = discord.Embed(color = self.bot.bot_color)
		try:
			message = await channel.fetch_message(record["message_id"])
			embed.description = f"[{record['reminder'] or 'Reminder'}]({message.jump_url})"
		except discord.NotFound:
			embed.description = record["reminder"] or "Reminder"
		embed.set_footer(text = "Reminder set")
		embed.timestamp = record["created_time"]
		try:
			await channel.send(user.mention, embed = embed)
		except discord.Forbidden:
			# TODO: Attempt to send without embed
			# TODO: Fall back to DM
			await self.bot.db.execute("UPDATE reminders.reminders SET failed = TRUE WHERE id = $1", record["id"])
		else:
			await self.bot.db.execute("UPDATE reminders.reminders SET reminded = TRUE WHERE id = $1", record["id"])

	@timer.before_loop
	async def before_timer(self):
		await self.initialize_database()
		await self.bot.wait_until_ready()

	@timer.after_loop
	async def after_timer(self):
		if self.restarting_timer:
			self.restarting_timer = False
		else:
			print(f"{self.bot.console_message_prefix}Reminders task cancelled @ {datetime.datetime.now().isoformat()}")

class RemindersMenu(Menu, menus.MenuPages):

	def __init__(self, records, per_page):
		super().__init__(RemindersSource(records, per_page), timeout = None, clear_reactions_after = True, check_embeds = True)

	async def send_initial_message(self, ctx, channel):
		message = await super().send_initial_message(ctx, channel)
		await ctx.bot.attempt_delete_message(ctx.message)
		return message

class RemindersSource(menus.ListPageSource):

	def __init__(self, records, per_page):
		super().__init__(records, per_page = per_page)

	async def format_page(self, menu, records):
		embed = discord.Embed(title = "Reminders", color = menu.bot.bot_color)
		embed.set_author(name = menu.ctx.author.display_name, icon_url = menu.ctx.author.avatar_url)
		if self.per_page == 1:
			records = [records]
		for record in records:
			value = record["reminder"] or "Reminder"
			if channel := menu.bot.get_channel(record["channel_id"]):
				try:
					message = await channel.fetch_message(record["message_id"])
					value = f"[{value}]({message.jump_url})"
				except discord.NotFound:
					pass
				value += f"\nIn {getattr(channel, 'mention', 'DMs')}"
			# TODO: Attempt to fetch channel?
			value += f"\nAt {record['remind_time'].isoformat(timespec = 'seconds').replace('+00:00', 'Z')}"
			embed.add_field(name = f"ID: {record['id']}", value = value)
		offset = menu.current_page * self.per_page
		start = offset + 1
		end = min(offset + self.per_page, len(self.entries))
		if start < end:
			embed.set_footer(text = f"Reminders {start} - {end} of {len(self.entries)}")
		elif start == end:
			embed.set_footer(text = f"Reminder {start} of {len(self.entries)}")
		else:
			return embed.set_footer(text = f"In response to: {menu.ctx.message.clean_content}")
		return {"content": f"In response to: `{menu.ctx.message.clean_content}`", "embed": embed}

<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Load incumbent candidate data scraped from the CAL-ACCESS website into OCD models.
"""
import re
from django.utils import timezone
from django.db.models import IntegerField
from django.db.models.functions import Cast
from calaccess_processed.management.commands import LoadOCDModelsCommand
from calaccess_processed.models.scraper import (
    IncumbentScrapedElection,
    ScrapedIncumbent,
)
from opencivicdata.models import Membership
from opencivicdata.models.elections import Election, Candidacy

class Command(LoadOCDModelsCommand):
    """
    Load incumbent candidate data scraped from the CAL-ACCESS website into OCD models.
    """
    help = 'Load incumbent candidate data scraped from the CAL-ACCESS website into OCD models'

    def handle(self, *args, **options):
        """
        Make it happen.
        """
        super(Command, self).handle(*args, **options)
        self.load()
        self.set_end_dates()
        if Candidacy.objects.exists():
            self.log('')
            self.set_incumbent_candidacies()
        self.success("Done!")

    def get_or_create_election(self, scraped_elec):
        """
        Get or create an OCD Election object using the IncumbentScrapedElection.

        Returns a tuple (Election object, created), where created is a boolean
        specifying whether a Election was created.
        """
        dt_obj = timezone.make_aware(
            timezone.datetime.combine(scraped_elec.date, timezone.datetime.min.time())
        )
        name = '{0} {1}'.format(dt_obj.year, scraped_elec.name)
        try:
            elec = Election.objects.get(start_time=dt_obj)
        except Election.DoesNotExist:
            # or make a new one
            elec = self.create_election(name, dt_obj)
            created = True
        else:
            created = False
            # if election already exists and is named 'SPECIAL' or 'RECALL'
            if 'SPECIAL' in elec.name.upper() or 'RECALL' in elec.name.upper():
                # and the matched election's name includes either 'GENERAL'
                # or 'PRIMARY'...
                if (
                    re.match(r'^\d{4} GENERAL$', scraped_elec.name) or
                    re.match(r'^\d{4} PRIMARY$', scraped_elec.name)
                ):
                    # update the name
                    elec.name = name
                    elec.save()
        return (elec, created)

    def load(self):
        """
        Load OCD Election, Membership and related models with data scraped from CAL-ACCESS website.
        """
        for scraped_elec in IncumbentScrapedElection.objects.all():
            ocd_elec = self.get_or_create_election(scraped_elec)[0]
            ocd_elec.sources.update_or_create(
                url=scraped_elec.url,
                note='Last scraped on {dt:%Y-%m-%d}'.format(
                    dt=scraped_elec.last_modified,
                )
            )

        for incumbent in ScrapedIncumbent.objects.all().order_by('-session'):
            # Get or create post
            post, post_created = self.get_or_create_post(
                incumbent.office_name,
            )
            if post_created and self.verbosity > 2:
                self.log('Created new Post: %s' % post.label)
            # Get or person
            person, person_created = self.get_or_create_person(
                incumbent.name,
                filer_id=incumbent.scraped_id,
            )
            if person_created and self.verbosity > 2:
                self.log('Created new Person: %s' % person.name)
            # Get or membership for post and person
            membership, membership_created = Membership.objects.get_or_create(
                person=person,
                post=post,
                role=post.role,
                organization=post.organization,
                person_name=person.sort_name,
            )
            if membership_created and self.verbosity > 2:
                self.log('Created new Membership: %s' % membership)
            # Handle start_date on membership
            if membership_created or membership.start_date == '':
                membership.start_date = incumbent.session
                membership.save()
            else:
                # increment start year down
                start = int(membership.start_date)
                if start > incumbent.session:
                    membership.start_date = incumbent.session
                    membership.save()

        return

    def set_end_dates(self):
        """
        Set the end_date for each Membership based on the start_date of each successor.
        """
        for member in Membership.objects.all():
            # Each member's end year should be the start year of their successor.
            # Successor is the member in the same post with the earliest
            # start year greater than the incumbent's start year
            successor_q = Membership.objects.exclude(
                start_date='',
            ).annotate(
                start_year=Cast('start_date', IntegerField()),
            ).filter(
                start_year__gt=int(member.start_date),
                post=member.post,
            ).order_by('start_year')

            if successor_q.exists():
                member.end_date = int(successor_q[0].start_date)
                member.save()
        return

    def set_incumbent_candidacies(self):
        """
        Set is_incumbent for candidacies within each member's start/end years.
        """
        # For every one of the members candidacies for the office where
        # the election of the contest happens after the start year
        # but before the end year, if it exists, mark as incumbent
        for member in Membership.objects.all():
            candidacies_q = member.person.candidacies.filter(
                contest__election__start_time__year__gt=int(member.start_date),
                post=member.post,
            ).exclude(is_incumbent=True)
            # Handle blank member end_date values
            if member.end_date != '':
                member_end_year = int(member.end_date)
                candidacies_q = candidacies_q.filter(contest__election__start_time__year__lte=member_end_year)

            if candidacies_q.exists():
                rows = candidacies_q.update(is_incumbent=True)
                self.log(
                    ' {0} identified as incumbent in {1} contests'.format(
                        member.person.name,
                        rows,
                    )
                )
        return
<commit_msg>Typo in comment<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
Load incumbent candidate data scraped from the CAL-ACCESS website into OCD models.
"""
import re
from django.utils import timezone
from django.db.models import IntegerField
from django.db.models.functions import Cast
from calaccess_processed.management.commands import LoadOCDModelsCommand
from calaccess_processed.models.scraper import (
    IncumbentScrapedElection,
    ScrapedIncumbent,
)
from opencivicdata.models import Membership
from opencivicdata.models.elections import Election, Candidacy

class Command(LoadOCDModelsCommand):
    """
    Load incumbent candidate data scraped from the CAL-ACCESS website into OCD models.
    """
    help = 'Load incumbent candidate data scraped from the CAL-ACCESS website into OCD models'

    def handle(self, *args, **options):
        """
        Make it happen.
        """
        super(Command, self).handle(*args, **options)
        self.load()
        self.set_end_dates()
        if Candidacy.objects.exists():
            self.log('')
            self.set_incumbent_candidacies()
        self.success("Done!")

    def get_or_create_election(self, scraped_elec):
        """
        Get or create an OCD Election object using the IncumbentScrapedElection.

        Returns a tuple (Election object, created), where created is a boolean
        specifying whether a Election was created.
        """
        dt_obj = timezone.make_aware(
            timezone.datetime.combine(scraped_elec.date, timezone.datetime.min.time())
        )
        name = '{0} {1}'.format(dt_obj.year, scraped_elec.name)
        try:
            elec = Election.objects.get(start_time=dt_obj)
        except Election.DoesNotExist:
            # or make a new one
            elec = self.create_election(name, dt_obj)
            created = True
        else:
            created = False
            # if election already exists and is named 'SPECIAL' or 'RECALL'
            if 'SPECIAL' in elec.name.upper() or 'RECALL' in elec.name.upper():
                # and the matched election's name includes either 'GENERAL'
                # or 'PRIMARY'...
                if (
                    re.match(r'^\d{4} GENERAL$', scraped_elec.name) or
                    re.match(r'^\d{4} PRIMARY$', scraped_elec.name)
                ):
                    # update the name
                    elec.name = name
                    elec.save()
        return (elec, created)

    def load(self):
        """
        Load OCD Election, Membership and related models with data scraped from CAL-ACCESS website.
        """
        for scraped_elec in IncumbentScrapedElection.objects.all():
            ocd_elec = self.get_or_create_election(scraped_elec)[0]
            ocd_elec.sources.update_or_create(
                url=scraped_elec.url,
                note='Last scraped on {dt:%Y-%m-%d}'.format(
                    dt=scraped_elec.last_modified,
                )
            )

        for incumbent in ScrapedIncumbent.objects.all().order_by('-session'):
            # Get or create post
            post, post_created = self.get_or_create_post(
                incumbent.office_name,
            )
            if post_created and self.verbosity > 2:
                self.log('Created new Post: %s' % post.label)
            # Get or person
            person, person_created = self.get_or_create_person(
                incumbent.name,
                filer_id=incumbent.scraped_id,
            )
            if person_created and self.verbosity > 2:
                self.log('Created new Person: %s' % person.name)
            # Get or membership for post and person
            membership, membership_created = Membership.objects.get_or_create(
                person=person,
                post=post,
                role=post.role,
                organization=post.organization,
                person_name=person.sort_name,
            )
            if membership_created and self.verbosity > 2:
                self.log('Created new Membership: %s' % membership)
            # Handle start_date on membership
            if membership_created or membership.start_date == '':
                membership.start_date = incumbent.session
                membership.save()
            else:
                # increment start year down
                start = int(membership.start_date)
                if start > incumbent.session:
                    membership.start_date = incumbent.session
                    membership.save()

        return

    def set_end_dates(self):
        """
        Set the end_date for each Membership based on the start_date of each successor.
        """
        for member in Membership.objects.all():
            # Each member's end year should be the start year of their successor.
            # Successor is the member in the same post with the earliest
            # start year greater than the incumbent's start year
            successor_q = Membership.objects.exclude(
                start_date='',
            ).annotate(
                start_year=Cast('start_date', IntegerField()),
            ).filter(
                start_year__gt=int(member.start_date),
                post=member.post,
            ).order_by('start_year')

            if successor_q.exists():
                member.end_date = int(successor_q[0].start_date)
                member.save()
        return

    def set_incumbent_candidacies(self):
        """
        Set is_incumbent for candidacies within each member's start/end years.
        """
        # For every one of the member's candidacies for the office where
        # the election of the contest happens after the start year
        # but before the end year, if it exists, mark as incumbent
        for member in Membership.objects.all():
            candidacies_q = member.person.candidacies.filter(
                contest__election__start_time__year__gt=int(member.start_date),
                post=member.post,
            ).exclude(is_incumbent=True)
            # Handle blank member end_date values
            if member.end_date != '':
                member_end_year = int(member.end_date)
                candidacies_q = candidacies_q.filter(contest__election__start_time__year__lte=member_end_year)

            if candidacies_q.exists():
                rows = candidacies_q.update(is_incumbent=True)
                self.log(
                    ' {0} identified as incumbent in {1} contests'.format(
                        member.person.name,
                        rows,
                    )
                )
        return
<|endoftext|>