<commit_before>from __future__ import division, print_function, absolute_import

import itertools

import numpy as np
from numpy.testing import (TestCase, dec, assert_, run_module_suite)
from scipy.special._testutils import FuncData
from scipy.special import smirnov, smirnovi

_rtol = 1e-10

class TestSmirnov(TestCase):
    def test_nan(self):
        assert_(np.isnan(smirnov(1, np.nan)))

    def test_basic(self):
        dataset = [(1, 0.1, 0.9),
                   (1, 0.875, 0.125),
                   (2, 0.875, 0.125 * 0.125),
                   (3, 0.875, 0.125 * 0.125 * 0.125)]

        dataset = np.asarray(dataset)
        FuncData(smirnov, dataset, (0, 1), 2, rtol=_rtol).check()

    def test_x_equals_0(self):
        dataset = [(n, 0, 1) for n in itertools.chain(range(2, 20), range(1010, 1020))]
        dataset = np.asarray(dataset)
        FuncData(smirnov, dataset, (0, 1), 2, rtol=_rtol).check()

    def test_x_equals_1(self):
        dataset = [(n, 1, 0) for n in itertools.chain(range(2, 20), range(1010, 1020))]
        dataset = np.asarray(dataset)
        FuncData(smirnov, dataset, (0, 1), 2, rtol=_rtol).check()

    def test_x_equals_0point5(self):
        dataset = [(1, 0.5, 0.5),
                   (2, 0.5, 0.25),
                   (3, 0.5, 0.166666666667),
                   (4, 0.5, 0.09375),
                   (5, 0.5, 0.056),
                   (6, 0.5, 0.0327932098765),
                   (7, 0.5, 0.0191958707681),
                   (8, 0.5, 0.0112953186035),
                   (9, 0.5, 0.00661933257355),
                   (10, 0.5, 0.003888705)]

        dataset = np.asarray(dataset)
        FuncData(smirnov, dataset, (0, 1), 2, rtol=_rtol).check()

    def test_n_equals_1(self):
        x = np.linspace(0, 1, 101, endpoint=True)
        dataset = np.stack([[1]*len(x), x, 1-x], axis=1)
        # dataset = np.asarray(dataset)
        FuncData(smirnov, dataset, (0, 1), 2, rtol=_rtol).check()

    def test_n_equals_2(self):
        x = np.linspace(0.5, 1, 101, endpoint=True)
        p = np.power(1-x, 2)
        n = np.array([2] * len(x))
        dataset = np.stack([n, x, p], axis=1)
        # dataset = np.asarray(dataset)
        FuncData(smirnov, dataset, (0, 1), 2, rtol=_rtol).check()

    def test_n_equals_3(self):
        x = np.linspace(0.7, 1, 31, endpoint=True)
        p = np.power(1-x, 3)
        n = np.array([3] * len(x))
        dataset = np.stack([n, x, p], axis=1)
        # dataset = np.asarray(dataset)
        FuncData(smirnov, dataset, (0, 1), 2, rtol=_rtol).check()

    def test_n_large(self):
        # test for large values of n
        # Probabilities should go down as n goes up
        x = 0.4
        pvals = np.array([smirnov(n, x) for n in range(400, 1100, 20)])
        dfs = np.diff(pvals)
        assert_(np.all(dfs <= 0), msg='Not all diffs negative %s' % dfs)

        dataset = [(1000, 1 - 1.0/2000, np.power(2000.0, -1000))]
        dataset = np.asarray(dataset)
        FuncData(smirnov, dataset, (0, 1), 2, rtol=_rtol).check()

        # Check asymptotic behaviour
        dataset = [(n, 1.0 / np.sqrt(n), np.exp(-2)) for n in range(1000, 5000, 1000)]
        dataset = np.asarray(dataset)
        FuncData(smirnov, dataset, (0, 1), 2, rtol=.05).check()

class TestSmirnovi(TestCase):
    def test_nan(self):
        assert_(np.isnan(smirnovi(1, np.nan)))

    @dec.knownfailureif(True, "test fails; smirnovi() is not always accurate")
    def test_basic(self):
        dataset = [(1, 0.4, 0.6),
                   (1, 0.6, 0.4),
                   (1, 0.99, 0.01),
                   (1, 0.01, 0.99),
                   (2, 0.125 * 0.125, 0.875),
                   (3, 0.125 * 0.125 * 0.125, 0.875),
                   (10, 1.0 / 16 ** 10, 1 - 1.0 / 16)]

        dataset = np.asarray(dataset)
        FuncData(smirnovi, dataset, (0, 1), 2, rtol=_rtol).check()

    @dec.knownfailureif(True, "test fails; smirnovi(_,0) is not accurate")
    def test_x_equals_0(self):
        dataset = [(n, 0, 1) for n in itertools.chain(range(2, 20), range(1010, 1020))]
        dataset = np.asarray(dataset)
        FuncData(smirnovi, dataset, (0, 1), 2, rtol=_rtol).check()

    def test_x_equals_1(self):
        dataset = [(n, 1, 0) for n in itertools.chain(range(2, 20), range(1010, 1020))]
        dataset = np.asarray(dataset)
        FuncData(smirnovi, dataset, (0, 1), 2, rtol=_rtol).check()

    @dec.knownfailureif(True, "test fails; smirnovi(1,) is not accurate")
    def test_n_equals_1(self):
        pp = np.linspace(0, 1, 101, endpoint=True)
        dataset = [(1, p, 1-p) for p in pp]
        dataset = np.asarray(dataset)
        FuncData(smirnovi, dataset, (0, 1), 2, rtol=_rtol).check()

    @dec.knownfailureif(True, "test fails; smirnovi(2,_) is not accurate")
    def test_n_equals_2(self):
        x = np.linspace(0.5, 1, 101, endpoint=True)
        p = np.power(1-x, 2)
        n = np.array([2] * len(x))
        dataset = np.stack([n, p, x], axis=1)
        # dataset = np.asarray(dataset)
        FuncData(smirnovi, dataset, (0, 1), 2, rtol=_rtol).check()

    @dec.knownfailureif(True, "test fails; smirnovi(3,_) is not accurate")
    def test_n_equals_3(self):
        x = np.linspace(0.7, 1, 31, endpoint=True)
        p = np.power(1-x, 3)
        n = np.array([3] * len(x))
        dataset = np.stack([n, p, x], axis=1)
        # dataset = np.asarray(dataset)
        FuncData(smirnovi, dataset, (0, 1), 2, rtol=_rtol).check()

    @dec.knownfailureif(True, "test fails; smirnovi(_,_) is not accurate")
    def test_round_trip(self):
        def _sm_smi(n, p):
            return smirnov(n, smirnovi(n, p))

        dataset = [(1, 0.4, 0.4),
                   (1, 0.6, 0.6),
                   (2, 0.875, 0.875),
                   (3, 0.875, 0.875),
                   (3, 0.125, 0.125),
                   (10, 0.999, 0.999),
                   (10, 0.0001, 0.0001)]

        dataset = np.asarray(dataset)
        FuncData(_sm_smi, dataset, (0, 1), 2, rtol=_rtol).check()

    def test_x_equals_0point5(self):
        dataset = [(1, 0.5, 0.5),
                   (2, 0.5, 0.366025403784),
                   (2, 0.25, 0.5),
                   (3, 0.5, 0.297156508177),
                   (4, 0.5, 0.255520481121),
                   (5, 0.5, 0.234559536069),
                   (6, 0.5, 0.21715965898),
                   (7, 0.5, 0.202722580034),
                   (8, 0.5, 0.190621765256),
                   (9, 0.5, 0.180363501362),
                   (10, 0.5, 0.17157867006)]

        dataset = np.asarray(dataset)
        FuncData(smirnovi, dataset, (0, 1), 2, rtol=_rtol).check()

if __name__ == "__main__":
    run_module_suite()
<commit_msg>Replaced np.stack with np.column_stack for numpy 1.8<commit_after>from __future__ import division, print_function, absolute_import

import itertools

import numpy as np
from numpy.testing import (TestCase, dec, assert_, run_module_suite)
from scipy.special._testutils import FuncData
from scipy.special import smirnov, smirnovi

_rtol = 1e-10

class TestSmirnov(TestCase):
    def test_nan(self):
        assert_(np.isnan(smirnov(1, np.nan)))

    def test_basic(self):
        dataset = [(1, 0.1, 0.9),
                   (1, 0.875, 0.125),
                   (2, 0.875, 0.125 * 0.125),
                   (3, 0.875, 0.125 * 0.125 * 0.125)]

        dataset = np.asarray(dataset)
        FuncData(smirnov, dataset, (0, 1), 2, rtol=_rtol).check()

    def test_x_equals_0(self):
        dataset = [(n, 0, 1) for n in itertools.chain(range(2, 20), range(1010, 1020))]
        dataset = np.asarray(dataset)
        FuncData(smirnov, dataset, (0, 1), 2, rtol=_rtol).check()

    def test_x_equals_1(self):
        dataset = [(n, 1, 0) for n in itertools.chain(range(2, 20), range(1010, 1020))]
        dataset = np.asarray(dataset)
        FuncData(smirnov, dataset, (0, 1), 2, rtol=_rtol).check()

    def test_x_equals_0point5(self):
        dataset = [(1, 0.5, 0.5),
                   (2, 0.5, 0.25),
                   (3, 0.5, 0.166666666667),
                   (4, 0.5, 0.09375),
                   (5, 0.5, 0.056),
                   (6, 0.5, 0.0327932098765),
                   (7, 0.5, 0.0191958707681),
                   (8, 0.5, 0.0112953186035),
                   (9, 0.5, 0.00661933257355),
                   (10, 0.5, 0.003888705)]

        dataset = np.asarray(dataset)
        FuncData(smirnov, dataset, (0, 1), 2, rtol=_rtol).check()

    def test_n_equals_1(self):
        x = np.linspace(0, 1, 101, endpoint=True)
        dataset = np.column_stack([[1]*len(x), x, 1-x])
        # dataset = np.asarray(dataset)
        FuncData(smirnov, dataset, (0, 1), 2, rtol=_rtol).check()

    def test_n_equals_2(self):
        x = np.linspace(0.5, 1, 101, endpoint=True)
        p = np.power(1-x, 2)
        n = np.array([2] * len(x))
        dataset = np.column_stack([n, x, p])
        # dataset = np.asarray(dataset)
        FuncData(smirnov, dataset, (0, 1), 2, rtol=_rtol).check()

    def test_n_equals_3(self):
        x = np.linspace(0.7, 1, 31, endpoint=True)
        p = np.power(1-x, 3)
        n = np.array([3] * len(x))
        dataset = np.column_stack([n, x, p])
        # dataset = np.asarray(dataset)
        FuncData(smirnov, dataset, (0, 1), 2, rtol=_rtol).check()

    def test_n_large(self):
        # test for large values of n
        # Probabilities should go down as n goes up
        x = 0.4
        pvals = np.array([smirnov(n, x) for n in range(400, 1100, 20)])
        dfs = np.diff(pvals)
        assert_(np.all(dfs <= 0), msg='Not all diffs negative %s' % dfs)

        dataset = [(1000, 1 - 1.0/2000, np.power(2000.0, -1000))]
        dataset = np.asarray(dataset)
        FuncData(smirnov, dataset, (0, 1), 2, rtol=_rtol).check()

        # Check asymptotic behaviour
        dataset = [(n, 1.0 / np.sqrt(n), np.exp(-2)) for n in range(1000, 5000, 1000)]
        dataset = np.asarray(dataset)
        FuncData(smirnov, dataset, (0, 1), 2, rtol=.05).check()

