<commit_before>'''
Created on Mar 2, 2013

@author: fryckbos
'''
import logging
LOGGER = logging.getLogger("openmotics")

import time
import traceback
from datetime import datetime
from threading import Thread

import power_api

class TimeKeeper:
    """ The TimeKeeper keeps track of time and sets the day or night mode on the power modules. """

    def __init__(self, power_communicator, power_controller, period):
        self.__power_communicator = power_communicator
        self.__power_controller = power_controller
        self.__period = period

        self.__mode = {}

        self.__thread = None
        self.__stop = False

    def start(self):
        """ Start the background thread of the TimeKeeper. """
        if self.__thread == None:
            LOGGER.info("Starting Power TimeKeeper")
            self.__stop = False
            self.__thread = Thread(target=self.__run, name="TimeKeeper thread")
            self.__thread.daemon = True
            self.__thread.start()
        else:
            raise Exception("TimeKeeper thread already running.")

    def stop(self):
        """ Stop the background thread in the TimeKeeper. """
        if self.__thread != None:
            self.__stop = True
        else:
            raise Exception("TimeKeeper thread not running.")

    def __run(self):
        """ Code for the background thread. """
        while not self.__stop:
            date = datetime.now()
            for module in self.__power_controller.get_power_modules().values():
                daynight = []
                for i in range(8):
                    if self.is_day_time(module['times%d' % i], date):
                        daynight.append(power_api.DAY)
                    else:
                        daynight.append(power_api.NIGHT)

                self.__set_mode(module['address'], daynight)

            time.sleep(self.__period)

        self.__thread = None

    def is_day_time(self, times, date):
        """ Check if a date is in day time. """
        if times == None:
            times = [ 0 for _ in range(14) ]
        else:
            times = map(lambda time: int(time.replace(":", "")), times.split(","))

        day_of_week = date.weekday() # 0 = Monday, 6 = Sunday
        current_time = date.hour * 100 + date.minute

        start = times[day_of_week * 2]
        stop = times[day_of_week * 2 + 1]

        return current_time >= start and current_time < stop

    def __set_mode(self, address, bytes):
        """ Set the power modules mode. """
        try:
            if address not in self.__mode or self.__mode[address] != bytes:
                LOGGER.info("Setting day/night mode to " + str(bytes))
                self.__power_communicator.do_command(address, power_api.set_day_night(), *bytes)
                self.__mode[address] = bytes
        except:
            ## Got an exception, we'll just try again later
            print "Exception while setting mode for power modules."
            traceback.print_exc()
<commit_msg>Improved error handling on TimeKeeper.<commit_after>'''
Created on Mar 2, 2013

@author: fryckbos
'''
import logging
LOGGER = logging.getLogger("openmotics")

import time
import traceback
from datetime import datetime
from threading import Thread

import power_api

class TimeKeeper:
    """ The TimeKeeper keeps track of time and sets the day or night mode on the power modules. """

    def __init__(self, power_communicator, power_controller, period):
        self.__power_communicator = power_communicator
        self.__power_controller = power_controller
        self.__period = period

        self.__mode = {}

        self.__thread = None
        self.__stop = False

    def start(self):
        """ Start the background thread of the TimeKeeper. """
        if self.__thread == None:
            LOGGER.info("Starting TimeKeeper")
            self.__stop = False
            self.__thread = Thread(target=self.__run, name="TimeKeeper thread")
            self.__thread.daemon = True
            self.__thread.start()
        else:
            raise Exception("TimeKeeper thread already running.")

    def stop(self):
        """ Stop the background thread in the TimeKeeper. """
        if self.__thread != None:
            self.__stop = True
        else:
            raise Exception("TimeKeeper thread not running.")

    def __run(self):
        """ Code for the background thread. """
        while not self.__stop:
            try:
                self.__run_once()
            except:
                LOGGER.exception("Exception in TimeKeeper")

            time.sleep(self.__period)

        LOGGER.info("Stopped TimeKeeper")
        self.__thread = None

    def __run_once(self):
        """ One run of the background thread. """
        date = datetime.now()
        for module in self.__power_controller.get_power_modules().values():
            daynight = []
            for i in range(8):
                if self.is_day_time(module['times%d' % i], date):
                    daynight.append(power_api.DAY)
                else:
                    daynight.append(power_api.NIGHT)

            self.__set_mode(module['address'], daynight)

    def is_day_time(self, times, date):
        """ Check if a date is in day time. """
        if times == None:
            times = [ 0 for _ in range(14) ]
        else:
            times = map(lambda time: int(time.replace(":", "")), times.split(","))

        day_of_week = date.weekday() # 0 = Monday, 6 = Sunday
        current_time = date.hour * 100 + date.minute

        start = times[day_of_week * 2]
        stop = times[day_of_week * 2 + 1]

        return current_time >= start and current_time < stop

    def __set_mode(self, address, bytes):
        """ Set the power modules mode. """
        if address not in self.__mode or self.__mode[address] != bytes:
            LOGGER.info("Setting day/night mode to " + str(bytes))
            self.__power_communicator.do_command(address, power_api.set_day_night(), *bytes)
            self.__mode[address] = bytes
<|endoftext|>


<commit_before>#!/usr/bin/env python

# Copyright (c) 2013, Carnegie Mellon University
# All rights reserved.
# Authors: Pras Velagapudi <pkv@cs.cmu.edu>
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# - Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
# - Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
# - Neither the name of Carnegie Mellon University nor the names of its
#   contributors may be used to endorse or promote products derived from this
#   software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
import logging
import time
import itertools
import numpy
import openravepy
from base import (BasePlanner, PlanningMethod, PlanningError,
                  UnsupportedPlanningError)

logger = logging.getLogger('prpy.planning.tsr')

class TSRPlanner(BasePlanner):
    def __init__(self, delegate_planner=None):
        super(TSRPlanner, self).__init__()
        self.delegate_planner = delegate_planner

    def __str__(self):
        return 'TSRPlanner'

    @PlanningMethod
    def PlanToIK(self, robot, goal_pose, **kw_args):
        from ..tsr import TSR, TSRChain

        tsr = TSR(T0_w=goal_pose, Tw_e=numpy.eye(4), Bw=numpy.zeros((6,2)))
        tsr_chain = TSRChain(sample_goal=True, TSR=tsr)

        return self.PlanToTSR(robot, [tsr_chain], **kw_args)

    @PlanningMethod
    def PlanToTSR(self, robot, tsrchains, tsr_timeout=2.0,
                  num_attempts=3, chunk_size=1, ranker=None,
                  max_deviation=2*numpy.pi, **kw_args):
        """
        Plan to a desired TSR set using a-priori goal sampling.  This planner
        samples a fixed number of goals from the specified TSRs up-front, then
        uses another planner's PlanToConfiguration (chunk_size = 1) or
        PlanToConfigurations (chunk_size > 1) to plan to the resulting affine
        transformations.

        This planner will return failure if the provided TSR chains require
        any constraint other than goal sampling.

        @param robot the robot whose active manipulator will be used
        @param tsrchains a list of TSR chains that define a goal set
        @param num_attempts the maximum number of planning attempts to make
        @param chunk_size the number of sampled goals to use per planning call
        @param tsr_timeout the maximum time to spend sampling goal TSR chains
        @param ranker an IK ranking function to use over the IK solutions
        @param max_deviation the maximum per-joint deviation from current pose
                             that can be considered a valid sample.
        @return traj a trajectory that satisfies the specified TSR chains
        """
        # Delegate to robot.planner by default.
        delegate_planner = self.delegate_planner or robot.planner

        # Plan using the active manipulator.
        with robot.GetEnv():
            manipulator = robot.GetActiveManipulator()

            # Distance from current configuration is default ranking.
            if ranker is None:
                from ..ik_ranking import NominalConfiguration
                ranker = NominalConfiguration(manipulator.GetArmDOFValues(),
                                              max_deviation=max_deviation)

        # Test for tsrchains that cannot be handled.
        for tsrchain in tsrchains:
            if tsrchain.sample_start or tsrchain.constrain:
                raise UnsupportedPlanningError(
                    'Cannot handle start or trajectory-wide TSR constraints.')
        tsrchains = [t for t in tsrchains if t.sample_goal]

        # Create an iterator that cycles through each TSR chain.
        tsr_cycler = itertools.cycle(tsrchains)

        # Create an iterator that cycles TSR chains until the timelimit.
        tsr_timelimit = time.time() + tsr_timeout
        tsr_sampler = itertools.takewhile(
            lambda v: time.time() < tsr_timelimit, tsr_cycler)

        # Sample a list of TSR poses and collate valid IK solutions.
        from openravepy import (IkFilterOptions,
                                IkParameterization,
                                IkParameterizationType)
        ik_solutions = []
        for tsrchain in tsr_sampler:
            ik_param = IkParameterization(
                tsrchain.sample(), IkParameterizationType.Transform6D)
            ik_solution = manipulator.FindIKSolutions(
                ik_param, IkFilterOptions.CheckEnvCollisions)
            if ik_solution.shape[0] > 0:
                ik_solutions.append(ik_solution)

        if len(ik_solutions) == 0:
            raise PlanningError('No collision-free IK solutions at goal TSRs.')

        # Sort the IK solutions in ascending order by the costs returned by the
        # ranker. Lower cost solutions are better and infinite cost solutions
        # are assumed to be infeasible.
        ik_solutions = numpy.vstack(ik_solutions)
        scores = ranker(robot, ik_solutions)
        valid_idxs = ~numpy.isposinf(scores)
        valid_scores = scores[valid_idxs]
        valid_solutions = ik_solutions[valid_idxs, :]
        ranked_indices = numpy.argsort(valid_scores)
        ranked_ik_solutions = valid_solutions[ranked_indices, :]

        # Group the IK solutions into sets of the specified size
        # (plan for each set of IK solutions together).
        ranked_ik_solution_sets = [
            ranked_ik_solutions[i:i+chunk_size, :]
            for i in range(0, ranked_ik_solutions.shape[0], chunk_size)
        ]

        num_attempts = min(len(ranked_ik_solution_sets), num_attempts)
        ik_set_list = enumerate(ranked_ik_solution_sets[:num_attempts])

        # Configure the robot to use the active manipulator for planning.
        p = openravepy.KinBody.SaveParameters
        with robot.CreateRobotStateSaver(p.ActiveDOF):
            robot.SetActiveDOFs(manipulator.GetArmIndices())

            # Try planning to each solution set in descending cost order.
            for i, ik_set in ik_set_list:
                try:
                    if ik_set.shape[0] > 1:
                        traj = delegate_planner.PlanToConfigurations(
                            robot, ik_set)
                    else:
                        traj = delegate_planner.PlanToConfiguration(
                            robot, ik_set[0])

                    logger.info('Planned to IK solution set %d of %d.',
                                i + 1, num_attempts)
                    return traj
                except PlanningError as e:
                    logger.warning(
                        'Planning to IK solution set %d of %d failed: %s',
                        i + 1, num_attempts, e)

        # If none of the planning attempts succeeded, report failure.
        raise PlanningError(
            'Planning to the top {:d} of {:d} IK solution sets failed.'
            .format(num_attempts, len(ranked_ik_solution_sets)))
