<commit_before># -*- coding: utf-8 -*-
"""
test_priority
~~~~~~~~~~~~~

Tests for the Priority trees
"""
from __future__ import division

import itertools

from hypothesis import given
from hypothesis.strategies import integers, lists, tuples

import priority

try:
    from functools import reduce
except ImportError:
    pass

try:
    from math import gcd
except ImportError:
    from fractions import gcd

STREAMS_AND_WEIGHTS = lists(
    elements=tuples(
        integers(min_value=1), integers(min_value=1, max_value=255)
    ),
    unique_by=lambda x: x[0],
)

class TestStream(object):
    def test_stream_repr(self):
        """
        The stream representation renders according to the README.
        """
        s = priority.Stream(stream_id=80, weight=16)
        assert repr(s) == "Stream<id=80, weight=16>"

class TestPriorityTree(object):
    @given(STREAMS_AND_WEIGHTS)
    def test_period_of_repetition(self, streams_and_weights):
        """
        The period of repetition of a priority sequence is given by the
        formula: sum(weights) / gcd(weights). Once that many values have been
        pulled out, the sequence should repeat identically.
        """
        p = priority.PriorityTree()
        weights = []

        for stream, weight in streams_and_weights:
            p.insert_stream(stream_id=stream, weight=weight)
            weights.append(weight)

        period = sum(weights)

        # Pop off the first n elements, which will always be evenly
        # distributed.
        print(streams_and_weights)
        for _ in weights:
            print("    " + str(next(p)))

        pattern = [next(p) for _ in range(period)]
        print("    " + str(pattern))
        pattern = itertools.cycle(pattern)

        for i in range(period * 20):
            assert next(p) == next(pattern), i
<commit_msg>Better test docstrings.<commit_after># -*- coding: utf-8 -*-
"""
test_priority
~~~~~~~~~~~~~

Tests for the Priority trees
"""
from __future__ import division

import itertools

from hypothesis import given
from hypothesis.strategies import integers, lists, tuples

import priority

try:
    from functools import reduce
except ImportError:
    pass

try:
    from math import gcd
except ImportError:
    from fractions import gcd

STREAMS_AND_WEIGHTS = lists(
    elements=tuples(
        integers(min_value=1), integers(min_value=1, max_value=255)
    ),
    unique_by=lambda x: x[0],
)

class TestStream(object):
    def test_stream_repr(self):
        """
        The stream representation renders according to the README.
        """
        s = priority.Stream(stream_id=80, weight=16)
        assert repr(s) == "Stream<id=80, weight=16>"

class TestPriorityTreeOutput(object):
    """
    These tests use Hypothesis to attempt to bound the output of iterating over
    the priority tree. In particular, their goal is to ensure that the output
    of the tree is "good enough": that it meets certain requirements on
    fairness and equidistribution.
    """
    @given(STREAMS_AND_WEIGHTS)
    def test_period_of_repetition(self, streams_and_weights):
        """
        The period of repetition of a priority sequence is given by the sum of
        the weights of the streams. Once that many values have been pulled out
        the sequence repeats identically.
        """
        p = priority.PriorityTree()
        weights = []

        for stream, weight in streams_and_weights:
            p.insert_stream(stream_id=stream, weight=weight)
            weights.append(weight)

        period = sum(weights)

        # Pop off the first n elements, which will always be evenly
        # distributed.
        print(streams_and_weights)
        for _ in weights:
            print("    " + str(next(p)))

        pattern = [next(p) for _ in range(period)]
        print("    " + str(pattern))
        pattern = itertools.cycle(pattern)

        for i in range(period * 20):
            assert next(p) == next(pattern), i
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
from helpers import create_monomial_exponents2

import math
import numpy
import numpy.testing
import pytest
import quadrature
import sympy

import os
import matplotlib as mpl
if 'DISPLAY' not in os.environ:
    # headless mode, for remote executions (and travis)
    mpl.use('Agg')
from matplotlib import pyplot as plt

def _integrate_exact(f, triangle):
    #
    # Note that
    #
    #     \int_T f(x) dx = \int_T0 |J(xi)| f(P(xi)) dxi
    #
    # with
    #
    #     P(xi) = x0 * (1-xi[0]-xi[1]) + x1 * xi[0] + x2 * xi[1].
    #
    # and T0 being the reference triangle [(0.0, 0.0), (1.0, 0.0), (0.0,
    # 1.0)].
    # The determinant of the transformation matrix J equals twice the volume of
    # the triangle. (See, e.g.,
    # <http://math2.uncc.edu/~shaodeng/TEACHING/math5172/Lectures/Lect_15.PDF>).
    #
    xi = sympy.DeferredVector('xi')
    x_xi = \
        + triangle[0] * (1 - xi[0] - xi[1]) \
        + triangle[1] * xi[0] \
        + triangle[2] * xi[1]
    abs_det_J = 2 * quadrature.triangle.volume(triangle)
    exact = sympy.integrate(
        sympy.integrate(abs_det_J * f(x_xi), (xi[1], 0, 1-xi[0])),
        (xi[0], 0, 1)
        )
    return float(exact)

def _integrate_monomial_over_standard_triangle(k):
    '''The integral of monomials over the standard triangle is given by

    \int_T x_0^k0 * x1^k1 = (k0!*k1!) / (2+k0+k1)!,

    see, e.g.,
    A set of symmetric quadrature rules on triangles and tetrahedra,
    Linbo Zhang, Tao Cui and Hui Liu,
    Journal of Computational Mathematics,
    Vol. 27, No. 1 (January 2009), pp. 89-96
    '''
    # exp-log to account for large values in numerator and denominator
    return math.exp(
        math.fsum([math.lgamma(kk+1) for kk in k])
        - math.lgamma(3 + sum(k))
        )

@pytest.mark.parametrize('scheme', [
    quadrature.triangle.Centroid(),
    quadrature.triangle.Vertex(),
    quadrature.triangle.SevenPoint(),
    quadrature.triangle.HammerMarloweStroud(1),
    quadrature.triangle.HammerMarloweStroud(2),
    quadrature.triangle.HammerMarloweStroud(3),
    quadrature.triangle.HammerMarloweStroud(4),
    quadrature.triangle.HammerMarloweStroud(5),
    quadrature.triangle.NewtonCotesClosed(1),
    quadrature.triangle.NewtonCotesClosed(2),
    quadrature.triangle.NewtonCotesClosed(3),
    quadrature.triangle.NewtonCotesClosed(4),
    quadrature.triangle.NewtonCotesClosed(5),
    quadrature.triangle.NewtonCotesOpen(0),
    quadrature.triangle.NewtonCotesOpen(1),
    quadrature.triangle.NewtonCotesOpen(2),
    quadrature.triangle.NewtonCotesOpen(3),
    quadrature.triangle.NewtonCotesOpen(4),
    quadrature.triangle.NewtonCotesOpen(5),
    quadrature.triangle.Strang(1),
    quadrature.triangle.Strang(2),
    quadrature.triangle.Strang(3),
    quadrature.triangle.Strang(4),
    quadrature.triangle.Strang(5),
    quadrature.triangle.Strang(6),
    quadrature.triangle.Strang(7),
    quadrature.triangle.Strang(8),
    quadrature.triangle.Strang(9),
    quadrature.triangle.Strang(10),
    quadrature.triangle.LynessJespersen(1),
    quadrature.triangle.LynessJespersen(2),
    quadrature.triangle.LynessJespersen(3),
    quadrature.triangle.LynessJespersen(4),
    quadrature.triangle.LynessJespersen(5),
    quadrature.triangle.LynessJespersen(6),
    quadrature.triangle.LynessJespersen(7),
    quadrature.triangle.LynessJespersen(8),
    quadrature.triangle.LynessJespersen(9),
    quadrature.triangle.LynessJespersen(10),
    quadrature.triangle.LynessJespersen(11),
    quadrature.triangle.LynessJespersen(12),
    quadrature.triangle.LynessJespersen(13),
    quadrature.triangle.LynessJespersen(14),
    quadrature.triangle.LynessJespersen(15),
    quadrature.triangle.LynessJespersen(16),
    quadrature.triangle.LynessJespersen(17),
    quadrature.triangle.LynessJespersen(18),
    quadrature.triangle.LynessJespersen(19),
    quadrature.triangle.LynessJespersen(20),
    quadrature.triangle.LynessJespersen(21),
    quadrature.triangle.Hillion(1),
    quadrature.triangle.Hillion(2),
    quadrature.triangle.Hillion(3),
    quadrature.triangle.Hillion(4),
    quadrature.triangle.Hillion(5),
    quadrature.triangle.LaursenGellert('1'),
    quadrature.triangle.LaursenGellert('2a'),
    quadrature.triangle.LaursenGellert('2b'),
    quadrature.triangle.LaursenGellert('3'),
    quadrature.triangle.LaursenGellert('4'),
    quadrature.triangle.LaursenGellert('5'),
    quadrature.triangle.LaursenGellert('6'),
    quadrature.triangle.LaursenGellert('7'),
    quadrature.triangle.LaursenGellert('8'),
    quadrature.triangle.LaursenGellert('9'),
    quadrature.triangle.LaursenGellert('10'),
    quadrature.triangle.LaursenGellert('11'),
    quadrature.triangle.LaursenGellert('12'),
    quadrature.triangle.LaursenGellert('13'),
    quadrature.triangle.LaursenGellert('14'),
    quadrature.triangle.LaursenGellert('15a'),
    quadrature.triangle.LaursenGellert('15b'),
    quadrature.triangle.Cubtri(),
    quadrature.triangle.Triex(19),
    quadrature.triangle.Triex(28),
    quadrature.triangle.Dunavant(1),
    quadrature.triangle.Dunavant(2),
    quadrature.triangle.Dunavant(3),
    quadrature.triangle.Dunavant(4),
    quadrature.triangle.Dunavant(5),
    quadrature.triangle.Dunavant(6),
    quadrature.triangle.Dunavant(7),
    quadrature.triangle.Dunavant(8),
    quadrature.triangle.Dunavant(9),
    quadrature.triangle.Dunavant(10),
    quadrature.triangle.Dunavant(11),
    quadrature.triangle.Dunavant(12),
    quadrature.triangle.Dunavant(13),
    quadrature.triangle.Dunavant(14),
    quadrature.triangle.Dunavant(15),
    quadrature.triangle.Dunavant(16),
    quadrature.triangle.Dunavant(17),
    quadrature.triangle.Dunavant(18),
    quadrature.triangle.Dunavant(19),
    quadrature.triangle.Dunavant(20),
    quadrature.triangle.CoolsHaegemans(1),
    quadrature.triangle.CoolsHaegemans(2),
    quadrature.triangle.BerntsenEspelid(1),
    quadrature.triangle.BerntsenEspelid(2),
    quadrature.triangle.BerntsenEspelid(3),
    quadrature.triangle.BerntsenEspelid(4),
    quadrature.triangle.LiuVinokur(1),
    quadrature.triangle.LiuVinokur(2),
    quadrature.triangle.LiuVinokur(3),
    quadrature.triangle.LiuVinokur(4),
    quadrature.triangle.LiuVinokur(5),
    quadrature.triangle.LiuVinokur(6),
    quadrature.triangle.LiuVinokur(7),
    quadrature.triangle.LiuVinokur(8),
    quadrature.triangle.LiuVinokur(9),
    quadrature.triangle.LiuVinokur(10),
    quadrature.triangle.LiuVinokur(11),
    quadrature.triangle.LiuVinokur(12),
    quadrature.triangle.LiuVinokur(13),
    quadrature.triangle.WandzuraXiao(1),
    quadrature.triangle.WandzuraXiao(2),
    quadrature.triangle.WandzuraXiao(3),
    quadrature.triangle.WandzuraXiao(4),
    quadrature.triangle.WandzuraXiao(5),
    quadrature.triangle.WandzuraXiao(6),
    quadrature.triangle.TaylorWingateBos(1),
    quadrature.triangle.TaylorWingateBos(2),
    quadrature.triangle.TaylorWingateBos(4),
    quadrature.triangle.TaylorWingateBos(5),
    quadrature.triangle.TaylorWingateBos(8),
    quadrature.triangle.ZhangCuiLiu(1),
    quadrature.triangle.ZhangCuiLiu(2),
    quadrature.triangle.ZhangCuiLiu(3),
    quadrature.triangle.XiaoGimbutas(1),
    quadrature.triangle.XiaoGimbutas(2),
    quadrature.triangle.XiaoGimbutas(3),
    quadrature.triangle.XiaoGimbutas(4),
    quadrature.triangle.XiaoGimbutas(5),
    quadrature.triangle.XiaoGimbutas(6),
    quadrature.triangle.XiaoGimbutas(7),
    quadrature.triangle.XiaoGimbutas(8),
    quadrature.triangle.XiaoGimbutas(9),
    quadrature.triangle.XiaoGimbutas(10),
    quadrature.triangle.XiaoGimbutas(11),
    quadrature.triangle.XiaoGimbutas(12),
    quadrature.triangle.XiaoGimbutas(13),
    ])