class TestSmirnovi(TestCase):
    def test_nan(self):
        assert_(np.isnan(smirnovi(1, np.nan)))

    @dec.knownfailureif(True, "test fails; smirnovi() is not always accurate")
    def test_basic(self):
        dataset = [(1, 0.4, 0.6),
                   (1, 0.6, 0.4),
                   (1, 0.99, 0.01),
                   (1, 0.01, 0.99),
                   (2, 0.125 * 0.125, 0.875),
                   (3, 0.125 * 0.125 * 0.125, 0.875),
                   (10, 1.0 / 16 ** 10, 1 - 1.0 / 16)]

        dataset = np.asarray(dataset)
        FuncData(smirnovi, dataset, (0, 1), 2, rtol=_rtol).check()

    @dec.knownfailureif(True, "test fails; smirnovi(_,0) is not accurate")
    def test_x_equals_0(self):
        dataset = [(n, 0, 1) for n in itertools.chain(range(2, 20), range(1010, 1020))]
        dataset = np.asarray(dataset)
        FuncData(smirnovi, dataset, (0, 1), 2, rtol=_rtol).check()

    def test_x_equals_1(self):
        dataset = [(n, 1, 0) for n in itertools.chain(range(2, 20), range(1010, 1020))]
        dataset = np.asarray(dataset)
        FuncData(smirnovi, dataset, (0, 1), 2, rtol=_rtol).check()

    @dec.knownfailureif(True, "test fails; smirnovi(1,) is not accurate")
    def test_n_equals_1(self):
        pp = np.linspace(0, 1, 101, endpoint=True)
        dataset = [(1, p, 1-p) for p in pp]
        dataset = np.asarray(dataset)
        FuncData(smirnovi, dataset, (0, 1), 2, rtol=_rtol).check()

    @dec.knownfailureif(True, "test fails; smirnovi(2,_) is not accurate")
    def test_n_equals_2(self):
        x = np.linspace(0.5, 1, 101, endpoint=True)
        p = np.power(1-x, 2)
        n = np.array([2] * len(x))
        dataset = np.column_stack([n, p, x])
        # dataset = np.asarray(dataset)
        FuncData(smirnovi, dataset, (0, 1), 2, rtol=_rtol).check()

    @dec.knownfailureif(True, "test fails; smirnovi(3,_) is not accurate")
    def test_n_equals_3(self):
        x = np.linspace(0.7, 1, 31, endpoint=True)
        p = np.power(1-x, 3)
        n = np.array([3] * len(x))
        dataset = np.column_stack([n, p, x])
        # dataset = np.asarray(dataset)
        FuncData(smirnovi, dataset, (0, 1), 2, rtol=_rtol).check()

    @dec.knownfailureif(True, "test fails; smirnovi(_,_) is not accurate")
    def test_round_trip(self):
        def _sm_smi(n, p):
            return smirnov(n, smirnovi(n, p))

        dataset = [(1, 0.4, 0.4),
                   (1, 0.6, 0.6),
                   (2, 0.875, 0.875),
                   (3, 0.875, 0.875),
                   (3, 0.125, 0.125),
                   (10, 0.999, 0.999),
                   (10, 0.0001, 0.0001)]

        dataset = np.asarray(dataset)
        FuncData(_sm_smi, dataset, (0, 1), 2, rtol=_rtol).check()

    def test_x_equals_0point5(self):
        dataset = [(1, 0.5, 0.5),
                   (2, 0.5, 0.366025403784),
                   (2, 0.25, 0.5),
                   (3, 0.5, 0.297156508177),
                   (4, 0.5, 0.255520481121),
                   (5, 0.5, 0.234559536069),
                   (6, 0.5, 0.21715965898),
                   (7, 0.5, 0.202722580034),
                   (8, 0.5, 0.190621765256),
                   (9, 0.5, 0.180363501362),
                   (10, 0.5, 0.17157867006)]

        dataset = np.asarray(dataset)
        FuncData(smirnovi, dataset, (0, 1), 2, rtol=_rtol).check()

if __name__ == "__main__":
    run_module_suite()
<|endoftext|>


<commit_before># pylint: disable=eval-used,bare-except,wildcard-import,unused-wildcard-import
"""
Tree.py
"""

import copy as cp
import numpy as np

from eugene.Primatives import *
from eugene.Node import Node, random_node

# indent branching [level:num] value (height) - {node_num|complexity}
DISPLAY_NODE_STR = '%s%s[%s:%s] %s (%s) - {%s|%s}'

# @profile
def random_tree(max_level=20, min_level=1, current_level=0):
    """generate a random tree of random nodes"""
    return Tree(random_node(max_level, min_level, current_level))

class Tree(object):
    """
    Defines a tree of nodes, with functions to operate on tree and subtrees
    """

    def __init__(self, nodes=None, subtree=False):
        self.nodes = nodes
        if not subtree:
            self.nodes.set_nums()

    @property
    def height(self):
        """
        Return the number of levels in the tree.
        """
        return self.nodes.height + 1

    @property
    def node_num(self):
        """
        Return the number of nodes in the tree.
        """
        return self.nodes.node_num

    @property
    def leaf_num(self):
        """
        Return the number of leaves in the tree.
        """
        return self.nodes.leaf_num

    @property
    def edge_num(self):
        """
        Return the number of edges in the tree.
        """
        return self.nodes.edge_num

    @property
    def complexity(self):
        """
        Return the complexity of the tree (sum of nodes in tree and each subtree).
        """
        return self.nodes.complexity

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        return str(self.nodes)

    # @profile
    def evaluate(self):
        """evaluate expression stored in tree.
        """
        try:
            result = np.array(eval(compile(self.__str__(), '', 'eval')))
        except:
            result = np.array(np.nan)
        return result

    # @profile
    def get_node(self, n=0):
        """
        Return a node from the tree.
        """

        # search tree until node number is found and take sub tree
        if self.nodes.num == n:
            return cp.deepcopy(self.nodes)
        elif len(self.nodes.children) > 0:
            for c in self.nodes.children:
                cn = Tree(c, subtree=True).get_node(n)
                if cn:
                    return cn
        else:
            return None

    # @profile
    def set_node(self, n=0, node=None):
        """
        Set a node in the tree.
        """

        # search tree until node number is found, and store sub tree
        if self.nodes.num == n:
            self.nodes = node
        else:
            self.nodes.children = tuple([Tree(c, subtree=True).set_node(n, node) for c in self.nodes.children])

        # rebase the numbers of the Tree
        self.nodes.set_nums()
        return self.nodes

    # @profile
    def list_edges(self):
        """
        Get edges of tree.
        """

        # get list of tuple edges between nodes e.g. [(n1,n2),(n1,n3)...]
        edges = [(self.nodes.value, c.value if len(c.children) > 0 else c.value) for c in self.nodes.children]
        children_nodes = [Tree(c, subtree=True).list_edges() for c in self.nodes.children if len(c.children) > 0]
        for i in xrange(len(children_nodes)):
            edges += children_nodes[i]
        return edges

    # @profile
    def list_nodes(self):
        """
        Return nodes of tree.
        """

        # get list of nodes
        node_list = []
        node_list.append(self.nodes.value)
        # add children
        node_list.extend([c.value for c in self.nodes.children if len(c.children) == 0])
        # add children's children
        grand_children = [Tree(c, subtree=True).list_nodes() for c in self.nodes.children if len(c.children) > 0]
        node_list.extend([node for grand_child in grand_children for node in grand_child])

        return node_list

    # @profile
    def prune(self):
        """
        Go thru nodes and remove or replace dead / constant branches (subtrees).
        """

        # create subtree
        sub_tree = Tree(self.nodes, subtree=True)
        # check if the tree contains a variable
        contains_variable = any([n in VARIABLES for n in sub_tree.list_nodes()])
        # evaluate subtree for inefficiencies
        sub_eval = sub_tree.evaluate()
        # check is evaluation exactly equals one of the variables
        equals_variable = [v for v in VARIABLES if np.array(Tree(Node(v)).evaluate() == sub_eval).all()]

        # if subtree of node does not contain variable, it must be constant
        if not contains_variable:
            self.nodes.value = sub_eval
            self.nodes.children = ()
        # if subtree contains a variable, but evaluates to exactly the variable, replace with variable
        elif equals_variable:
            self.nodes.value = equals_variable[0]
            self.nodes.children = ()
        # can't make more effiecient
        else:
            for child in self.nodes.children:
                Tree(child, subtree=True).prune()

        # rebase the numbers of the Tree
        self.nodes.set_nums()
        return self.nodes

    def display(self, level=0, level_list=None, tree_string='', stdout=True):
        """
        Display helper. Returns variable and prints to stdout (default stdout=True).
        """
        level_list = level_list if level_list else []

        if level == 0:
            node_str = DISPLAY_NODE_STR % (
                '',
                '',
                self.nodes.level,
                self.nodes.num,
                self.nodes.value,
                self.nodes.height,
                self.nodes.node_num,
                self.nodes.complexity
            )
        else:
            branching = '\\-' if level_list[-1] == '      ' else '|-'
            indent = '    ' + ''.join(level_list[:-1])
            node_str = DISPLAY_NODE_STR % (
                indent,
                branching,
                self.nodes.level,
                self.nodes.num,
                self.nodes.value,
                self.nodes.height,
                self.nodes.node_num,
                self.nodes.complexity
            )
        # Place this node of the tree
        if stdout:
            print node_str
        tree_string += node_str
        # Append children nodes
        for i, child in enumerate(self.nodes.children):
            child_string = Tree(child, subtree=True).display(
                level + 1, level_list + ['      ' if i == len(self.nodes.children) - 1 else '|     '],
                stdout=stdout
            )
            tree_string += '\n'
            tree_string += child_string
        return tree_string
<commit_msg>title case comments<commit_after># pylint: disable=eval-used,bare-except,wildcard-import,unused-wildcard-import
"""
Tree.py
"""

import copy as cp
import numpy as np

from eugene.Primatives import *
from eugene.Node import Node, random_node

# indent branching [level:num] value (height) - {node_num|complexity}
DISPLAY_NODE_STR = '%s%s[%s:%s] %s (%s) - {%s|%s}'

# @profile
def random_tree(max_level=20, min_level=1, current_level=0):
    """generate a random tree of random nodes"""
    return Tree(random_node(max_level, min_level, current_level))