<commit_msg>Disabled PlanToIK on TSRPlanner.<commit_after>#!/usr/bin/env python

# Copyright (c) 2013, Carnegie Mellon University
# All rights reserved.
# Authors: Pras Velagapudi <pkv@cs.cmu.edu>
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# - Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
# - Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
# - Neither the name of Carnegie Mellon University nor the names of its
#   contributors may be used to endorse or promote products derived from this
#   software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
import logging
import time
import itertools
import numpy
import openravepy
from base import (BasePlanner, PlanningMethod, PlanningError,
                  UnsupportedPlanningError)

logger = logging.getLogger('prpy.planning.tsr')

class TSRPlanner(BasePlanner):
    def __init__(self, delegate_planner=None):
        super(TSRPlanner, self).__init__()
        self.delegate_planner = delegate_planner

    def __str__(self):
        return 'TSRPlanner'

    # TODO: Temporarily disabled based on Pras' feedback.
    """
    @PlanningMethod
    def PlanToIK(self, robot, goal_pose, **kw_args):
        from ..tsr import TSR, TSRChain

        tsr = TSR(T0_w=goal_pose, Tw_e=numpy.eye(4), Bw=numpy.zeros((6,2)))
        tsr_chain = TSRChain(sample_goal=True, TSR=tsr)

        return self.PlanToTSR(robot, [tsr_chain], **kw_args)
    """

    @PlanningMethod
    def PlanToTSR(self, robot, tsrchains, tsr_timeout=2.0,
                  num_attempts=3, chunk_size=1, ranker=None,
                  max_deviation=2*numpy.pi, **kw_args):
        """
        Plan to a desired TSR set using a-priori goal sampling.  This planner
        samples a fixed number of goals from the specified TSRs up-front, then
        uses another planner's PlanToConfiguration (chunk_size = 1) or
        PlanToConfigurations (chunk_size > 1) to plan to the resulting affine
        transformations.

        This planner will return failure if the provided TSR chains require
        any constraint other than goal sampling.

        @param robot the robot whose active manipulator will be used
        @param tsrchains a list of TSR chains that define a goal set
        @param num_attempts the maximum number of planning attempts to make
        @param chunk_size the number of sampled goals to use per planning call
        @param tsr_timeout the maximum time to spend sampling goal TSR chains
        @param ranker an IK ranking function to use over the IK solutions
        @param max_deviation the maximum per-joint deviation from current pose
                             that can be considered a valid sample.
        @return traj a trajectory that satisfies the specified TSR chains
        """
        # Delegate to robot.planner by default.
        delegate_planner = self.delegate_planner or robot.planner

        # Plan using the active manipulator.
        with robot.GetEnv():
            manipulator = robot.GetActiveManipulator()

            # Distance from current configuration is default ranking.
            if ranker is None:
                from ..ik_ranking import NominalConfiguration
                ranker = NominalConfiguration(manipulator.GetArmDOFValues(),
                                              max_deviation=max_deviation)

        # Test for tsrchains that cannot be handled.
        for tsrchain in tsrchains:
            if tsrchain.sample_start or tsrchain.constrain:
                raise UnsupportedPlanningError(
                    'Cannot handle start or trajectory-wide TSR constraints.')
        tsrchains = [t for t in tsrchains if t.sample_goal]

        # Create an iterator that cycles through each TSR chain.
        tsr_cycler = itertools.cycle(tsrchains)

        # Create an iterator that cycles TSR chains until the timelimit.
        tsr_timelimit = time.time() + tsr_timeout
        tsr_sampler = itertools.takewhile(
            lambda v: time.time() < tsr_timelimit, tsr_cycler)

        # Sample a list of TSR poses and collate valid IK solutions.
        from openravepy import (IkFilterOptions,
                                IkParameterization,
                                IkParameterizationType)
        ik_solutions = []
        for tsrchain in tsr_sampler:
            ik_param = IkParameterization(
                tsrchain.sample(), IkParameterizationType.Transform6D)
            ik_solution = manipulator.FindIKSolutions(
                ik_param, IkFilterOptions.CheckEnvCollisions)
            if ik_solution.shape[0] > 0:
                ik_solutions.append(ik_solution)

        if len(ik_solutions) == 0:
            raise PlanningError('No collision-free IK solutions at goal TSRs.')

        # Sort the IK solutions in ascending order by the costs returned by the
        # ranker. Lower cost solutions are better and infinite cost solutions
        # are assumed to be infeasible.
        ik_solutions = numpy.vstack(ik_solutions)
        scores = ranker(robot, ik_solutions)
        valid_idxs = ~numpy.isposinf(scores)
        valid_scores = scores[valid_idxs]
        valid_solutions = ik_solutions[valid_idxs, :]
        ranked_indices = numpy.argsort(valid_scores)
        ranked_ik_solutions = valid_solutions[ranked_indices, :]

        # Group the IK solutions into sets of the specified size
        # (plan for each set of IK solutions together).
        ranked_ik_solution_sets = [
            ranked_ik_solutions[i:i+chunk_size, :]
            for i in range(0, ranked_ik_solutions.shape[0], chunk_size)
        ]

        num_attempts = min(len(ranked_ik_solution_sets), num_attempts)
        ik_set_list = enumerate(ranked_ik_solution_sets[:num_attempts])

        # Configure the robot to use the active manipulator for planning.
        p = openravepy.KinBody.SaveParameters
        with robot.CreateRobotStateSaver(p.ActiveDOF):
            robot.SetActiveDOFs(manipulator.GetArmIndices())

            # Try planning to each solution set in descending cost order.
            for i, ik_set in ik_set_list:
                try:
                    if ik_set.shape[0] > 1:
                        traj = delegate_planner.PlanToConfigurations(
                            robot, ik_set)
                    else:
                        traj = delegate_planner.PlanToConfiguration(
                            robot, ik_set[0])

                    logger.info('Planned to IK solution set %d of %d.',
                                i + 1, num_attempts)
                    return traj
                except PlanningError as e:
                    logger.warning(
                        'Planning to IK solution set %d of %d failed: %s',
                        i + 1, num_attempts, e)

        # If none of the planning attempts succeeded, report failure.
        raise PlanningError(
            'Planning to the top {:d} of {:d} IK solution sets failed.'
            .format(num_attempts, len(ranked_ik_solution_sets)))
<|endoftext|>


<commit_before>
from .publisher import Publisher
from .models import ModelPublisher
<commit_msg>Sort imports<commit_after>
from .models import ModelPublisher
from .publisher import Publisher
<|endoftext|>


<commit_before>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import numpy as np
import scipy.linalg as linalg
import sympy as sym
import itertools as it

from globalopts import opts as gopts

import settings
import linprog.linprog as linprog

from IPython import embed

import err
import utils as U

class SolverFaliure(Exception):
    pass

def part_constraints(partition_trace):
    """constraints due to partitions of the pwa model

    Parameters
    ----------
    pwa_trace : pwa trace form the bmc module

    Returns
    -------
    cons corresponding to Cx - d <= 0

    Notes
    ------
    """
    C_, d_ = [], []
    # for each sub_model
    for p in partition_trace:
        C_.append(p.C)
        d_.append(p.d)

    C = linalg.block_diag(*C_)
    d = np.hstack(d_)

    assert(C.shape[0] == d.shape[0])

    return C, d

def dyn_constraints(models, vars_grouped_by_states):
    """constraints due to dynamics of the pwa model

    Parameters
    ----------
    models : list of pwa models
    vars_grouped_by_states: variables for state vectors grouped
    together. e.g. for a 2dim/state vector (x1,x2):

        k=0 (init)     k=1       k=2
        [(v0, v1),  (v2, v3), (v4, v5), ...]

    Returns
    -------
    polynomial constraints corresponding to x' \in p(x) + [error.l, error.h]

    Notes
    ------
    """

    # for each sub_model
    # x' >= p(x) + e.l
    # x' <= p(x) + e.h

    dyn_cons = []

    for m, (Vars, next_vars) in zip(models, U.pairwise(vars_grouped_by_states)):
        for p, el, eh, x, x_ in zip(m.poly, m.error.l, m.error.h, Vars, next_vars):
            p.truncate_coeffs(gopts.bmc_prec)
            #assert(len(p.vars) == ndimx)
            old2new_var_map = {v: v_ for v, v_ in zip(p.vars, Vars)}
            tpoly = p.subs_vars(old2new_var_map)
            # xi' <= p(x) + ehi
            cons_ub = x_ - (tpoly.as_expr() + eh)
            # xi' >= p(x) + eli
            cons_lb = tpoly.as_expr() + el - x_
            dyn_cons.append(cons_ub <= 0)
            dyn_cons.append(cons_lb <= 0)

    return dyn_cons