def test_scheme(scheme):
    # Test integration until we get to a polynomial degree `d` that can no
    # longer be integrated exactly. The scheme's degree is `d-1`.
    triangle = numpy.array([
        [0.0, 0.0],
        [1.0, 0.0],
        [0.0, 1.0]
        ])
    success = True
    degree = 0
    max_degree = scheme.degree + 1
    while success:
        for k in create_monomial_exponents2(degree):
            def poly(x):
                return x[0]**k[0] * x[1]**k[1]
            # exact_val = _integrate_exact(poly, triangle)
            exact_val = _integrate_monomial_over_standard_triangle(k)
            val = quadrature.triangle.integrate(
                    poly, triangle, scheme
                    )
            if abs(exact_val - val) > 1.0e-10:
                success = False
                break
        if not success:
            break
        if degree >= max_degree:
            break
        degree += 1
    assert degree-1 >= scheme.degree
    # numpy.testing.assert_equal(degree-1, scheme.degree)
    return

def test_show():
    triangle = numpy.array([
        [numpy.cos(0.5*numpy.pi), numpy.sin(0.5*numpy.pi)],
        [numpy.cos(7.0/6.0*numpy.pi), numpy.sin(7.0/6.0*numpy.pi)],
        [numpy.cos(11.0/6.0*numpy.pi), numpy.sin(11.0/6.0*numpy.pi)],
        ])
    quadrature.triangle.show(
        triangle,
        # quadrature.triangle.Centroid()
        # quadrature.triangle.Vertex()
        # quadrature.triangle.SevenPoint()
        # quadrature.triangle.Strang(9)
        quadrature.triangle.Dunavant(20)
        # quadrature.triangle.CoolsHaegemans(2)
        )
    return

if __name__ == '__main__':
    test_show()
    plt.show()
    scheme = quadrature.triangle.CoolsHaegemans(2)
    test_scheme(scheme)
<commit_msg>more compact test listing<commit_after># -*- coding: utf-8 -*-
#
from helpers import create_monomial_exponents2

import math
import numpy
import numpy.testing
import pytest
import quadrature
import sympy

import os
import matplotlib as mpl
if 'DISPLAY' not in os.environ:
    # headless mode, for remote executions (and travis)
    mpl.use('Agg')
from matplotlib import pyplot as plt

def _integrate_exact(f, triangle):
    #
    # Note that
    #
    #     \int_T f(x) dx = \int_T0 |J(xi)| f(P(xi)) dxi
    #
    # with
    #
    #     P(xi) = x0 * (1-xi[0]-xi[1]) + x1 * xi[0] + x2 * xi[1].
    #
    # and T0 being the reference triangle [(0.0, 0.0), (1.0, 0.0), (0.0,
    # 1.0)].
    # The determinant of the transformation matrix J equals twice the volume of
    # the triangle. (See, e.g.,
    # <http://math2.uncc.edu/~shaodeng/TEACHING/math5172/Lectures/Lect_15.PDF>).
    #
    xi = sympy.DeferredVector('xi')
    x_xi = \
        + triangle[0] * (1 - xi[0] - xi[1]) \
        + triangle[1] * xi[0] \
        + triangle[2] * xi[1]
    abs_det_J = 2 * quadrature.triangle.volume(triangle)
    exact = sympy.integrate(
        sympy.integrate(abs_det_J * f(x_xi), (xi[1], 0, 1-xi[0])),
        (xi[0], 0, 1)
        )
    return float(exact)

def _integrate_monomial_over_standard_triangle(k):
    '''The integral of monomials over the standard triangle is given by

    \int_T x_0^k0 * x1^k1 = (k0!*k1!) / (2+k0+k1)!,

    see, e.g.,
    A set of symmetric quadrature rules on triangles and tetrahedra,
    Linbo Zhang, Tao Cui and Hui Liu,
    Journal of Computational Mathematics,
    Vol. 27, No. 1 (January 2009), pp. 89-96
    '''
    # exp-log to account for large values in numerator and denominator
    return math.exp(
        math.fsum([math.lgamma(kk+1) for kk in k])
        - math.lgamma(3 + sum(k))
        )

@pytest.mark.parametrize(
    'scheme',
    [quadrature.triangle.Centroid()]
    + [quadrature.triangle.Vertex()]
    + [quadrature.triangle.SevenPoint()]
    + [quadrature.triangle.HammerMarloweStroud(k) for k in range(1, 6)]
    + [quadrature.triangle.NewtonCotesClosed(k) for k in range(1, 6)]
    + [quadrature.triangle.NewtonCotesOpen(k) for k in range(6)]
    + [quadrature.triangle.Strang(k) for k in range(1, 11)]
    + [quadrature.triangle.LynessJespersen(k) for k in range(1, 22)]
    + [quadrature.triangle.Hillion(k) for k in range(1, 6)]
    + [quadrature.triangle.LaursenGellert(key) for key in [
        '1', '2a', '2b', '3', '4', '5', '6', '7', '8', '9', '10', '11', '12',
        '13', '14', '15a', '15b'
        ]]
    + [quadrature.triangle.Cubtri()]
    + [quadrature.triangle.Triex(19), quadrature.triangle.Triex(28)]
    + [quadrature.triangle.Dunavant(k) for k in range(1, 21)]
    + [quadrature.triangle.CoolsHaegemans(k) for k in [1, 2]]
    + [quadrature.triangle.BerntsenEspelid(k) for k in range(1, 5)]
    + [quadrature.triangle.LiuVinokur(k) for k in range(1, 14)]
    + [quadrature.triangle.WandzuraXiao(k) for k in range(1, 7)]
    + [quadrature.triangle.TaylorWingateBos(k) for k in [1, 2, 4, 5, 8]]
    + [quadrature.triangle.ZhangCuiLiu(k) for k in [1, 2, 3]]
    + [quadrature.triangle.XiaoGimbutas(k) for k in range(1, 14)]
    )
def test_scheme(scheme):
    # Test integration until we get to a polynomial degree `d` that can no
    # longer be integrated exactly. The scheme's degree is `d-1`.
    triangle = numpy.array([
        [0.0, 0.0],
        [1.0, 0.0],
        [0.0, 1.0]
        ])
    success = True
    degree = 0
    max_degree = scheme.degree + 1
    while success:
        for k in create_monomial_exponents2(degree):
            def poly(x):
                return x[0]**k[0] * x[1]**k[1]
            # exact_val = _integrate_exact(poly, triangle)
            exact_val = _integrate_monomial_over_standard_triangle(k)
            val = quadrature.triangle.integrate(
                    poly, triangle, scheme
                    )
            if abs(exact_val - val) > 1.0e-10:
                success = False
                break
        if not success:
            break
        if degree >= max_degree:
            break
        degree += 1
    assert degree-1 >= scheme.degree
    # numpy.testing.assert_equal(degree-1, scheme.degree)
    return

def test_show():
    triangle = numpy.array([
        [numpy.cos(0.5*numpy.pi), numpy.sin(0.5*numpy.pi)],
        [numpy.cos(7.0/6.0*numpy.pi), numpy.sin(7.0/6.0*numpy.pi)],
        [numpy.cos(11.0/6.0*numpy.pi), numpy.sin(11.0/6.0*numpy.pi)],
        ])
    quadrature.triangle.show(
        triangle,
        # quadrature.triangle.Centroid()
        # quadrature.triangle.Vertex()
        # quadrature.triangle.SevenPoint()
        # quadrature.triangle.Strang(9)
        quadrature.triangle.Dunavant(20)
        # quadrature.triangle.CoolsHaegemans(2)
        )
    return

if __name__ == '__main__':
    test_show()
    plt.show()
    scheme = quadrature.triangle.CoolsHaegemans(2)
    test_scheme(scheme)
<|endoftext|>


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

import os
import sys
import re
import unittest
from cgi import escape as html_escape

sys.path.append('../lib/')
import xml_util

text_fixtures = os.path.join(os.curdir, 'fixtures/text/')
with open(text_fixtures+'accented_characters.txt') as f:
    accent_file = f.read()

class Test_xml_util(unittest.TestCase):

    def setUp(self):
        pass

    def test_flatten(self):
        testlist = [ [1,4,7], [2,5,8], [3,6,9] ]
        reslist = xml_util.flatten(testlist)
        goallist = [ [1,2,3], [4,5,6], [7,8,9] ]
        self.assertTrue(reslist == goallist, \
            "{0}\nshould be\n{1}".format(reslist,goallist))

    def test_extend_padding(self):
        testlist = [ [1,2,3], [4,5], [5,6,7,8] ]
        reslist = xml_util.extend_padding(testlist,0)
        goallist = [ [1,2,3,0], [4,5,0,0], [5,6,7,8] ]
        self.assertTrue(reslist == goallist, \
            "{0}\nshould be\n{1}".format(reslist,goallist))

    def test_extend_padding_string(self):
        testlist = [ ['a','b','c'], ['d'] ]
        reslist = xml_util.extend_padding(testlist)
        goallist = [ ['a','b','c'], ['d','',''] ]
        self.assertTrue(reslist == goallist, \
            "{0}\nshould be\n{1}".format(reslist,goallist))

    def test_flatten_with_extend(self):
        testlist = [ [1,4,7], [2,5,8], [3,6] ]
        testlist = xml_util.extend_padding(testlist,0)
        reslist = xml_util.flatten(testlist)
        goallist = [ [1,2,3], [4,5,6], [7,8,0] ]
        self.assertTrue(reslist == goallist, \
            "{0}\nshould be\n{1}".format(reslist,goallist))

    def test_flatten_with_extend_multiple(self):
        testlist = [ [1,4,7], [2], [3,6] ]
        testlist = xml_util.extend_padding(testlist,0)
        reslist = xml_util.flatten(testlist)
        goallist = [ [1,2,3], [4,0,6], [7,0,0] ]
        self.assertTrue(reslist == goallist, \
            "{0}\nshould be\n{1}".format(reslist,goallist))

    def test_escape_html_nosub(self):
        teststring = "<tag1> ampersand here: & </tag1>"
        resstring = xml_util.escape_html_nosub(teststring)
        goalstring = html_escape(teststring)
        self.assertTrue(resstring == goalstring, \
            "{0}\nshould be\n{1}".format(resstring,goalstring))

    def test_escape_html_nosub2(self):
        substart = ""
        subend = ""
        teststring = "<escape & skip sub tags>"
        resstring = xml_util.escape_html_nosub(substart+teststring+subend)
        goalstring = substart+html_escape(teststring)+subend
        self.assertTrue(resstring == goalstring, \
            "{0}\nshould be\n{1}".format(resstring,goalstring))

    def test_normalize_utf8_defaultstring(self):
        # this is a PYTHON DEFAULT string consisting of the characters supported by unicode
        teststring_normal = """!@#$%^&*()_+-=QWERTYqwerty<>,.:";'?/{}[]|\\"""
        resstring = xml_util.normalize_utf8(teststring_normal)
        self.assertTrue(teststring_normal == resstring, \
            "{0}\nshould be\n{1}".format(resstring, teststring_normal))

    def test_normalize_utf8_unicodestring(self):
        # this is a UNICODE string consisting of the characters supported by unicode
        teststring_normal = unicode("""!@#$%^&*()_+-=QWERTYqwerty<>,.:";'?/{}[]|\\""")
        resstring = xml_util.normalize_utf8(teststring_normal)
        self.assertTrue(teststring_normal == resstring, \
            "{0}\nshould be\n{1}".format(resstring, teststring_normal))

    def test_normalize_utf8_withaccents(self):
        teststring = accent_file
        resstring = xml_util.normalize_utf8(teststring)
        self.assertTrue(teststring.decode('utf-8') == resstring, resstring)