class Tree(object):
    """
    Defines a tree of nodes, with functions to operate on tree and subtrees
    """

    def __init__(self, nodes=None, subtree=False):
        self.nodes = nodes
        if not subtree:
            self.nodes.set_nums()

    @property
    def height(self):
        """
        Return the number of levels in the tree.
        """
        return self.nodes.height + 1

    @property
    def node_num(self):
        """
        Return the number of nodes in the tree.
        """
        return self.nodes.node_num

    @property
    def leaf_num(self):
        """
        Return the number of leaves in the tree.
        """
        return self.nodes.leaf_num

    @property
    def edge_num(self):
        """
        Return the number of edges in the tree.
        """
        return self.nodes.edge_num

    @property
    def complexity(self):
        """
        Return the complexity of the tree (sum of nodes in tree and each subtree).
        """
        return self.nodes.complexity

    def __repr__(self):
        return self.__str__()

    def __str__(self):
        return str(self.nodes)

    # @profile
    def evaluate(self):
        """evaluate expression stored in tree.
        """
        try:
            result = np.array(eval(compile(self.__str__(), '', 'eval')))
        except:
            result = np.array(np.nan)
        return result

    # @profile
    def get_node(self, n=0):
        """
        Return a node from the tree.
        """

        # Search tree until node number is found and take sub tree
        if self.nodes.num == n:
            return cp.deepcopy(self.nodes)
        elif len(self.nodes.children) > 0:
            for c in self.nodes.children:
                cn = Tree(c, subtree=True).get_node(n)
                if cn:
                    return cn
        else:
            return None

    # @profile
    def set_node(self, n=0, node=None):
        """
        Set a node in the tree.
        """

        # Search tree until node number is found, and store sub tree
        if self.nodes.num == n:
            self.nodes = node
        else:
            self.nodes.children = tuple([Tree(c, subtree=True).set_node(n, node) for c in self.nodes.children])

        # Rebase the numbers of the Tree
        self.nodes.set_nums()
        return self.nodes

    # @profile
    def list_edges(self):
        """
        Get edges of tree.
        """

        # Get list of tuple edges between nodes e.g. [(n1,n2),(n1,n3)...]
        edges = [(self.nodes.value, c.value if len(c.children) > 0 else c.value) for c in self.nodes.children]
        children_nodes = [Tree(c, subtree=True).list_edges() for c in self.nodes.children if len(c.children) > 0]
        for i in xrange(len(children_nodes)):
            edges += children_nodes[i]
        return edges

    # @profile
    def list_nodes(self):
        """
        Return nodes of tree.
        """

        # Get list of nodes
        node_list = []
        node_list.append(self.nodes.value)
        # Add children
        node_list.extend([c.value for c in self.nodes.children if len(c.children) == 0])
        # Add children's children
        grand_children = [Tree(c, subtree=True).list_nodes() for c in self.nodes.children if len(c.children) > 0]
        node_list.extend([node for grand_child in grand_children for node in grand_child])

        return node_list

    # @profile
    def prune(self):
        """
        Go thru nodes and remove or replace dead / constant branches (subtrees).
        """

        # Create subtree
        sub_tree = Tree(self.nodes, subtree=True)
        # Check if the tree contains a variable
        contains_variable = any([n in VARIABLES for n in sub_tree.list_nodes()])
        # Evaluate subtree for inefficiencies
        sub_eval = sub_tree.evaluate()
        # Check is evaluation exactly equals one of the variables
        equals_variable = [v for v in VARIABLES if np.array(Tree(Node(v)).evaluate() == sub_eval).all()]

        # If subtree of node does not contain variable, it must be constant
        if not contains_variable:
            self.nodes.value = sub_eval
            self.nodes.children = ()
        # If subtree contains a variable, but evaluates to exactly the variable, replace with variable
        elif equals_variable:
            self.nodes.value = equals_variable[0]
            self.nodes.children = ()
        # Can't make more effiecient
        else:
            for child in self.nodes.children:
                Tree(child, subtree=True).prune()

        # Rebase the numbers of the Tree
        self.nodes.set_nums()
        return self.nodes

    def display(self, level=0, level_list=None, tree_string='', stdout=True):
        """
        Display helper. Returns variable and prints to stdout (default stdout=True).
        """
        level_list = level_list if level_list else []

        if level == 0:
            node_str = DISPLAY_NODE_STR % (
                '',
                '',
                self.nodes.level,
                self.nodes.num,
                self.nodes.value,
                self.nodes.height,
                self.nodes.node_num,
                self.nodes.complexity
            )
        else:
            branching = '\\-' if level_list[-1] == '      ' else '|-'
            indent = '    ' + ''.join(level_list[:-1])
            node_str = DISPLAY_NODE_STR % (
                indent,
                branching,
                self.nodes.level,
                self.nodes.num,
                self.nodes.value,
                self.nodes.height,
                self.nodes.node_num,
                self.nodes.complexity
            )
        # Place this node of the tree
        if stdout:
            print node_str
        tree_string += node_str
        # Append children nodes
        for i, child in enumerate(self.nodes.children):
            child_string = Tree(child, subtree=True).display(
                level + 1, level_list + ['      ' if i == len(self.nodes.children) - 1 else '|     '],
                stdout=stdout
            )
            tree_string += '\n'
            tree_string += child_string
        return tree_string
<|endoftext|>


<commit_before>import sys
import time
import random
import math
from generate_random_operations import *

def getOperationMaximums(args):
	maximums = [ -1, -1, -1, -1 ]
	maxIndex = min(4, len(args))
	for i in range(maxIndex):
		maximums[i] = int(args[i])
	return maximums

def parsePointStr(pointStr):
	fields = pointStr.split()
	return [ float(x) for x in fields ]

def generateDatasetOperation(numDimensions, cumulativeOpProbs, dataset):
	# Randomly choose operation
	val = random.random()
	if val >= 0 and val < cumulativeOpProbs[0]:
		opType = "I"
	elif val >= cumulativeOpProbs[0] and val < cumulativeOpProbs[1]:
		opType = "D"
	elif val >= cumulativeOpProbs[1] and val < cumulativeOpProbs[2]:
		opType = "U"
	elif val >= cumulativeOpProbs[2] and val <= cumulativeOpProbs[3]:
		opType = "P"
	else:
		raise RuntimeError("Unknown random value generated in generateRandomOperation()")

	return generateDatasetOperationWithType(numDimensions, opType, dataset)

def generateDatasetOperationWithType(numDimensions, opType, dataset):
	# Choose random point from dataset
	randomIndex = random.randint(0, len(dataset) - 1)
	point = list(dataset[randomIndex])
	# Select another point if the operation is an Update
	if opType == "U":
		randomIndex = random.randint(0, len(dataset) - 1)
		point += dataset[randomIndex]

	return (opType, point)

if __name__ == "__main__":
	# Seed random number generator
	random.seed(time.time())

	# Parse command line arguments
	if len(sys.argv) < 3:
		sys.exit("Usage: {} <datasetFilename> [all|random|equal] {{<numOperations> <insertProbability> <deleteProbability> <updateProbability> <pqueryProbability> <maxInsert> <maxDelete> <maxUpdate> <maxPQuery>}}".format(sys.argv[0]))
	datasetFilename = sys.argv[1]
	listType = sys.argv[2]
	if not listType in [ "all", "random", "equal" ]:
		sys.exit("Invalid list type '{}' given".format(listType))
	numOperations = 0
	if listType != "all":
		if len(sys.argv) < 4:
			sys.exit("The number of operations to generate must be given if 'all' mode is node used")
		else:
			numOperations = int(sys.argv[3])

	operationProbabilities = getOperationProbabilities(sys.argv[4:])
	operationProbabilities = normaliseProbabilities(operationProbabilities)
	cumulativeOpProbs = computeCumulativeProbabilities(operationProbabilities)

	operationMaximums = getOperationMaximums(sys.argv[8:])

	# Load specified dataset into memory
	numDimensions = 0
	dataset = []
	with open(datasetFilename, "r") as f:
		headerFields = f.readline().split()
		if len(headerFields) < 2:
			sys.exit("Dataset header file does not have both #dimensions and #points")
		numDimensions, numDatasetPoints = int(headerFields[0]), int(headerFields[1])

		dataset = [ parsePointStr(line) for line in f.readlines() ]
		dataset = dataset[:numDatasetPoints]

	# If generate mode is 'all', then all points are inserted, queries, then deleted in bulk
	if listType == "all":
		print("{} {}".format(numDimensions, len(dataset) * 3))
		for p in dataset:
			print("I {}".format( pointToStr(p) ))
		for p in dataset:
			print("P {}".format( pointToStr(p) ))
		for p in dataset:
			print("D {}".format( pointToStr(p) ))
	# If 'random', randomly generate operations that use points FROM THE DATASET using given probabilities
	elif listType == "random":
		print("{} {}".format(numDimensions, numOperations))
		for i in range(numOperations):
			op = generateDatasetOperation(numDimensions, cumulativeOpProbs, dataset)
			print(operationToStr(op))
	elif listType == "equal":
		print("{} {}".format(numDimensions, numOperations))

		# Compute how many of each operation to generate
		opsPerType = int(numOperations / 4) + 1
		# Generate the same amount of operations for each type
		operations = [ generateDatasetOperationWithType(numDimensions, "I", dataset) for i in range(opsPerType) ]
		operations += [ generateDatasetOperationWithType(numDimensions, "D", dataset) for i in range(opsPerType) ]
		operations += [ generateDatasetOperationWithType(numDimensions, "U", dataset) for i in range(opsPerType) ]
		operations += [ generateDatasetOperationWithType(numDimensions, "P", dataset) for i in range(opsPerType) ]
		operations = operations[:numOperations]
		# Shuffle all operations
		random.shuffle(operations)
		# Output all opewrations
		for op in operations:
			print(operationToStr(op))<commit_msg>Modified operation list generator so it can generate op type specific lists that run through all the points in the dataset (i.e. non-random)<commit_after>import sys
import time
import random
import math
from generate_random_operations import *
from generate_random_dataset import generatePoint

def getOperationMaximums(args):
	maximums = [ -1, -1, -1, -1 ]
	maxIndex = min(4, len(args))
	for i in range(maxIndex):
		maximums[i] = int(args[i])
	return maximums

def parsePointStr(pointStr):
	fields = pointStr.split()
	return [ float(x) for x in fields ]

def generateDatasetOperation(numDimensions, cumulativeOpProbs, dataset):
	# Randomly choose operation
	val = random.random()
	if val >= 0 and val < cumulativeOpProbs[0]:
		opType = "I"
	elif val >= cumulativeOpProbs[0] and val < cumulativeOpProbs[1]:
		opType = "D"
	elif val >= cumulativeOpProbs[1] and val < cumulativeOpProbs[2]:
		opType = "U"
	elif val >= cumulativeOpProbs[2] and val <= cumulativeOpProbs[3]:
		opType = "P"
	else:
		raise RuntimeError("Unknown random value generated in generateRandomOperation()")

	return generateDatasetOperationWithType(numDimensions, opType, dataset)

def generateDatasetOperationWithType(numDimensions, opType, dataset):
	# Choose random point from dataset
	randomIndex = random.randint(0, len(dataset) - 1)
	point = list(dataset[randomIndex])
	# Select another point if the operation is an Update
	if opType == "U":
		randomIndex = random.randint(0, len(dataset) - 1)
		point += dataset[randomIndex]

	return (opType, point)

def printSingleOperationTypeList(opType, numDimensions, dataset):
	print("{} {}".format(numDimensions, len(dataset)))
	for p in dataset:
		op = (opType, list(p))
		print(operationToStr(op))