def param_constraints(num_dims, prop, num_partitions):
    """constraints due to the initial set, final set and ci/pi

    Parameters
    ----------
    prop :

    Returns
    -------

    Notes
    ------
    """
    #trace_len = len(pwa_trace)
    iA, ib = prop.init_cons.poly()
    fA, fb = prop.final_cons.poly()

    assert(iA.shape == fA.shape)

    # find the dim(x) using any submodel's b vector
    # This is the dim of total variables: dimX = dimW
    #dimX = pwa_trace[0].m.b.size

    dimX = num_dims.x + num_dims.pi
    dimW = num_dims.pi

    # pad iA, ib, fA, fb with 0's to accomodate w/pi
    # Remember: the state vector for pwa is: [x]
    #                                        [w]
    # Of course this much code is not required, and the intended
    # operation can be done much succintly. But this is much easier to
    # understand IMO.

    # append 0s for wi: A -> [A 0 .. 0]
    padding_scheme = ((0, 0), (0, dimW))

    iA = np.pad(iA, padding_scheme, 'constant')
    fA = np.pad(fA, padding_scheme, 'constant')

    # number of rows and coloumns of iA/fA
    nr, nc = iA.shape
    num_cons = 2 * nr

    # Add 1, because trace captures transitions,
    # and num_states = num_trans + 1
    #A = np.zeros((num_cons, dimX * (trace_len)))
    A = np.zeros((num_cons, dimX * (num_partitions)))
    A[0:nr, 0:nc] = iA
    A[-nr:, -nc:] = fA
    b = np.hstack((ib, fb))
    #print(A)
    #print(b)
    # Each constraint expression in A has a value in b
    assert(A.shape[0] == b.size)
    return A, b

def truncate(*args):
    assert(isinstance(gopts.bmc_prec, int))
    prec = gopts.bmc_prec

    # Round off to the same amount as the bmc query
    def arr2str(n): return '{n:0.{p}f}'.format(n=n, p=prec)
    trunc_array = np.vectorize(lambda x: np.float(arr2str(x)))

    return (trunc_array(X) for X in args)

def pwatrace2cons(pwa_trace, num_dims, prop):

    # find the dim(x) using any submodel's b vector
    m = pwa_trace.models[0]
    ndimx = m.error.l.size

    nvars = (len(pwa_trace.models)+1) * ndimx
    all_vars = sym.var(','.join(('x{}'.format(i) for i in range(nvars))))
    vars_grouped_by_state = zip(*[all_vars[i::ndimx] for i in range(ndimx)])

    C, d = part_constraints(pwa_trace.partitions)
    dyn_cons = dyn_constraints(pwa_trace.models, vars_grouped_by_state)

    P, q = param_constraints(num_dims, prop, len(pwa_trace.partitions))

    C, d = truncate(C, d)
    P, q = truncate(P, q)
    part_cons_ = np.dot(C, all_vars) - d
    part_cons = (c <= 0 for c in part_cons_)

    param_cons_ = np.dot(P, all_vars) - q
    param_cons = (c <= 0 for c in param_cons_)

    all_cons = it.chain(part_cons, dyn_cons, param_cons)
    #all_cons = (part_cons.tolist() + dyn_cons + param_cons.tolist())

    return all_cons, all_vars

def optsoln2x(x, trace_len):
    '''converts solution of lp: a vector to a concrete trace: numpy
    array'''
    x = np.array(x)
    assert(x.ndim == 1)
    num_vars = len(x)//trace_len
    return np.reshape(x, (trace_len, num_vars))

def feasible(num_dims, prop, pwa_trace, solver=gopts.opt_engine):
    import nonlinprog.z3opt as z3opt
    import nonlinprog.scipyopt as scipyopt
    cons, Vars = pwatrace2cons(pwa_trace, num_dims, prop)
    cons = list(cons)

    #nvars = num_dims.x + num_dims.pi

    #A_ub, b_ub = truncate(A_ub, b_ub)
    obj = 0

    #err.warn_severe('faking output of optimizer')
    #res = True
    #varval_map = {v: 0 for v in Vars}
    # TODO: Make choice of opt engine
    #res, varval_map = z3opt.nlinprog(obj, cons, Vars)
    res, varval_map = scipyopt.nlinprog(obj, cons, Vars)

    ret_val = optsoln2x([varval_map[v] for v in Vars], len(pwa_trace)) if res else None
    print(ret_val)
    #embed()
    return ret_val

    #return lpsoln2x(varval_map, len(pwa_trace)) if res else None

def overapprox_x0(num_dims, prop, pwa_trace, solver=gopts.opt_engine):
    raise NotImplementedError
<commit_msg>added back overapproxx0<commit_after>from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

import numpy as np
import scipy.linalg as linalg
import sympy as sym
import itertools as it

from globalopts import opts as gopts

import settings
import nonlinprog.nonlinprog as nonlinprog

from IPython import embed

import err
import utils as U
from constraints import IntervalCons

class SolverFaliure(Exception):
    pass

def part_constraints(partition_trace):
    """constraints due to partitions of the pwa model

    Parameters
    ----------
    pwa_trace : pwa trace form the bmc module

    Returns
    -------
    cons corresponding to Cx - d <= 0

    Notes
    ------
    """
    C_, d_ = [], []
    # for each sub_model
    for p in partition_trace:
        C_.append(p.C)
        d_.append(p.d)

    C = linalg.block_diag(*C_)
    d = np.hstack(d_)

    assert(C.shape[0] == d.shape[0])

    return C, d

def dyn_constraints(models, vars_grouped_by_state):
    """constraints due to dynamics of the pwa model

    Parameters
    ----------
    models : list of pwa models
    vars_grouped_by_state: variables for state vectors grouped
    together. e.g. for a 2dim/state vector (x1,x2):

        k=0 (init)     k=1       k=2
        [(v0, v1),  (v2, v3), (v4, v5), ...]

    Returns
    -------
    polynomial constraints corresponding to x' \in p(x) + [error.l, error.h]

    Notes
    ------
    """

    # for each sub_model
    # x' >= p(x) + e.l
    # x' <= p(x) + e.h

    dyn_cons = []

    for m, (Vars, next_vars) in zip(models, U.pairwise(vars_grouped_by_state)):
        for p, el, eh, x, x_ in zip(m.poly, m.error.l, m.error.h, Vars, next_vars):
            p.truncate_coeffs(gopts.bmc_prec)
            #assert(len(p.vars) == ndimx)
            old2new_var_map = {v: v_ for v, v_ in zip(p.vars, Vars)}
            tpoly = p.subs_vars(old2new_var_map)
            # xi' <= p(x) + ehi
            cons_ub = x_ - (tpoly.as_expr() + eh)
            # xi' >= p(x) + eli
            cons_lb = tpoly.as_expr() + el - x_
            dyn_cons.append(cons_ub <= 0)
            dyn_cons.append(cons_lb <= 0)

    return dyn_cons

def param_constraints(num_dims, prop, num_partitions):
    """constraints due to the initial set, final set and ci/pi

    Parameters
    ----------
    prop :

    Returns
    -------

    Notes
    ------
    """
    #trace_len = len(pwa_trace)
    iA, ib = prop.init_cons.poly()
    fA, fb = prop.final_cons.poly()

    assert(iA.shape == fA.shape)

    # find the dim(x) using any submodel's b vector
    # This is the dim of total variables: dimX = dimW
    #dimX = pwa_trace[0].m.b.size

    dimX = num_dims.x + num_dims.pi
    dimW = num_dims.pi

    # pad iA, ib, fA, fb with 0's to accomodate w/pi
    # Remember: the state vector for pwa is: [x]
    #                                        [w]
    # Of course this much code is not required, and the intended
    # operation can be done much succintly. But this is much easier to
    # understand IMO.

    # append 0s for wi: A -> [A 0 .. 0]
    padding_scheme = ((0, 0), (0, dimW))

    iA = np.pad(iA, padding_scheme, 'constant')
    fA = np.pad(fA, padding_scheme, 'constant')

    # number of rows and coloumns of iA/fA
    nr, nc = iA.shape
    num_cons = 2 * nr

    # Add 1, because trace captures transitions,
    # and num_states = num_trans + 1
    #A = np.zeros((num_cons, dimX * (trace_len)))
    A = np.zeros((num_cons, dimX * (num_partitions)))
    A[0:nr, 0:nc] = iA
    A[-nr:, -nc:] = fA
    b = np.hstack((ib, fb))
    #print(A)
    #print(b)
    # Each constraint expression in A has a value in b
    assert(A.shape[0] == b.size)
    return A, b

def truncate(*args):
    assert(isinstance(gopts.bmc_prec, int))
    prec = gopts.bmc_prec

    # Round off to the same amount as the bmc query
    def arr2str(n): return '{n:0.{p}f}'.format(n=n, p=prec)
    trunc_array = np.vectorize(lambda x: np.float(arr2str(x)))

    return (trunc_array(X) for X in args)

def pwatrace2cons(pwa_trace, num_dims, prop):
    """pwatrace2cons

    Parameters
    ----------
    pwa_trace :
    num_dims :
    prop :

    Returns
    -------
    all_cons: All constraints in the form g(x) <= 0
    all_vars: List of vars, ordered against the trace
    vars_grouped_by_state:

    Notes
    ------
    """

    # find the dim(x) using any submodel's b vector
    m = pwa_trace.models[0]
    ndimx = m.error.l.size

    nvars = (len(pwa_trace.models)+1) * ndimx
    all_vars = sym.var(','.join(('x{}'.format(i) for i in range(nvars))))
    vars_grouped_by_state = zip(*[all_vars[i::ndimx] for i in range(ndimx)])

    C, d = part_constraints(pwa_trace.partitions)
    dyn_cons = dyn_constraints(pwa_trace.models, vars_grouped_by_state)

    P, q = param_constraints(num_dims, prop, len(pwa_trace.partitions))

    C, d = truncate(C, d)
    P, q = truncate(P, q)
    part_cons_ = np.dot(C, all_vars) - d
    part_cons = (c <= 0 for c in part_cons_)

    param_cons_ = np.dot(P, all_vars) - q
    param_cons = (c <= 0 for c in param_cons_)

    all_cons = it.chain(part_cons, dyn_cons, param_cons)
    #all_cons = (part_cons.tolist() + dyn_cons + param_cons.tolist())

    return all_cons, all_vars, vars_grouped_by_state