unittest.main()
<commit_msg>remove whitespace<commit_after>#!/usr/bin/env python

import os
import sys
import re
import unittest
from cgi import escape as html_escape

sys.path.append('../lib/')
import xml_util

text_fixtures = os.path.join(os.curdir, 'fixtures/text/')
with open(text_fixtures+'accented_characters.txt') as f:
    accent_file = f.read()

class Test_xml_util(unittest.TestCase):

    def setUp(self):
        pass

    def test_flatten(self):
        testlist = [ [1,4,7], [2,5,8], [3,6,9] ]
        reslist = xml_util.flatten(testlist)
        goallist = [ [1,2,3], [4,5,6], [7,8,9] ]
        self.assertTrue(reslist == goallist, \
            "{0}\nshould be\n{1}".format(reslist,goallist))

    def test_extend_padding(self):
        testlist = [ [1,2,3], [4,5], [5,6,7,8] ]
        reslist = xml_util.extend_padding(testlist,0)
        goallist = [ [1,2,3,0], [4,5,0,0], [5,6,7,8] ]
        self.assertTrue(reslist == goallist, \
            "{0}\nshould be\n{1}".format(reslist,goallist))

    def test_extend_padding_string(self):
        testlist = [ ['a','b','c'], ['d'] ]
        reslist = xml_util.extend_padding(testlist)
        goallist = [ ['a','b','c'], ['d','',''] ]
        self.assertTrue(reslist == goallist, \
            "{0}\nshould be\n{1}".format(reslist,goallist))

    def test_flatten_with_extend(self):
        testlist = [ [1,4,7], [2,5,8], [3,6] ]
        testlist = xml_util.extend_padding(testlist,0)
        reslist = xml_util.flatten(testlist)
        goallist = [ [1,2,3], [4,5,6], [7,8,0] ]
        self.assertTrue(reslist == goallist, \
            "{0}\nshould be\n{1}".format(reslist,goallist))

    def test_flatten_with_extend_multiple(self):
        testlist = [ [1,4,7], [2], [3,6] ]
        testlist = xml_util.extend_padding(testlist,0)
        reslist = xml_util.flatten(testlist)
        goallist = [ [1,2,3], [4,0,6], [7,0,0] ]
        self.assertTrue(reslist == goallist, \
            "{0}\nshould be\n{1}".format(reslist,goallist))

    def test_escape_html_nosub(self):
        teststring = "<tag1> ampersand here: & </tag1>"
        resstring = xml_util.escape_html_nosub(teststring)
        goalstring = html_escape(teststring)
        self.assertTrue(resstring == goalstring, \
            "{0}\nshould be\n{1}".format(resstring,goalstring))

    def test_escape_html_nosub2(self):
        substart = ""
        subend = ""
        teststring = "<escape & skip sub tags>"
        resstring = xml_util.escape_html_nosub(substart+teststring+subend)
        goalstring = substart+html_escape(teststring)+subend
        self.assertTrue(resstring == goalstring, \
            "{0}\nshould be\n{1}".format(resstring,goalstring))

    def test_normalize_utf8_defaultstring(self):
        # this is a PYTHON DEFAULT string consisting of the characters supported by unicode
        teststring_normal = """!@#$%^&*()_+-=QWERTYqwerty<>,.:";'?/{}[]|\\"""
        resstring = xml_util.normalize_utf8(teststring_normal)
        self.assertTrue(teststring_normal == resstring, \
            "{0}\nshould be\n{1}".format(resstring, teststring_normal))

    def test_normalize_utf8_unicodestring(self):
        # this is a UNICODE string consisting of the characters supported by unicode
        teststring_normal = unicode("""!@#$%^&*()_+-=QWERTYqwerty<>,.:";'?/{}[]|\\""")
        resstring = xml_util.normalize_utf8(teststring_normal)
        self.assertTrue(teststring_normal == resstring, \
            "{0}\nshould be\n{1}".format(resstring, teststring_normal))

    def test_normalize_utf8_withaccents(self):
        teststring = accent_file
        resstring = xml_util.normalize_utf8(teststring)
        self.assertTrue(teststring.decode('utf-8') == resstring, resstring)

unittest.main()
<|endoftext|>


<commit_before>import unittest
import numpy as np

import pgmpy.tests.help_functions as hf
from pgmpy.models import DynamicBayesianNetwork
from pgmpy.factors import TabularCPD

class TestDynamicBayesianNetworkCreation(unittest.TestCase):
    def setUp(self):
        self.network = DynamicBayesianNetwork()

    def test_add_single_node(self):
        self.network.add_node('a')
        self.assertListEqual(self.network.nodes(), ['a'])

    def test_add_multiple_nodes(self):
        self.network.add_nodes_from(['a', 'b', 'c'])
        self.assertListEqual(sorted(self.network.nodes()), ['a', 'b', 'c'])

    def test_add_single_edge_with_timeslice(self):
        self.network.add_edge(('a', 0), ('b', 0))
        self.assertListEqual(sorted(self.network.edges()), [(('a', 0), ('b', 0)), (('a', 1), ('b', 1))])
        self.assertListEqual(sorted(self.network.nodes()), ['a', 'b'])

    def test_add_edge_with_different_number_timeslice(self):
        self.network.add_edge(('a', 2), ('b', 2))
        self.assertListEqual(sorted(self.network.edges()), [(('a', 0), ('b', 0)), (('a', 1), ('b', 1))])

    def test_add_edge_going_backward(self):
        self.assertRaises(ValueError, self.network.add_edge, ('a', 1), ('b', 0))

    def test_add_edge_with_farther_timeslice(self):
        self.assertRaises(ValueError, self.network.add_edge, ('a', 2), ('b', 4))

    def test_add_edge_with_self_loop(self):
        self.assertRaises(ValueError, self.network.add_edge, ('a', 0), ('a', 0))

    def test_add_edge_with_varying_length(self):
        self.assertRaises(ValueError, self.network.add_edge, ('a', 1, 1), ('b', 2))
        self.assertRaises(ValueError, self.network.add_edge, ('b', 2), ('a', 2, 3))

    def test_add_edge_with_closed_path(self):
        self.assertRaises(ValueError, self.network.add_edges_from,
                          [(('a', 0), ('b', 0)), (('b', 0), ('c', 0)), (('c', 0), ('a', 0))])

    def test_add_single_edge_without_timeslice(self):
        self.assertRaises(ValueError, self.network.add_edge, 'a', 'b')

    def test_add_single_edge_with_incorrect_timeslice(self):
        self.assertRaises(ValueError, self.network.add_edge, ('a', 'b'), ('b', 'c'))

    def test_add_multiple_edges(self):
        self.network.add_edges_from([(('a', 0), ('b', 0)), (('a', 0), ('a', 1)), (('b', 0), ('b', 1))])
        self.assertListEqual(sorted(self.network.edges()),
                             [(('a', 0), ('a', 1)), (('a', 0), ('b', 0)), (('a', 1), ('b', 1)), (('b', 0), ('b', 1))])

    def tearDown(self):
        del self.network

class TestDynamicBayesianNetworkMethods(unittest.TestCase):
    def setUp(self):
        self.network = DynamicBayesianNetwork()
        self.grade_cpd = TabularCPD(('G', 0), 3, [[0.3, 0.05, 0.9, 0.5],
                                             [0.4, 0.25, 0.08, 0.3],
                                             [0.3, 0.7, 0.2, 0.2]], [('D', 0), ('I', 0)], [2, 2])
        self.d_i_cpd = TabularCPD(('D', 1), 2, [[0.6, 0.3], [0.4, 0.7]], [('D', 0)], 2)
        self.diff_cpd = TabularCPD(('D', 0), 2, [[0.6, 0.4]])
        self.intel_cpd = TabularCPD(('I', 0), 2, [[0.7, 0.3]])
        self.i_i_cpd = TabularCPD(('I', 1), 2, [[0.5, 0.4], [0.5, 0.6]], [('I', 0)], 2)
        self.grade_1_cpd = TabularCPD(('G', 1), 3, [[0.3, 0.05, 0.9, 0.5],
                                             [0.4, 0.25, 0.08, 0.3],
                                             [0.3, 0.7, 0.2, 0.2]], [('D', 1), ('I', 1)], [2, 2])

    def test_get_intra_and_inter_edges(self):
        self.network.add_edges_from([(('a', 0), ('b', 0)), (('a', 0), ('a', 1)), (('b', 0), ('b', 1))])
        self.assertListEqual(sorted(self.network.get_intra_edges()), [(('a', 0), ('b', 0))])
        self.assertListEqual(sorted(self.network.get_intra_edges(1)), [(('a', 1), ('b', 1))])
        self.assertRaises(ValueError, self.network.get_intra_edges, -1)
        self.assertRaises(TypeError, self.network.get_intra_edges, '-')
        self.assertListEqual(sorted(self.network.get_inter_edges()), [(('a', 0), ('a', 1)), (('b', 0), ('b', 1))])

    def test_get_interface_nodes(self):
        self.network.add_edges_from(
            [(('D', 0), ('G', 0)), (('I', 0), ('G', 0)), (('D', 0), ('D', 1)), (('I', 0), ('I', 1))])
        self.assertListEqual(sorted(self.network.get_interface_nodes()), [('D', 0), ('I',0)])
        self.assertRaises(ValueError, self.network.get_interface_nodes, -1)
        self.assertRaises(TypeError, self.network.get_interface_nodes, '-')

    def test_get_slice_nodes(self):
        self.network.add_edges_from(
            [(('D', 0), ('G', 0)), (('I', 0), ('G', 0)), (('D', 0), ('D', 1)), (('I', 0), ('I', 1))])
        self.assertListEqual(sorted(self.network.get_slice_nodes()), [('D', 0), ('G', 0), ('I', 0)])
        self.assertListEqual(sorted(self.network.get_slice_nodes(1)), [('D', 1), ('G', 1), ('I', 1)])
        self.assertRaises(ValueError, self.network.get_slice_nodes, -1)
        self.assertRaises(TypeError, self.network.get_slice_nodes, '-')

    def test_add_single_cpds(self):
        self.network.add_edges_from([(('D', 0), ('G', 0)), (('I', 0), ('G', 0))])
        self.network.add_cpds(self.grade_cpd)
        self.assertListEqual(self.network.get_cpds(), [self.grade_cpd])

    def test_get_cpds(self):
        self.network.add_edges_from(
            [(('D', 0), ('G', 0)), (('I', 0), ('G', 0)), (('D', 0), ('D', 1)), (('I', 0), ('I', 1))])
        self.network.add_cpds(self.grade_cpd, self.d_i_cpd, self.diff_cpd, self.intel_cpd, self.i_i_cpd)
        self.network.initialize_initial_state()
        self.assertEqual(set(self.network.get_cpds()), set([self.diff_cpd, self.intel_cpd, self.grade_cpd]))
        self.assertEqual(self.network.get_cpds(time_slice=1)[0].variable, ('G', 1))

    def test_add_multiple_cpds(self):
        self.network.add_edges_from(
            [(('D', 0), ('G', 0)), (('I', 0), ('G', 0)), (('D', 0), ('D', 1)), (('I', 0), ('I', 1))])
        self.network.add_cpds(self.grade_cpd, self.d_i_cpd, self.diff_cpd, self.intel_cpd, self.i_i_cpd)
        self.assertEqual(self.network.get_cpds(('G', 0)).variable, ('G', 0))
        self.assertEqual(self.network.get_cpds(('D', 1)).variable, ('D', 1))
        self.assertEqual(self.network.get_cpds(('D', 0)).variable, ('D', 0))
        self.assertEqual(self.network.get_cpds(('I', 0)).variable, ('I', 0))
        self.assertEqual(self.network.get_cpds(('I', 1)).variable, ('I', 1))

    def test_initialize_initial_state(self):

        self.network.add_nodes_from(['D', 'G', 'I', 'S', 'L'])
        self.network.add_edges_from(
            [(('D', 0), ('G', 0)), (('I', 0), ('G', 0)), (('D', 0), ('D', 1)), (('I', 0), ('I', 1))])
        self.network.add_cpds(self.grade_cpd, self.d_i_cpd, self.diff_cpd, self.intel_cpd, self.i_i_cpd)
        self.network.initialize_initial_state()
        self.assertEqual(len(self.network.cpds), 6)
        self.assertEqual(self.network.get_cpds(('G', 1)).variable, ('G', 1))

    def test_moralize(self):
        self.network.add_edges_from(([(('D',0), ('G',0)), (('I',0), ('G',0))]))
        moral_graph = self.network.moralize()
        self.assertListEqual(hf.recursive_sorted(moral_graph.edges()),
                             [[('D', 0), ('G', 0)], [('D', 0), ('I', 0)],
                              [('D', 1), ('G', 1)], [('D', 1), ('I', 1)],
                              [('G', 0), ('I', 0)], [('G', 1), ('I', 1)]])

    def tearDown(self):
        del self.network