if __name__ == "__main__":
	# Seed random number generator
	random.seed(time.time())

	# Parse command line arguments
	if len(sys.argv) < 3:
		sys.exit("Usage: {} <datasetFilename> [all|random|equal|insert|delete|update|pquery] {{<numOperations> <insertProbability> <deleteProbability> <updateProbability> <pqueryProbability> <maxInsert> <maxDelete> <maxUpdate> <maxPQuery>}}".format(sys.argv[0]))
	datasetFilename = sys.argv[1]
	listType = sys.argv[2]
	if not listType in [ "all", "random", "equal", "insert", "delete", "update", "pquery" ]:
		sys.exit("Invalid list type '{}' given".format(listType))
	numOperations = 0
	if listType in ["random", "equal"]:
		if len(sys.argv) < 4:
			sys.exit("The number of operations to generate must be given if 'all' mode is node used")
		else:
			numOperations = int(sys.argv[3])

	operationProbabilities = getOperationProbabilities(sys.argv[4:])
	operationProbabilities = normaliseProbabilities(operationProbabilities)
	cumulativeOpProbs = computeCumulativeProbabilities(operationProbabilities)

	operationMaximums = getOperationMaximums(sys.argv[8:])

	# Load specified dataset into memory
	numDimensions = 0
	dataset = []
	with open(datasetFilename, "r") as f:
		headerFields = f.readline().split()
		if len(headerFields) < 2:
			sys.exit("Dataset header file does not have both #dimensions and #points")
		numDimensions, numDatasetPoints = int(headerFields[0]), int(headerFields[1])

		dataset = [ parsePointStr(line) for line in f.readlines() ]
		dataset = dataset[:numDatasetPoints]

	# If generate mode is 'all', then all points are inserted, queries, then deleted in bulk
	if listType == "all":
		print("{} {}".format(numDimensions, len(dataset) * 3))
		for p in dataset:
			print("I {}".format( pointToStr(p) ))
		for p in dataset:
			print("P {}".format( pointToStr(p) ))
		for p in dataset:
			print("D {}".format( pointToStr(p) ))
	# If 'random', randomly generate operations that use points FROM THE DATASET using given probabilities
	elif listType == "random":
		print("{} {}".format(numDimensions, numOperations))
		for i in range(numOperations):
			op = generateDatasetOperation(numDimensions, cumulativeOpProbs, dataset)
			print(operationToStr(op))
	elif listType == "equal":
		print("{} {}".format(numDimensions, numOperations))

		# Compute how many of each operation to generate
		opsPerType = int(numOperations / 4) + 1
		# Generate the same amount of operations for each type
		operations = [ generateDatasetOperationWithType(numDimensions, "I", dataset) for i in range(opsPerType) ]
		operations += [ generateDatasetOperationWithType(numDimensions, "D", dataset) for i in range(opsPerType) ]
		operations += [ generateDatasetOperationWithType(numDimensions, "U", dataset) for i in range(opsPerType) ]
		operations += [ generateDatasetOperationWithType(numDimensions, "P", dataset) for i in range(opsPerType) ]
		operations = operations[:numOperations]
		# Shuffle all operations
		random.shuffle(operations)
		# Output all opewrations
		for op in operations:
			print(operationToStr(op))
	# Generate olists which have a single type of operation, each operation being for a single point in the dataset
	elif listType == "insert":
		printSingleOperationTypeList("I", numDimensions, dataset)
	elif listType == "delete":
		printSingleOperationTypeList("D", numDimensions, dataset)
	elif listType == "update":
		print("{} {}".format(numDimensions, len(dataset)))
		for p in dataset:
			updateArgs = list(p)
			# Generate random point to update the given point with
			updateArgs += generatePoint(numDimensions)
			op = ("U", updateArgs)
			print(operationToStr(op))
	elif listType == "pquery":
		printSingleOperationTypeList("P", numDimensions, dataset)<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# ***** BEGIN LICENSE BLOCK *****
# Copyright (C) 2012  Hayaki Saito <user@zuse.jp>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
# ***** END LICENSE BLOCK *****

import os, termios, pty, signal, fcntl, struct, select, errno
import codecs
try:
    from cStringIO import StringIO
except:
    try:
        from StringIO import StringIO
    except:
        from io import StringIO

from interface import * # terminal filter framework interface

BUFFER_SIZE=2048

################################################################################
#
# Scanner implementation
#
class DefaultScanner(Scanner):
    ''' scan input stream and iterate characters '''
    __data = None

    def assign(self, value, termenc):
        self.__data = unicode(value, termenc, 'ignore')

    def __iter__(self):
        for x in self.__data:
            yield ord(x)

################################################################################
#
# Simple Parser implementation
#
class SimpleParser(Parser):
    ''' simple parser, don't parse ESC/CSI/string seqneces '''
    def __init__(self):
        pass

    def parse(self, context):
        for c in context:
            context.dispatch_char(c)

################################################################################
#
# Default Parser implementation
#
STATE_GROUND = 0
STATE_ESC = 1
STATE_ESC_FINAL = 2

STATE_CSI_PARAMETER = 3
STATE_CSI_INTERMEDIATE = 4
STATE_CSI_FINAL = 5

STATE_OSC = 6
STATE_STR = 7

class DefaultParser(Parser):
    ''' parse ESC/CSI/string seqneces '''

    def __init__(self):
        self.__parse_state = STATE_GROUND
        self.__csi_parameter = []
        self.__csi_intermediate = []
        self.__esc_prefix = []
        self.__str = []
        self.__str_prefix = None
        self.__str_esc_state = False

    def parse(self, context):
        for c in context:
            if self.__parse_state == STATE_OSC:
                # parse control string
                if c == 0x1b:
                    self.__str_esc_state = True
                elif c == 0x5c and self.__str_esc_state == True:
                    context.dispatch_control_string(self.__str_prefix, self.__str)
                    self.__parse_state = STATE_GROUND
                elif c == 0x07:
                    context.dispatch_control_string(self.__str_prefix, self.__str)
                    self.__parse_state = STATE_GROUND
                elif c < 0x20:
                    self.__parse_state = STATE_GROUND
                else:
                    self.__str.append(c)

            elif self.__parse_state == STATE_STR:
                # parse control string
                if c == 0x1b:
                    self.__str_esc_state = True
                elif c == 0x5c and self.__str_esc_state == True:
                    context.dispatch_control_string(self.__str_prefix, self.__str)
                    self.__parse_state = STATE_GROUND
                elif c < 0x20:
                    self.__parse_state = STATE_GROUND
                else:
                    self.__str.append(c)

            elif c == 0x1b: # ESC
                self.__esc_prefix = []
                self.__parse_state = STATE_ESC

            elif c < 0x20 or c == 0x7f: # control character
                context.dispatch_char(c)
                self.__parse_state = STATE_GROUND

            elif self.__parse_state == STATE_ESC:
                if c == 0x5b: # [
                    self.__csi_parameter = []
                    self.__csi_intermediate = []
                    self.__parse_state = STATE_CSI_PARAMETER
                elif c == 0x5d: # ]
                    self.__str_esc_state = False
                    self.__str = []
                    self.__str_prefix = c
                    self.__parse_state = STATE_OSC
                elif c == 0x50 or c == 0x58 or c == 0x5e or c == 0x5f:
                    # P(DCS) or X(SOS) or ^(PM) or _(APC)
                    self.__str_esc_state = False
                    self.__str = []
                    self.__str_prefix = c
                    self.__parse_state = STATE_STR
                elif 0x20 <= c and c <= 0x2f: # SP to /
                    self.__esc_prefix.append(c)
                    self.__parse_state = STATE_ESC_FINAL
                else:
                    context.dispatch_esc(self.__esc_prefix, c)
                    self.__parse_state = STATE_GROUND

            elif self.__parse_state == STATE_ESC_FINAL:
                context.dispatch_esc(self.__esc_prefix, c)
                self.__parse_state = STATE_GROUND

            elif self.__parse_state == STATE_CSI_PARAMETER:
                # parse control sequence
                #
                # CSI P ... P I ... I F
                #     ^
                if 0x30 <= c and c <= 0x3f: # parameter, 0 to ?
                    self.__csi_parameter.append(c)
                elif 0x20 <= c and c <= 0x2f: # intermediate, SP to /
                    self.__csi_intermediate.append(c)
                    self.__parse_state = STATE_CSI_INTERMEDIATE
                elif 0x40 <= c and c <= 0x7e: # Final byte, @ to ~
                    context.dispatch_csi(self.__csi_parameter,
                                         self.__csi_intermediate,
                                         c)
                    self.__parse_state = STATE_GROUND
                else:
                    self.__parse_state = STATE_GROUND

            elif self.__parse_state == STATE_CSI_INTERMEDIATE:
                # parse control sequence
                #
                # CSI P ... P I ... I F
                #             ^
                if 0x20 <= c and c <= 0x2f: # intermediate, SP to /
                    self.__csi_intermediate.append(c)
                    self.__parse_state = STATE_CSI_INTERMEDIATE
                elif 0x40 <= c and c <= 0x7e: # Final byte, @ to ~
                    context.dispatch_csi(self.__csi_parameter,
                                         self.__csi_intermediate,
                                         c)
                    self.__parse_state = STATE_GROUND
                else:
                    self.__parse_state = STATE_GROUND

            else:
                context.dispatch_char(c)

################################################################################
#
# Handler implementation
#
class DefaultHandler(EventObserver):
    ''' default handler, pass through all ESC/CSI/string seqnceses '''
    def __init__(self):
        pass

# EventObserver
    def handle_csi(self, context, parameter, intermediate, final):
        context.write(0x1b) # ESC
        context.write(0x5b) # [
        for c in parameter:
            context.write(c)
        for c in intermediate:
            context.write(c)
        context.write(final)

    def handle_esc(self, context, prefix, final):
        context.write(0x1b) # ESC
        for c in prefix:
            context.write(c)
        context.write(final)

    def handle_control_string(self, context, prefix, value):
        context.write(0x1b) # ESC
        context.write(prefix)
        for c in value:
            context.write(c)
        context.write(0x1b) # ESC
        context.write(0x5c) # \

    def handle_char(self, context, c):
        if c < 0x20 or c == 0x7f:
            context.write(c)
        else:
            context.write(c)

################################################################################
#
# Dispatcher implementation
#
class ParseContext(OutputStream, EventDispatcher):

    def __init__(self,
                 termenc = 'UTF-8',
                 scanner = DefaultScanner(),
                 handler = DefaultHandler()):
        self.__termenc = termenc
        self.__scanner = scanner
        self.__handler = handler
        self.__output = codecs.getwriter(termenc)(StringIO())

    def __iter__(self):
        return self.__scanner.__iter__()

    def writestring(self, data):
        self.__output.write(data)

    def assign(self, data):
        self.__scanner.assign(data, self.__termenc)
        self.__output.truncate(0)

# OutputStream
    def write(self, c):
        if c < 0x80:
            self.__output.write(chr(c))
        else:
            try:
                self.__output.write(unichr(c))
            except:
                self.__output.write('?')

    def getvalue(self):
        return self.__output.getvalue()

# EventDispatcher
    def dispatch_esc(self, prefix, final):
        self.__handler.handle_esc(self, prefix, final)

    def dispatch_csi(self, parameter, intermediate, final):
        self.__handler.handle_csi(self, parameter, intermediate, final)

    def dispatch_control_string(self, prefix, value):
        self.__handler.handle_control_string(self, prefix, value)

    def dispatch_char(self, c):
        self.__handler.handle_char(self, c)