def optsoln2x(x, trace_len):
    '''converts solution of lp: a vector to a concrete trace: numpy
    array'''
    x = np.array(x)
    assert(x.ndim == 1)
    num_vars = len(x)//trace_len
    return np.reshape(x, (trace_len, num_vars))

def feasible(num_dims, prop, pwa_trace, solver=gopts.opt_engine):
    cons, Vars, vars_grouped_by_state = pwatrace2cons(pwa_trace, num_dims, prop)
    cons = list(cons)

    #nvars = num_dims.x + num_dims.pi

    #A_ub, b_ub = truncate(A_ub, b_ub)
    def return_zero(x):
        return 0
    obj = return_zero

    #err.warn_severe('faking output of optimizer')
    #res = True
    #varval_map = {v: 0 for v in Vars}
    # TODO: Make choice of opt engine
    #res, varval_map = z3opt.nlinprog(obj, cons, Vars)

    #res, varval_map = nlpfun(solver)(obj, cons, Vars)
    #ret_val = optsoln2x([varval_map[v] for v in Vars], len(pwa_trace)) if res else None
    res = nlpfun(solver)(obj, cons, Vars)
    ret_val = optsoln2x(res.x, len(pwa_trace)) if res.success else None
    print(ret_val)
    #embed()
    return ret_val

    #return lpsoln2x(varval_map, len(pwa_trace)) if res else None

#TODO: getting lpsolver every time a linprog is executed!
#@memoize
def nlpfun(solver):
    return nonlinprog.factory(solver)