<commit_msg>another change in test case<commit_after>import unittest
import numpy as np

import pgmpy.tests.help_functions as hf
from pgmpy.models import DynamicBayesianNetwork
from pgmpy.factors import TabularCPD

class TestDynamicBayesianNetworkCreation(unittest.TestCase):
    def setUp(self):
        self.network = DynamicBayesianNetwork()

    def test_add_single_node(self):
        self.network.add_node('a')
        self.assertListEqual(self.network.nodes(), ['a'])

    def test_add_multiple_nodes(self):
        self.network.add_nodes_from(['a', 'b', 'c'])
        self.assertListEqual(sorted(self.network.nodes()), ['a', 'b', 'c'])

    def test_add_single_edge_with_timeslice(self):
        self.network.add_edge(('a', 0), ('b', 0))
        self.assertListEqual(sorted(self.network.edges()), [(('a', 0), ('b', 0)), (('a', 1), ('b', 1))])
        self.assertListEqual(sorted(self.network.nodes()), ['a', 'b'])

    def test_add_edge_with_different_number_timeslice(self):
        self.network.add_edge(('a', 2), ('b', 2))
        self.assertListEqual(sorted(self.network.edges()), [(('a', 0), ('b', 0)), (('a', 1), ('b', 1))])

    def test_add_edge_going_backward(self):
        self.assertRaises(ValueError, self.network.add_edge, ('a', 1), ('b', 0))

    def test_add_edge_with_farther_timeslice(self):
        self.assertRaises(ValueError, self.network.add_edge, ('a', 2), ('b', 4))

    def test_add_edge_with_self_loop(self):
        self.assertRaises(ValueError, self.network.add_edge, ('a', 0), ('a', 0))

    def test_add_edge_with_varying_length(self):
        self.assertRaises(ValueError, self.network.add_edge, ('a', 1, 1), ('b', 2))
        self.assertRaises(ValueError, self.network.add_edge, ('b', 2), ('a', 2, 3))

    def test_add_edge_with_closed_path(self):
        self.assertRaises(ValueError, self.network.add_edges_from,
                          [(('a', 0), ('b', 0)), (('b', 0), ('c', 0)), (('c', 0), ('a', 0))])

    def test_add_single_edge_without_timeslice(self):
        self.assertRaises(ValueError, self.network.add_edge, 'a', 'b')

    def test_add_single_edge_with_incorrect_timeslice(self):
        self.assertRaises(ValueError, self.network.add_edge, ('a', 'b'), ('b', 'c'))

    def test_add_multiple_edges(self):
        self.network.add_edges_from([(('a', 0), ('b', 0)), (('a', 0), ('a', 1)), (('b', 0), ('b', 1))])
        self.assertListEqual(sorted(self.network.edges()),
                             [(('a', 0), ('a', 1)), (('a', 0), ('b', 0)), (('a', 1), ('b', 1)), (('b', 0), ('b', 1))])

    def tearDown(self):
        del self.network

class TestDynamicBayesianNetworkMethods(unittest.TestCase):
    def setUp(self):
        self.network = DynamicBayesianNetwork()
        self.grade_cpd = TabularCPD(('G', 0), 3, [[0.3, 0.05, 0.9, 0.5],
                                             [0.4, 0.25, 0.08, 0.3],
                                             [0.3, 0.7, 0.2, 0.2]], [('D', 0), ('I', 0)], [2, 2])
        self.d_i_cpd = TabularCPD(('D', 1), 2, [[0.6, 0.3], [0.4, 0.7]], [('D', 0)], 2)
        self.diff_cpd = TabularCPD(('D', 0), 2, [[0.6, 0.4]])
        self.intel_cpd = TabularCPD(('I', 0), 2, [[0.7, 0.3]])
        self.i_i_cpd = TabularCPD(('I', 1), 2, [[0.5, 0.4], [0.5, 0.6]], [('I', 0)], 2)
        self.grade_1_cpd = TabularCPD(('G', 1), 3, [[0.3, 0.05, 0.9, 0.5],
                                             [0.4, 0.25, 0.08, 0.3],
                                             [0.3, 0.7, 0.2, 0.2]], [('D', 1), ('I', 1)], [2, 2])

    def test_get_intra_and_inter_edges(self):
        self.network.add_edges_from([(('a', 0), ('b', 0)), (('a', 0), ('a', 1)), (('b', 0), ('b', 1))])
        self.assertListEqual(sorted(self.network.get_intra_edges()), [(('a', 0), ('b', 0))])
        self.assertListEqual(sorted(self.network.get_intra_edges(1)), [(('a', 1), ('b', 1))])
        self.assertRaises(ValueError, self.network.get_intra_edges, -1)
        self.assertRaises(ValueError, self.network.get_intra_edges, '-')
        self.assertListEqual(sorted(self.network.get_inter_edges()), [(('a', 0), ('a', 1)), (('b', 0), ('b', 1))])

    def test_get_interface_nodes(self):
        self.network.add_edges_from(
            [(('D', 0), ('G', 0)), (('I', 0), ('G', 0)), (('D', 0), ('D', 1)), (('I', 0), ('I', 1))])
        self.assertListEqual(sorted(self.network.get_interface_nodes()), [('D', 0), ('I',0)])
        self.assertRaises(ValueError, self.network.get_interface_nodes, -1)
        self.assertRaises(ValueError, self.network.get_interface_nodes, '-')

    def test_get_slice_nodes(self):
        self.network.add_edges_from(
            [(('D', 0), ('G', 0)), (('I', 0), ('G', 0)), (('D', 0), ('D', 1)), (('I', 0), ('I', 1))])
        self.assertListEqual(sorted(self.network.get_slice_nodes()), [('D', 0), ('G', 0), ('I', 0)])
        self.assertListEqual(sorted(self.network.get_slice_nodes(1)), [('D', 1), ('G', 1), ('I', 1)])
        self.assertRaises(ValueError, self.network.get_slice_nodes, -1)
        self.assertRaises(ValueError, self.network.get_slice_nodes, '-')

    def test_add_single_cpds(self):
        self.network.add_edges_from([(('D', 0), ('G', 0)), (('I', 0), ('G', 0))])
        self.network.add_cpds(self.grade_cpd)
        self.assertListEqual(self.network.get_cpds(), [self.grade_cpd])

    def test_get_cpds(self):
        self.network.add_edges_from(
            [(('D', 0), ('G', 0)), (('I', 0), ('G', 0)), (('D', 0), ('D', 1)), (('I', 0), ('I', 1))])
        self.network.add_cpds(self.grade_cpd, self.d_i_cpd, self.diff_cpd, self.intel_cpd, self.i_i_cpd)
        self.network.initialize_initial_state()
        self.assertEqual(set(self.network.get_cpds()), set([self.diff_cpd, self.intel_cpd, self.grade_cpd]))
        self.assertEqual(self.network.get_cpds(time_slice=1)[0].variable, ('G', 1))

    def test_add_multiple_cpds(self):
        self.network.add_edges_from(
            [(('D', 0), ('G', 0)), (('I', 0), ('G', 0)), (('D', 0), ('D', 1)), (('I', 0), ('I', 1))])
        self.network.add_cpds(self.grade_cpd, self.d_i_cpd, self.diff_cpd, self.intel_cpd, self.i_i_cpd)
        self.assertEqual(self.network.get_cpds(('G', 0)).variable, ('G', 0))
        self.assertEqual(self.network.get_cpds(('D', 1)).variable, ('D', 1))
        self.assertEqual(self.network.get_cpds(('D', 0)).variable, ('D', 0))
        self.assertEqual(self.network.get_cpds(('I', 0)).variable, ('I', 0))
        self.assertEqual(self.network.get_cpds(('I', 1)).variable, ('I', 1))

    def test_initialize_initial_state(self):

        self.network.add_nodes_from(['D', 'G', 'I', 'S', 'L'])
        self.network.add_edges_from(
            [(('D', 0), ('G', 0)), (('I', 0), ('G', 0)), (('D', 0), ('D', 1)), (('I', 0), ('I', 1))])
        self.network.add_cpds(self.grade_cpd, self.d_i_cpd, self.diff_cpd, self.intel_cpd, self.i_i_cpd)
        self.network.initialize_initial_state()
        self.assertEqual(len(self.network.cpds), 6)
        self.assertEqual(self.network.get_cpds(('G', 1)).variable, ('G', 1))

    def test_moralize(self):
        self.network.add_edges_from(([(('D',0), ('G',0)), (('I',0), ('G',0))]))
        moral_graph = self.network.moralize()
        self.assertListEqual(hf.recursive_sorted(moral_graph.edges()),
                             [[('D', 0), ('G', 0)], [('D', 0), ('I', 0)],
                              [('D', 1), ('G', 1)], [('D', 1), ('I', 1)],
                              [('G', 0), ('I', 0)], [('G', 1), ('I', 1)]])

    def tearDown(self):
        del self.network
<|endoftext|>


<commit_before>"""
Module :mod:`hmtk.parsers.catalogue.csv_catalogue_parser`
implements :class:`CsvCatalogueParser`.
"""
import csv
import numpy as np

from hmtk.seismicity.catalogue import Catalogue
from hmtk.parsers.catalogue.base import BaseCatalogueParser

class CsvCatalogueParser(BaseCatalogueParser):
    """CSV Catalogue Parser Class
    """

    def read_file(self):
        """
        """
        filedata = open(self.input_file, 'rt')
        catalogue = Catalogue()
        # Reading the data file
        data = csv.DictReader(filedata)
        # Parsing the data content
        for irow, row in enumerate(data):
            if irow == 0:
                valid_key_list = self._header_check(row.keys(),
                    catalogue.TOTAL_ATTRIBUTE_LIST)
            for key in valid_key_list:
                if key in catalogue.FLOAT_ATTRIBUTE_LIST:
                    catalogue.data[key] = self._float_check(
                                                        catalogue.data[key],
                                                        row[key])
                elif key in catalogue.INT_ATTRIBUTE_LIST:
                    catalogue.data[key] = self._int_check(
                                                        catalogue.data[key],
                                                        row[key])
                else:
                    catalogue.data[key].append(row[key])
        return catalogue

    def _header_check(self, input_keys, catalogue_keys):
        valid_key_list = []
        for element in input_keys:
            print element
            if element in catalogue_keys:
                valid_key_list.append(element)
            else:
                print 'Catalogue Attribute %s is not a recognised catalogue key'\
                    % element
        return valid_key_list

    def _float_check(self, attribute_array, value):
        '''Checks if value is valid float, appends to array if valid, appends
        nan if not'''
        value = value.strip(' ')
        if value:
            attribute_array = np.hstack([attribute_array, float(value)])
        else:
            attribute_array = np.hstack([attribute_array, np.nan])
        return attribute_array

    def _int_check(self, attribute_array, value):
        '''Checks if value is valid integer, appends to array if valid, appends
        nan if not'''
        value = value.strip(' ')
        if value:
            attribute_array = np.hstack([attribute_array, int(value)])
            print int(value)
        else:
            attribute_array = np.hstack([attribute_array, np.nan])
        return attribute_array