################################################################################
#
# Settings
#
class Settings:

    def __init__(self,
                 command,
                 term,
                 lang,
                 termenc,
                 stdin,
                 stdout,
                 inputscanner=DefaultScanner(),
                 inputparser=DefaultParser(),
                 inputhandler=DefaultHandler(),
                 outputscanner=DefaultScanner(),
                 outputparser=DefaultParser(),
                 outputhandler=DefaultHandler()):
        self.command = command
        self.term = term
        self.lang = lang
        self.termenc = termenc
        self.stdin = stdin
        self.stdout = stdout
        self.inputscanner = inputscanner
        self.inputparser = inputparser
        self.inputhandler = inputhandler
        self.outputscanner = outputscanner
        self.outputparser = outputparser
        self.outputhandler = outputhandler

################################################################################
#
# DefaultPTY
#
class DefaultPTY(PTY):

    __stdin_fileno = None
    __backup = None
    __master = None

    def __init__(self, settings):
        self.__stdin_fileno = settings.stdin.fileno()
        self.__backup = termios.tcgetattr(self.__stdin_fileno)
        new = termios.tcgetattr(self.__stdin_fileno)
        new[0] &= ~(termios.IGNBRK | termios.BRKINT | termios.PARMRK
                  | termios.ISTRIP | termios.INLCR | termios. IGNCR
                  | termios.ICRNL | termios.IXON)
        new[3] = new[3] &~ (termios.ECHO | termios.ICANON)

        vdisable = os.fpathconf(self.__stdin_fileno, 'PC_VDISABLE')

        new[6][termios.VINTR] = vdisable     # Ctrl-C
        new[6][termios.VREPRINT] = vdisable  # Ctrl-R
        new[6][termios.VSTART] = vdisable    # Ctrl-Q
        new[6][termios.VSTOP] = vdisable     # Ctrl-S
        new[6][termios.VLNEXT] = vdisable    # Ctrl-V
        new[6][termios.VWERASE] = vdisable   # Ctrl-W
        new[6][termios.VKILL] = vdisable     # Ctrl-X
        new[6][termios.VSUSP] = vdisable     # Ctrl-Z
        new[6][termios.VQUIT] = vdisable     # Ctrl-\

        termios.tcsetattr(self.__stdin_fileno, termios.TCSANOW, new)
        pid, master = pty.fork()
        if not pid:
            os.environ['TERM'] = settings.term
            os.environ['LANG'] = settings.lang
            os.execlp('/bin/sh',
                      '/bin/sh', '-c',
                      'cd $HOME && exec %s' % settings.command)

        self.__pid = pid
        self.__master = master
        signal.signal(signal.SIGWINCH, lambda no, frame: self.fitsize())

        # call resize once
        self.fitsize()

    def __del__(self):
        termios.tcsetattr(self.__stdin_fileno, termios.TCSANOW, self.__backup)

    def fitsize(self):
         winsize = fcntl.ioctl(self.__stdin_fileno, termios.TIOCGWINSZ, 'hhhh')
         height, width = struct.unpack('hh', winsize)
         self.resize(height, width)

    def resize(self, height, width):
         winsize = struct.pack('HHHH', height, width, 0, 0)
         fcntl.ioctl(self.__master, termios.TIOCSWINSZ, winsize)
         # notify Application process that terminal size has been changed.
         os.kill(self.__pid, signal.SIGWINCH)
         return height, width

    def read(self):
        return os.read(self.__master, BUFFER_SIZE)

    def write(self, data):
        os.write(self.__master, data)

    def drive(self):
        master = self.__master
        stdin_fileno = self.__stdin_fileno
        rfds = [stdin_fileno, master]
        wfds = []
        xfds = [stdin_fileno, master]
        try:
            while True:
                try:
                    rfd, wfd, xfd = select.select(rfds, wfds, xfds)
                    if xfd:
                        break
                    for fd in rfd:
                        if fd == stdin_fileno:
                            data = os.read(stdin_fileno, BUFFER_SIZE)
                            if data:
                                yield data, None
                        elif fd == master:
                            data = self.read()
                            if data:
                                yield None, data
                except OSError, e:
                    no, msg = e
                    if no == errno.EIO:
                        return
                except select.error, e:
                    no, msg = e
                    if no == errno.EINTR:
                        pass
                    else:
                        raise
        finally:
            os.close(master)

################################################################################
#
# Session
#
class Session:

    def start(self, settings):

        tty = DefaultPTY(settings)

        input_context = ParseContext(termenc=settings.termenc,
                                     scanner=settings.inputscanner,
                                     handler=settings.inputhandler)
        input_parser = settings.inputparser

        output_context = ParseContext(termenc=settings.termenc,
                                      scanner=settings.outputscanner,
                                      handler=settings.outputhandler)
        output_parser = settings.outputparser

        stdout = settings.stdout
        for idata, odata in tty.drive():
            if idata:
                input_context.assign(idata)
                input_parser.parse(input_context)
                tty.write(input_context.getvalue())
            if odata:
                output_context.assign(odata)
                output_parser.parse(output_context)
                stdout.write(output_context.getvalue())
                stdout.flush()

''' main '''
if __name__ == '__main__':
    pass

<commit_msg>Fix VDSUSP (Ctrl-y) issue<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# ***** BEGIN LICENSE BLOCK *****
# Copyright (C) 2012  Hayaki Saito <user@zuse.jp>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
# ***** END LICENSE BLOCK *****

import os, termios, pty, signal, fcntl, struct, select, errno
import codecs
try:
    from cStringIO import StringIO
except:
    try:
        from StringIO import StringIO
    except:
        from io import StringIO

from interface import * # terminal filter framework interface

BUFFER_SIZE=2048

################################################################################
#
# Scanner implementation
#
class DefaultScanner(Scanner):
    ''' scan input stream and iterate characters '''
    __data = None

    def assign(self, value, termenc):
        self.__data = unicode(value, termenc, 'ignore')

    def __iter__(self):
        for x in self.__data:
            yield ord(x)

################################################################################
#
# Simple Parser implementation
#
class SimpleParser(Parser):
    ''' simple parser, don't parse ESC/CSI/string seqneces '''
    def __init__(self):
        pass

    def parse(self, context):
        for c in context:
            context.dispatch_char(c)

################################################################################
#
# Default Parser implementation
#
STATE_GROUND = 0
STATE_ESC = 1
STATE_ESC_FINAL = 2

STATE_CSI_PARAMETER = 3
STATE_CSI_INTERMEDIATE = 4
STATE_CSI_FINAL = 5

STATE_OSC = 6
STATE_STR = 7

class DefaultParser(Parser):
    ''' parse ESC/CSI/string seqneces '''

    def __init__(self):
        self.__parse_state = STATE_GROUND
        self.__csi_parameter = []
        self.__csi_intermediate = []
        self.__esc_prefix = []
        self.__str = []
        self.__str_prefix = None
        self.__str_esc_state = False

    def parse(self, context):
        for c in context:
            if self.__parse_state == STATE_OSC:
                # parse control string
                if c == 0x1b:
                    self.__str_esc_state = True
                elif c == 0x5c and self.__str_esc_state == True:
                    context.dispatch_control_string(self.__str_prefix, self.__str)
                    self.__parse_state = STATE_GROUND
                elif c == 0x07:
                    context.dispatch_control_string(self.__str_prefix, self.__str)
                    self.__parse_state = STATE_GROUND
                elif c < 0x20:
                    self.__parse_state = STATE_GROUND
                else:
                    self.__str.append(c)

            elif self.__parse_state == STATE_STR:
                # parse control string
                if c == 0x1b:
                    self.__str_esc_state = True
                elif c == 0x5c and self.__str_esc_state == True:
                    context.dispatch_control_string(self.__str_prefix, self.__str)
                    self.__parse_state = STATE_GROUND
                elif c < 0x20:
                    self.__parse_state = STATE_GROUND
                else:
                    self.__str.append(c)

            elif c == 0x1b: # ESC
                self.__esc_prefix = []
                self.__parse_state = STATE_ESC

            elif c < 0x20 or c == 0x7f: # control character
                context.dispatch_char(c)
                self.__parse_state = STATE_GROUND

            elif self.__parse_state == STATE_ESC:
                if c == 0x5b: # [
                    self.__csi_parameter = []
                    self.__csi_intermediate = []
                    self.__parse_state = STATE_CSI_PARAMETER
                elif c == 0x5d: # ]
                    self.__str_esc_state = False
                    self.__str = []
                    self.__str_prefix = c
                    self.__parse_state = STATE_OSC
                elif c == 0x50 or c == 0x58 or c == 0x5e or c == 0x5f:
                    # P(DCS) or X(SOS) or ^(PM) or _(APC)
                    self.__str_esc_state = False
                    self.__str = []
                    self.__str_prefix = c
                    self.__parse_state = STATE_STR
                elif 0x20 <= c and c <= 0x2f: # SP to /
                    self.__esc_prefix.append(c)
                    self.__parse_state = STATE_ESC_FINAL
                else:
                    context.dispatch_esc(self.__esc_prefix, c)
                    self.__parse_state = STATE_GROUND

            elif self.__parse_state == STATE_ESC_FINAL:
                context.dispatch_esc(self.__esc_prefix, c)
                self.__parse_state = STATE_GROUND

            elif self.__parse_state == STATE_CSI_PARAMETER:
                # parse control sequence
                #
                # CSI P ... P I ... I F
                #     ^
                if 0x30 <= c and c <= 0x3f: # parameter, 0 to ?
                    self.__csi_parameter.append(c)
                elif 0x20 <= c and c <= 0x2f: # intermediate, SP to /
                    self.__csi_intermediate.append(c)
                    self.__parse_state = STATE_CSI_INTERMEDIATE
                elif 0x40 <= c and c <= 0x7e: # Final byte, @ to ~
                    context.dispatch_csi(self.__csi_parameter,
                                         self.__csi_intermediate,
                                         c)
                    self.__parse_state = STATE_GROUND
                else:
                    self.__parse_state = STATE_GROUND

            elif self.__parse_state == STATE_CSI_INTERMEDIATE:
                # parse control sequence
                #
                # CSI P ... P I ... I F
                #             ^
                if 0x20 <= c and c <= 0x2f: # intermediate, SP to /
                    self.__csi_intermediate.append(c)
                    self.__parse_state = STATE_CSI_INTERMEDIATE
                elif 0x40 <= c and c <= 0x7e: # Final byte, @ to ~
                    context.dispatch_csi(self.__csi_parameter,
                                         self.__csi_intermediate,
                                         c)
                    self.__parse_state = STATE_GROUND
                else:
                    self.__parse_state = STATE_GROUND

            else:
                context.dispatch_char(c)

################################################################################
#
# Handler implementation
#
class DefaultHandler(EventObserver):
    ''' default handler, pass through all ESC/CSI/string seqnceses '''
    def __init__(self):
        pass

# EventObserver
    def handle_csi(self, context, parameter, intermediate, final):
        context.write(0x1b) # ESC
        context.write(0x5b) # [
        for c in parameter:
            context.write(c)
        for c in intermediate:
            context.write(c)
        context.write(final)

    def handle_esc(self, context, prefix, final):
        context.write(0x1b) # ESC
        for c in prefix:
            context.write(c)
        context.write(final)

    def handle_control_string(self, context, prefix, value):
        context.write(0x1b) # ESC
        context.write(prefix)
        for c in value:
            context.write(c)
        context.write(0x1b) # ESC
        context.write(0x5c) # \

    def handle_char(self, context, c):
        if c < 0x20 or c == 0x7f:
            context.write(c)
        else:
            context.write(c)