#TODO: ugly function
def overapprox_x0(num_dims, prop, pwa_trace, solver=gopts.opt_engine):
    cons, Vars, vars_grouped_by_state = pwatrace2cons(pwa_trace, num_dims, prop)
    cons = list(cons)

    num_opt_vars = len(Vars)
    nvars = num_dims.x + num_dims.pi

    #directions = [(1, 0), (-1, 0), (0, 1), (0, -1)]
    I = np.eye(nvars)
    directions = np.vstack((I, -I))
    left_over_vars = num_opt_vars - nvars
    directions_ext = np.pad(directions, [(0, 0), (0, left_over_vars)], 'constant')

    var_array = np.array(Vars)
    lambdafied = tuple(
            sym.lambdify(Vars, np.dot(direction, var_array), str('numpy')) for direction in directions_ext)
    obj_f = tuple(lambda x, l=l: l(*x) for l in lambdafied)

    x_arr = np.array(
            sym.symbols(
                ['x{}'.format(i) for i in range(nvars)]
                ))

    res = solve_mult_opt(nlpfun(solver), obj_f, cons, Vars)

    l = len(res)
    assert(l % 2 == 0)
    min_res, max_res = res[:l//2], res[l//2:]

    ranges = []

    for di, rmin, rmax in zip(directions, min_res, max_res):
        if (rmin.success and rmax.success):
            print('{} \in [{}, {}]'.format(np.dot(di, x_arr), rmin.fun, -rmax.fun))
            ranges.append([rmin.fun, -rmax.fun])
        else:
            if settings.debug:
                print('LP failed')
                print('rmin status:', rmin.status)
                print('rmax status:', rmax.status)
            return None

    r = np.asarray(ranges)
    # due to numerical errors, the interval can be malformed
    try:
        ret_val = IntervalCons(r[:, 0], r[:, 1])
    except ValueError:
        err.warn('linprog fp failure: Malformed Interval! Please fix.')
        return None

    return ret_val

# raise an exception as soon as the first lp fails...better than
# solving all directions when the problem in infeasible
def solve_mult_opt(nlp_fun, directions_ext, cons, Vars):
    res = []
    for obj in directions_ext:
        # A_ub and b_ub do not change in the for loop...can do a warm
        # restart? or some kind of caching?
        ret_val = nlp_fun(obj, cons, Vars)
        if not ret_val.success:
            raise RuntimeError('lp solver failed')
        else:
            res.append(ret_val)
    return res
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

"""

Node-Link JSON
--------------
This module contains IO functions for interconversion between BEL graphs and Node-Link JSON

"""

import json

from networkx.readwrite.json_graph import node_link_data, node_link_graph

from .utils import ensure_version
from ..constants import GRAPH_ANNOTATION_LIST
from ..struct import BELGraph
from ..utils import list2tuple

__all__ = [
    'to_json',
    'to_json_file',
    'to_jsons',
    'from_json',
    'from_json_file',
    'from_jsons',
]

def to_json(graph):
    """Converts this graph to a Node-Link JSON object

    :param BELGraph graph: A BEL graph
    :return: A Node-Link JSON object representing the given graph
    :rtype: dict
    """
    graph_json_dict = node_link_data(graph)
    graph_json_dict['graph'][GRAPH_ANNOTATION_LIST] = {
        k: list(sorted(v))
        for k, v in graph_json_dict['graph'].get(GRAPH_ANNOTATION_LIST, {}).items()
    }
    return graph_json_dict

def to_json_file(graph, file):
    """Writes this graph as Node-Link JSON to a file

    :param BELGraph graph: A BEL graph
    :param file file: A write-supporting file or file-like object
    """
    graph_json_dict = to_json(graph)
    json.dump(graph_json_dict, file, ensure_ascii=False)

def to_jsons(graph):
    """Dumps this graph as a Node-Link JSON object to a string

    :param BELGraph graph: A BEL graph
    :return: A string representation of the Node-Link JSON produced for this graph by :func:`pybel.to_json`
    :rtype: str
    """
    graph_json_str = to_json(graph)
    return json.dumps(graph_json_str, ensure_ascii=False)

def from_json(graph_json_dict, check_version=True):
    """Builds a graph from Node-Link JSON Object

    :param dict graph_json_dict: A JSON dictionary representing a graph
    :param bool check_version: Checks if the graph was produced by this version of PyBEL
    :return: A BEL graph
    :rtype: BELGraph
    """
    for i, node in enumerate(graph_json_dict['nodes']):
        graph_json_dict['nodes'][i]['id'] = list2tuple(graph_json_dict['nodes'][i]['id'])

    graph = node_link_graph(graph_json_dict, directed=True, multigraph=True)
    graph = BELGraph(data=graph)
    return ensure_version(graph, check_version=check_version)

def from_json_file(file, check_version=True):
    """Builds a graph from the Node-Link JSON contained in the given file

    :param file file: A readable file or file-like
    :param bool check_version: Checks if the graph was produced by this version of PyBEL
    :return: A BEL graph
    :rtype: BELGraph
    """
    graph_json_dict = json.load(file)
    return from_json(graph_json_dict, check_version=check_version)

def from_jsons(graph_json_str, check_version=True):
    """Reads a BEL graph from a Node-Link JSON string

    :param str graph_json_str: A Node-Link JSON string produced by PyBEL
    :param bool check_version: Checks if the graph was produced by this version of PyBEL
    :return: A BEL graph
    :rtype: BELGraph
    """
    graph_json_dict = json.loads(graph_json_str)
    return from_json(graph_json_dict, check_version=check_version)
<commit_msg>Remove default getter<commit_after># -*- coding: utf-8 -*-

"""

Node-Link JSON
--------------
This module contains IO functions for interconversion between BEL graphs and Node-Link JSON

"""

import json

from networkx.readwrite.json_graph import node_link_data, node_link_graph

from .utils import ensure_version
from ..constants import GRAPH_ANNOTATION_LIST
from ..struct import BELGraph
from ..utils import list2tuple

__all__ = [
    'to_json',
    'to_json_file',
    'to_jsons',
    'from_json',
    'from_json_file',
    'from_jsons',
]

def to_json(graph):
    """Converts this graph to a Node-Link JSON object

    :param BELGraph graph: A BEL graph
    :return: A Node-Link JSON object representing the given graph
    :rtype: dict
    """
    graph_json_dict = node_link_data(graph)
    graph_json_dict['graph'][GRAPH_ANNOTATION_LIST] = {
        k: list(sorted(v))
        for k, v in graph_json_dict['graph'][GRAPH_ANNOTATION_LIST].items()
    }
    return graph_json_dict

def to_json_file(graph, file):
    """Writes this graph as Node-Link JSON to a file

    :param BELGraph graph: A BEL graph
    :param file file: A write-supporting file or file-like object
    """
    graph_json_dict = to_json(graph)
    json.dump(graph_json_dict, file, ensure_ascii=False)

def to_jsons(graph):
    """Dumps this graph as a Node-Link JSON object to a string

    :param BELGraph graph: A BEL graph
    :return: A string representation of the Node-Link JSON produced for this graph by :func:`pybel.to_json`
    :rtype: str
    """
    graph_json_str = to_json(graph)
    return json.dumps(graph_json_str, ensure_ascii=False)

def from_json(graph_json_dict, check_version=True):
    """Builds a graph from Node-Link JSON Object

    :param dict graph_json_dict: A JSON dictionary representing a graph
    :param bool check_version: Checks if the graph was produced by this version of PyBEL
    :return: A BEL graph
    :rtype: BELGraph
    """
    for i, node in enumerate(graph_json_dict['nodes']):
        graph_json_dict['nodes'][i]['id'] = list2tuple(graph_json_dict['nodes'][i]['id'])

    graph = node_link_graph(graph_json_dict, directed=True, multigraph=True)
    graph = BELGraph(data=graph)
    return ensure_version(graph, check_version=check_version)

def from_json_file(file, check_version=True):
    """Builds a graph from the Node-Link JSON contained in the given file

    :param file file: A readable file or file-like
    :param bool check_version: Checks if the graph was produced by this version of PyBEL
    :return: A BEL graph
    :rtype: BELGraph
    """
    graph_json_dict = json.load(file)
    return from_json(graph_json_dict, check_version=check_version)

def from_jsons(graph_json_str, check_version=True):
    """Reads a BEL graph from a Node-Link JSON string

    :param str graph_json_str: A Node-Link JSON string produced by PyBEL
    :param bool check_version: Checks if the graph was produced by this version of PyBEL
    :return: A BEL graph
    :rtype: BELGraph
    """
    graph_json_dict = json.loads(graph_json_str)
    return from_json(graph_json_dict, check_version=check_version)
<|endoftext|>


<commit_before>#!/usr/bin/env python3

import abc
import argparse
import json
import logging
import logging.config
import logging.handlers
import os
import socket
import sys
import typing

import dns.name
import dns.rdataclass
import dns.rdatatype
import dns.rdtypes.IN.A
import dns.rdtypes.IN.AAAA
import dns.resolver
import dns.tsigkeyring
import dns.update
import netifaces

class Platform(metaclass=abc.ABCMeta):
    """
    Encapsulates knowledge about a specific operating system.
    """

    @abc.abstractmethod
    def getName(self) -> str:
        """Return the name as shown in os.name."""
        pass

    @abc.abstractmethod
    def getPermanentIPv6Addresses(self, addresses: typing.List[str]) -> typing.List[str]:
        """
        Return only the IPv6 addresses that are permanent (i.e. not generated
        by RFC4941 privacy extensions).

        For IPv6, most NICs will have multiple addresses: one permanent and a
        bunch of temporary generated by SLAAC privacy extensions. We should
        include only the permanent address, not the temporary ones.

        This is platform-specific because netifaces does not return any
        metadata about the address indicating whether it is permanent or
        temporary (and the addresses themselves are indistinguishable).
        However, specific operating systems appear to have a consistent
        ordering as to whether permanent addresses are returned before or after
        temporary addresses.
        """
        pass

    @abc.abstractmethod
    def getDefaultConfigFilename(self) -> str:
        """Return the default location of the configuration file."""
        pass

    @abc.abstractmethod
    def getDefaultCacheFilename(self) -> str:
        """Return the default location of the cache file."""
        pass

    @abc.abstractmethod
    def platformSpecificSetup(self) -> None:
        """
        Do any work that is specific to this platform for initializing the
        program.
        """
        pass

class POSIXPlatform(Platform):
    def getName(self) -> str:
        return "posix"

    def getPermanentIPv6Addresses(self, addresses: typing.List[str]) -> typing.List[str]:
        # Linux returns permanent addresses last.
        if len(addresses) == 0:
            return []
        else:
            return [addresses[-1]]

    def getDefaultConfigFilename(self) -> str:
        return "/etc/pydyndns.conf"

    def getDefaultCacheFilename(self) -> str:
        return "/run/pydyndns.cache"

    def platformSpecificSetup(self) -> None:
        pass

class WindowsPlatform(Platform):
    def getName(self) -> str:
        return "nt"

    def getPermanentIPv6Addresses(self, addresses: typing.List[str]) -> typing.List[str]:
        # Windows returns permanent addresses first.
        if len(addresses) == 0:
            return []
        else:
            return [addresses[0]]

    def getDefaultConfigFilename(self) -> str:
        return os.path.join(os.path.dirname(os.path.abspath(__file__)), "pydyndns.conf")

    def getDefaultCacheFilename(self) -> str:
        localAppData = os.environ.get("LOCALAPPDATA")
        if localAppData is None:
            localAppData = os.path.join(os.path.expanduser("~"), "AppData", "Local")
        return os.path.join(localAppData, "Temp", "pydyndns.cache")

    def platformSpecificSetup(self) -> None:
        # Python’s NTEventLogHandler class unconditionally tries to add the
        # event source to the Windows registry. This fails when running as a
        # low-privileged account. If somebody had already added the event
        # source then logging events using that source could work, but
        # NTEventLogHandler doesn’t bother trying that, it just fails in
        # construction if the registration fails.  Work around this by
        # swallowing exceptions during source registration.
        try:
            import pywintypes
            import win32evtlogutil
            oldAddSourceToRegistry = win32evtlogutil.AddSourceToRegistry
            def replacement(appname: str, dllname: str, logtype: str) -> None:
                try:
                    oldAddSourceToRegistry(appname, dllname, logtype)
                except pywintypes.error:
                    pass
            win32evtlogutil.AddSourceToRegistry = replacement
        except ImportError as exp:
            # Guess Win32 extensions are not installed.
            pass

class UnknownPlatform(Platform):
    def getName(self) -> str:
        return "unknown"

    def getPermanentIPv6Addresses(self, addresses: typing.List[str]) -> typing.List[str]:
        # No idea what the convention is on this platform, so just return all
        # of them.
        return addresses

    def getDefaultConfigFilename(self) -> str:
        return "pydyndns.conf"

    def getDefaultCacheFilename(self) -> str:
        return "pydyndns.cache"

    def platformSpecificSetup(self) -> None:
        pass

class Family(metaclass=abc.ABCMeta):
    """
    Encapsulates knowledge about a specific address family.
    """

    @abc.abstractmethod
    def getName(self) -> str:
        """Return the name used as a cache key for addresses in this family."""
        pass

    @abc.abstractmethod
    def getNetIFacesConstant(self) -> int:
        """Return the numeric ID used as a key in netifaces’ output."""
        pass

    @abc.abstractmethod
    def addAddressToUpdate(self, update: dns.update.Update, hostPart: dns.name.Name, ttl: int, address: str) -> None:
        """Add an address in this family to a DNS update request."""
        pass

    @abc.abstractmethod
    def filterAddressList(self, addresses: typing.Iterable[str]) -> typing.List[str]:
        """
        Return only those addresses that are useful, e.g. not loopback,
        link-local, temporary, or other special addresses that should not be
        registered.
        """
        pass

class IPv4(Family):
    def getName(self) -> str:
        return "ipv4"

    def getNetIFacesConstant(self) -> int:
        return netifaces.AF_INET

    def addAddressToUpdate(self, update: dns.update.Update, hostPart: dns.name.Name, ttl: int, address: str) -> None:
        update.add(hostPart, ttl, dns.rdtypes.IN.A.A(dns.rdataclass.IN, dns.rdatatype.A, address))

    def filterAddressList(self, addresses: typing.Iterable[str]) -> typing.List[str]:
        # For IPv4 most NICs have only one address. It’s not clear that there
        # are any specific rules about how multiple addresses ought to be
        # handled. Just include all of them that are acceptable.
        return [x for x in addresses if self.includeAddress(x)]

    def includeAddress(self, address: str) -> bool:
        parts = [int(part) for part in address.split(".")]
        if parts[0] == 127:
            return False # Loopback address
        elif parts[0] >= 240:
            return False # Multicast or reserved address
        return True

class IPv6(Family):
    def __init__(self, platform: Platform, config: typing.Mapping[typing.Any, typing.Any]):
        self._platform = platform
        self._config = config

    def getName(self) -> str:
        return "ipv6"

    def getNetIFacesConstant(self) -> int:
        return netifaces.AF_INET6

    def addAddressToUpdate(self, update: dns.update.Update, hostPart: dns.name.Name, ttl: int, address: str) -> None:
        update.add(hostPart, ttl, dns.rdtypes.IN.AAAA.AAAA(dns.rdataclass.IN, dns.rdatatype.AAAA, address))

    def filterAddressList(self, addresses: typing.Iterable[str]) -> typing.List[str]:
        return self._platform.getPermanentIPv6Addresses([x for x in addresses if self.includeAddress(x)])

    def includeAddress(self, address: str) -> bool:
        first_word = int(address.split(":")[0] or "0", 16)
        second_word = int(address.split(":")[1] or "0", 16)
        if first_word == 0x0000:
            return False # Unspecified, local, or IPv6-mapped address
        elif first_word == 0x0100:
            return False # Discard address
        elif (first_word & 0xFE00) == 0xFC00:
            return False # Unique local address
        elif first_word == 0xFE80:
            return False # Link-local address
        elif (first_word & 0xFF00) == 0xFF00:
            return False # Multicast address
        elif ((first_word == 0x2001) and (second_word == 0x0000)) and not self._config["teredo"]:
            return False # Teredo address
        return True

def run(platform: Platform, args: argparse.Namespace, config: typing.Mapping[typing.Any, typing.Any], logger: logging.Logger) -> None:
    """
    Run the program.

    platform -- an instance of a subclass of Platform
    args -- a module containing parsed command-line arguments
    config -- a dict containing the parsed configuration file
    logger -- a logger to log messages to
    """
    # Decide which families to use.
    families: typing.List[Family] = []
    if config["ipv4"]:
        families.append(IPv4())
    if config["ipv6"]["enable"]:
        families.append(IPv6(platform, config["ipv6"]))
    if not families:
        logger.error("No address families are enabled.")
        return

    # Grab the TTL from the config file.
    ttl = int(config["ttl"])

    # Decide which cache file to use, if any.
    cacheFile: typing.Optional[str]
    if isinstance(config["cache"], str):
        cacheFile = config["cache"]
    elif config["cache"] == True:
        cacheFile = platform.getDefaultCacheFilename()
    else:
        cacheFile = None
    if cacheFile is None:
        logger.debug("Using no cache file.")
    else:
        logger.debug("Using cache file %s.", cacheFile)

    # Wipe the cache file if in force mode. Doing this, rather than just
    # unconditionally updating right now, means that if this update fails, the
    # cache file will not be written and the next update will also be
    # unconditional, which is a more useful behaviour for force (you can run in
    # force mode once and be sure that at least one update will happen
    # successfully before we stop trying).
    if args.force and cacheFile is not None:
        logger.debug("Wiping cache due to --force.")
        try:
            os.remove(cacheFile)
        except OSError:
            pass

    # Load the cache file, if any.
    if cacheFile is not None:
        try:
            with open(cacheFile, "r") as fp:
                cache = json.load(fp)
        except (KeyboardInterrupt, SystemExit):
            raise
        except:
            cache = None
    else:
        cache = None

    # Rip apart my hostname.
    fqdn = dns.name.from_text(socket.getfqdn())
    zone = fqdn.parent()
    hostPart = fqdn.relativize(zone)

    # Find which nameserver we should talk to using an SOA query.
    resp = dns.resolver.query(zone, dns.rdatatype.SOA)
    if len(resp.rrset) != 1:
        raise Exception(f"Got {len(resp.rrset)} SOA records for zone {zone}, expected 1.")
    server = resp.rrset[0].mname.to_text(omit_final_dot=True)
    logger.debug("Using nameserver %s.", server)

    # Find my addresses.
    addresses: typing.Dict[str, typing.List[str]] = {family.getName(): [] for family in families}
    for interface in (args.interface or netifaces.interfaces()):
        for family in families:
            ifAddresses = netifaces.ifaddresses(interface).get(family.getNetIFacesConstant(), [])
            addresses[family.getName()] += family.filterAddressList([addr["addr"] for addr in ifAddresses])
    for family in families:
        addresses[family.getName()].sort()

    # Get the hostname and addresses most recently sent from the cache.
    if cache:
        last_hostname = cache.get("hostname")
        last_addresses = cache.get("addresses")
    else:
        last_hostname = None
        last_addresses = None

    # Check if the current hostname and addresses are the same as the last one.
    if fqdn.to_text() == last_hostname and addresses == last_addresses:
        logger.info("Eliding DNS record update for %s to %s as cache says addresses have not changed.", fqdn, addresses)
    else:
        # Construct the DNS update.
        logger.info("Updating DNS record for %s to %s.", fqdn, addresses)
        update = dns.update.Update(zone)
        update.delete(hostPart)
        for family in families:
            for address in addresses[family.getName()]:
                family.addAddressToUpdate(update, hostPart, ttl, address)
        if "tsig" in config:
            knownAlgorithms = {
                "hmac-md5": dns.tsig.HMAC_MD5,
                "hmac-sha1": dns.tsig.HMAC_SHA1,
                "hmac-sha224": dns.tsig.HMAC_SHA224,
                "hmac-sha256": dns.tsig.HMAC_SHA256,
                "hmac-sha384": dns.tsig.HMAC_SHA384,
                "hmac-sha512": dns.tsig.HMAC_SHA512,
            }
            if config["tsig"]["algorithm"] not in knownAlgorithms:
                raise ValueError(f"TSIG algorithm {config['tsig']['algorithm']} not recognized.")
            tsigAlgorithm = knownAlgorithms[config["tsig"]["algorithm"]]
            tsigRing = dns.tsigkeyring.from_text({config["tsig"]["keyname"]: config["tsig"]["key"]})
            update.use_tsig(keyring=tsigRing, algorithm=tsigAlgorithm)
            logger.debug("Update will be authenticated with TSIG %s.", tsigAlgorithm)
        else:
            logger.debug("Update will be unauthenticated.")

        # Resolve the nameserver hostname from the SOA record to one or more IP
        # addresses, and try to send an update to each one in turn until one
        # succeeds or they all fail.
        server_addresses = socket.getaddrinfo(server, "domain", type=socket.SOCK_STREAM)
        errors = []
        for (_, _, _, _, sockaddr) in server_addresses:
            try:
                # Send the update.
                server_address = sockaddr[0]
                logger.debug("Sending update to DNS server at %s.", server_address)
                resp = dns.query.tcp(update, where=server_address, timeout=30)
                if resp.rcode() != dns.rcode.NOERROR:
                    raise Exception("Update failed with rcode {resp.rrcode()}.")

                # This update was successful, so no need to try the rest of the
                # nameserver’s addresses.
                break
            except (OSError, dns.exception.DNSException) as exp:
                # Hold onto the error, but don’t report it yet—try the next
                # address instead.
                errors.append(exp)
        else:
            # All the nameserver’s addresses failed. Bail out.
            raise Exception("Unable to contact any nameservers: " + "; ".join(f"{address[4][0]}: {error}" for (address, error) in zip(server_addresses, errors)))

        # Update the cache to remember that we did the update.
        if cacheFile is not None:
            with open(cacheFile, "w") as fp:
                json.dump({"hostname": fqdn.to_text(), "addresses": addresses}, fp, ensure_ascii=False, allow_nan=False)

def main() -> None:
    # Choose a platform.
    platform: Platform = UnknownPlatform()
    for i in (POSIXPlatform(), WindowsPlatform()):
        if i.getName() == os.name:
            platform = i
    platform.platformSpecificSetup()

    # Parse command-line arguments.
    parser = argparse.ArgumentParser(description="Dynamically update DNS records.")
    parser.add_argument("-c", "--config", default=platform.getDefaultConfigFilename(), type=str, help=f"which configuration file to read (default: {platform.getDefaultConfigFilename()})", metavar="FILE")
    parser.add_argument("-f", "--force", action="store_true", help="update even if cache says unnecessary")
    parser.add_argument("interface", nargs="*", help="the name of an interface whose address(es) to register (default: all interfaces)")
    args = parser.parse_args()

    # Load configuration file.
    with open(args.config, "r") as configFile:
        config = json.load(configFile)

    # Configure logging.
    if "logging" in config:
        logging.config.dictConfig(config["logging"])
    else:
        logging.basicConfig(level=logging.DEBUG)
        logging.warn("No logging section in config file.")
    logging.captureWarnings(True)

    # Run the program.
    try:
        run(platform, args, config, logging.getLogger("pydyndns"))
    except (KeyboardInterrupt, SystemExit):
        raise
    except:
        logging.getLogger("pydyndns").error("Unhandled exception", exc_info=True)
<commit_msg>Make AF constants more type-checking-friendly<commit_after>#!/usr/bin/env python3

import abc
import argparse
import json
import logging
import logging.config
import logging.handlers
import os
import socket
import sys
import typing

import dns.name
import dns.rdataclass
import dns.rdatatype
import dns.rdtypes.IN.A
import dns.rdtypes.IN.AAAA
import dns.resolver
import dns.tsigkeyring
import dns.update
import netifaces

class Platform(metaclass=abc.ABCMeta):
    """
    Encapsulates knowledge about a specific operating system.
    """

    @abc.abstractmethod
    def getName(self) -> str:
        """Return the name as shown in os.name."""
        pass

    @abc.abstractmethod
    def getPermanentIPv6Addresses(self, addresses: typing.List[str]) -> typing.List[str]:
        """
        Return only the IPv6 addresses that are permanent (i.e. not generated
        by RFC4941 privacy extensions).

        For IPv6, most NICs will have multiple addresses: one permanent and a
        bunch of temporary generated by SLAAC privacy extensions. We should
        include only the permanent address, not the temporary ones.

        This is platform-specific because netifaces does not return any
        metadata about the address indicating whether it is permanent or
        temporary (and the addresses themselves are indistinguishable).
        However, specific operating systems appear to have a consistent
        ordering as to whether permanent addresses are returned before or after
        temporary addresses.
        """
        pass

    @abc.abstractmethod
    def getDefaultConfigFilename(self) -> str:
        """Return the default location of the configuration file."""
        pass

    @abc.abstractmethod
    def getDefaultCacheFilename(self) -> str:
        """Return the default location of the cache file."""
        pass

    @abc.abstractmethod
    def platformSpecificSetup(self) -> None:
        """
        Do any work that is specific to this platform for initializing the
        program.
        """
        pass

class POSIXPlatform(Platform):
    def getName(self) -> str:
        return "posix"

    def getPermanentIPv6Addresses(self, addresses: typing.List[str]) -> typing.List[str]:
        # Linux returns permanent addresses last.
        if len(addresses) == 0:
            return []
        else:
            return [addresses[-1]]

    def getDefaultConfigFilename(self) -> str:
        return "/etc/pydyndns.conf"

    def getDefaultCacheFilename(self) -> str:
        return "/run/pydyndns.cache"

    def platformSpecificSetup(self) -> None:
        pass

class WindowsPlatform(Platform):
    def getName(self) -> str:
        return "nt"

    def getPermanentIPv6Addresses(self, addresses: typing.List[str]) -> typing.List[str]:
        # Windows returns permanent addresses first.
        if len(addresses) == 0:
            return []
        else:
            return [addresses[0]]

    def getDefaultConfigFilename(self) -> str:
        return os.path.join(os.path.dirname(os.path.abspath(__file__)), "pydyndns.conf")

    def getDefaultCacheFilename(self) -> str:
        localAppData = os.environ.get("LOCALAPPDATA")
        if localAppData is None:
            localAppData = os.path.join(os.path.expanduser("~"), "AppData", "Local")
        return os.path.join(localAppData, "Temp", "pydyndns.cache")

    def platformSpecificSetup(self) -> None:
        # Python’s NTEventLogHandler class unconditionally tries to add the
        # event source to the Windows registry. This fails when running as a
        # low-privileged account. If somebody had already added the event
        # source then logging events using that source could work, but
        # NTEventLogHandler doesn’t bother trying that, it just fails in
        # construction if the registration fails.  Work around this by
        # swallowing exceptions during source registration.
        try:
            import pywintypes
            import win32evtlogutil
            oldAddSourceToRegistry = win32evtlogutil.AddSourceToRegistry
            def replacement(appname: str, dllname: str, logtype: str) -> None:
                try:
                    oldAddSourceToRegistry(appname, dllname, logtype)
                except pywintypes.error:
                    pass
            win32evtlogutil.AddSourceToRegistry = replacement
        except ImportError as exp:
            # Guess Win32 extensions are not installed.
            pass

class UnknownPlatform(Platform):
    def getName(self) -> str:
        return "unknown"

    def getPermanentIPv6Addresses(self, addresses: typing.List[str]) -> typing.List[str]:
        # No idea what the convention is on this platform, so just return all
        # of them.
        return addresses

    def getDefaultConfigFilename(self) -> str:
        return "pydyndns.conf"

    def getDefaultCacheFilename(self) -> str:
        return "pydyndns.cache"

    def platformSpecificSetup(self) -> None:
        pass

class Family(metaclass=abc.ABCMeta):
    """
    Encapsulates knowledge about a specific address family.
    """

    @abc.abstractmethod
    def getName(self) -> str:
        """Return the name used as a cache key for addresses in this family."""
        pass

    @abc.abstractmethod
    def getNetIFacesConstant(self) -> int:
        """Return the numeric ID used as a key in netifaces’ output."""
        pass

    @abc.abstractmethod
    def addAddressToUpdate(self, update: dns.update.Update, hostPart: dns.name.Name, ttl: int, address: str) -> None:
        """Add an address in this family to a DNS update request."""
        pass

    @abc.abstractmethod
    def filterAddressList(self, addresses: typing.Iterable[str]) -> typing.List[str]:
        """
        Return only those addresses that are useful, e.g. not loopback,
        link-local, temporary, or other special addresses that should not be
        registered.
        """
        pass

class IPv4(Family):
    def getName(self) -> str:
        return "ipv4"

    def getNetIFacesConstant(self) -> int:
        return int(socket.AF_INET)

    def addAddressToUpdate(self, update: dns.update.Update, hostPart: dns.name.Name, ttl: int, address: str) -> None:
        update.add(hostPart, ttl, dns.rdtypes.IN.A.A(dns.rdataclass.IN, dns.rdatatype.A, address))

    def filterAddressList(self, addresses: typing.Iterable[str]) -> typing.List[str]:
        # For IPv4 most NICs have only one address. It’s not clear that there
        # are any specific rules about how multiple addresses ought to be
        # handled. Just include all of them that are acceptable.
        return [x for x in addresses if self.includeAddress(x)]

    def includeAddress(self, address: str) -> bool:
        parts = [int(part) for part in address.split(".")]
        if parts[0] == 127:
            return False # Loopback address
        elif parts[0] >= 240:
            return False # Multicast or reserved address
        return True

class IPv6(Family):
    def __init__(self, platform: Platform, config: typing.Mapping[typing.Any, typing.Any]):
        self._platform = platform
        self._config = config

    def getName(self) -> str:
        return "ipv6"

    def getNetIFacesConstant(self) -> int:
        return int(socket.AF_INET6)

    def addAddressToUpdate(self, update: dns.update.Update, hostPart: dns.name.Name, ttl: int, address: str) -> None:
        update.add(hostPart, ttl, dns.rdtypes.IN.AAAA.AAAA(dns.rdataclass.IN, dns.rdatatype.AAAA, address))

    def filterAddressList(self, addresses: typing.Iterable[str]) -> typing.List[str]:
        return self._platform.getPermanentIPv6Addresses([x for x in addresses if self.includeAddress(x)])

    def includeAddress(self, address: str) -> bool:
        first_word = int(address.split(":")[0] or "0", 16)
        second_word = int(address.split(":")[1] or "0", 16)
        if first_word == 0x0000:
            return False # Unspecified, local, or IPv6-mapped address
        elif first_word == 0x0100:
            return False # Discard address
        elif (first_word & 0xFE00) == 0xFC00:
            return False # Unique local address
        elif first_word == 0xFE80:
            return False # Link-local address
        elif (first_word & 0xFF00) == 0xFF00:
            return False # Multicast address
        elif ((first_word == 0x2001) and (second_word == 0x0000)) and not self._config["teredo"]:
            return False # Teredo address
        return True

def run(platform: Platform, args: argparse.Namespace, config: typing.Mapping[typing.Any, typing.Any], logger: logging.Logger) -> None:
    """
    Run the program.

    platform -- an instance of a subclass of Platform
    args -- a module containing parsed command-line arguments
    config -- a dict containing the parsed configuration file
    logger -- a logger to log messages to
    """
    # Decide which families to use.
    families: typing.List[Family] = []
    if config["ipv4"]:
        families.append(IPv4())
    if config["ipv6"]["enable"]:
        families.append(IPv6(platform, config["ipv6"]))
    if not families:
        logger.error("No address families are enabled.")
        return

    # Grab the TTL from the config file.
    ttl = int(config["ttl"])

    # Decide which cache file to use, if any.
    cacheFile: typing.Optional[str]
    if isinstance(config["cache"], str):
        cacheFile = config["cache"]
    elif config["cache"] == True:
        cacheFile = platform.getDefaultCacheFilename()
    else:
        cacheFile = None
    if cacheFile is None:
        logger.debug("Using no cache file.")
    else:
        logger.debug("Using cache file %s.", cacheFile)

    # Wipe the cache file if in force mode. Doing this, rather than just
    # unconditionally updating right now, means that if this update fails, the
    # cache file will not be written and the next update will also be
    # unconditional, which is a more useful behaviour for force (you can run in
    # force mode once and be sure that at least one update will happen
    # successfully before we stop trying).
    if args.force and cacheFile is not None:
        logger.debug("Wiping cache due to --force.")
        try:
            os.remove(cacheFile)
        except OSError:
            pass

    # Load the cache file, if any.
    if cacheFile is not None:
        try:
            with open(cacheFile, "r") as fp:
                cache = json.load(fp)
        except (KeyboardInterrupt, SystemExit):
            raise
        except:
            cache = None
    else:
        cache = None

    # Rip apart my hostname.
    fqdn = dns.name.from_text(socket.getfqdn())
    zone = fqdn.parent()
    hostPart = fqdn.relativize(zone)

    # Find which nameserver we should talk to using an SOA query.
    resp = dns.resolver.query(zone, dns.rdatatype.SOA)
    if len(resp.rrset) != 1:
        raise Exception(f"Got {len(resp.rrset)} SOA records for zone {zone}, expected 1.")
    server = resp.rrset[0].mname.to_text(omit_final_dot=True)
    logger.debug("Using nameserver %s.", server)

    # Find my addresses.
    addresses: typing.Dict[str, typing.List[str]] = {family.getName(): [] for family in families}
    for interface in (args.interface or netifaces.interfaces()):
        for family in families:
            ifAddresses = netifaces.ifaddresses(interface).get(family.getNetIFacesConstant(), [])
            addresses[family.getName()] += family.filterAddressList([addr["addr"] for addr in ifAddresses])
    for family in families:
        addresses[family.getName()].sort()

    # Get the hostname and addresses most recently sent from the cache.
    if cache:
        last_hostname = cache.get("hostname")
        last_addresses = cache.get("addresses")
    else:
        last_hostname = None
        last_addresses = None

    # Check if the current hostname and addresses are the same as the last one.
    if fqdn.to_text() == last_hostname and addresses == last_addresses:
        logger.info("Eliding DNS record update for %s to %s as cache says addresses have not changed.", fqdn, addresses)
    else:
        # Construct the DNS update.
        logger.info("Updating DNS record for %s to %s.", fqdn, addresses)
        update = dns.update.Update(zone)
        update.delete(hostPart)
        for family in families:
            for address in addresses[family.getName()]:
                family.addAddressToUpdate(update, hostPart, ttl, address)
        if "tsig" in config:
            knownAlgorithms = {
                "hmac-md5": dns.tsig.HMAC_MD5,
                "hmac-sha1": dns.tsig.HMAC_SHA1,
                "hmac-sha224": dns.tsig.HMAC_SHA224,
                "hmac-sha256": dns.tsig.HMAC_SHA256,
                "hmac-sha384": dns.tsig.HMAC_SHA384,
                "hmac-sha512": dns.tsig.HMAC_SHA512,
            }
            if config["tsig"]["algorithm"] not in knownAlgorithms:
                raise ValueError(f"TSIG algorithm {config['tsig']['algorithm']} not recognized.")
            tsigAlgorithm = knownAlgorithms[config["tsig"]["algorithm"]]
            tsigRing = dns.tsigkeyring.from_text({config["tsig"]["keyname"]: config["tsig"]["key"]})
            update.use_tsig(keyring=tsigRing, algorithm=tsigAlgorithm)
            logger.debug("Update will be authenticated with TSIG %s.", tsigAlgorithm)
        else:
            logger.debug("Update will be unauthenticated.")

        # Resolve the nameserver hostname from the SOA record to one or more IP
        # addresses, and try to send an update to each one in turn until one
        # succeeds or they all fail.
        server_addresses = socket.getaddrinfo(server, "domain", type=socket.SOCK_STREAM)
        errors = []
        for (_, _, _, _, sockaddr) in server_addresses:
            try:
                # Send the update.
                server_address = sockaddr[0]
                logger.debug("Sending update to DNS server at %s.", server_address)
                resp = dns.query.tcp(update, where=server_address, timeout=30)
                if resp.rcode() != dns.rcode.NOERROR:
                    raise Exception("Update failed with rcode {resp.rrcode()}.")

                # This update was successful, so no need to try the rest of the
                # nameserver’s addresses.
                break
            except (OSError, dns.exception.DNSException) as exp:
                # Hold onto the error, but don’t report it yet—try the next
                # address instead.
                errors.append(exp)
        else:
            # All the nameserver’s addresses failed. Bail out.
            raise Exception("Unable to contact any nameservers: " + "; ".join(f"{address[4][0]}: {error}" for (address, error) in zip(server_addresses, errors)))

        # Update the cache to remember that we did the update.
        if cacheFile is not None:
            with open(cacheFile, "w") as fp:
                json.dump({"hostname": fqdn.to_text(), "addresses": addresses}, fp, ensure_ascii=False, allow_nan=False)

def main() -> None:
    # Choose a platform.
    platform: Platform = UnknownPlatform()
    for i in (POSIXPlatform(), WindowsPlatform()):
        if i.getName() == os.name:
            platform = i
    platform.platformSpecificSetup()

    # Parse command-line arguments.
    parser = argparse.ArgumentParser(description="Dynamically update DNS records.")
    parser.add_argument("-c", "--config", default=platform.getDefaultConfigFilename(), type=str, help=f"which configuration file to read (default: {platform.getDefaultConfigFilename()})", metavar="FILE")
    parser.add_argument("-f", "--force", action="store_true", help="update even if cache says unnecessary")
    parser.add_argument("interface", nargs="*", help="the name of an interface whose address(es) to register (default: all interfaces)")
    args = parser.parse_args()

    # Load configuration file.
    with open(args.config, "r") as configFile:
        config = json.load(configFile)

    # Configure logging.
    if "logging" in config:
        logging.config.dictConfig(config["logging"])
    else:
        logging.basicConfig(level=logging.DEBUG)
        logging.warn("No logging section in config file.")
    logging.captureWarnings(True)

    # Run the program.
    try:
        run(platform, args, config, logging.getLogger("pydyndns"))
    except (KeyboardInterrupt, SystemExit):
        raise
    except:
        logging.getLogger("pydyndns").error("Unhandled exception", exc_info=True)
<|endoftext|>


<commit_before># This file is part of the pyMOR project (http://www.pymor.org).
# Copyright Holders: Rene Milk, Stephan Rave, Felix Schindler
# License: BSD 2-Clause License (http://opensource.org/licenses/BSD-2-Clause)

"""Utilities for colorized log output.
via http://stackoverflow.com/questions/384076/how-can-i-make-the-python-logging-output-to-be-colored
Can not be moved because it's needed to be imported in the root __init__.py OR ELSE
"""
from __future__ import absolute_import, division, print_function
import curses
import logging
import os
import time

from pymor.core.defaults import defaults

BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)