<commit_msg>Cleaned csv catalogue parser<commit_after>"""
Module :mod:`hmtk.parsers.catalogue.csv_catalogue_parser`
implements :class:`CsvCatalogueParser`.
"""
import csv
import numpy as np

from hmtk.seismicity.catalogue import Catalogue
from hmtk.parsers.catalogue.base import BaseCatalogueParser

class CsvCatalogueParser(BaseCatalogueParser):
    """CSV Catalogue Parser Class
    """

    def read_file(self):
        """
        """
        filedata = open(self.input_file, 'rt')
        catalogue = Catalogue()
        # Reading the data file
        data = csv.DictReader(filedata)
        # Parsing the data content
        for irow, row in enumerate(data):
            if irow == 0:
                valid_key_list = self._header_check(row.keys(),
                    catalogue.TOTAL_ATTRIBUTE_LIST)
            for key in valid_key_list:
                if key in catalogue.FLOAT_ATTRIBUTE_LIST:
                    catalogue.data[key] = self._float_check(
                                                        catalogue.data[key],
                                                        row[key])
                elif key in catalogue.INT_ATTRIBUTE_LIST:
                    catalogue.data[key] = self._int_check(
                                                        catalogue.data[key],
                                                        row[key])
                else:
                    catalogue.data[key].append(row[key])
        return catalogue

    def _header_check(self, input_keys, catalogue_keys):
        valid_key_list = []
        for element in input_keys:
            if element in catalogue_keys:
                valid_key_list.append(element)
            else:
                print 'Catalogue Attribute %s is not a recognised catalogue key'\
                    % element
        return valid_key_list

    def _float_check(self, attribute_array, value):
        '''Checks if value is valid float, appends to array if valid, appends
        nan if not'''
        value = value.strip(' ')
        if value:
            attribute_array = np.hstack([attribute_array, float(value)])
        else:
            attribute_array = np.hstack([attribute_array, np.nan])
        return attribute_array

    def _int_check(self, attribute_array, value):
        '''Checks if value is valid integer, appends to array if valid, appends
        nan if not'''
        value = value.strip(' ')
        if value:
            attribute_array = np.hstack([attribute_array, int(value)])
        else:
            attribute_array = np.hstack([attribute_array, np.nan])
        return attribute_array
<|endoftext|>


<commit_before>from django.db import models

from systextil.models.base import HDoc001

class ContaEstoque(models.Model):
    conta_estoque = models.IntegerField(
        primary_key=True,
        verbose_name='Código')
    descr_ct_estoque = models.CharField(
        max_length=100,
        verbose_name='Descrição')

    def __str__(self):
        return '{}-{}'.format(self.conta_estoque, self.descr_ct_estoque)

    class Meta:
        managed = False
        app_label = 'systextil'
        db_table = "BASI_150"
        verbose_name = "Conta de estoque"

class TipoContaEstoqueManager(models.Manager):
    def get_queryset(self):
        return super(
            TipoContaEstoqueManager, self).get_queryset().filter(tipo=420)

class TipoContaEstoque(HDoc001):
    objects = TipoContaEstoqueManager()

    def __str__(self):
        return '{}-{}'.format(self.codigo, self.descricao)

    class Meta:
        proxy = True
        managed = False
        app_label = 'systextil'
        db_table = "HDOC_001"
        verbose_name = "Tipo de conta de estoque"

class TipoTransacao(models.Model):
    codigo_transacao = models.IntegerField(
        primary_key=True,
        verbose_name='Código da transação')
    descricao = models.CharField(
        max_length=20,
        verbose_name='Descrição')

    def __str__(self):
        return '{} - {}'.format(self.codigo_transacao, self.descricao)

    class Meta:
        managed = False
        app_label = 'systextil'
        db_table = "ESTQ_005"
        verbose_name = "Código de transação de estoque"
<commit_msg>Adicionado campo entrada_saida da tabela TipoTransacao<commit_after>from django.db import models

from systextil.models.base import HDoc001

class ContaEstoque(models.Model):
    conta_estoque = models.IntegerField(
        primary_key=True,
        verbose_name='Código')
    descr_ct_estoque = models.CharField(
        max_length=100,
        verbose_name='Descrição')

    def __str__(self):
        return '{}-{}'.format(self.conta_estoque, self.descr_ct_estoque)

    class Meta:
        managed = False
        app_label = 'systextil'
        db_table = "BASI_150"
        verbose_name = "Conta de estoque"

class TipoContaEstoqueManager(models.Manager):
    def get_queryset(self):
        return super(
            TipoContaEstoqueManager, self).get_queryset().filter(tipo=420)

class TipoContaEstoque(HDoc001):
    objects = TipoContaEstoqueManager()

    def __str__(self):
        return '{}-{}'.format(self.codigo, self.descricao)

    class Meta:
        proxy = True
        managed = False
        app_label = 'systextil'
        db_table = "HDOC_001"
        verbose_name = "Tipo de conta de estoque"

class TipoTransacao(models.Model):
    codigo_transacao = models.IntegerField(
        primary_key=True,
        verbose_name='Código da transação')
    descricao = models.CharField(
        max_length=20,
        verbose_name='Descrição')
    entrada_saida = models.CharField(
        max_length=1,
        verbose_name='Entrada/Saída')

    def __str__(self):
        return (f'{self.codigo_transacao} ({self.entrada_saida}) - '
                f'{self.descricao}')

    class Meta:
        managed = False
        app_label = 'systextil'
        db_table = "ESTQ_005"
        verbose_name = "Código de transação de estoque"
<|endoftext|>


<commit_before>"""CoolMasterNet platform to control of CoolMasteNet Climate Devices."""

import logging

import voluptuous as vol

from homeassistant.components.climate import ClimateDevice, PLATFORM_SCHEMA
from homeassistant.components.climate.const import (
    HVAC_MODE_OFF, HVAC_MODE_AUTO, HVAC_MODE_COOL, HVAC_MODE_DRY,
    HVAC_MODE_FAN_ONLY, HVAC_MODE_HEAT, SUPPORT_FAN_MODE,
    SUPPORT_TARGET_TEMPERATURE)
from homeassistant.const import (
    ATTR_TEMPERATURE, CONF_HOST, CONF_PORT, TEMP_CELSIUS, TEMP_FAHRENHEIT)
import homeassistant.helpers.config_validation as cv

SUPPORT_FLAGS = (SUPPORT_TARGET_TEMPERATURE | SUPPORT_FAN_MODE)

DEFAULT_PORT = 10102

AVAILABLE_MODES = [HVAC_MODE_OFF, HVAC_MODE_HEAT, HVAC_MODE_COOL,
                   HVAC_MODE_DRY, HVAC_MODE_AUTO, HVAC_MODE_FAN_ONLY]

CM_TO_HA_STATE = {
    'heat': HVAC_MODE_HEAT,
    'cool': HVAC_MODE_COOL,
    'auto': HVAC_MODE_AUTO,
    'dry': HVAC_MODE_DRY,
    'fan': HVAC_MODE_FAN_ONLY,
}

HA_STATE_TO_CM = {value: key for key, value in CM_TO_HA_STATE.items()}

FAN_MODES = ['low', 'med', 'high', 'auto']

CONF_SUPPORTED_MODES = 'supported_modes'
PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
    vol.Required(CONF_HOST): cv.string,
    vol.Optional(CONF_PORT, default=DEFAULT_PORT): cv.port,
    vol.Optional(CONF_SUPPORTED_MODES, default=AVAILABLE_MODES):
        vol.All(cv.ensure_list, [vol.In(AVAILABLE_MODES)]),
})

_LOGGER = logging.getLogger(__name__)

def _build_entity(device, supported_modes):
    _LOGGER.debug("Found device %s", device.uid)
    return CoolmasterClimate(device, supported_modes)

def setup_platform(hass, config, add_entities, discovery_info=None):
    """Set up the CoolMasterNet climate platform."""
    from pycoolmasternet import CoolMasterNet

    supported_modes = config.get(CONF_SUPPORTED_MODES)
    host = config[CONF_HOST]
    port = config[CONF_PORT]
    cool = CoolMasterNet(host, port=port)
    devices = cool.devices()

    all_devices = [_build_entity(device, supported_modes)
                   for device in devices]

    add_entities(all_devices, True)

class CoolmasterClimate(ClimateDevice):
    """Representation of a coolmaster climate device."""

    def __init__(self, device, supported_modes):
        """Initialize the climate device."""
        self._device = device
        self._uid = device.uid
        self._hvac_modes = supported_modes
        self._hvac_mode = None
        self._target_temperature = None
        self._current_temperature = None
        self._current_fan_mode = None
        self._current_operation = None
        self._on = None
        self._unit = None

    def update(self):
        """Pull state from CoolMasterNet."""
        status = self._device.status
        self._target_temperature = status['thermostat']
        self._current_temperature = status['temperature']
        self._current_fan_mode = status['fan_speed']
        self._on = status['is_on']

        device_mode = status['mode']
        if self._on:
            self._hvac_mode = CM_TO_HA_STATE[device_mode]
        else:
            self._hvac_mode = HVAC_MODE_OFF

        if status['unit'] == 'celsius':
            self._unit = TEMP_CELSIUS
        else:
            self._unit = TEMP_FAHRENHEIT

    @property
    def unique_id(self):
        """Return unique ID for this device."""
        return self._uid

    @property
    def supported_features(self):
        """Return the list of supported features."""
        return SUPPORT_FLAGS

    @property
    def name(self):
        """Return the name of the climate device."""
        return self.unique_id

    @property
    def temperature_unit(self):
        """Return the unit of measurement."""
        return self._unit

    @property
    def current_temperature(self):
        """Return the current temperature."""
        return self._current_temperature

    @property
    def target_temperature(self):
        """Return the temperature we are trying to reach."""
        return self._target_temperature

    @property
    def hvac_mode(self):
        """Return hvac target hvac state."""
        return self._hvac_mode

    @property
    def hvac_modes(self):
        """Return the list of available operation modes."""
        return self._hvac_modes

    @property
    def fan_mode(self):
        """Return the fan setting."""
        return self._current_fan_mode

    @property
    def fan_modes(self):
        """Return the list of available fan modes."""
        return FAN_MODES

    def set_temperature(self, **kwargs):
        """Set new target temperatures."""
        temp = kwargs.get(ATTR_TEMPERATURE)
        if temp is not None:
            _LOGGER.debug("Setting temp of %s to %s", self.unique_id,
                          str(temp))
            self._device.set_thermostat(str(temp))

    def set_fan_mode(self, fan_mode):
        """Set new fan mode."""
        _LOGGER.debug("Setting fan mode of %s to %s", self.unique_id,
                      fan_mode)
        self._device.set_fan_speed(fan_mode)

    def set_hvac_mode(self, hvac_mode):
        """Set new operation mode."""
        _LOGGER.debug("Setting operation mode of %s to %s", self.unique_id,
                      hvac_mode)

        if hvac_mode == HVAC_MODE_OFF:
            self._device.turn_off()
        else:
            self._device.set_mode(HA_STATE_TO_CM[hvac_mode])
            self._device.turn_on()