################################################################################
#
# Dispatcher implementation
#
class ParseContext(OutputStream, EventDispatcher):

    def __init__(self,
                 termenc = 'UTF-8',
                 scanner = DefaultScanner(),
                 handler = DefaultHandler()):
        self.__termenc = termenc
        self.__scanner = scanner
        self.__handler = handler
        self.__output = codecs.getwriter(termenc)(StringIO())

    def __iter__(self):
        return self.__scanner.__iter__()

    def writestring(self, data):
        self.__output.write(data)

    def assign(self, data):
        self.__scanner.assign(data, self.__termenc)
        self.__output.truncate(0)

# OutputStream
    def write(self, c):
        if c < 0x80:
            self.__output.write(chr(c))
        else:
            try:
                self.__output.write(unichr(c))
            except:
                self.__output.write('?')

    def getvalue(self):
        return self.__output.getvalue()

# EventDispatcher
    def dispatch_esc(self, prefix, final):
        self.__handler.handle_esc(self, prefix, final)

    def dispatch_csi(self, parameter, intermediate, final):
        self.__handler.handle_csi(self, parameter, intermediate, final)

    def dispatch_control_string(self, prefix, value):
        self.__handler.handle_control_string(self, prefix, value)

    def dispatch_char(self, c):
        self.__handler.handle_char(self, c)

################################################################################
#
# Settings
#
class Settings:

    def __init__(self,
                 command,
                 term,
                 lang,
                 termenc,
                 stdin,
                 stdout,
                 inputscanner=DefaultScanner(),
                 inputparser=DefaultParser(),
                 inputhandler=DefaultHandler(),
                 outputscanner=DefaultScanner(),
                 outputparser=DefaultParser(),
                 outputhandler=DefaultHandler()):
        self.command = command
        self.term = term
        self.lang = lang
        self.termenc = termenc
        self.stdin = stdin
        self.stdout = stdout
        self.inputscanner = inputscanner
        self.inputparser = inputparser
        self.inputhandler = inputhandler
        self.outputscanner = outputscanner
        self.outputparser = outputparser
        self.outputhandler = outputhandler

################################################################################
#
# DefaultPTY
#
class DefaultPTY(PTY):

    __stdin_fileno = None
    __backup = None
    __master = None

    def __init__(self, settings):
        self.__stdin_fileno = settings.stdin.fileno()
        self.__backup = termios.tcgetattr(self.__stdin_fileno)
        new = termios.tcgetattr(self.__stdin_fileno)
        new[0] &= ~(termios.IGNBRK | termios.BRKINT | termios.PARMRK
                  | termios.ISTRIP | termios.INLCR | termios. IGNCR
                  | termios.ICRNL | termios.IXON)
        new[3] = new[3] &~ (termios.ECHO | termios.ICANON)

        vdisable = os.fpathconf(self.__stdin_fileno, 'PC_VDISABLE')

        new[6][termios.VINTR] = vdisable     # Ctrl-C
        new[6][termios.VREPRINT] = vdisable  # Ctrl-R
        new[6][termios.VSTART] = vdisable    # Ctrl-Q
        new[6][termios.VSTOP] = vdisable     # Ctrl-S
        new[6][termios.VLNEXT] = vdisable    # Ctrl-V
        new[6][termios.VWERASE] = vdisable   # Ctrl-W
        new[6][termios.VKILL] = vdisable     # Ctrl-X
        new[6][termios.VSUSP] = vdisable     # Ctrl-Z
        new[6][termios.VQUIT] = vdisable     # Ctrl-\

        VDSUSP = 11
        new[6][VDSUSP] = vdisable    # Ctrl-Y

        termios.tcsetattr(self.__stdin_fileno, termios.TCSANOW, new)
        pid, master = pty.fork()
        if not pid:
            os.environ['TERM'] = settings.term
            os.environ['LANG'] = settings.lang
            os.execlp('/bin/sh',
                      '/bin/sh', '-c',
                      'cd $HOME && exec %s' % settings.command)

        self.__pid = pid
        self.__master = master
        signal.signal(signal.SIGWINCH, lambda no, frame: self.fitsize())

        # call resize once
        self.fitsize()

    def __del__(self):
        termios.tcsetattr(self.__stdin_fileno, termios.TCSANOW, self.__backup)

    def fitsize(self):
         winsize = fcntl.ioctl(self.__stdin_fileno, termios.TIOCGWINSZ, 'hhhh')
         height, width = struct.unpack('hh', winsize)
         self.resize(height, width)

    def resize(self, height, width):
         winsize = struct.pack('HHHH', height, width, 0, 0)
         fcntl.ioctl(self.__master, termios.TIOCSWINSZ, winsize)
         # notify Application process that terminal size has been changed.
         os.kill(self.__pid, signal.SIGWINCH)
         return height, width

    def read(self):
        return os.read(self.__master, BUFFER_SIZE)

    def write(self, data):
        os.write(self.__master, data)

    def drive(self):
        master = self.__master
        stdin_fileno = self.__stdin_fileno
        rfds = [stdin_fileno, master]
        wfds = []
        xfds = [stdin_fileno, master]
        try:
            while True:
                try:
                    rfd, wfd, xfd = select.select(rfds, wfds, xfds)
                    if xfd:
                        break
                    for fd in rfd:
                        if fd == stdin_fileno:
                            data = os.read(stdin_fileno, BUFFER_SIZE)
                            if data:
                                yield data, None
                        elif fd == master:
                            data = self.read()
                            if data:
                                yield None, data
                except OSError, e:
                    no, msg = e
                    if no == errno.EIO:
                        return
                except select.error, e:
                    no, msg = e
                    if no == errno.EINTR:
                        pass
                    else:
                        raise
        finally:
            os.close(master)

################################################################################
#
# Session
#
class Session:

    def start(self, settings):

        tty = DefaultPTY(settings)

        input_context = ParseContext(termenc=settings.termenc,
                                     scanner=settings.inputscanner,
                                     handler=settings.inputhandler)
        input_parser = settings.inputparser

        output_context = ParseContext(termenc=settings.termenc,
                                      scanner=settings.outputscanner,
                                      handler=settings.outputhandler)
        output_parser = settings.outputparser

        stdout = settings.stdout
        for idata, odata in tty.drive():
            if idata:
                input_context.assign(idata)
                input_parser.parse(input_context)
                tty.write(input_context.getvalue())
            if odata:
                output_context.assign(odata)
                output_parser.parse(output_context)
                stdout.write(output_context.getvalue())
                stdout.flush()

''' main '''
if __name__ == '__main__':
    pass

<|endoftext|>


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

# Author : Himanshu Shekhar < https://github.com/himanshub16 > (2017)
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""
Checks for broken links in file names specified as command line parameters.

There are a ton of a solutions available for validating URLs in string using
regex, but less for searching, of which very few are accurate.
This snippet is intended to just do the required work, and avoid complexities.
Django Validator has pretty good regex for validation, but we have to find
urls instead of validating them. (REFERENCES [7])
There's always room for improvement.

Method:
* Match URLs using regex (REFERENCES [1]])
* Some URLs need to be fixed, as they have < (or) > due to inefficient regex.
* Remove duplicates (because regex is not 100% efficient as of now).
* Check validity of URL, using HEAD request. (HEAD to save bandwidth)
  Uses requests module for others are painful to use. REFERENCES[9]
  Handles redirects, http, https, ftp as well.

REFERENCES:
Using [1] with some modificatons for including ftp
[1] http://stackoverflow.com/a/6883094/5163807
[2] http://stackoverflow.com/a/31952097/5163807
[3] http://daringfireball.net/2010/07/improved_regex_for_matching_urls
[4] https://mathiasbynens.be/demo/url-regex
[5] https://github.com/django/django/blob/master/django/core/validators.py
[6] https://data.iana.org/TLD/tlds-alpha-by-domain.txt
[7] https://codereview.stackexchange.com/questions/19663/http-url-validating
[8] https://developer.mozilla.org/en-US/docs/Web/HTTP/Methods/HEAD
[9] http://docs.python-requests.org/

"""

from __future__ import print_function

import os
import re
import sys

import requests

HERE = os.path.abspath(os.path.dirname(__file__))

REGEX = r'(?:http|ftp|https)?://' \
        r'(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\(\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+'

def get_urls(filename):
    """Extracts all URLs available in specified filename
    """
    # fname = os.path.abspath(os.path.join(HERE, filename))
    # expecting absolute path
    fname = os.path.abspath(filename)
    text = ''
    with open(fname) as f:
        text = f.read()

    urls = re.findall(REGEX, text)
    # remove duplicates, list for sets are not iterable
    urls = list(set(urls))
    # correct urls which are between < and/or >
    i = 0
    while i < len(urls):
        urls[i] = re.sub("[<>]", '', urls[i])
        i += 1

    return urls

def validate_url(url):
    """Validate the URL by attempting an HTTP connection.
    Makes an HTTP-HEAD request for each URL.
    Uses requests module.
    """
    try:
        res = requests.head(url)
        return res.ok
    except requests.exceptions.RequestException:
        return False

def main():
    """Main function
    """
    files = sys.argv[1:]
    fails = []
    for fname in files:
        urls = get_urls(fname)
        i = 0
        last = len(urls)
        for url in urls:
            i += 1
            if not validate_url(url):
                fails.append((url, fname))
            sys.stdout.write("\r " +
                             fname + " : " + str(i) + " / " + str(last))
            sys.stdout.flush()

    print()
    if len(fails) == 0:
        print("All URLs are valid. Cheers!")
    else:
        print("Total :", len(fails), "fails!")
        print("Writing failed urls to fails.txt")
        with open("../../fails.txt", 'w') as f:
            for fail in fails:
                f.write(fail[1] + ' : ' + fail[0] + os.linesep)
            f.write('-' * 20)
            f.write(os.linesep*2)

if __name__ == '__main__':
    main()
<commit_msg>add message if no files are provided<commit_after>#!/usr/bin/env python

# Author : Himanshu Shekhar < https://github.com/himanshub16 > (2017)
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

"""
Checks for broken links in file names specified as command line parameters.

There are a ton of a solutions available for validating URLs in string using
regex, but less for searching, of which very few are accurate.
This snippet is intended to just do the required work, and avoid complexities.
Django Validator has pretty good regex for validation, but we have to find
urls instead of validating them. (REFERENCES [7])
There's always room for improvement.

Method:
* Match URLs using regex (REFERENCES [1]])
* Some URLs need to be fixed, as they have < (or) > due to inefficient regex.
* Remove duplicates (because regex is not 100% efficient as of now).
* Check validity of URL, using HEAD request. (HEAD to save bandwidth)
  Uses requests module for others are painful to use. REFERENCES[9]
  Handles redirects, http, https, ftp as well.

REFERENCES:
Using [1] with some modificatons for including ftp
[1] http://stackoverflow.com/a/6883094/5163807
[2] http://stackoverflow.com/a/31952097/5163807
[3] http://daringfireball.net/2010/07/improved_regex_for_matching_urls
[4] https://mathiasbynens.be/demo/url-regex
[5] https://github.com/django/django/blob/master/django/core/validators.py
[6] https://data.iana.org/TLD/tlds-alpha-by-domain.txt
[7] https://codereview.stackexchange.com/questions/19663/http-url-validating
[8] https://developer.mozilla.org/en-US/docs/Web/HTTP/Methods/HEAD
[9] http://docs.python-requests.org/