# The background is set with 40 plus the number of the color, and the foreground with 30
# These are the sequences needed to get colored output
RESET_SEQ = "\033[0m"
COLOR_SEQ = "\033[1;%dm"
BOLD_SEQ = "\033[1m"
COLORS = {
    'WARNING':  YELLOW,
    'INFO':     GREEN,
    'DEBUG':    BLUE,
    'CRITICAL': MAGENTA,
    'ERROR':    RED
}

FORMAT = '%(asctime)s$BOLD%(levelname)s|$BOLD%(name)s$RESET: %(message)s'
MAX_HIERACHY_LEVEL = 3

start_time = time.time()

def formatter_message(message, use_color):
    if use_color:
        message = message.replace("$RESET", RESET_SEQ).replace("$BOLD", BOLD_SEQ)
    else:
        message = message.replace("$RESET", "").replace("$BOLD", "")
    return message

class ColoredFormatter(logging.Formatter):
    """A logging.Formatter that inserts tty control characters to color
    loglevel keyword output. Coloring can be disabled by setting the
    `PYMOR_COLORS_DISABLE` environment variable to `1`.
    """

    def __init__(self):
        disable_colors = int(os.environ.get('PYMOR_COLORS_DISABLE', 0)) == 1
        if disable_colors:
            self.use_color = False
        else:
            try:
                curses.setupterm()
                self.use_color = curses.tigetnum("colors") > 1
            except Exception:
                self.use_color = False

        super(ColoredFormatter, self).__init__(formatter_message(FORMAT, self.use_color))

    def formatTime(self, record, datefmt=None):
        elapsed = int(time.time() - start_time)
        days, remainder = divmod(elapsed, 86400)
        hours, remainder = divmod(remainder, 3600)
        minutes, seconds = divmod(remainder, 60)
        if days:
            return '{}d {:02}:{:02}:{:02}'.format(days, hours, minutes, seconds)
        elif hours:
            return '{:02}:{:02}:{:02}'.format(hours, minutes, seconds)
        else:
            return '{:02}:{:02}'.format(minutes, seconds)

    def format(self, record):
        if not record.msg:
            return ''
        tokens = record.name.split('.')
        record.name = '.'.join(tokens[1:MAX_HIERACHY_LEVEL])
        if len(tokens) > MAX_HIERACHY_LEVEL - 1:
            record.name += '.' + tokens[-1]
        levelname = record.levelname
        if self.use_color and levelname in COLORS.keys():
            if levelname is 'INFO':
                levelname_color = RESET_SEQ
            else:
                levelname_color = RESET_SEQ + '|' + COLOR_SEQ % (30 + COLORS[levelname]) + levelname + RESET_SEQ
            record.levelname = levelname_color
        elif levelname is 'INFO':
            record.levelname = ''
        return logging.Formatter.format(self, record)