<commit_msg>Add turn on/off to coolmaster (#25097)<commit_after>"""CoolMasterNet platform to control of CoolMasteNet Climate Devices."""

import logging

import voluptuous as vol

from homeassistant.components.climate import ClimateDevice, PLATFORM_SCHEMA
from homeassistant.components.climate.const import (
    HVAC_MODE_OFF, HVAC_MODE_AUTO, HVAC_MODE_COOL, HVAC_MODE_DRY,
    HVAC_MODE_FAN_ONLY, HVAC_MODE_HEAT, SUPPORT_FAN_MODE,
    SUPPORT_TARGET_TEMPERATURE)
from homeassistant.const import (
    ATTR_TEMPERATURE, CONF_HOST, CONF_PORT, TEMP_CELSIUS, TEMP_FAHRENHEIT)
import homeassistant.helpers.config_validation as cv

SUPPORT_FLAGS = (SUPPORT_TARGET_TEMPERATURE | SUPPORT_FAN_MODE)

DEFAULT_PORT = 10102

AVAILABLE_MODES = [HVAC_MODE_OFF, HVAC_MODE_HEAT, HVAC_MODE_COOL,
                   HVAC_MODE_DRY, HVAC_MODE_AUTO, HVAC_MODE_FAN_ONLY]

CM_TO_HA_STATE = {
    'heat': HVAC_MODE_HEAT,
    'cool': HVAC_MODE_COOL,
    'auto': HVAC_MODE_AUTO,
    'dry': HVAC_MODE_DRY,
    'fan': HVAC_MODE_FAN_ONLY,
}

HA_STATE_TO_CM = {value: key for key, value in CM_TO_HA_STATE.items()}

FAN_MODES = ['low', 'med', 'high', 'auto']

CONF_SUPPORTED_MODES = 'supported_modes'
PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend({
    vol.Required(CONF_HOST): cv.string,
    vol.Optional(CONF_PORT, default=DEFAULT_PORT): cv.port,
    vol.Optional(CONF_SUPPORTED_MODES, default=AVAILABLE_MODES):
        vol.All(cv.ensure_list, [vol.In(AVAILABLE_MODES)]),
})

_LOGGER = logging.getLogger(__name__)

def _build_entity(device, supported_modes):
    _LOGGER.debug("Found device %s", device.uid)
    return CoolmasterClimate(device, supported_modes)

def setup_platform(hass, config, add_entities, discovery_info=None):
    """Set up the CoolMasterNet climate platform."""
    from pycoolmasternet import CoolMasterNet

    supported_modes = config.get(CONF_SUPPORTED_MODES)
    host = config[CONF_HOST]
    port = config[CONF_PORT]
    cool = CoolMasterNet(host, port=port)
    devices = cool.devices()

    all_devices = [_build_entity(device, supported_modes)
                   for device in devices]

    add_entities(all_devices, True)

class CoolmasterClimate(ClimateDevice):
    """Representation of a coolmaster climate device."""

    def __init__(self, device, supported_modes):
        """Initialize the climate device."""
        self._device = device
        self._uid = device.uid
        self._hvac_modes = supported_modes
        self._hvac_mode = None
        self._target_temperature = None
        self._current_temperature = None
        self._current_fan_mode = None
        self._current_operation = None
        self._on = None
        self._unit = None

    def update(self):
        """Pull state from CoolMasterNet."""
        status = self._device.status
        self._target_temperature = status['thermostat']
        self._current_temperature = status['temperature']
        self._current_fan_mode = status['fan_speed']
        self._on = status['is_on']

        device_mode = status['mode']
        if self._on:
            self._hvac_mode = CM_TO_HA_STATE[device_mode]
        else:
            self._hvac_mode = HVAC_MODE_OFF

        if status['unit'] == 'celsius':
            self._unit = TEMP_CELSIUS
        else:
            self._unit = TEMP_FAHRENHEIT

    @property
    def unique_id(self):
        """Return unique ID for this device."""
        return self._uid

    @property
    def supported_features(self):
        """Return the list of supported features."""
        return SUPPORT_FLAGS

    @property
    def name(self):
        """Return the name of the climate device."""
        return self.unique_id

    @property
    def temperature_unit(self):
        """Return the unit of measurement."""
        return self._unit

    @property
    def current_temperature(self):
        """Return the current temperature."""
        return self._current_temperature

    @property
    def target_temperature(self):
        """Return the temperature we are trying to reach."""
        return self._target_temperature

    @property
    def hvac_mode(self):
        """Return hvac target hvac state."""
        return self._hvac_mode

    @property
    def hvac_modes(self):
        """Return the list of available operation modes."""
        return self._hvac_modes

    @property
    def fan_mode(self):
        """Return the fan setting."""
        return self._current_fan_mode

    @property
    def fan_modes(self):
        """Return the list of available fan modes."""
        return FAN_MODES

    def set_temperature(self, **kwargs):
        """Set new target temperatures."""
        temp = kwargs.get(ATTR_TEMPERATURE)
        if temp is not None:
            _LOGGER.debug("Setting temp of %s to %s", self.unique_id,
                          str(temp))
            self._device.set_thermostat(str(temp))

    def set_fan_mode(self, fan_mode):
        """Set new fan mode."""
        _LOGGER.debug("Setting fan mode of %s to %s", self.unique_id,
                      fan_mode)
        self._device.set_fan_speed(fan_mode)

    def set_hvac_mode(self, hvac_mode):
        """Set new operation mode."""
        _LOGGER.debug("Setting operation mode of %s to %s", self.unique_id,
                      hvac_mode)

        if hvac_mode == HVAC_MODE_OFF:
            self.turn_off()
        else:
            self._device.set_mode(HA_STATE_TO_CM[hvac_mode])
            self.turn_on()

    def turn_on(self):
        """Turn on."""
        _LOGGER.debug("Turning %s on", self.unique_id)
        self._device.turn_on()

    def turn_off(self):
        """Turn off."""
        _LOGGER.debug("Turning %s off", self.unique_id)
        self._device.turn_off()
<|endoftext|>


<commit_before>"""Sensor support for Wireless Sensor Tags platform."""
import logging

import voluptuous as vol

from homeassistant.components.sensor import PLATFORM_SCHEMA, SensorEntity
from homeassistant.const import CONF_MONITORED_CONDITIONS
from homeassistant.core import callback
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.dispatcher import async_dispatcher_connect

from . import DOMAIN as WIRELESSTAG_DOMAIN, SIGNAL_TAG_UPDATE, WirelessTagBaseSensor

_LOGGER = logging.getLogger(__name__)

SENSOR_TEMPERATURE = "temperature"
SENSOR_HUMIDITY = "humidity"
SENSOR_MOISTURE = "moisture"
SENSOR_LIGHT = "light"

SENSOR_TYPES = [SENSOR_TEMPERATURE, SENSOR_HUMIDITY, SENSOR_MOISTURE, SENSOR_LIGHT]

PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
    {
        vol.Required(CONF_MONITORED_CONDITIONS, default=[]): vol.All(
            cv.ensure_list, [vol.In(SENSOR_TYPES)]
        )
    }
)

def setup_platform(hass, config, add_entities, discovery_info=None):
    """Set up the sensor platform."""
    platform = hass.data.get(WIRELESSTAG_DOMAIN)
    sensors = []
    tags = platform.tags
    for tag in tags.values():
        for sensor_type in config.get(CONF_MONITORED_CONDITIONS):
            if sensor_type in tag.allowed_sensor_types:
                sensors.append(
                    WirelessTagSensor(platform, tag, sensor_type, hass.config)
                )

    add_entities(sensors, True)

class WirelessTagSensor(WirelessTagBaseSensor, SensorEntity):
    """Representation of a Sensor."""

    def __init__(self, api, tag, sensor_type, config):
        """Initialize a WirelessTag sensor."""
        super().__init__(api, tag)

        self._sensor_type = sensor_type
        self._name = self._tag.name

        # I want to see entity_id as:
        # sensor.wirelesstag_bedroom_temperature
        # and not as sensor.bedroom for temperature and
        # sensor.bedroom_2 for humidity
        self._entity_id = (
            f"sensor.{WIRELESSTAG_DOMAIN}_{self.underscored_name}_{self._sensor_type}"
        )

    async def async_added_to_hass(self):
        """Register callbacks."""
        self.async_on_remove(
            async_dispatcher_connect(
                self.hass,
                SIGNAL_TAG_UPDATE.format(self.tag_id, self.tag_manager_mac),
                self._update_tag_info_callback,
            )
        )

    @property
    def entity_id(self):
        """Overridden version."""
        return self._entity_id

    @property
    def underscored_name(self):
        """Provide name savvy to be used in entity_id name of self."""
        return self.name.lower().replace(" ", "_")

    @property
    def native_value(self):
        """Return the state of the sensor."""
        return self._state

    @property
    def device_class(self):
        """Return the class of the sensor."""
        return self._sensor_type

    @property
    def native_unit_of_measurement(self):
        """Return the unit of measurement."""
        return self._sensor.unit

    @property
    def principal_value(self):
        """Return sensor current value."""
        return self._sensor.value

    @property
    def _sensor(self):
        """Return tag sensor entity."""
        return self._tag.sensor[self._sensor_type]

    @callback
    def _update_tag_info_callback(self, new_tag):
        """Handle push notification sent by tag manager."""
        _LOGGER.debug("Entity to update state: %s with new tag: %s", self, new_tag)
        self._tag = new_tag
        self._state = self.updated_state_value()
        self.async_write_ha_state()