"""

from __future__ import print_function

import os
import re
import sys

import requests

HERE = os.path.abspath(os.path.dirname(__file__))

REGEX = r'(?:http|ftp|https)?://' \
        r'(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\(\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+'

def get_urls(filename):
    """Extracts all URLs available in specified filename
    """
    # fname = os.path.abspath(os.path.join(HERE, filename))
    # expecting absolute path
    fname = os.path.abspath(filename)
    text = ''
    with open(fname) as f:
        text = f.read()

    urls = re.findall(REGEX, text)
    # remove duplicates, list for sets are not iterable
    urls = list(set(urls))
    # correct urls which are between < and/or >
    i = 0
    while i < len(urls):
        urls[i] = re.sub("[<>]", '', urls[i])
        i += 1

    return urls

def validate_url(url):
    """Validate the URL by attempting an HTTP connection.
    Makes an HTTP-HEAD request for each URL.
    Uses requests module.
    """
    try:
        res = requests.head(url)
        return res.ok
    except requests.exceptions.RequestException:
        return False

def main():
    """Main function
    """
    files = sys.argv[1:]

    if not files:
        return sys.exit("usage: %s <FILES...>" % __name__)
    fails = []
    for fname in files:
        urls = get_urls(fname)
        i = 0
        last = len(urls)
        for url in urls:
            i += 1
            if not validate_url(url):
                fails.append((url, fname))
            sys.stdout.write("\r " +
                             fname + " : " + str(i) + " / " + str(last))
            sys.stdout.flush()

    print()
    if len(fails) == 0:
        print("All links are valid. Cheers!")
    else:
        print("Total :", len(fails), "fails!")
        print("Writing failed urls to fails.txt")
        with open("../../fails.txt", 'w') as f:
            for fail in fails:
                f.write(fail[1] + ' : ' + fail[0] + os.linesep)
            f.write('-' * 20)
            f.write(os.linesep*2)

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before>from __future__ import absolute_import
from datetime import datetime, date, time
from couchdbkit.ext.django.schema import *
from corehq.apps.case import const
from dimagi.utils import parsing
from couchdbkit.schema.properties_proxy import SchemaListProperty
import logging
from datetime import datetime, date, time
from couchdbkit.ext.django.schema import *
from corehq.apps.case import const
from dimagi.utils import parsing
from couchdbkit.schema.properties_proxy import SchemaListProperty
import logging

"""
Couch models for commcare cases.

For details on casexml check out:
http://bitbucket.org/javarosa/javarosa/wiki/casexml
"""

class CaseBase(Document):
    """
    Base class for cases and referrals.
    """
    opened_on = DateTimeProperty()
    modified_on = DateTimeProperty()
    type = StringProperty()
    closed = BooleanProperty(default=False)
    closed_on = DateTimeProperty()

    class Meta:
        app_label = 'case'

class CommCareCaseAction(Document):
    """
    An atomic action on a case. Either a create, update, or close block in
    the xml.
    """
    action_type = StringProperty()
    date = DateTimeProperty()
    visit_date = DateProperty()

    @classmethod
    def from_action_block(cls, action, date, visit_date, action_block):
        if not action in const.CASE_ACTIONS:
            raise ValueError("%s not a valid case action!")

        action = CommCareCaseAction(action_type=action, date=date, visit_date=visit_date)

        # a close block can come without anything inside.
        # if this is the case don't bother trying to post
        # process anything
        if isinstance(action_block, basestring):
            return action

        for item in action_block:
            #if item not in const.CASE_TAGS:
            action[item] = action_block[item]
        return action

    @classmethod
    def new_create_action(cls, date=None):
        """
        Get a new create action
        """
        if not date: date = datetime.utcnow()
        return CommCareCaseAction(action_type=const.CASE_ACTION_CLOSE,
                                  date=date, visit_date=date.date(),
                                  opened_on=date)

    @classmethod
    def new_close_action(cls, date=None):
        """
        Get a new close action
        """
        if not date: date = datetime.utcnow()
        return CommCareCaseAction(action_type=const.CASE_ACTION_CLOSE,
                                  date=date, visit_date=date.date(),
                                  closed_on=date)

    class Meta:
        app_label = 'case'

class Referral(CaseBase):
    """
    A referral, taken from casexml.
    """

    # Referrals have top-level couch guids, but this id is important
    # to the phone, so we keep it here.  This is _not_ globally unique
    # but case_id/referral_id/type should be.
    # (in our world: case_id/referral_id/type)
    referral_id = StringProperty()
    followup_on = DateTimeProperty()
    outcome = StringProperty()

    class Meta:
        app_label = 'case'

    def __unicode__(self):
        return ("%s:%s" % (self.type, self.referral_id))

    def apply_updates(self, date, referral_block):
        if not const.REFERRAL_ACTION_UPDATE in referral_block:
            logging.warn("No update action found in referral block, nothing to be applied")
            return

        update_block = referral_block[const.REFERRAL_ACTION_UPDATE]
        if not self.type == update_block[const.REFERRAL_TAG_TYPE]:
            raise logging.warn("Tried to update from a block with a mismatched type!")
            return

        if date > self.modified_on:
            self.modified_on = date

        if const.REFERRAL_TAG_FOLLOWUP_DATE in referral_block:
            self.followup_on = parsing.string_to_datetime(referral_block[const.REFERRAL_TAG_FOLLOWUP_DATE])

        if const.REFERRAL_TAG_DATE_CLOSED in update_block:
            self.closed = True
            self.closed_on = parsing.string_to_datetime(update_block[const.REFERRAL_TAG_DATE_CLOSED])

    @classmethod
    def from_block(cls, date, block):
        """
        Create referrals from a block of processed data (a dictionary)
        """
        if not const.REFERRAL_ACTION_OPEN in block:
            raise ValueError("No open tag found in referral block!")
        id = block[const.REFERRAL_TAG_ID]
        follow_date = parsing.string_to_datetime(block[const.REFERRAL_TAG_FOLLOWUP_DATE])
        open_block = block[const.REFERRAL_ACTION_OPEN]
        types = open_block[const.REFERRAL_TAG_TYPES].split(" ")

        ref_list = []
        for type in types:
            ref = Referral(referral_id=id, followup_on=follow_date,
                            type=type, opened_on=date, modified_on=date,
                            closed=False)
            ref_list.append(ref)

        # there could be a single update block that closes a referral
        # that we just opened.  not sure why this would happen, but
        # we'll support it.
        if const.REFERRAL_ACTION_UPDATE in block:
            update_block = block[const.REFERRAL_ACTION_UPDATE]
            for ref in ref_list:
                if ref.type == update_block[const.REFERRAL_TAG_TYPE]:
                    ref.apply_updates(date, block)

        return ref_list

class CommCareCase(CaseBase):
    """
    A case, taken from casexml.  This represents the latest
    representation of the case - the result of playing all
    the actions in sequence.
    """

    case_id = StringProperty()

    external_id = StringProperty()
    user_id = StringProperty()

    referrals = SchemaListProperty(Referral)
    actions = SchemaListProperty(CommCareCaseAction)
    name = StringProperty()
    followup_type = StringProperty()

    # date the case actually starts, before this won't be sent to phone.
    # this is for missed appointments, which don't start until the appointment
    # is actually missed
    start_date = DateProperty()
    activation_date = DateProperty() # date the phone triggers it active
    due_date = DateProperty()        # date the phone thinks it's due

    class Meta:
        app_label = 'case'

    def __unicode__(self):
        return "CommCareCase: %s (%s)" % (self.case_id, self.get_id)

    def get_version_token(self):
        """
        A unique token for this version.
        """
        # in theory since case ids are unique and modification dates get updated
        # upon any change, this is all we need
        return "%(case_id)s::%(date_modified)s" % (self.case_id, self.date_modified)

    def is_started(self, since=None):
        """
        Whether the case has started (since a date, or today).
        """
        if since is None:
            since = date.today()
        return self.start_date <= since if self.start_date else True

    @classmethod
    def from_doc(cls, case_block):
        """
        Create a case object from a case block.
        """
        if not const.CASE_ACTION_CREATE in case_block:
            raise ValueError("No create tag found in case block!")

        # create case from required fields in the case/create block
        create_block = case_block[const.CASE_ACTION_CREATE]
        id = case_block[const.CASE_TAG_ID]
        opened_on = parsing.string_to_datetime(case_block[const.CASE_TAG_MODIFIED])

        # create block
        type = create_block[const.CASE_TAG_TYPE_ID]
        name = create_block[const.CASE_TAG_NAME]
        external_id = create_block[const.CASE_TAG_EXTERNAL_ID]
        user_id = create_block[const.CASE_TAG_USER_ID] if const.CASE_TAG_USER_ID in create_block else ""
        create_action = CommCareCaseAction.from_action_block(const.CASE_ACTION_CREATE,
                                                             opened_on, opened_on.date(),
                                                             create_block)

        case = CommCareCase(case_id=id, opened_on=opened_on, modified_on=opened_on,
                     type=type, name=name, user_id=user_id, external_id=external_id,
                     closed=False, actions=[create_action,])

        # apply initial updates, referrals and such, if present
        case.update_from_block(case_block)
        return case

    @classmethod
    def get_by_case_id(cls, id):
        return cls.view(const.VIEW_BY_CASE_ID, key=id).one()

    def update_from_block(self, case_block, visit_date=None):

        mod_date = parsing.string_to_datetime(case_block[const.CASE_TAG_MODIFIED])
        if mod_date > self.modified_on:
            self.modified_on = mod_date

        # you can pass in a visit date, to override the udpate/close action dates
        if not visit_date:
            visit_date = mod_date.date()

        if const.CASE_ACTION_UPDATE in case_block:

            update_block = case_block[const.CASE_ACTION_UPDATE]
            update_action = CommCareCaseAction.from_action_block(const.CASE_ACTION_UPDATE,
                                                                 mod_date, visit_date,
                                                                 update_block)
            self.apply_updates(update_action)
            self.actions.append(update_action)

        if const.CASE_ACTION_CLOSE in case_block:
            close_block = case_block[const.CASE_ACTION_CLOSE]
            close_action = CommCareCaseAction.from_action_block(const.CASE_ACTION_CLOSE,
                                                                mod_date, visit_date,
                                                                close_block)
            self.apply_close(close_action)
            self.actions.append(close_action)

        if const.REFERRAL_TAG in case_block:
            referral_block = case_block[const.REFERRAL_TAG]
            if const.REFERRAL_ACTION_OPEN in referral_block:
                referrals = Referral.from_block(mod_date, referral_block)
                # for some reason extend doesn't work.  disconcerting
                # self.referrals.extend(referrals)
                for referral in referrals:
                    self.referrals.append(referral)
            elif const.REFERRAL_ACTION_UPDATE in referral_block:
                found = False
                update_block = referral_block[const.REFERRAL_ACTION_UPDATE]
                for ref in self.referrals:
                    if ref.type == update_block[const.REFERRAL_TAG_TYPE]:
                        ref.apply_updates(mod_date, referral_block)
                        found = True
                if not found:
                    logging.error(("Tried to update referral type %s for referral %s in case %s "
                                   "but it didn't exist! Nothing will be done about this.") % \
                                   update_block[const.REFERRAL_TAG_TYPE],
                                   referral_block[const.REFERRAL_TAG_ID],
                                   self.case_id)

    def apply_updates(self, update_action):
        """
        Applies updates to a case
        """
        if const.CASE_TAG_TYPE_ID in update_action:
            self.type = update_action[const.CASE_TAG_TYPE_ID]
        if const.CASE_TAG_NAME in update_action:
            self.name = update_action[const.CASE_TAG_NAME]
        if const.CASE_TAG_DATE_OPENED in update_action:
            self.opened_on = update_action[const.CASE_TAG_DATE_OPENED]
        for item in update_action.dynamic_properties():
            print "%s: %s"  %(item, update_action[item])
            if item not in const.CASE_TAGS:
                self[item] = update_action[item]

    def apply_close(self, close_action):
        self.closed = True
        self.closed_on = datetime.combine(close_action.visit_date, time())

import corehq.apps.case.signals<commit_msg>remove printout<commit_after>from __future__ import absolute_import
from datetime import datetime, date, time
from couchdbkit.ext.django.schema import *
from corehq.apps.case import const
from dimagi.utils import parsing
from couchdbkit.schema.properties_proxy import SchemaListProperty
import logging
from datetime import datetime, date, time
from couchdbkit.ext.django.schema import *
from corehq.apps.case import const
from dimagi.utils import parsing
from couchdbkit.schema.properties_proxy import SchemaListProperty
import logging

"""
Couch models for commcare cases.