def getLogger(module, level=None, filename=None, handler_cls=logging.StreamHandler):
    module = 'pymor' if module == '__main__' else module
    logger = logging.getLogger(module)
    streamhandler = handler_cls()
    streamformatter = ColoredFormatter()
    streamhandler.setFormatter(streamformatter)
    logger.handlers = [streamhandler]
    logger.propagate = False
    if level:
        logger.setLevel(level)
    return logger

class DummyLogger(object):

    __slots__ = []

    def nop(self, *args, **kwargs):
        return None

    propagate = False
    debug = nop
    info = nop
    warn = nop
    warning = nop
    error = nop
    critical = nop
    log = nop
    exception = nop

    def isEnabledFor(sefl, lvl):
        return False

    def getEffectiveLevel(self):
        return None

    def getChild(self):
        return self

dummy_logger = DummyLogger()

@defaults('levels')
def set_log_levels(levels={'pymor': 'WARN',
                           'pymor.core': 'WARN'}):
    """Set log levels for pyMOR's logging facility.

    Parameters
    ----------
        Dict of log levels. Keys names of loggers (see :func:`logging.getLogger`),
        values are the log levels to set for the loggers of the given names
        (see :meth:`~logging.Logger.setLevel`).
    """
    for k, v in levels.items():
        getLogger(k).setLevel(v)