<commit_msg>Add wirelesstag ambient temperature sensor for Outdoor Probe (#60243)<commit_after>"""Sensor support for Wireless Sensor Tags platform."""
from __future__ import annotations

import logging

import voluptuous as vol

from homeassistant.components.sensor import (
    PLATFORM_SCHEMA,
    STATE_CLASS_MEASUREMENT,
    SensorEntity,
    SensorEntityDescription,
)
from homeassistant.const import (
    CONF_MONITORED_CONDITIONS,
    DEVICE_CLASS_HUMIDITY,
    DEVICE_CLASS_ILLUMINANCE,
    DEVICE_CLASS_TEMPERATURE,
)
from homeassistant.core import callback
import homeassistant.helpers.config_validation as cv
from homeassistant.helpers.dispatcher import async_dispatcher_connect

from . import DOMAIN as WIRELESSTAG_DOMAIN, SIGNAL_TAG_UPDATE, WirelessTagBaseSensor

_LOGGER = logging.getLogger(__name__)

SENSOR_TEMPERATURE = "temperature"
SENSOR_AMBIENT_TEMPERATURE = "ambient_temperature"
SENSOR_HUMIDITY = "humidity"
SENSOR_MOISTURE = "moisture"
SENSOR_LIGHT = "light"

SENSOR_TYPES: dict[str, SensorEntityDescription] = {
    SENSOR_TEMPERATURE: SensorEntityDescription(
        key=SENSOR_TEMPERATURE,
        device_class=DEVICE_CLASS_TEMPERATURE,
        state_class=STATE_CLASS_MEASUREMENT,
    ),
    SENSOR_AMBIENT_TEMPERATURE: SensorEntityDescription(
        key=SENSOR_AMBIENT_TEMPERATURE,
        device_class=DEVICE_CLASS_TEMPERATURE,
        state_class=STATE_CLASS_MEASUREMENT,
    ),
    SENSOR_HUMIDITY: SensorEntityDescription(
        key=SENSOR_HUMIDITY,
        device_class=DEVICE_CLASS_HUMIDITY,
        state_class=STATE_CLASS_MEASUREMENT,
    ),
    SENSOR_MOISTURE: SensorEntityDescription(
        key=SENSOR_MOISTURE,
        device_class=SENSOR_MOISTURE,
        state_class=STATE_CLASS_MEASUREMENT,
    ),
    SENSOR_LIGHT: SensorEntityDescription(
        key=SENSOR_LIGHT,
        device_class=DEVICE_CLASS_ILLUMINANCE,
        state_class=STATE_CLASS_MEASUREMENT,
    ),
}

SENSOR_KEYS: list[str] = list(SENSOR_TYPES)

PLATFORM_SCHEMA = PLATFORM_SCHEMA.extend(
    {
        vol.Required(CONF_MONITORED_CONDITIONS, default=[]): vol.All(
            cv.ensure_list, [vol.In(SENSOR_KEYS)]
        )
    }
)

def setup_platform(hass, config, add_entities, discovery_info=None):
    """Set up the sensor platform."""
    platform = hass.data.get(WIRELESSTAG_DOMAIN)
    sensors = []
    tags = platform.tags
    for tag in tags.values():
        for key in config[CONF_MONITORED_CONDITIONS]:
            if key not in tag.allowed_sensor_types:
                continue
            description = SENSOR_TYPES[key]
            sensors.append(WirelessTagSensor(platform, tag, description))

    add_entities(sensors, True)

class WirelessTagSensor(WirelessTagBaseSensor, SensorEntity):
    """Representation of a Sensor."""

    entity_description: SensorEntityDescription

    def __init__(self, api, tag, description):
        """Initialize a WirelessTag sensor."""
        super().__init__(api, tag)

        self._sensor_type = description.key
        self.entity_description = description
        self._name = self._tag.name

        # I want to see entity_id as:
        # sensor.wirelesstag_bedroom_temperature
        # and not as sensor.bedroom for temperature and
        # sensor.bedroom_2 for humidity
        self._entity_id = (
            f"sensor.{WIRELESSTAG_DOMAIN}_{self.underscored_name}_{self._sensor_type}"
        )

    async def async_added_to_hass(self):
        """Register callbacks."""
        self.async_on_remove(
            async_dispatcher_connect(
                self.hass,
                SIGNAL_TAG_UPDATE.format(self.tag_id, self.tag_manager_mac),
                self._update_tag_info_callback,
            )
        )

    @property
    def entity_id(self):
        """Overridden version."""
        return self._entity_id

    @property
    def underscored_name(self):
        """Provide name savvy to be used in entity_id name of self."""
        return self.name.lower().replace(" ", "_")

    @property
    def native_value(self):
        """Return the state of the sensor."""
        return self._state

    @property
    def native_unit_of_measurement(self):
        """Return the unit of measurement."""
        return self._sensor.unit

    @property
    def principal_value(self):
        """Return sensor current value."""
        return self._sensor.value

    @property
    def _sensor(self):
        """Return tag sensor entity."""
        return self._tag.sensor[self._sensor_type]

    @callback
    def _update_tag_info_callback(self, new_tag):
        """Handle push notification sent by tag manager."""
        _LOGGER.debug("Entity to update state: %s with new tag: %s", self, new_tag)
        self._tag = new_tag
        self._state = self.updated_state_value()
        self.async_write_ha_state()
<|endoftext|>


<commit_before>"""Device discovery functions for Zigbee Home Automation."""
from __future__ import annotations

from collections import Counter
from collections.abc import Callable
import logging

from homeassistant import const as ha_const
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.dispatcher import (
    async_dispatcher_connect,
    async_dispatcher_send,
)
from homeassistant.helpers.entity_registry import async_entries_for_device

from . import const as zha_const, registries as zha_regs, typing as zha_typing
from .. import (  # noqa: F401 pylint: disable=unused-import,
    alarm_control_panel,
    binary_sensor,
    button,
    climate,
    cover,
    device_tracker,
    fan,
    light,
    lock,
    number,
    select,
    sensor,
    siren,
    switch,
)
from .channels import base

_LOGGER = logging.getLogger(__name__)

@callback
async def async_add_entities(
    _async_add_entities: Callable,
    entities: list[
        tuple[
            zha_typing.ZhaEntityType,
            tuple[str, zha_typing.ZhaDeviceType, list[zha_typing.ChannelType]],
        ]
    ],
    update_before_add: bool = True,
) -> None:
    """Add entities helper."""
    if not entities:
        return
    to_add = [ent_cls.create_entity(*args) for ent_cls, args in entities]
    to_add = [entity for entity in to_add if entity is not None]
    _async_add_entities(to_add, update_before_add=update_before_add)
    entities.clear()

class ProbeEndpoint:
    """All discovered channels and entities of an endpoint."""

    def __init__(self):
        """Initialize instance."""
        self._device_configs = {}

    @callback
    def discover_entities(self, channel_pool: zha_typing.ChannelPoolType) -> None:
        """Process an endpoint on a zigpy device."""
        self.discover_by_device_type(channel_pool)
        self.discover_multi_entities(channel_pool)
        self.discover_by_cluster_id(channel_pool)
        zha_regs.ZHA_ENTITIES.clean_up()

    @callback
    def discover_by_device_type(self, channel_pool: zha_typing.ChannelPoolType) -> None:
        """Process an endpoint on a zigpy device."""

        unique_id = channel_pool.unique_id

        component = self._device_configs.get(unique_id, {}).get(ha_const.CONF_TYPE)
        if component is None:
            ep_profile_id = channel_pool.endpoint.profile_id
            ep_device_type = channel_pool.endpoint.device_type
            component = zha_regs.DEVICE_CLASS[ep_profile_id].get(ep_device_type)

        if component and component in zha_const.PLATFORMS:
            channels = channel_pool.unclaimed_channels()
            entity_class, claimed = zha_regs.ZHA_ENTITIES.get_entity(
                component, channel_pool.manufacturer, channel_pool.model, channels
            )
            if entity_class is None:
                return
            channel_pool.claim_channels(claimed)
            channel_pool.async_new_entity(component, entity_class, unique_id, claimed)

    @callback
    def discover_by_cluster_id(self, channel_pool: zha_typing.ChannelPoolType) -> None:
        """Process an endpoint on a zigpy device."""

        items = zha_regs.SINGLE_INPUT_CLUSTER_DEVICE_CLASS.items()
        single_input_clusters = {
            cluster_class: match
            for cluster_class, match in items
            if not isinstance(cluster_class, int)
        }
        remaining_channels = channel_pool.unclaimed_channels()
        for channel in remaining_channels:
            if channel.cluster.cluster_id in zha_regs.CHANNEL_ONLY_CLUSTERS:
                channel_pool.claim_channels([channel])
                continue

            component = zha_regs.SINGLE_INPUT_CLUSTER_DEVICE_CLASS.get(
                channel.cluster.cluster_id
            )
            if component is None:
                for cluster_class, match in single_input_clusters.items():
                    if isinstance(channel.cluster, cluster_class):
                        component = match
                        break

            self.probe_single_cluster(component, channel, channel_pool)

        # until we can get rid off registries
        self.handle_on_off_output_cluster_exception(channel_pool)

    @staticmethod
    def probe_single_cluster(
        component: str,
        channel: zha_typing.ChannelType,
        ep_channels: zha_typing.ChannelPoolType,
    ) -> None:
        """Probe specified cluster for specific component."""
        if component is None or component not in zha_const.PLATFORMS:
            return
        channel_list = [channel]
        unique_id = f"{ep_channels.unique_id}-{channel.cluster.cluster_id}"

        entity_class, claimed = zha_regs.ZHA_ENTITIES.get_entity(
            component, ep_channels.manufacturer, ep_channels.model, channel_list
        )
        if entity_class is None:
            return
        ep_channels.claim_channels(claimed)
        ep_channels.async_new_entity(component, entity_class, unique_id, claimed)

    def handle_on_off_output_cluster_exception(
        self, ep_channels: zha_typing.ChannelPoolType
    ) -> None:
        """Process output clusters of the endpoint."""

        profile_id = ep_channels.endpoint.profile_id
        device_type = ep_channels.endpoint.device_type
        if device_type in zha_regs.REMOTE_DEVICE_TYPES.get(profile_id, []):
            return

        for cluster_id, cluster in ep_channels.endpoint.out_clusters.items():
            component = zha_regs.SINGLE_OUTPUT_CLUSTER_DEVICE_CLASS.get(
                cluster.cluster_id
            )
            if component is None:
                continue

            channel_class = zha_regs.ZIGBEE_CHANNEL_REGISTRY.get(
                cluster_id, base.ZigbeeChannel
            )
            channel = channel_class(cluster, ep_channels)
            self.probe_single_cluster(component, channel, ep_channels)

    @staticmethod
    @callback
    def discover_multi_entities(channel_pool: zha_typing.ChannelPoolType) -> None:
        """Process an endpoint on and discover multiple entities."""

        ep_profile_id = channel_pool.endpoint.profile_id
        ep_device_type = channel_pool.endpoint.device_type
        cmpt_by_dev_type = zha_regs.DEVICE_CLASS[ep_profile_id].get(ep_device_type)
        remaining_channels = channel_pool.unclaimed_channels()

        matches, claimed = zha_regs.ZHA_ENTITIES.get_multi_entity(
            channel_pool.manufacturer, channel_pool.model, remaining_channels
        )

        channel_pool.claim_channels(claimed)
        for component, ent_n_chan_list in matches.items():
            for entity_and_channel in ent_n_chan_list:
                _LOGGER.debug(
                    "'%s' component -> '%s' using %s",
                    component,
                    entity_and_channel.entity_class.__name__,
                    [ch.name for ch in entity_and_channel.claimed_channel],
                )
        for component, ent_n_chan_list in matches.items():
            for entity_and_channel in ent_n_chan_list:
                if component == cmpt_by_dev_type:
                    # for well known device types, like thermostats we'll take only 1st class
                    channel_pool.async_new_entity(
                        component,
                        entity_and_channel.entity_class,
                        channel_pool.unique_id,
                        entity_and_channel.claimed_channel,
                    )
                    break
                first_ch = entity_and_channel.claimed_channel[0]
                channel_pool.async_new_entity(
                    component,
                    entity_and_channel.entity_class,
                    f"{channel_pool.unique_id}-{first_ch.cluster.cluster_id}",
                    entity_and_channel.claimed_channel,
                )

    def initialize(self, hass: HomeAssistant) -> None:
        """Update device overrides config."""
        zha_config = hass.data[zha_const.DATA_ZHA].get(zha_const.DATA_ZHA_CONFIG, {})
        if overrides := zha_config.get(zha_const.CONF_DEVICE_CONFIG):
            self._device_configs.update(overrides)

class GroupProbe:
    """Determine the appropriate component for a group."""

    def __init__(self):
        """Initialize instance."""
        self._hass = None
        self._unsubs = []

    def initialize(self, hass: HomeAssistant) -> None:
        """Initialize the group probe."""
        self._hass = hass
        self._unsubs.append(
            async_dispatcher_connect(
                hass, zha_const.SIGNAL_GROUP_ENTITY_REMOVED, self._reprobe_group
            )
        )

    def cleanup(self):
        """Clean up on when zha shuts down."""
        for unsub in self._unsubs[:]:
            unsub()
            self._unsubs.remove(unsub)

    @callback
    def _reprobe_group(self, group_id: int) -> None:
        """Reprobe a group for entities after its members change."""
        zha_gateway = self._hass.data[zha_const.DATA_ZHA][zha_const.DATA_ZHA_GATEWAY]
        if (zha_group := zha_gateway.groups.get(group_id)) is None:
            return
        self.discover_group_entities(zha_group)

    @callback
    def discover_group_entities(self, group: zha_typing.ZhaGroupType) -> None:
        """Process a group and create any entities that are needed."""
        # only create a group entity if there are 2 or more members in a group
        if len(group.members) < 2:
            _LOGGER.debug(
                "Group: %s:0x%04x has less than 2 members - skipping entity discovery",
                group.name,
                group.group_id,
            )
            return

        entity_domains = GroupProbe.determine_entity_domains(self._hass, group)

        if not entity_domains:
            return

        zha_gateway = self._hass.data[zha_const.DATA_ZHA][zha_const.DATA_ZHA_GATEWAY]
        for domain in entity_domains:
            entity_class = zha_regs.ZHA_ENTITIES.get_group_entity(domain)
            if entity_class is None:
                continue
            self._hass.data[zha_const.DATA_ZHA][domain].append(
                (
                    entity_class,
                    (
                        group.get_domain_entity_ids(domain),
                        f"{domain}_zha_group_0x{group.group_id:04x}",
                        group.group_id,
                        zha_gateway.coordinator_zha_device,
                    ),
                )
            )
        async_dispatcher_send(self._hass, zha_const.SIGNAL_ADD_ENTITIES)

    @staticmethod
    def determine_entity_domains(
        hass: HomeAssistant, group: zha_typing.ZhaGroupType
    ) -> list[str]:
        """Determine the entity domains for this group."""
        entity_domains: list[str] = []
        zha_gateway = hass.data[zha_const.DATA_ZHA][zha_const.DATA_ZHA_GATEWAY]
        all_domain_occurrences = []
        for member in group.members:
            if member.device.is_coordinator:
                continue
            entities = async_entries_for_device(
                zha_gateway.ha_entity_registry,
                member.device.device_id,
                include_disabled_entities=True,
            )
            all_domain_occurrences.extend(
                [
                    entity.domain
                    for entity in entities
                    if entity.domain in zha_regs.GROUP_ENTITY_DOMAINS
                ]
            )
        if not all_domain_occurrences:
            return entity_domains
        # get all domains we care about if there are more than 2 entities of this domain
        counts = Counter(all_domain_occurrences)
        entity_domains = [domain[0] for domain in counts.items() if domain[1] >= 2]
        _LOGGER.debug(
            "The entity domains are: %s for group: %s:0x%04x",
            entity_domains,
            group.name,
            group.group_id,
        )
        return entity_domains

PROBE = ProbeEndpoint()
GROUP_PROBE = GroupProbe()
<commit_msg>Add zha typing [core.discovery] (1) (#68359)<commit_after>"""Device discovery functions for Zigbee Home Automation."""
from __future__ import annotations

from collections import Counter
from collections.abc import Callable
import logging
from typing import TYPE_CHECKING

from homeassistant import const as ha_const
from homeassistant.core import HomeAssistant, callback
from homeassistant.helpers.dispatcher import (
    async_dispatcher_connect,
    async_dispatcher_send,
)
from homeassistant.helpers.entity_platform import AddEntitiesCallback
from homeassistant.helpers.entity_registry import async_entries_for_device
from homeassistant.helpers.typing import ConfigType

from . import const as zha_const, registries as zha_regs
from .. import (  # noqa: F401 pylint: disable=unused-import,
    alarm_control_panel,
    binary_sensor,
    button,
    climate,
    cover,
    device_tracker,
    fan,
    light,
    lock,
    number,
    select,
    sensor,
    siren,
    switch,
)
from .channels import base

if TYPE_CHECKING:
    from ..entity import ZhaEntity
    from .channels import ChannelPool
    from .device import ZHADevice
    from .gateway import ZHAGateway
    from .group import ZHAGroup

_LOGGER = logging.getLogger(__name__)

@callback
async def async_add_entities(
    _async_add_entities: AddEntitiesCallback,
    entities: list[
        tuple[
            type[ZhaEntity],
            tuple[str, ZHADevice, list[base.ZigbeeChannel]],
        ]
    ],
    update_before_add: bool = True,
) -> None:
    """Add entities helper."""
    if not entities:
        return
    to_add = [ent_cls.create_entity(*args) for ent_cls, args in entities]
    entities_to_add = [entity for entity in to_add if entity is not None]
    _async_add_entities(entities_to_add, update_before_add=update_before_add)
    entities.clear()

class ProbeEndpoint:
    """All discovered channels and entities of an endpoint."""

    def __init__(self) -> None:
        """Initialize instance."""
        self._device_configs: ConfigType = {}

    @callback
    def discover_entities(self, channel_pool: ChannelPool) -> None:
        """Process an endpoint on a zigpy device."""
        self.discover_by_device_type(channel_pool)
        self.discover_multi_entities(channel_pool)
        self.discover_by_cluster_id(channel_pool)
        zha_regs.ZHA_ENTITIES.clean_up()

    @callback
    def discover_by_device_type(self, channel_pool: ChannelPool) -> None:
        """Process an endpoint on a zigpy device."""

        unique_id = channel_pool.unique_id

        component: str | None = self._device_configs.get(unique_id, {}).get(
            ha_const.CONF_TYPE
        )
        if component is None:
            ep_profile_id = channel_pool.endpoint.profile_id
            ep_device_type = channel_pool.endpoint.device_type
            component = zha_regs.DEVICE_CLASS[ep_profile_id].get(ep_device_type)

        if component and component in zha_const.PLATFORMS:
            channels = channel_pool.unclaimed_channels()
            entity_class, claimed = zha_regs.ZHA_ENTITIES.get_entity(
                component, channel_pool.manufacturer, channel_pool.model, channels
            )
            if entity_class is None:
                return
            channel_pool.claim_channels(claimed)
            channel_pool.async_new_entity(component, entity_class, unique_id, claimed)

    @callback
    def discover_by_cluster_id(self, channel_pool: ChannelPool) -> None:
        """Process an endpoint on a zigpy device."""

        items = zha_regs.SINGLE_INPUT_CLUSTER_DEVICE_CLASS.items()
        single_input_clusters = {
            cluster_class: match
            for cluster_class, match in items
            if not isinstance(cluster_class, int)
        }
        remaining_channels = channel_pool.unclaimed_channels()
        for channel in remaining_channels:
            if channel.cluster.cluster_id in zha_regs.CHANNEL_ONLY_CLUSTERS:
                channel_pool.claim_channels([channel])
                continue

            component = zha_regs.SINGLE_INPUT_CLUSTER_DEVICE_CLASS.get(
                channel.cluster.cluster_id
            )
            if component is None:
                for cluster_class, match in single_input_clusters.items():
                    if isinstance(channel.cluster, cluster_class):
                        component = match
                        break

            self.probe_single_cluster(component, channel, channel_pool)

        # until we can get rid off registries
        self.handle_on_off_output_cluster_exception(channel_pool)

    @staticmethod
    def probe_single_cluster(
        component: str,
        channel: base.ZigbeeChannel,
        ep_channels: ChannelPool,
    ) -> None:
        """Probe specified cluster for specific component."""
        if component is None or component not in zha_const.PLATFORMS:
            return
        channel_list = [channel]
        unique_id = f"{ep_channels.unique_id}-{channel.cluster.cluster_id}"

        entity_class, claimed = zha_regs.ZHA_ENTITIES.get_entity(
            component, ep_channels.manufacturer, ep_channels.model, channel_list
        )
        if entity_class is None:
            return
        ep_channels.claim_channels(claimed)
        ep_channels.async_new_entity(component, entity_class, unique_id, claimed)

    def handle_on_off_output_cluster_exception(self, ep_channels: ChannelPool) -> None:
        """Process output clusters of the endpoint."""

        profile_id = ep_channels.endpoint.profile_id
        device_type = ep_channels.endpoint.device_type
        if device_type in zha_regs.REMOTE_DEVICE_TYPES.get(profile_id, []):
            return

        for cluster_id, cluster in ep_channels.endpoint.out_clusters.items():
            component = zha_regs.SINGLE_OUTPUT_CLUSTER_DEVICE_CLASS.get(
                cluster.cluster_id
            )
            if component is None:
                continue

            channel_class = zha_regs.ZIGBEE_CHANNEL_REGISTRY.get(
                cluster_id, base.ZigbeeChannel
            )
            channel = channel_class(cluster, ep_channels)
            self.probe_single_cluster(component, channel, ep_channels)

    @staticmethod
    @callback
    def discover_multi_entities(channel_pool: ChannelPool) -> None:
        """Process an endpoint on and discover multiple entities."""

        ep_profile_id = channel_pool.endpoint.profile_id
        ep_device_type = channel_pool.endpoint.device_type
        cmpt_by_dev_type = zha_regs.DEVICE_CLASS[ep_profile_id].get(ep_device_type)
        remaining_channels = channel_pool.unclaimed_channels()

        matches, claimed = zha_regs.ZHA_ENTITIES.get_multi_entity(
            channel_pool.manufacturer, channel_pool.model, remaining_channels
        )

        channel_pool.claim_channels(claimed)
        for component, ent_n_chan_list in matches.items():
            for entity_and_channel in ent_n_chan_list:
                _LOGGER.debug(
                    "'%s' component -> '%s' using %s",
                    component,
                    entity_and_channel.entity_class.__name__,
                    [ch.name for ch in entity_and_channel.claimed_channel],
                )
        for component, ent_n_chan_list in matches.items():
            for entity_and_channel in ent_n_chan_list:
                if component == cmpt_by_dev_type:
                    # for well known device types, like thermostats we'll take only 1st class
                    channel_pool.async_new_entity(
                        component,
                        entity_and_channel.entity_class,
                        channel_pool.unique_id,
                        entity_and_channel.claimed_channel,
                    )
                    break
                first_ch = entity_and_channel.claimed_channel[0]
                channel_pool.async_new_entity(
                    component,
                    entity_and_channel.entity_class,
                    f"{channel_pool.unique_id}-{first_ch.cluster.cluster_id}",
                    entity_and_channel.claimed_channel,
                )

    def initialize(self, hass: HomeAssistant) -> None:
        """Update device overrides config."""
        zha_config: ConfigType = hass.data[zha_const.DATA_ZHA].get(
            zha_const.DATA_ZHA_CONFIG, {}
        )
        if overrides := zha_config.get(zha_const.CONF_DEVICE_CONFIG):
            self._device_configs.update(overrides)

class GroupProbe:
    """Determine the appropriate component for a group."""

    _hass: HomeAssistant

    def __init__(self) -> None:
        """Initialize instance."""
        self._unsubs: list[Callable[[], None]] = []

    def initialize(self, hass: HomeAssistant) -> None:
        """Initialize the group probe."""
        self._hass = hass
        self._unsubs.append(
            async_dispatcher_connect(
                hass, zha_const.SIGNAL_GROUP_ENTITY_REMOVED, self._reprobe_group
            )
        )

    def cleanup(self) -> None:
        """Clean up on when zha shuts down."""
        for unsub in self._unsubs[:]:
            unsub()
            self._unsubs.remove(unsub)

    @callback
    def _reprobe_group(self, group_id: int) -> None:
        """Reprobe a group for entities after its members change."""
        zha_gateway: ZHAGateway = self._hass.data[zha_const.DATA_ZHA][
            zha_const.DATA_ZHA_GATEWAY
        ]
        if (zha_group := zha_gateway.groups.get(group_id)) is None:
            return
        self.discover_group_entities(zha_group)

    @callback
    def discover_group_entities(self, group: ZHAGroup) -> None:
        """Process a group and create any entities that are needed."""
        # only create a group entity if there are 2 or more members in a group
        if len(group.members) < 2:
            _LOGGER.debug(
                "Group: %s:0x%04x has less than 2 members - skipping entity discovery",
                group.name,
                group.group_id,
            )
            return

        entity_domains = GroupProbe.determine_entity_domains(self._hass, group)

        if not entity_domains:
            return

        zha_gateway: ZHAGateway = self._hass.data[zha_const.DATA_ZHA][
            zha_const.DATA_ZHA_GATEWAY
        ]
        for domain in entity_domains:
            entity_class = zha_regs.ZHA_ENTITIES.get_group_entity(domain)
            if entity_class is None:
                continue
            self._hass.data[zha_const.DATA_ZHA][domain].append(
                (
                    entity_class,
                    (
                        group.get_domain_entity_ids(domain),
                        f"{domain}_zha_group_0x{group.group_id:04x}",
                        group.group_id,
                        zha_gateway.coordinator_zha_device,
                    ),
                )
            )
        async_dispatcher_send(self._hass, zha_const.SIGNAL_ADD_ENTITIES)

    @staticmethod
    def determine_entity_domains(hass: HomeAssistant, group: ZHAGroup) -> list[str]:
        """Determine the entity domains for this group."""
        entity_domains: list[str] = []
        zha_gateway: ZHAGateway = hass.data[zha_const.DATA_ZHA][
            zha_const.DATA_ZHA_GATEWAY
        ]
        all_domain_occurrences = []
        for member in group.members:
            if member.device.is_coordinator:
                continue
            entities = async_entries_for_device(
                zha_gateway.ha_entity_registry,
                member.device.device_id,
                include_disabled_entities=True,
            )
            all_domain_occurrences.extend(
                [
                    entity.domain
                    for entity in entities
                    if entity.domain in zha_regs.GROUP_ENTITY_DOMAINS
                ]
            )
        if not all_domain_occurrences:
            return entity_domains
        # get all domains we care about if there are more than 2 entities of this domain
        counts = Counter(all_domain_occurrences)
        entity_domains = [domain[0] for domain in counts.items() if domain[1] >= 2]
        _LOGGER.debug(
            "The entity domains are: %s for group: %s:0x%04x",
            entity_domains,
            group.name,
            group.group_id,
        )
        return entity_domains

PROBE = ProbeEndpoint()
GROUP_PROBE = GroupProbe()
<|endoftext|>