For details on casexml check out:
http://bitbucket.org/javarosa/javarosa/wiki/casexml
"""

class CaseBase(Document):
    """
    Base class for cases and referrals.
    """
    opened_on = DateTimeProperty()
    modified_on = DateTimeProperty()
    type = StringProperty()
    closed = BooleanProperty(default=False)
    closed_on = DateTimeProperty()

    class Meta:
        app_label = 'case'

class CommCareCaseAction(Document):
    """
    An atomic action on a case. Either a create, update, or close block in
    the xml.
    """
    action_type = StringProperty()
    date = DateTimeProperty()
    visit_date = DateProperty()

    @classmethod
    def from_action_block(cls, action, date, visit_date, action_block):
        if not action in const.CASE_ACTIONS:
            raise ValueError("%s not a valid case action!")

        action = CommCareCaseAction(action_type=action, date=date, visit_date=visit_date)

        # a close block can come without anything inside.
        # if this is the case don't bother trying to post
        # process anything
        if isinstance(action_block, basestring):
            return action

        for item in action_block:
            #if item not in const.CASE_TAGS:
            action[item] = action_block[item]
        return action

    @classmethod
    def new_create_action(cls, date=None):
        """
        Get a new create action
        """
        if not date: date = datetime.utcnow()
        return CommCareCaseAction(action_type=const.CASE_ACTION_CLOSE,
                                  date=date, visit_date=date.date(),
                                  opened_on=date)

    @classmethod
    def new_close_action(cls, date=None):
        """
        Get a new close action
        """
        if not date: date = datetime.utcnow()
        return CommCareCaseAction(action_type=const.CASE_ACTION_CLOSE,
                                  date=date, visit_date=date.date(),
                                  closed_on=date)

    class Meta:
        app_label = 'case'

class Referral(CaseBase):
    """
    A referral, taken from casexml.
    """

    # Referrals have top-level couch guids, but this id is important
    # to the phone, so we keep it here.  This is _not_ globally unique
    # but case_id/referral_id/type should be.
    # (in our world: case_id/referral_id/type)
    referral_id = StringProperty()
    followup_on = DateTimeProperty()
    outcome = StringProperty()

    class Meta:
        app_label = 'case'

    def __unicode__(self):
        return ("%s:%s" % (self.type, self.referral_id))

    def apply_updates(self, date, referral_block):
        if not const.REFERRAL_ACTION_UPDATE in referral_block:
            logging.warn("No update action found in referral block, nothing to be applied")
            return

        update_block = referral_block[const.REFERRAL_ACTION_UPDATE]
        if not self.type == update_block[const.REFERRAL_TAG_TYPE]:
            raise logging.warn("Tried to update from a block with a mismatched type!")
            return

        if date > self.modified_on:
            self.modified_on = date

        if const.REFERRAL_TAG_FOLLOWUP_DATE in referral_block:
            self.followup_on = parsing.string_to_datetime(referral_block[const.REFERRAL_TAG_FOLLOWUP_DATE])

        if const.REFERRAL_TAG_DATE_CLOSED in update_block:
            self.closed = True
            self.closed_on = parsing.string_to_datetime(update_block[const.REFERRAL_TAG_DATE_CLOSED])

    @classmethod
    def from_block(cls, date, block):
        """
        Create referrals from a block of processed data (a dictionary)
        """
        if not const.REFERRAL_ACTION_OPEN in block:
            raise ValueError("No open tag found in referral block!")
        id = block[const.REFERRAL_TAG_ID]
        follow_date = parsing.string_to_datetime(block[const.REFERRAL_TAG_FOLLOWUP_DATE])
        open_block = block[const.REFERRAL_ACTION_OPEN]
        types = open_block[const.REFERRAL_TAG_TYPES].split(" ")

        ref_list = []
        for type in types:
            ref = Referral(referral_id=id, followup_on=follow_date,
                            type=type, opened_on=date, modified_on=date,
                            closed=False)
            ref_list.append(ref)

        # there could be a single update block that closes a referral
        # that we just opened.  not sure why this would happen, but
        # we'll support it.
        if const.REFERRAL_ACTION_UPDATE in block:
            update_block = block[const.REFERRAL_ACTION_UPDATE]
            for ref in ref_list:
                if ref.type == update_block[const.REFERRAL_TAG_TYPE]:
                    ref.apply_updates(date, block)

        return ref_list

class CommCareCase(CaseBase):
    """
    A case, taken from casexml.  This represents the latest
    representation of the case - the result of playing all
    the actions in sequence.
    """

    case_id = StringProperty()

    external_id = StringProperty()
    user_id = StringProperty()

    referrals = SchemaListProperty(Referral)
    actions = SchemaListProperty(CommCareCaseAction)
    name = StringProperty()
    followup_type = StringProperty()

    # date the case actually starts, before this won't be sent to phone.
    # this is for missed appointments, which don't start until the appointment
    # is actually missed
    start_date = DateProperty()
    activation_date = DateProperty() # date the phone triggers it active
    due_date = DateProperty()        # date the phone thinks it's due

    class Meta:
        app_label = 'case'

    def __unicode__(self):
        return "CommCareCase: %s (%s)" % (self.case_id, self.get_id)

    def get_version_token(self):
        """
        A unique token for this version.
        """
        # in theory since case ids are unique and modification dates get updated
        # upon any change, this is all we need
        return "%(case_id)s::%(date_modified)s" % (self.case_id, self.date_modified)

    def is_started(self, since=None):
        """
        Whether the case has started (since a date, or today).
        """
        if since is None:
            since = date.today()
        return self.start_date <= since if self.start_date else True

    @classmethod
    def from_doc(cls, case_block):
        """
        Create a case object from a case block.
        """
        if not const.CASE_ACTION_CREATE in case_block:
            raise ValueError("No create tag found in case block!")

        # create case from required fields in the case/create block
        create_block = case_block[const.CASE_ACTION_CREATE]
        id = case_block[const.CASE_TAG_ID]
        opened_on = parsing.string_to_datetime(case_block[const.CASE_TAG_MODIFIED])

        # create block
        type = create_block[const.CASE_TAG_TYPE_ID]
        name = create_block[const.CASE_TAG_NAME]
        external_id = create_block[const.CASE_TAG_EXTERNAL_ID]
        user_id = create_block[const.CASE_TAG_USER_ID] if const.CASE_TAG_USER_ID in create_block else ""
        create_action = CommCareCaseAction.from_action_block(const.CASE_ACTION_CREATE,
                                                             opened_on, opened_on.date(),
                                                             create_block)

        case = CommCareCase(case_id=id, opened_on=opened_on, modified_on=opened_on,
                     type=type, name=name, user_id=user_id, external_id=external_id,
                     closed=False, actions=[create_action,])

        # apply initial updates, referrals and such, if present
        case.update_from_block(case_block)
        return case

    @classmethod
    def get_by_case_id(cls, id):
        return cls.view(const.VIEW_BY_CASE_ID, key=id).one()

    def update_from_block(self, case_block, visit_date=None):

        mod_date = parsing.string_to_datetime(case_block[const.CASE_TAG_MODIFIED])
        if mod_date > self.modified_on:
            self.modified_on = mod_date

        # you can pass in a visit date, to override the udpate/close action dates
        if not visit_date:
            visit_date = mod_date.date()

        if const.CASE_ACTION_UPDATE in case_block:

            update_block = case_block[const.CASE_ACTION_UPDATE]
            update_action = CommCareCaseAction.from_action_block(const.CASE_ACTION_UPDATE,
                                                                 mod_date, visit_date,
                                                                 update_block)
            self.apply_updates(update_action)
            self.actions.append(update_action)

        if const.CASE_ACTION_CLOSE in case_block:
            close_block = case_block[const.CASE_ACTION_CLOSE]
            close_action = CommCareCaseAction.from_action_block(const.CASE_ACTION_CLOSE,
                                                                mod_date, visit_date,
                                                                close_block)
            self.apply_close(close_action)
            self.actions.append(close_action)

        if const.REFERRAL_TAG in case_block:
            referral_block = case_block[const.REFERRAL_TAG]
            if const.REFERRAL_ACTION_OPEN in referral_block:
                referrals = Referral.from_block(mod_date, referral_block)
                # for some reason extend doesn't work.  disconcerting
                # self.referrals.extend(referrals)
                for referral in referrals:
                    self.referrals.append(referral)
            elif const.REFERRAL_ACTION_UPDATE in referral_block:
                found = False
                update_block = referral_block[const.REFERRAL_ACTION_UPDATE]
                for ref in self.referrals:
                    if ref.type == update_block[const.REFERRAL_TAG_TYPE]:
                        ref.apply_updates(mod_date, referral_block)
                        found = True
                if not found:
                    logging.error(("Tried to update referral type %s for referral %s in case %s "
                                   "but it didn't exist! Nothing will be done about this.") % \
                                   update_block[const.REFERRAL_TAG_TYPE],
                                   referral_block[const.REFERRAL_TAG_ID],
                                   self.case_id)

    def apply_updates(self, update_action):
        """
        Applies updates to a case
        """
        if const.CASE_TAG_TYPE_ID in update_action:
            self.type = update_action[const.CASE_TAG_TYPE_ID]
        if const.CASE_TAG_NAME in update_action:
            self.name = update_action[const.CASE_TAG_NAME]
        if const.CASE_TAG_DATE_OPENED in update_action:
            self.opened_on = update_action[const.CASE_TAG_DATE_OPENED]
        for item in update_action.dynamic_properties():
            if item not in const.CASE_TAGS:
                self[item] = update_action[item]

    def apply_close(self, close_action):
        self.closed = True
        self.closed_on = datetime.combine(close_action.visit_date, time())

import corehq.apps.case.signals<|endoftext|>