<commit_msg>[core.logger] set pyMOR-wide default log level to INFO.<commit_after># This file is part of the pyMOR project (http://www.pymor.org).
# Copyright Holders: Rene Milk, Stephan Rave, Felix Schindler
# License: BSD 2-Clause License (http://opensource.org/licenses/BSD-2-Clause)

"""Utilities for colorized log output.
via http://stackoverflow.com/questions/384076/how-can-i-make-the-python-logging-output-to-be-colored
Can not be moved because it's needed to be imported in the root __init__.py OR ELSE
"""
from __future__ import absolute_import, division, print_function
import curses
import logging
import os
import time

from pymor.core.defaults import defaults

BLACK, RED, GREEN, YELLOW, BLUE, MAGENTA, CYAN, WHITE = range(8)

# The background is set with 40 plus the number of the color, and the foreground with 30
# These are the sequences needed to get colored output
RESET_SEQ = "\033[0m"
COLOR_SEQ = "\033[1;%dm"
BOLD_SEQ = "\033[1m"
COLORS = {
    'WARNING':  YELLOW,
    'INFO':     GREEN,
    'DEBUG':    BLUE,
    'CRITICAL': MAGENTA,
    'ERROR':    RED
}

FORMAT = '%(asctime)s$BOLD%(levelname)s|$BOLD%(name)s$RESET: %(message)s'
MAX_HIERACHY_LEVEL = 3

start_time = time.time()

def formatter_message(message, use_color):
    if use_color:
        message = message.replace("$RESET", RESET_SEQ).replace("$BOLD", BOLD_SEQ)
    else:
        message = message.replace("$RESET", "").replace("$BOLD", "")
    return message

class ColoredFormatter(logging.Formatter):
    """A logging.Formatter that inserts tty control characters to color
    loglevel keyword output. Coloring can be disabled by setting the
    `PYMOR_COLORS_DISABLE` environment variable to `1`.
    """

    def __init__(self):
        disable_colors = int(os.environ.get('PYMOR_COLORS_DISABLE', 0)) == 1
        if disable_colors:
            self.use_color = False
        else:
            try:
                curses.setupterm()
                self.use_color = curses.tigetnum("colors") > 1
            except Exception:
                self.use_color = False

        super(ColoredFormatter, self).__init__(formatter_message(FORMAT, self.use_color))

    def formatTime(self, record, datefmt=None):
        elapsed = int(time.time() - start_time)
        days, remainder = divmod(elapsed, 86400)
        hours, remainder = divmod(remainder, 3600)
        minutes, seconds = divmod(remainder, 60)
        if days:
            return '{}d {:02}:{:02}:{:02}'.format(days, hours, minutes, seconds)
        elif hours:
            return '{:02}:{:02}:{:02}'.format(hours, minutes, seconds)
        else:
            return '{:02}:{:02}'.format(minutes, seconds)

    def format(self, record):
        if not record.msg:
            return ''
        tokens = record.name.split('.')
        record.name = '.'.join(tokens[1:MAX_HIERACHY_LEVEL])
        if len(tokens) > MAX_HIERACHY_LEVEL - 1:
            record.name += '.' + tokens[-1]
        levelname = record.levelname
        if self.use_color and levelname in COLORS.keys():
            if levelname is 'INFO':
                levelname_color = RESET_SEQ
            else:
                levelname_color = RESET_SEQ + '|' + COLOR_SEQ % (30 + COLORS[levelname]) + levelname + RESET_SEQ
            record.levelname = levelname_color
        elif levelname is 'INFO':
            record.levelname = ''
        return logging.Formatter.format(self, record)

def getLogger(module, level=None, filename=None, handler_cls=logging.StreamHandler):
    module = 'pymor' if module == '__main__' else module
    logger = logging.getLogger(module)
    streamhandler = handler_cls()
    streamformatter = ColoredFormatter()
    streamhandler.setFormatter(streamformatter)
    logger.handlers = [streamhandler]
    logger.propagate = False
    if level:
        logger.setLevel(level)
    return logger

class DummyLogger(object):

    __slots__ = []

    def nop(self, *args, **kwargs):
        return None

    propagate = False
    debug = nop
    info = nop
    warn = nop
    warning = nop
    error = nop
    critical = nop
    log = nop
    exception = nop

    def isEnabledFor(sefl, lvl):
        return False

    def getEffectiveLevel(self):
        return None

    def getChild(self):
        return self

dummy_logger = DummyLogger()

@defaults('levels')
def set_log_levels(levels={'pymor': 'INFO'}):
    """Set log levels for pyMOR's logging facility.

    Parameters
    ----------
        Dict of log levels. Keys names of loggers (see :func:`logging.getLogger`),
        values are the log levels to set for the loggers of the given names
        (see :meth:`~logging.Logger.setLevel`).
    """
    for k, v in levels.items():
        getLogger(k).setLevel(v)
<|endoftext|>