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

# Import python libs
import os
import tempfile
import textwrap

# Import Salt Testing libs
from salttesting import TestCase
from salttesting.helpers import ensure_in_syspath
from salttesting.mock import MagicMock

ensure_in_syspath('../../')

# Import Salt libs
from salt.modules import file as filemod
from salt.modules import config as configmod
from salt.modules import cmdmod
from salt.exceptions import CommandExecutionError

filemod.__salt__ = {
    'config.manage_mode': configmod.manage_mode,
    'cmd.run': cmdmod.run,
    'cmd.run_all': cmdmod.run_all
}
filemod.__opts__ = {'test': False}

SED_CONTENT = """test
some
content
/var/lib/foo/app/test
here
"""

class FileReplaceTestCase(TestCase):
    MULTILINE_STRING = textwrap.dedent('''\
        Lorem ipsum dolor sit amet, consectetur adipiscing elit. Nam rhoncus
        enim ac bibendum vulputate. Etiam nibh velit, placerat ac auctor in,
        lacinia a turpis. Nulla elit elit, ornare in sodales eu, aliquam sit
        amet nisl.

        Fusce ac vehicula lectus. Vivamus justo nunc, pulvinar in ornare nec,
        sollicitudin id sem. Pellentesque sed ipsum dapibus, dapibus elit id,
        malesuada nisi.

        Lorem ipsum dolor sit amet, consectetur adipiscing elit. Donec
        venenatis tellus eget massa facilisis, in auctor ante aliquet. Sed nec
        cursus metus. Curabitur massa urna, vehicula id porttitor sed, lobortis
        quis leo.
        ''')

    def setUp(self):
        self.tfile = tempfile.NamedTemporaryFile(delete=False)
        self.tfile.write(self.MULTILINE_STRING)
        self.tfile.close()

    def tearDown(self):
        os.remove(self.tfile.name)

    def test_replace(self):
        filemod.replace(self.tfile.name, r'Etiam', 'Salticus', backup=False)

        with open(self.tfile.name, 'rb') as fp:
            self.assertIn('Salticus', fp.read())

    def test_backup(self):
        fext = '.bak'
        bak_file = '{0}{1}'.format(self.tfile.name, fext)

        filemod.replace(self.tfile.name, r'Etiam', 'Salticus', backup=fext)

        self.assertTrue(os.path.exists(bak_file))
        os.unlink(bak_file)

    def test_nobackup(self):
        fext = '.bak'
        bak_file = '{0}{1}'.format(self.tfile.name, fext)

        filemod.replace(self.tfile.name, r'Etiam', 'Salticus', backup=False)

        self.assertFalse(os.path.exists(bak_file))

    def test_dry_run(self):
        before_ctime = os.stat(self.tfile.name).st_mtime
        filemod.replace(self.tfile.name, r'Etiam', 'Salticus', dry_run=True)
        after_ctime = os.stat(self.tfile.name).st_mtime

        self.assertEqual(before_ctime, after_ctime)

    def test_show_changes(self):
        ret = filemod.replace(self.tfile.name,
                              r'Etiam', 'Salticus',
                              show_changes=True)

        self.assertTrue(ret.startswith('---'))  # looks like a diff

    def test_noshow_changes(self):
        ret = filemod.replace(self.tfile.name,
                              r'Etiam', 'Salticus',
                              show_changes=False)

        self.assertIsInstance(ret, bool)

    def test_re_str_flags(self):
        # upper- & lower-case
        filemod.replace(self.tfile.name,
                        r'Etiam', 'Salticus',
                        flags=['MULTILINE', 'ignorecase'])

    def test_re_int_flags(self):
        filemod.replace(self.tfile.name, r'Etiam', 'Salticus', flags=10)

class FileBlockReplaceTestCase(TestCase):
    MULTILINE_STRING = textwrap.dedent('''\
        Lorem ipsum dolor sit amet, consectetur adipiscing elit. Nam rhoncus
        enim ac bibendum vulputate. Etiam nibh velit, placerat ac auctor in,
        lacinia a turpis. Nulla elit elit, ornare in sodales eu, aliquam sit
        amet nisl.

        Fusce ac vehicula lectus. Vivamus justo nunc, pulvinar in ornare nec,
        sollicitudin id sem. Pellentesque sed ipsum dapibus, dapibus elit id,
        malesuada nisi.

        first part of start line // START BLOCK : part of start line not removed
        to be removed
        first part of end line // END BLOCK : part of end line not removed

        #-- START BLOCK UNFINISHED

        #-- START BLOCK 1
        old content part 1
        old content part 2
        #-- END BLOCK 1

        Lorem ipsum dolor sit amet, consectetur adipiscing elit. Donec
        venenatis tellus eget massa facilisis, in auctor ante aliquet. Sed nec
        cursus metus. Curabitur massa urna, vehicula id porttitor sed, lobortis
        quis leo.
        ''')

    def setUp(self):
        self.tfile = tempfile.NamedTemporaryFile(delete=False,
                                                 prefix='blockrepltmp')
        self.tfile.write(self.MULTILINE_STRING)
        self.tfile.close()
        manage_mode_mock = MagicMock()
        filemod.__salt__['config.manage_mode'] = manage_mode_mock

    def tearDown(self):
        os.remove(self.tfile.name)

    def test_replace_multiline(self):
        new_multiline_content = (
            "Who's that then?\nWell, how'd you become king,"
            "then?\nWe found them. I'm not a witch.\nWe shall"
            "say 'Ni' again to you, if you do not appease us."
        )
        filemod.blockreplace(self.tfile.name,
                             '#-- START BLOCK 1',
                             '#-- END BLOCK 1',
                             new_multiline_content,
                             backup=False)

        with open(self.tfile.name, 'rb') as fp:
            filecontent = fp.read()
        self.assertIn('#-- START BLOCK 1'
                      + "\n" + new_multiline_content
                      + "\n"
                      + '#-- END BLOCK 1', filecontent)
        self.assertNotIn('old content part 1', filecontent)
        self.assertNotIn('old content part 2', filecontent)

    def test_replace_append(self):
        new_content = "Well, I didn't vote for you."

        self.assertRaises(
            CommandExecutionError,
            filemod.blockreplace,
            self.tfile.name,
            '#-- START BLOCK 2',
            '#-- END BLOCK 2',
            new_content,
            append_if_not_found=False,
            backup=False
        )
        with open(self.tfile.name, 'rb') as fp:
            self.assertNotIn('#-- START BLOCK 2'
                             + "\n" + new_content + "\n"
                             + '#-- END BLOCK 2', fp.read())

        filemod.blockreplace(self.tfile.name,
                             '#-- START BLOCK 2',
                             '#-- END BLOCK 2',
                             new_content,
                             backup=False,
                             append_if_not_found=True)

        with open(self.tfile.name, 'rb') as fp:
            self.assertIn('#-- START BLOCK 2'
                          + "\n" + new_content
                          + "\n" + '#-- END BLOCK 2', fp.read())

    def test_replace_prepend(self):
        new_content = "Well, I didn't vote for you."

        self.assertRaises(
            CommandExecutionError,
            filemod.blockreplace,
            self.tfile.name,
            '#-- START BLOCK 2',
            '#-- END BLOCK 2',
            new_content,
            prepend_if_not_found=False,
            backup=False
        )
        with open(self.tfile.name, 'rb') as fp:
            self.assertNotIn(
                '#-- START BLOCK 2' + "\n"
                + new_content + "\n" + '#-- END BLOCK 2',
                fp.read())

        filemod.blockreplace(self.tfile.name,
                             '#-- START BLOCK 2', '#-- END BLOCK 2',
                             new_content,
                             backup=False,
                             prepend_if_not_found=True)

        with open(self.tfile.name, 'rb') as fp:
            self.assertTrue(
                fp.read().startswith(
                    '#-- START BLOCK 2'
                    + "\n" + new_content
                    + "\n" + '#-- END BLOCK 2'))

    def test_replace_partial_marked_lines(self):
        filemod.blockreplace(self.tfile.name,
                             '// START BLOCK',
                             '// END BLOCK',
                             'new content 1',
                             backup=False)

        with open(self.tfile.name, 'rb') as fp:
            filecontent = fp.read()
        self.assertIn('new content 1', filecontent)
        self.assertNotIn('to be removed', filecontent)
        self.assertIn('first part of start line', filecontent)
        self.assertIn('first part of end line', filecontent)
        self.assertIn('part of start line not removed', filecontent)
        self.assertIn('part of end line not removed', filecontent)

    def test_backup(self):
        fext = '.bak'
        bak_file = '{0}{1}'.format(self.tfile.name, fext)

        filemod.blockreplace(
            self.tfile.name,
            '// START BLOCK', '// END BLOCK', 'new content 2',
            backup=fext)

        self.assertTrue(os.path.exists(bak_file))
        os.unlink(bak_file)
        self.assertFalse(os.path.exists(bak_file))

        fext = '.bak'
        bak_file = '{0}{1}'.format(self.tfile.name, fext)

        filemod.blockreplace(self.tfile.name,
                             '// START BLOCK', '// END BLOCK', 'new content 3',
                             backup=False)

        self.assertFalse(os.path.exists(bak_file))

    def test_no_modifications(self):
        filemod.blockreplace(self.tfile.name,
                             '// START BLOCK', '// END BLOCK',
                             'new content 4',
                             backup=False)
        before_ctime = os.stat(self.tfile.name).st_mtime
        filemod.blockreplace(self.tfile.name,
                             '// START BLOCK',
                             '// END BLOCK',
                             'new content 4',
                             backup=False)
        after_ctime = os.stat(self.tfile.name).st_mtime

        self.assertEqual(before_ctime, after_ctime)

    def test_dry_run(self):
        before_ctime = os.stat(self.tfile.name).st_mtime
        filemod.blockreplace(self.tfile.name,
                             '// START BLOCK',
                             '// END BLOCK',
                             'new content 5',
                             dry_run=True)
        after_ctime = os.stat(self.tfile.name).st_mtime

        self.assertEqual(before_ctime, after_ctime)

    def test_show_changes(self):
        ret = filemod.blockreplace(self.tfile.name,
                                   '// START BLOCK',
                                   '// END BLOCK',
                                   'new content 6',
                                   backup=False,
                                   show_changes=True)

        self.assertTrue(ret.startswith('---'))  # looks like a diff

        ret = filemod.blockreplace(self.tfile.name,
                                   '// START BLOCK',
                                   '// END BLOCK',
                                   'new content 7',
                                   backup=False,
                                   show_changes=False)

        self.assertIsInstance(ret, bool)

    def test_unfinished_block_exception(self):
        self.assertRaises(
            CommandExecutionError,
            filemod.blockreplace,
            self.tfile.name,
            '#-- START BLOCK UNFINISHED',
            '#-- END BLOCK UNFINISHED',
            'foobar',
            backup=False
        )

class FileModuleTestCase(TestCase):
    def test_sed_limit_escaped(self):
        with tempfile.NamedTemporaryFile() as tfile:
            tfile.write(SED_CONTENT)
            tfile.seek(0, 0)

            path = tfile.name
            before = '/var/lib/foo'
            after = ''
            limit = '^{0}'.format(before)

            filemod.sed(path, before, after, limit=limit)

            with open(path, 'rb') as newfile:
                self.assertEqual(
                    SED_CONTENT.replace(before, ''),
                    newfile.read()
                )

    def test_append_newline_at_eof(self):
        '''
        Check that file.append works consistently on files with and without
        newlines at end of file.
        '''
        # File ending with a newline
        with tempfile.NamedTemporaryFile() as tfile:
            tfile.write('foo\n')
            tfile.flush()
            filemod.append(tfile.name, 'bar')
            with open(tfile.name) as tfile2:
                self.assertEqual(tfile2.read(), 'foo\nbar\n')
        # File not ending with a newline
        with tempfile.NamedTemporaryFile() as tfile:
            tfile.write('foo')
            tfile.flush()
            filemod.append(tfile.name, 'bar')
            with open(tfile.name) as tfile2:
                self.assertEqual(tfile2.read(), 'foo\nbar\n')
        # A newline should not be added in empty files
        with tempfile.NamedTemporaryFile() as tfile:
            filemod.append(tfile.name, 'bar')
            with open(tfile.name) as tfile2:
                self.assertEqual(tfile2.read(), 'bar\n')

if __name__ == '__main__':
    from integration import run_tests
    run_tests(FileModuleTestCase,
              FileReplaceTestCase,
              FileBlockReplaceTestCase,
              needs_daemon=False)
<commit_msg>whitespace<commit_after># -*- coding: utf-8 -*-

# Import python libs
import os
import tempfile
import textwrap

# Import Salt Testing libs
from salttesting import TestCase
from salttesting.helpers import ensure_in_syspath
from salttesting.mock import MagicMock

ensure_in_syspath('../../')

# Import Salt libs
from salt.modules import file as filemod
from salt.modules import config as configmod
from salt.modules import cmdmod
from salt.exceptions import CommandExecutionError

filemod.__salt__ = {
    'config.manage_mode': configmod.manage_mode,
    'cmd.run': cmdmod.run,
    'cmd.run_all': cmdmod.run_all
}
filemod.__opts__ = {'test': False}

SED_CONTENT = """test
some
content
/var/lib/foo/app/test
here
"""

class FileReplaceTestCase(TestCase):
    MULTILINE_STRING = textwrap.dedent('''\
        Lorem ipsum dolor sit amet, consectetur adipiscing elit. Nam rhoncus
        enim ac bibendum vulputate. Etiam nibh velit, placerat ac auctor in,
        lacinia a turpis. Nulla elit elit, ornare in sodales eu, aliquam sit
        amet nisl.

        Fusce ac vehicula lectus. Vivamus justo nunc, pulvinar in ornare nec,
        sollicitudin id sem. Pellentesque sed ipsum dapibus, dapibus elit id,
        malesuada nisi.

        Lorem ipsum dolor sit amet, consectetur adipiscing elit. Donec
        venenatis tellus eget massa facilisis, in auctor ante aliquet. Sed nec
        cursus metus. Curabitur massa urna, vehicula id porttitor sed, lobortis
        quis leo.
        ''')

    def setUp(self):
        self.tfile = tempfile.NamedTemporaryFile(delete=False)
        self.tfile.write(self.MULTILINE_STRING)
        self.tfile.close()

    def tearDown(self):
        os.remove(self.tfile.name)

    def test_replace(self):
        filemod.replace(self.tfile.name, r'Etiam', 'Salticus', backup=False)

        with open(self.tfile.name, 'rb') as fp:
            self.assertIn('Salticus', fp.read())

    def test_backup(self):
        fext = '.bak'
        bak_file = '{0}{1}'.format(self.tfile.name, fext)

        filemod.replace(self.tfile.name, r'Etiam', 'Salticus', backup=fext)

        self.assertTrue(os.path.exists(bak_file))
        os.unlink(bak_file)

    def test_nobackup(self):
        fext = '.bak'
        bak_file = '{0}{1}'.format(self.tfile.name, fext)

        filemod.replace(self.tfile.name, r'Etiam', 'Salticus', backup=False)

        self.assertFalse(os.path.exists(bak_file))

    def test_dry_run(self):
        before_ctime = os.stat(self.tfile.name).st_mtime
        filemod.replace(self.tfile.name, r'Etiam', 'Salticus', dry_run=True)
        after_ctime = os.stat(self.tfile.name).st_mtime

        self.assertEqual(before_ctime, after_ctime)

    def test_show_changes(self):
        ret = filemod.replace(self.tfile.name,
                              r'Etiam', 'Salticus',
                              show_changes=True)

        self.assertTrue(ret.startswith('---'))  # looks like a diff

    def test_noshow_changes(self):
        ret = filemod.replace(self.tfile.name,
                              r'Etiam', 'Salticus',
                              show_changes=False)

        self.assertIsInstance(ret, bool)

    def test_re_str_flags(self):
        # upper- & lower-case
        filemod.replace(self.tfile.name,
                        r'Etiam', 'Salticus',
                        flags=['MULTILINE', 'ignorecase'])

    def test_re_int_flags(self):
        filemod.replace(self.tfile.name, r'Etiam', 'Salticus', flags=10)

class FileBlockReplaceTestCase(TestCase):
    MULTILINE_STRING = textwrap.dedent('''\
        Lorem ipsum dolor sit amet, consectetur adipiscing elit. Nam rhoncus
        enim ac bibendum vulputate. Etiam nibh velit, placerat ac auctor in,
        lacinia a turpis. Nulla elit elit, ornare in sodales eu, aliquam sit
        amet nisl.

        Fusce ac vehicula lectus. Vivamus justo nunc, pulvinar in ornare nec,
        sollicitudin id sem. Pellentesque sed ipsum dapibus, dapibus elit id,
        malesuada nisi.

        first part of start line // START BLOCK : part of start line not removed
        to be removed
        first part of end line // END BLOCK : part of end line not removed

        #-- START BLOCK UNFINISHED

        #-- START BLOCK 1
        old content part 1
        old content part 2
        #-- END BLOCK 1

        Lorem ipsum dolor sit amet, consectetur adipiscing elit. Donec
        venenatis tellus eget massa facilisis, in auctor ante aliquet. Sed nec
        cursus metus. Curabitur massa urna, vehicula id porttitor sed, lobortis
        quis leo.
        ''')

    def setUp(self):
        self.tfile = tempfile.NamedTemporaryFile(delete=False,
                                                 prefix='blockrepltmp')
        self.tfile.write(self.MULTILINE_STRING)
        self.tfile.close()
        manage_mode_mock = MagicMock()
        filemod.__salt__['config.manage_mode'] = manage_mode_mock

    def tearDown(self):
        os.remove(self.tfile.name)

    def test_replace_multiline(self):
        new_multiline_content = (
            "Who's that then?\nWell, how'd you become king,"
            "then?\nWe found them. I'm not a witch.\nWe shall"
            "say 'Ni' again to you, if you do not appease us."
        )
        filemod.blockreplace(self.tfile.name,
                             '#-- START BLOCK 1',
                             '#-- END BLOCK 1',
                             new_multiline_content,
                             backup=False)

        with open(self.tfile.name, 'rb') as fp:
            filecontent = fp.read()
        self.assertIn('#-- START BLOCK 1'
                      + "\n" + new_multiline_content
                      + "\n"
                      + '#-- END BLOCK 1', filecontent)
        self.assertNotIn('old content part 1', filecontent)
        self.assertNotIn('old content part 2', filecontent)

    def test_replace_append(self):
        new_content = "Well, I didn't vote for you."

        self.assertRaises(
            CommandExecutionError,
            filemod.blockreplace,
            self.tfile.name,
            '#-- START BLOCK 2',
            '#-- END BLOCK 2',
            new_content,
            append_if_not_found=False,
            backup=False
        )
        with open(self.tfile.name, 'rb') as fp:
            self.assertNotIn('#-- START BLOCK 2'
                             + "\n" + new_content + "\n"
                             + '#-- END BLOCK 2', fp.read())

        filemod.blockreplace(self.tfile.name,
                             '#-- START BLOCK 2',
                             '#-- END BLOCK 2',
                             new_content,
                             backup=False,
                             append_if_not_found=True)

        with open(self.tfile.name, 'rb') as fp:
            self.assertIn('#-- START BLOCK 2'
                          + "\n" + new_content
                          + "\n" + '#-- END BLOCK 2', fp.read())

    def test_replace_prepend(self):
        new_content = "Well, I didn't vote for you."

        self.assertRaises(
            CommandExecutionError,
            filemod.blockreplace,
            self.tfile.name,
            '#-- START BLOCK 2',
            '#-- END BLOCK 2',
            new_content,
            prepend_if_not_found=False,
            backup=False
        )
        with open(self.tfile.name, 'rb') as fp:
            self.assertNotIn(
                '#-- START BLOCK 2' + "\n"
                + new_content + "\n" + '#-- END BLOCK 2',
                fp.read())

        filemod.blockreplace(self.tfile.name,
                             '#-- START BLOCK 2', '#-- END BLOCK 2',
                             new_content,
                             backup=False,
                             prepend_if_not_found=True)

        with open(self.tfile.name, 'rb') as fp:
            self.assertTrue(
                fp.read().startswith(
                    '#-- START BLOCK 2'
                    + "\n" + new_content
                    + "\n" + '#-- END BLOCK 2'))

    def test_replace_partial_marked_lines(self):
        filemod.blockreplace(self.tfile.name,
                             '// START BLOCK',
                             '// END BLOCK',
                             'new content 1',
                             backup=False)

        with open(self.tfile.name, 'rb') as fp:
            filecontent = fp.read()
        self.assertIn('new content 1', filecontent)
        self.assertNotIn('to be removed', filecontent)
        self.assertIn('first part of start line', filecontent)
        self.assertIn('first part of end line', filecontent)
        self.assertIn('part of start line not removed', filecontent)
        self.assertIn('part of end line not removed', filecontent)

    def test_backup(self):
        fext = '.bak'
        bak_file = '{0}{1}'.format(self.tfile.name, fext)

        filemod.blockreplace(
            self.tfile.name,
            '// START BLOCK', '// END BLOCK', 'new content 2',
            backup=fext)

        self.assertTrue(os.path.exists(bak_file))
        os.unlink(bak_file)
        self.assertFalse(os.path.exists(bak_file))

        fext = '.bak'
        bak_file = '{0}{1}'.format(self.tfile.name, fext)

        filemod.blockreplace(self.tfile.name,
                             '// START BLOCK', '// END BLOCK', 'new content 3',
                             backup=False)

        self.assertFalse(os.path.exists(bak_file))

    def test_no_modifications(self):
        filemod.blockreplace(self.tfile.name,
                             '// START BLOCK', '// END BLOCK',
                             'new content 4',
                             backup=False)
        before_ctime = os.stat(self.tfile.name).st_mtime
        filemod.blockreplace(self.tfile.name,
                             '// START BLOCK',
                             '// END BLOCK',
                             'new content 4',
                             backup=False)
        after_ctime = os.stat(self.tfile.name).st_mtime

        self.assertEqual(before_ctime, after_ctime)

    def test_dry_run(self):
        before_ctime = os.stat(self.tfile.name).st_mtime
        filemod.blockreplace(self.tfile.name,
                             '// START BLOCK',
                             '// END BLOCK',
                             'new content 5',
                             dry_run=True)
        after_ctime = os.stat(self.tfile.name).st_mtime

        self.assertEqual(before_ctime, after_ctime)

    def test_show_changes(self):
        ret = filemod.blockreplace(self.tfile.name,
                                   '// START BLOCK',
                                   '// END BLOCK',
                                   'new content 6',
                                   backup=False,
                                   show_changes=True)

        self.assertTrue(ret.startswith('---'))  # looks like a diff

        ret = filemod.blockreplace(self.tfile.name,
                                   '// START BLOCK',
                                   '// END BLOCK',
                                   'new content 7',
                                   backup=False,
                                   show_changes=False)

        self.assertIsInstance(ret, bool)

    def test_unfinished_block_exception(self):
        self.assertRaises(
            CommandExecutionError,
            filemod.blockreplace,
            self.tfile.name,
            '#-- START BLOCK UNFINISHED',
            '#-- END BLOCK UNFINISHED',
            'foobar',
            backup=False
        )

class FileModuleTestCase(TestCase):
    def test_sed_limit_escaped(self):
        with tempfile.NamedTemporaryFile() as tfile:
            tfile.write(SED_CONTENT)
            tfile.seek(0, 0)

            path = tfile.name
            before = '/var/lib/foo'
            after = ''
            limit = '^{0}'.format(before)

            filemod.sed(path, before, after, limit=limit)

            with open(path, 'rb') as newfile:
                self.assertEqual(
                    SED_CONTENT.replace(before, ''),
                    newfile.read()
                )

    def test_append_newline_at_eof(self):
        '''
        Check that file.append works consistently on files with and without
        newlines at end of file.
        '''
        # File ending with a newline
        with tempfile.NamedTemporaryFile() as tfile:
            tfile.write('foo\n')
            tfile.flush()
            filemod.append(tfile.name, 'bar')
            with open(tfile.name) as tfile2:
                self.assertEqual(tfile2.read(), 'foo\nbar\n')
        # File not ending with a newline
        with tempfile.NamedTemporaryFile() as tfile:
            tfile.write('foo')
            tfile.flush()
            filemod.append(tfile.name, 'bar')
            with open(tfile.name) as tfile2:
                self.assertEqual(tfile2.read(), 'foo\nbar\n')
        # A newline should not be added in empty files
        with tempfile.NamedTemporaryFile() as tfile:
            filemod.append(tfile.name, 'bar')
            with open(tfile.name) as tfile2:
                self.assertEqual(tfile2.read(), 'bar\n')

if __name__ == '__main__':
    from integration import run_tests
    run_tests(FileModuleTestCase,
              FileReplaceTestCase,
              FileBlockReplaceTestCase,
              needs_daemon=False)
<|endoftext|>


<commit_before>import unittest

from zipper import Zipper

# Tests adapted from `problem-specifications//canonical-data.json` @ v1.0.0

def bt(value, left, right):
    return {
        'value': value,
        'left': left,
        'right': right
        }

def leaf(value):
    return bt(value, None, None)

EMPTY_TREE = None

def create_trees():
    t1 = bt(1, bt(2, EMPTY_TREE, leaf(3)), leaf(4))
    t2 = bt(1, bt(5, EMPTY_TREE, leaf(3)), leaf(4))
    t3 = bt(1, bt(2, leaf(5), leaf(3)), leaf(4))
    t4 = bt(1, leaf(2), leaf(4))
    return (t1, t2, t3, t4)

class ZipperTest(unittest.TestCase):
    def test_data_is_retained(self):
        t1, _, _, _ = create_trees()
        zipper = Zipper.from_tree(t1)
        tree = zipper.to_tree()
        self.assertEqual(tree, t1)

    def test_left_and_right_value(self):
        t1, _, _, _ = create_trees()
        zipper = Zipper.from_tree(t1)
        self.assertEqual(zipper.left().right().value(), 3)

    def test_dead_end(self):
        t1, _, _, _ = create_trees()
        zipper = Zipper.from_tree(t1)
        self.assertIsNone(zipper.left().left())

    def test_tree_from_deep_focus(self):
        t1, _, _, _ = create_trees()
        zipper = Zipper.from_tree(t1)
        self.assertEqual(zipper.left().right().to_tree(), t1)

    def test_set_value(self):
        t1, t2, _, _ = create_trees()
        zipper = Zipper.from_tree(t1)
        updatedZipper = zipper.left().set_value(5)
        tree = updatedZipper.to_tree()
        self.assertEqual(tree, t2)

    def test_set_left_with_value(self):
        t1, _, t3, _ = create_trees()
        zipper = Zipper.from_tree(t1)
        updatedZipper = zipper.left().set_left(leaf(5))
        tree = updatedZipper.to_tree()
        self.assertEqual(tree, t3)

    def test_set_right_to_none(self):
        t1, _, _, t4 = create_trees()
        zipper = Zipper.from_tree(t1)
        updatedZipper = zipper.left().set_right(None)
        tree = updatedZipper.to_tree()
        self.assertEqual(tree, t4)

    def test_different_paths_to_same_zipper(self):
        t1, _, _, _ = create_trees()
        zipper = Zipper.from_tree(t1)
        self.assertEqual(zipper.left().up().right().to_tree(),
                         zipper.right().to_tree())

if __name__ == '__main__':
    unittest.main()
<commit_msg>zipper: update tests to v1.1.0 (#1337)<commit_after>import unittest

from zipper import Zipper

# Tests adapted from `problem-specifications//canonical-data.json` @ v1.1.0

class ZipperTest(unittest.TestCase):
    def bt(self, value, left, right):
        return {
            'value': value,
            'left': left,
            'right': right
        }

    def leaf(self, value):
        return self.bt(value, None, None)

    def create_trees(self):
        t1 = self.bt(1, self.bt(2, None, self.leaf(3)), self.leaf(4))
        t2 = self.bt(1, self.bt(5, None, self.leaf(3)), self.leaf(4))
        t3 = self.bt(1, self.bt(2, self.leaf(5), self.leaf(3)), self.leaf(4))
        t4 = self.bt(1, self.leaf(2), self.leaf(4))
        return (t1, t2, t3, t4)

    def test_data_is_retained(self):
        t1, _, _, _ = self.create_trees()
        zipper = Zipper.from_tree(t1)
        tree = zipper.to_tree()
        self.assertEqual(tree, t1)

    def test_left_and_right_value(self):
        t1, _, _, _ = self.create_trees()
        zipper = Zipper.from_tree(t1)
        self.assertEqual(zipper.left().right().value(), 3)

    def test_dead_end(self):
        t1, _, _, _ = self.create_trees()
        zipper = Zipper.from_tree(t1)
        self.assertIsNone(zipper.left().left())

    def test_tree_from_deep_focus(self):
        t1, _, _, _ = self.create_trees()
        zipper = Zipper.from_tree(t1)
        self.assertEqual(zipper.left().right().to_tree(), t1)

    def test_set_value(self):
        t1, t2, _, _ = self.create_trees()
        zipper = Zipper.from_tree(t1)
        updatedZipper = zipper.left().set_value(5)
        tree = updatedZipper.to_tree()
        self.assertEqual(tree, t2)

    def test_set_left_with_value(self):
        t1, _, t3, _ = self.create_trees()
        zipper = Zipper.from_tree(t1)
        updatedZipper = zipper.left().set_left(self.leaf(5))
        tree = updatedZipper.to_tree()
        self.assertEqual(tree, t3)

    def test_set_right_to_none(self):
        t1, _, _, t4 = self.create_trees()
        zipper = Zipper.from_tree(t1)
        updatedZipper = zipper.left().set_right(None)
        tree = updatedZipper.to_tree()
        self.assertEqual(tree, t4)

    def test_different_paths_to_same_zipper(self):
        t1, _, _, _ = self.create_trees()
        zipper = Zipper.from_tree(t1)
        self.assertEqual(zipper.left().up().right().to_tree(),
                         zipper.right().to_tree())

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


<commit_before>from __future__ import absolute_import, print_function

import zmq
from itertools import chain
from bisect import bisect
from operator import add
from time import sleep, time
from toolz import accumulate, topk, pluck
import uuid
from collections import defaultdict
from contextlib import contextmanager
from threading import Thread, Lock
from toolz import keymap
from datetime import datetime
from . import core
from .compatibility import Queue, Empty

context = zmq.Context()

tuple_sep = '-|-'

def log(*args):
    with open('log', 'a') as f:
        print(datetime.now(), *args, file=f)

log('Import zmq')

@contextmanager
def logduration(message, nbytes=None):
    start = time()
    try:
        yield
    finally:
        end = time()
        log(message, end - start)
        if nbytes:
            log("MB/s:", nbytes / (end - start) / 1e6)

@contextmanager
def logerrors():
    try:
        yield
    except Exception as e:
        log('Error!', str(e))
        raise

class Server(object):
    def __init__(self, path, address=None, available_memory=1e9):
        self.path = path
        self.inmem = defaultdict(list)
        self.lengths = defaultdict(lambda: 0)
        self.socket = context.socket(zmq.ROUTER)

        if address is None:
            address = 'ipc://server-%s' % str(uuid.uuid1())
        self.address = address
        self.socket.bind(self.address)
        with open(core.filename(self.path, '.address'), 'w') as f:
            f.write(self.address)

        self.available_memory=available_memory
        self.memory_usage = 0
        self.status = 'run'
        self._out_disk_buffer = Queue(maxsize=3)
        self._frozen_sockets = Queue()

        self._file_lock = core.lock(path)
        self._file_lock.acquire()
        self._lock = Lock()
        self._socket_lock = Lock()

    def start(self):
        log('Start server at', self.address)
        self._listen_thread = Thread(target=self.listen)
        self._listen_thread.start()
        self._write_to_disk_thread = Thread(target=self._write_to_disk)
        self._write_to_disk_thread.start()
        self._free_frozen_sockets_thread = Thread(target=self._free_frozen_sockets)
        self._free_frozen_sockets_thread.start()

    def listen(self):
        with logerrors():
            log(self.address, 'start listening')
            while self.status != 'closed':
                if not self.socket.poll(100):
                    continue

                with self._socket_lock:
                    payload = self.socket.recv_multipart()

                address, command, payload = payload[0], payload[1], payload[2:]
                log(self.address, 'Receive command', command)
                if command == b'close':
                    self.status = 'closed'
                    self.ack(address)
                    break

                elif command == b'put':
                    keys, values = payload[::2], payload[1::2]
                    keys = list(map(deserialize_key, keys))
                    data = dict(zip(keys, values))
                    self.put(data)
                    self.ack(address)

                elif command == b'get':
                    keys = list(map(deserialize_key, payload))
                    log('get', keys)
                    result = self.get(keys)
                    self.send_to_client(address, result)
                    self.ack(address, flow_control=False)

                elif command == b'syn':
                    log('syn-ack')
                    self.ack(address)

                else:
                    log("Unknown command", command)
                    raise ValueError("Unknown command: " + command)

    def send_to_client(self, address, result):
        if not isinstance(result, list):
            result = [result]
        with self._socket_lock:
            log('Server sends', address)
            self.socket.send_multipart([address] + result)

    def _write_to_disk(self):
        while self.status != 'closed':
            try:
                data = self._out_disk_buffer.get(timeout=1)
            except Empty:
                continue
            else:
                with self._lock:
                    nbytes = sum(map(len, data.values()))
                    with logduration("Write %d files" % len(data),
                                     nbytes=nbytes):
                        core.put(self.path, data, lock=False)
                self._out_disk_buffer.task_done()

    def ack(self, address, flow_control=True):
        if flow_control and self._out_disk_buffer.full():
            self._frozen_sockets.put(address)
        else:
            log('server sends ack')
            self.send_to_client(address, b'ack')

    def _free_frozen_sockets(self):
        while self.status != 'closed':
            try:
                data = self._frozen_sockets.get(timeout=1)
            except Empty:
                continue
            else:
                self._out_disk_buffer.join()
                while not self._frozen_sockets.empty():
                    self.ack(self._frozen_sockets.get())

    def put(self, data):
        for k, v in data.items():
            self.inmem[k].append(v)
            self.lengths[k] += len(v)
            self.memory_usage += len(v)

        if self.memory_usage > self.available_memory:
            keys = keys_to_flush(self.lengths, 0.25)
            self.flush(keys)

    def flush(self, keys=None, block=None):
        """ Flush keys to disk

        Parameters
        ----------

        keys: list or None
            list of keys to flush
        block: bool (defaults to None)
            Whether or not to block until all writing is complete

        If no keys are given then flush all keys
        """
        if keys is None:
            keys = list(self.lengths)
            if block is None:
                block = True
        assert isinstance(keys, (tuple, list))
        payload = dict((key, ''.join(self.inmem[key])) for key in keys)
        self._out_disk_buffer.put(payload)

        for key in keys:
            self.memory_usage -= self.lengths[key]
            del self.inmem[key]
            del self.lengths[key]

        if block:
            self._out_disk_buffer.join()

    def get(self, keys):
        self._out_disk_buffer.join()  # block until everything is written
        with self._lock:
            from_disk = core.get(self.path, keys, lock=False)
            result = [from_disk[i] + ''.join(self.inmem[k])
                          for i, k in enumerate(keys)]
        return result

    def close(self):
        self.status = 'closed'
        self._file_lock.release()

def keys_to_flush(lengths, fraction=0.4):
    """ Which keys to remove

    >>> lengths = {'a': 20, 'b': 10, 'c': 15, 'd': 15,
    ...            'e': 10, 'f': 25, 'g': 5}
    >>> keys_to_flush(lengths, 0.5)
    ['f', 'a']
    """
    tophalf = topk(max(len(lengths) / 2, 1),
                       lengths.items(),
                       key=1)
    total = sum(lengths.values())
    cutoff = max(1, bisect(list(accumulate(add, pluck(1, tophalf))),
                    total * fraction))
    return [k for k, v in tophalf[:cutoff]]

def create(path, **kwargs):
    core.create(path)
    server = Server(path, **kwargs)
    server.start()
    return server

sockets = dict()

def socket(path):
    try:
        return sockets[path]
    except KeyError:
        sock = context.socket(zmq.DEALER)
        lock = Lock()
        with open(core.filename(path, '.address')) as f:
            addr = f.read()
        sock.connect(addr)
        sockets[path] = (sock, lock)
        send(path, b'syn', [], ack_required=False)
        return sock, lock

def destroy(path, server=None):
    if server:
        server.close()
    else:
        send(path, b'close', [])
    del sockets[path]
    core.destroy(path)

def get(path, keys):
    keys = list(map(serialize_key, keys))
    return send(path, b'get', keys, recv=True)

def put(path, data):
    data = keymap(serialize_key, data)
    payload = list(chain.from_iterable(data.items()))
    send(path, b'put', payload)

def send(path, command, payload, recv=False, ack_required=True):
    sock, lock = socket(path)
    with lock:
        if ack_required:
            ack = sock.recv_multipart()
            assert ack == [b'ack']
        log('Client sends command', command)
        sock.send_multipart([command] + payload)
        if recv:
            result = sock.recv_multipart()
        else:
            result = None
    return result

ensure = core.ensure

@contextmanager
def partd(path=None, **kwargs):
    if path is None:
        path = 'tmp.partd'
    server = create(path, **kwargs)

    try:
        yield path, server
    finally:
        destroy(path, server)

def serialize_key(key):
    """

    >>> serialize_key('x')
    'x'
    >>> serialize_key(('a', 'b', 1))
    'a-|-b-|-1'
    """
    if isinstance(key, tuple):
        return tuple_sep.join(map(serialize_key, key))
    if isinstance(key, (bytes, str)):
        return key
    return str(key)

def deserialize_key(text):
    """

    >>> deserialize_key('x')
    'x'
    >>> deserialize_key('a-|-b-|-1')
    ('a', 'b', '1')
    """
    if tuple_sep in text:
        return tuple(text.split(tuple_sep))
    else:
        return text
<commit_msg>excessive logging.  Solve free address problem<commit_after>from __future__ import absolute_import, print_function

import zmq
from itertools import chain
from bisect import bisect
from operator import add
from time import sleep, time
from toolz import accumulate, topk, pluck
import uuid
from collections import defaultdict
from contextlib import contextmanager
from threading import Thread, Lock
from toolz import keymap
from datetime import datetime
from . import core
from .compatibility import Queue, Empty

context = zmq.Context()

tuple_sep = '-|-'

def log(*args):
    with open('log', 'a') as f:
        print(datetime.now(), *args, file=f)

log('Import zmq')

@contextmanager
def logduration(message, nbytes=None):
    start = time()
    try:
        yield
    finally:
        end = time()
        log(message, end - start)
        if nbytes:
            log("MB/s:", nbytes / (end - start) / 1e6)

@contextmanager
def logerrors():
    try:
        yield
    except Exception as e:
        log('Error!', str(e))
        raise

class Server(object):
    def __init__(self, path, address=None, available_memory=1e9):
        self.path = path
        self.inmem = defaultdict(list)
        self.lengths = defaultdict(lambda: 0)
        self.socket = context.socket(zmq.ROUTER)

        if address is None:
            address = 'ipc://server-%s' % str(uuid.uuid1())
        self.address = address
        self.socket.bind(self.address)
        with open(core.filename(self.path, '.address'), 'w') as f:
            f.write(self.address)

        self.available_memory=available_memory
        self.memory_usage = 0
        self.status = 'run'
        self._out_disk_buffer = Queue(maxsize=3)
        self._frozen_sockets = Queue()

        self._file_lock = core.lock(path)
        self._file_lock.acquire()
        self._lock = Lock()
        self._socket_lock = Lock()

    def start(self):
        log('Start server at', self.address)
        self._listen_thread = Thread(target=self.listen)
        self._listen_thread.start()
        self._write_to_disk_thread = Thread(target=self._write_to_disk)
        self._write_to_disk_thread.start()
        self._free_frozen_sockets_thread = Thread(target=self._free_frozen_sockets)
        self._free_frozen_sockets_thread.start()

    def listen(self):
        with logerrors():
            log(self.address, 'start listening')
            while self.status != 'closed':
                if not self.socket.poll(100):
                    continue

                with self._socket_lock:
                    payload = self.socket.recv_multipart()

                address, command, payload = payload[0], payload[1], payload[2:]
                if command == b'close':
                    log('Server closes')
                    self.status = 'closed'
                    self.ack(address)
                    break

                elif command == b'put':
                    keys, values = payload[::2], payload[1::2]
                    keys = list(map(deserialize_key, keys))
                    data = dict(zip(keys, values))
                    self.put(data)
                    self.ack(address)

                elif command == b'get':
                    keys = list(map(deserialize_key, payload))
                    log('get', keys)
                    result = self.get(keys)
                    self.send_to_client(address, result)
                    self.ack(address, flow_control=False)

                elif command == b'syn':
                    self.ack(address)

                else:
                    log("Unknown command", command)
                    raise ValueError("Unknown command: " + command)

    def send_to_client(self, address, result):
        if not isinstance(result, list):
            result = [result]
        with self._socket_lock:
            self.socket.send_multipart([address] + result)

    def _write_to_disk(self):
        while self.status != 'closed':
            try:
                data = self._out_disk_buffer.get(timeout=1)
            except Empty:
                continue
            else:
                with self._lock:
                    nbytes = sum(map(len, data.values()))
                    with logduration("Write %d files" % len(data),
                                     nbytes=nbytes):
                        core.put(self.path, data, lock=False)
                self._out_disk_buffer.task_done()

    def ack(self, address, flow_control=True):
        if flow_control and self._out_disk_buffer.full():
            log('Out disk buffer full - Flow control in effect',
                'Freezing address', address)
            self._frozen_sockets.put(address)
        else:
            log('Server sends ack')
            self.send_to_client(address, b'ack')

    def _free_frozen_sockets(self):
        while self.status != 'closed':
            try:
                data = self._frozen_sockets.get(timeout=1)
                self._frozen_sockets.put(data)  # put back in for the moment
            except Empty:
                continue
            else:
                log('Freeing frozen sockets, waiting on disk buffer to clear')
                self._out_disk_buffer.join()
                log('Disk buffer cleared, sending acks to %d sockets' %
                        self._frozen_sockets.qsize())
                while not self._frozen_sockets.empty():
                    addr = self._frozen_sockets.get()
                    log('Free', addr)
                    self.ack(addr)

    def put(self, data):
        total_mem = 0
        for k, v in data.items():
            self.inmem[k].append(v)
            self.lengths[k] += len(v)
            total_mem += len(v)
        self.memory_usage += total_mem

        log('Server puts %d keys' % len(data.keys()), 'of %d bytes' % total_mem)

        if self.memory_usage > self.available_memory:
            keys = keys_to_flush(self.lengths, 0.25)
            self.flush(keys)

    def flush(self, keys=None, block=None):
        """ Flush keys to disk

        Parameters
        ----------

        keys: list or None
            list of keys to flush
        block: bool (defaults to None)
            Whether or not to block until all writing is complete

        If no keys are given then flush all keys
        """
        if keys is None:
            keys = list(self.lengths)
            if block is None:
                block = True
        assert isinstance(keys, (tuple, list))
        payload = dict((key, ''.join(self.inmem[key])) for key in keys)
        log('Put data into out-disk-buffer', 'nkeys', len(keys))
        self._out_disk_buffer.put(payload)

        for key in keys:
            self.memory_usage -= self.lengths[key]
            del self.inmem[key]
            del self.lengths[key]

        if block:
            log('Blocking on out disk buffer from flush')
            self._out_disk_buffer.join()

    def get(self, keys):
        log('Server gets keys', keys)
        self._out_disk_buffer.join()  # block until everything is written
        with self._lock:
            from_disk = core.get(self.path, keys, lock=False)
            result = [from_disk[i] + ''.join(self.inmem[k])
                          for i, k in enumerate(keys)]
        return result

    def close(self):
        log('Server closes')
        self.status = 'closed'
        self._file_lock.release()

def keys_to_flush(lengths, fraction=0.4):
    """ Which keys to remove

    >>> lengths = {'a': 20, 'b': 10, 'c': 15, 'd': 15,
    ...            'e': 10, 'f': 25, 'g': 5}
    >>> keys_to_flush(lengths, 0.5)
    ['f', 'a']
    """
    tophalf = topk(max(len(lengths) / 2, 1),
                       lengths.items(),
                       key=1)
    total = sum(lengths.values())
    cutoff = max(1, bisect(list(accumulate(add, pluck(1, tophalf))),
                    total * fraction))
    return [k for k, v in tophalf[:cutoff]]

def create(path, **kwargs):
    core.create(path)
    server = Server(path, **kwargs)
    server.start()
    return server

sockets = dict()

def socket(path):
    try:
        return sockets[path]
    except KeyError:
        sock = context.socket(zmq.DEALER)
        lock = Lock()
        with open(core.filename(path, '.address')) as f:
            addr = f.read()
        sock.connect(addr)
        sockets[path] = (sock, lock)
        send(path, b'syn', [], ack_required=False)
        return sock, lock

def destroy(path, server=None):
    if server:
        server.close()
    else:
        send(path, b'close', [])
    del sockets[path]
    core.destroy(path)

def get(path, keys):
    log('Client gets', path, keys)
    keys = list(map(serialize_key, keys))
    return send(path, b'get', keys, recv=True)

def put(path, data):
    log('Client puts', path, str(len(data)) + ' keys')
    data = keymap(serialize_key, data)
    payload = list(chain.from_iterable(data.items()))
    send(path, b'put', payload)

def send(path, command, payload, recv=False, ack_required=True):
    sock, lock = socket(path)
    with lock:
        if ack_required:
            ack = sock.recv_multipart()
            assert ack == [b'ack']
        log('Client sends command', command)
        sock.send_multipart([command] + payload)
        if recv:
            result = sock.recv_multipart()
        else:
            result = None
    return result

ensure = core.ensure

@contextmanager
def partd(path=None, **kwargs):
    if path is None:
        path = 'tmp.partd'
    server = create(path, **kwargs)

    try:
        yield path, server
    finally:
        destroy(path, server)

def serialize_key(key):
    """

    >>> serialize_key('x')
    'x'
    >>> serialize_key(('a', 'b', 1))
    'a-|-b-|-1'
    """
    if isinstance(key, tuple):
        return tuple_sep.join(map(serialize_key, key))
    if isinstance(key, (bytes, str)):
        return key
    return str(key)

def deserialize_key(text):
    """

    >>> deserialize_key('x')
    'x'
    >>> deserialize_key('a-|-b-|-1')
    ('a', 'b', '1')
    """
    if tuple_sep in text:
        return tuple(text.split(tuple_sep))
    else:
        return text
<|endoftext|>


<commit_before>"""Metrics for estimating performance of drug discovery methods implemented in ODDT"""

from math import ceil
import numpy as np
from sklearn.metrics import roc_curve as roc, roc_auc_score as roc_auc, auc

__all__ = ['roc', 'auc', 'roc_auc', 'roc_log_auc', 'enrichment_factor', '_random_set']

def enrichment_factor(y_true, y_score, percentage=1, pos_label=None):
    """Computes enrichment factor for given percentage, i.e. EF_1% is enrichment factor for first percent of given samples.

    Parameters
    ----------
        y_true : array, shape=[n_samples]
            True binary labels, in range {0,1} or {-1,1}. If positive label is different than 1, it must be explicitly defined.

        y_score : array, shape=[n_samples]
            Scores for tested series of samples

        percentage : int or float
            The percentage for which EF is being calculated

        pos_label: int
            Positive label of samples (if other than 1)

    Returns
    -------
        ef : float
            Enrichment Factor for given percenage
    """
    if pos_label is None:
        pos_label = 1
    labels = y_true == pos_label
    # calculate fraction of positve labels
    n_perc = ceil(float(percentage)/100.*len(labels))
    return labels[:n_perc].sum()/n_perc*100

def roc_log_auc(y_true, y_score, pos_label=None, log_min=0.001, log_max=1.):
    """Computes area under semi-log ROC for random distribution.

    Parameters
    ----------
        y_true : array, shape=[n_samples]
            True binary labels, in range {0,1} or {-1,1}. If positive label is different than 1, it must be explicitly defined.

        y_score : array, shape=[n_samples]
            Scores for tested series of samples

        pos_label: int
            Positive label of samples (if other than 1)

        log_min : float (default=0.001)
            Minimum logarithm value for estimating AUC

        log_max : float (default=1.)
            Maximum logarithm value for estimating AUC.

    Returns
    -------
        auc : float
            semi-log ROC AUC
    """
    fpr, tpr, t = roc(y_true, y_score, pos_label=pos_label)
    idx = (fpr >= log_min) & (fpr <= log_max)
    log_fpr = 1-np.log10(fpr[idx])/np.log10(log_min)
    return auc(log_fpr, tpr[idx])

def random_roc_log_auc(log_min=0.001, log_max=1.):
    """Computes area under semi-log ROC for random distribution.

    Parameters
    ----------
        log_min : float (default=0.001)
            Minimum logarithm value for estimating AUC

        log_max : float (default=1.)
            Maximum logarithm value for estimating AUC.

    Returns
    -------
        auc : float
            semi-log ROC AUC for random distribution
    """
    return (log_max-log_min)/(np.log(10)*np.log10(log_max/log_min))
<commit_msg>Add MES<commit_after>"""Metrics for estimating performance of drug discovery methods implemented in ODDT"""

from math import ceil
import numpy as np
from sklearn.metrics import roc_curve as roc, roc_auc_score as roc_auc, auc, mean_squared_error as mse

__all__ = ['roc', 'auc', 'roc_auc', 'roc_log_auc', 'enrichment_factor', 'random_roc_log_auc', 'mse']

def enrichment_factor(y_true, y_score, percentage=1, pos_label=None):
    """Computes enrichment factor for given percentage, i.e. EF_1% is enrichment factor for first percent of given samples.

    Parameters
    ----------
        y_true : array, shape=[n_samples]
            True binary labels, in range {0,1} or {-1,1}. If positive label is different than 1, it must be explicitly defined.

        y_score : array, shape=[n_samples]
            Scores for tested series of samples

        percentage : int or float
            The percentage for which EF is being calculated

        pos_label: int
            Positive label of samples (if other than 1)

    Returns
    -------
        ef : float
            Enrichment Factor for given percenage
    """
    if pos_label is None:
        pos_label = 1
    labels = y_true == pos_label
    # calculate fraction of positve labels
    n_perc = ceil(float(percentage)/100.*len(labels))
    return labels[:n_perc].sum()/n_perc*100

def roc_log_auc(y_true, y_score, pos_label=None, log_min=0.001, log_max=1.):
    """Computes area under semi-log ROC for random distribution.

    Parameters
    ----------
        y_true : array, shape=[n_samples]
            True binary labels, in range {0,1} or {-1,1}. If positive label is different than 1, it must be explicitly defined.

        y_score : array, shape=[n_samples]
            Scores for tested series of samples

        pos_label: int
            Positive label of samples (if other than 1)

        log_min : float (default=0.001)
            Minimum logarithm value for estimating AUC

        log_max : float (default=1.)
            Maximum logarithm value for estimating AUC.

    Returns
    -------
        auc : float
            semi-log ROC AUC
    """
    fpr, tpr, t = roc(y_true, y_score, pos_label=pos_label)
    idx = (fpr >= log_min) & (fpr <= log_max)
    log_fpr = 1-np.log10(fpr[idx])/np.log10(log_min)
    return auc(log_fpr, tpr[idx])

def random_roc_log_auc(log_min=0.001, log_max=1.):
    """Computes area under semi-log ROC for random distribution.

    Parameters
    ----------
        log_min : float (default=0.001)
            Minimum logarithm value for estimating AUC

        log_max : float (default=1.)
            Maximum logarithm value for estimating AUC.

    Returns
    -------
        auc : float
            semi-log ROC AUC for random distribution
    """
    return (log_max-log_min)/(np.log(10)*np.log10(log_max/log_min))
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
'''
    tests.unit.utils.extend_test
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    Test the salt extend script, leave templates/test alone to keep this working!
'''

# Import python libs
from __future__ import absolute_import

import os
import shutil
from datetime import date

# Import salt libs
import salt.utils.extend
from salt.utils import fopen

# Import Salt Testing libs
from salttesting import TestCase
from salttesting.helpers import ensure_in_syspath
ensure_in_syspath('../../')

class ExtendTestCase(TestCase):
    def setUp(self):
        self.out = None

    def tearDown(self):
        if self.out is not None:
            if os.path.exists(self.out):
                shutil.rmtree(self.out, True)

    def test_run(self):
        out = salt.utils.extend.run('test', 'test', 'this description', '.', False)
        self.out = out
        year = date.today().strftime('%Y')
        self.assertTrue(os.path.exists(out))
        self.assertFalse(os.path.exists(os.path.join(out, 'template.yml')))
        self.assertTrue(os.path.exists(os.path.join(out, 'directory')))
        self.assertTrue(os.path.exists(os.path.join(out, 'directory', 'test.py')))
        with fopen(os.path.join(out, 'directory', 'test.py'), 'r') as test_f:
            self.assertEqual(test_f.read(), year)

if __name__ == '__main__':
    from unit import run_tests
    run_tests(ExtendTestCase, needs_daemon=False)<commit_msg>add newline to test<commit_after># -*- coding: utf-8 -*-
'''
    tests.unit.utils.extend_test
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~

    Test the salt extend script, leave templates/test alone to keep this working!
'''

# Import python libs
from __future__ import absolute_import

import os
import shutil
from datetime import date

# Import salt libs
import salt.utils.extend
from salt.utils import fopen

# Import Salt Testing libs
from salttesting import TestCase
from salttesting.helpers import ensure_in_syspath
ensure_in_syspath('../../')

class ExtendTestCase(TestCase):
    def setUp(self):
        self.out = None

    def tearDown(self):
        if self.out is not None:
            if os.path.exists(self.out):
                shutil.rmtree(self.out, True)

    def test_run(self):
        out = salt.utils.extend.run('test', 'test', 'this description', '.', False)
        self.out = out
        year = date.today().strftime('%Y')
        self.assertTrue(os.path.exists(out))
        self.assertFalse(os.path.exists(os.path.join(out, 'template.yml')))
        self.assertTrue(os.path.exists(os.path.join(out, 'directory')))
        self.assertTrue(os.path.exists(os.path.join(out, 'directory', 'test.py')))
        with fopen(os.path.join(out, 'directory', 'test.py'), 'r') as test_f:
            self.assertEqual(test_f.read(), year)

if __name__ == '__main__':
    from unit import run_tests
    run_tests(ExtendTestCase, needs_daemon=False)
<|endoftext|>


<commit_before>from merc import emitter

class Message(object):
  def as_params(self, client):
    raise NotImplementedError

  def emit(self, client, prefix):
    return emitter.emit_message(prefix, self.NAME, self.as_params(client))
<commit_msg>implement message truncation<commit_after>from merc import emitter

class Message(object):
  MAX_LENGTH = 510

  def as_params(self, client):
    raise NotImplementedError

  def emit(self, client, prefix):
    raw = emitter.emit_message(prefix, self.NAME, self.as_params(client))
    return raw.encode("utf-8")[:self.MAX_LENGTH].decode("utf-8", "ignore")
<|endoftext|>


<commit_before>
import numpy as np

from ._mie import dmiess_module as dmiess
from ._mie import bhmie_module as bhmie
from ._mie import dmilay_module as dmilay

#: Density of pure water, in kg/m^3
RHO_WATER = 1e3

__all__ = ["bhmie_scatter", "core_shell_scatter", "integrate_mode", ]

def bhmie_scatter(particle_radius, radiation_lambda, n_particle):
    """ Compute the scattering/absorption efficiency and asymmetry
    parameter for a homogenous particle.

    This function interfaces with the compiled Mie theory modules in order
    to determine the scattering parameters.

    Parameters
    ----------
    particle_radius : float
        The total particle radius (core + shell) in microns
    radiation_lambda : float
        Wavelength of incident radiation in microns
    n_particle : complex
        Complex refractive indices of the particle material

    Returns
    -------
    Qsca, Qabs, asym : floats
        The scattering efficiency, absorption efficiency, and asymmetry
        parameter for the specified particle

    """

    # Pass directly to Mie module
    Qext0, Qsca0, asym0 = bhmie.bhmie_driver(
        particle_radius,
        np.real(n_particle), np.imag(n_particle),
        radiation_lambda)

    # Post-process to properly set scattering and absorption efficiencies
    Qsca = np.min([Qsca0, Qext0])  # scattering efficiency
    Qabs = Qext0 - Qsca            # absorption efficiency
    asym = asym0

    return Qsca, Qabs, asym

def core_shell_scatter(particle_radius, core_fraction, radiation_lambda,
                       n_shell, n_core):
    """ Compute the scattering/absorption efficiency and asymmetry
    parameter for a heterogeneous, core-shell mixed particle.

    This function interfaces with the compiled Mie theory modules in order
    to determine the scattering parameters.

    Parameters
    ----------
    particle_radius : float
        The total particle radius (core + shell) in microns
    core_fraction : float
        The fraction of the particle comprised by its core, 0.0-1.0
    radiation_lambda : float
        Wavelength of incident radiation in microns
    n_shell, n_core : complex
        Complex refractive indices of the shell, and core respectively

    Returns
    -------
    Qsca, Qabs, asym : floats
        The scattering efficiency, absorption efficiency, and asymmetry
        parameter for the specified particle

    """

    # Pass directly to Mie module
    Qext0, Qsca0, asym0 = dmiess.dmiess_driver(
        particle_radius,
        core_fraction*particle_radius,
        n_shell,
        n_core,
        radiation_lambda)

    # Post-process to properly set scattering and absorption efficiencies
    Qsca = np.min([Qsca0, Qext0])  # scattering efficiency
    Qabs = Qext0 - Qsca            # absorption efficiency
    asym = asym0

    return Qsca, Qabs, asym

def integrate_mode(core_fraction, n_shell, n_core, radiation_lambda,
                   mode_radius, mode_sigma,
                   r_min=1e-3, r_max=100., nr=200):
    """ Integrate Mie theory calculations over a lognormal aerosol mode with
    homogeneous particle properties, weighting by size distribution.

    Parameters
    ----------
    core_fraction : float
        The fraction of the particle comprised by its core, 0.0-1.0
    n_shell, n_core : complex
        Complex refractive indices of the shell, and core respectively
    radiation_lambda : float
        Wavelength of incident radiation in microns
    mode_radius : float
        The geometric mean or mode radius of the aerosol size distribution, in
        microns
    mode_sigma : float
        The geometric standard deviation of the aerosol size distribution
    r_min, r_max : float (optional)
        The minimum and maximum particle radii to use in the integration, in
        microns
    nr : int (optional)
        The number of particle radii to use in the integration

    Returns
    -------
    mie_sca, mie_abs, mie_asym : floats
        Scattering efficiency, absorption efficiency, and asymmetry parameter
        integrated over an aerosol size distribution

    """

    # Generate the integration grid for the particle size distribution
    logr, dlogr = np.linspace(np.log(r_min), np.log(r_max), nr, retstep=True)
    radii = np.exp(logr)

    sumsca  = 0.0
    sumabs  = 0.0
    sumg    = 0.0
    volwet  = 0.0
    volcore = 0.0

    # Integration loop
    for i, radius in enumerate(radii):

        # Mie theory calculation
        Qsca, Qabs, asym = core_shell_scatter(
            radius, core_fraction, radiation_lambda, n_shell, n_core
        )

        # Compute weights and volumes for integral sum
        exparg   = np.log(radius / mode_radius)/np.log(mode_sigma)
        dsdlogr  = np.exp(-0.5*exparg**2)  # m^2/m3(air)] log-normal cross section area
        volwet  += (4./3.)*(radius*1e-6)*dsdlogr*dlogr  # [m3/m3(air)] wet volume
        volcore += (4./3.)*(core_fraction**3)*(radius*1e-6)*dsdlogr*dlogr

        sumabs += Qabs*dsdlogr*dlogr  # [m^2/m3(air)] absorption cross-section
        sumsca += Qsca*dsdlogr*dlogr  # [m^2/m3(air)] scattering cross-section
        sumg   += asym*Qsca*dsdlogr*dlogr

    mie_sca  = np.max([sumsca, 0.]) / (volwet * RHO_WATER)  # [m^2/kg] specific scattering cross-section
    mie_abs  = np.max([sumabs, 0.]) / (volwet * RHO_WATER)  # [m^2/kg] specific absorption cross-section
    mie_asym =                 sumg / sumsca                # average asymmetry

    return mie_sca, mie_abs, mie_asym
<commit_msg>bhmie_driver accepts a complex refractive index<commit_after>
import numpy as np

from ._mie import dmiess_module as dmiess
from ._mie import bhmie_module as bhmie
from ._mie import dmilay_module as dmilay

#: Density of pure water, in kg/m^3
RHO_WATER = 1e3

__all__ = ["bhmie_scatter", "core_shell_scatter", "integrate_mode", ]

def bhmie_scatter(particle_radius, radiation_lambda, n_particle):
    """ Compute the scattering/absorption efficiency and asymmetry
    parameter for a homogenous particle.

    This function interfaces with the compiled Mie theory modules in order
    to determine the scattering parameters.

    Parameters
    ----------
    particle_radius : float
        The total particle radius (core + shell) in microns
    radiation_lambda : float
        Wavelength of incident radiation in microns
    n_particle : complex
        Complex refractive indices of the particle material

    Returns
    -------
    Qsca, Qabs, asym : floats
        The scattering efficiency, absorption efficiency, and asymmetry
        parameter for the specified particle

    """

    # Pass directly to Mie module
    Qext0, Qsca0, asym0 = bhmie.bhmie_driver(particle_radius, n_particle,
                                             radiation_lambda)

    # Post-process to properly set scattering and absorption efficiencies
    Qsca = np.min([Qsca0, Qext0])  # scattering efficiency
    Qabs = Qext0 - Qsca            # absorption efficiency
    asym = asym0

    return Qsca, Qabs, asym

def core_shell_scatter(particle_radius, core_fraction, radiation_lambda,
                       n_shell, n_core):
    """ Compute the scattering/absorption efficiency and asymmetry
    parameter for a heterogeneous, core-shell mixed particle.

    This function interfaces with the compiled Mie theory modules in order
    to determine the scattering parameters.

    Parameters
    ----------
    particle_radius : float
        The total particle radius (core + shell) in microns
    core_fraction : float
        The fraction of the particle comprised by its core, 0.0-1.0
    radiation_lambda : float
        Wavelength of incident radiation in microns
    n_shell, n_core : complex
        Complex refractive indices of the shell, and core respectively

    Returns
    -------
    Qsca, Qabs, asym : floats
        The scattering efficiency, absorption efficiency, and asymmetry
        parameter for the specified particle

    """

    # Pass directly to Mie module
    Qext0, Qsca0, asym0 = dmiess.dmiess_driver(
        particle_radius,
        core_fraction*particle_radius,
        n_shell,
        n_core,
        radiation_lambda)

    # Post-process to properly set scattering and absorption efficiencies
    Qsca = np.min([Qsca0, Qext0])  # scattering efficiency
    Qabs = Qext0 - Qsca            # absorption efficiency
    asym = asym0

    return Qsca, Qabs, asym

def integrate_mode(core_fraction, n_shell, n_core, radiation_lambda,
                   mode_radius, mode_sigma,
                   r_min=1e-3, r_max=100., nr=200):
    """ Integrate Mie theory calculations over a lognormal aerosol mode with
    homogeneous particle properties, weighting by size distribution.

    Parameters
    ----------
    core_fraction : float
        The fraction of the particle comprised by its core, 0.0-1.0
    n_shell, n_core : complex
        Complex refractive indices of the shell, and core respectively
    radiation_lambda : float
        Wavelength of incident radiation in microns
    mode_radius : float
        The geometric mean or mode radius of the aerosol size distribution, in
        microns
    mode_sigma : float
        The geometric standard deviation of the aerosol size distribution
    r_min, r_max : float (optional)
        The minimum and maximum particle radii to use in the integration, in
        microns
    nr : int (optional)
        The number of particle radii to use in the integration

    Returns
    -------
    mie_sca, mie_abs, mie_asym : floats
        Scattering efficiency, absorption efficiency, and asymmetry parameter
        integrated over an aerosol size distribution

    """

    # Generate the integration grid for the particle size distribution
    logr, dlogr = np.linspace(np.log(r_min), np.log(r_max), nr, retstep=True)
    radii = np.exp(logr)

    sumsca  = 0.0
    sumabs  = 0.0
    sumg    = 0.0
    volwet  = 0.0
    volcore = 0.0

    # Integration loop
    for i, radius in enumerate(radii):

        # Mie theory calculation
        Qsca, Qabs, asym = core_shell_scatter(
            radius, core_fraction, radiation_lambda, n_shell, n_core
        )

        # Compute weights and volumes for integral sum
        exparg   = np.log(radius / mode_radius)/np.log(mode_sigma)
        dsdlogr  = np.exp(-0.5*exparg**2)  # m^2/m3(air)] log-normal cross section area
        volwet  += (4./3.)*(radius*1e-6)*dsdlogr*dlogr  # [m3/m3(air)] wet volume
        volcore += (4./3.)*(core_fraction**3)*(radius*1e-6)*dsdlogr*dlogr

        sumabs += Qabs*dsdlogr*dlogr  # [m^2/m3(air)] absorption cross-section
        sumsca += Qsca*dsdlogr*dlogr  # [m^2/m3(air)] scattering cross-section
        sumg   += asym*Qsca*dsdlogr*dlogr

    mie_sca  = np.max([sumsca, 0.]) / (volwet * RHO_WATER)  # [m^2/kg] specific scattering cross-section
    mie_abs  = np.max([sumabs, 0.]) / (volwet * RHO_WATER)  # [m^2/kg] specific absorption cross-section
    mie_asym =                 sumg / sumsca                # average asymmetry

    return mie_sca, mie_abs, mie_asym
<|endoftext|>


<commit_before>'''
Test the verification routines
'''

# Import Python libs
import getpass
import os
import sys
import stat
import shutil
import resource
import tempfile
import socket

# Import salt libs
try:
    import integration
except ImportError:
    if __name__ == '__main__':
        sys.path.insert(
            0, os.path.abspath(
                os.path.join(
                    os.path.dirname(__file__), '../../'
                )
            )
        )
    import integration
import salt.utils
from salt.utils.verify import (
    check_user,
    verify_env,
    verify_socket,
    zmq_version,
    check_max_open_files
)

# Import Salt Testing libs
from salttesting import skipIf, TestCase
from salttesting.helpers import TestsLoggingHandler

class TestVerify(TestCase):
    '''
    Verify module tests
    '''

    def test_zmq_verify(self):
        self.assertTrue(zmq_version())

    def test_zmq_verify_insufficient(self):
        import zmq
        zmq.__version__ = '2.1.0'
        self.assertFalse(zmq_version())

    def test_user(self):
        self.assertTrue(check_user(getpass.getuser()))

    def test_no_user(self):
        # Catch sys.stderr here since no logging is configured and
        # check_user WILL write to sys.stderr
        class FakeWriter(object):
            def __init__(self):
                self.output = ""

            def write(self, data):
                self.output += data
        stderr = sys.stderr
        writer = FakeWriter()
        sys.stderr = writer
        # Now run the test
        self.assertFalse(check_user('nouser'))
        # Restore sys.stderr
        sys.stderr = stderr
        if writer.output != 'CRITICAL: User not found: "nouser"\n':
            # If there's a different error catch, write it to sys.stderr
            sys.stderr.write(writer.output)

    @skipIf(sys.platform.startswith('win'), 'No verify_env Windows')
    def test_verify_env(self):
        root_dir = tempfile.mkdtemp()
        var_dir = os.path.join(root_dir, 'var', 'log', 'salt')
        verify_env([var_dir], getpass.getuser())
        self.assertTrue(os.path.exists(var_dir))
        dir_stat = os.stat(var_dir)
        self.assertEqual(dir_stat.st_uid, os.getuid())
        self.assertEqual(dir_stat.st_gid, os.getgid())
        self.assertEqual(dir_stat.st_mode & stat.S_IRWXU, stat.S_IRWXU)
        self.assertEqual(dir_stat.st_mode & stat.S_IRWXG, 40)
        self.assertEqual(dir_stat.st_mode & stat.S_IRWXO, 5)

    def test_verify_socket(self):
        self.assertTrue(verify_socket('', 18000, 18001))
        if socket.has_ipv6:
            # Only run if Python is built with IPv6 support; otherwise
            # this will just fail.
            self.assertTrue(verify_socket('::', 18000, 18001))

    @skipIf(os.environ.get('TRAVIS_PYTHON_VERSION', None) is not None,
            'Travis environment does not like too many open files')
    def test_max_open_files(self):

        with TestsLoggingHandler() as handler:
            logmsg_dbg = (
                'DEBUG:This salt-master instance has accepted {0} minion keys.'
            )
            logmsg_chk = (
                '{0}:The number of accepted minion keys({1}) should be lower '
                'than 1/4 of the max open files soft setting({2}). According '
                'to the system\'s hard limit, there\'s still a margin of {3} '
                'to raise the salt\'s max_open_files setting. Please consider '
                'raising this value.'
            )
            logmsg_crash = (
                '{0}:The number of accepted minion keys({1}) should be lower '
                'than 1/4 of the max open files soft setting({2}). '
                'salt-master will crash pretty soon! According to the '
                'system\'s hard limit, there\'s still a margin of {3} to '
                'raise the salt\'s max_open_files setting. Please consider '
                'raising this value.'
            )

            mof_s, mof_h = resource.getrlimit(resource.RLIMIT_NOFILE)
            tempdir = tempfile.mkdtemp(prefix='fake-keys')
            keys_dir = os.path.join(tempdir, 'minions')
            os.makedirs(keys_dir)

            mof_test = 256

            resource.setrlimit(resource.RLIMIT_NOFILE, (mof_test, mof_h))

            try:
                prev = 0
                for newmax, level in ((24, None), (66, 'INFO'),
                                      (127, 'WARNING'), (196, 'CRITICAL')):

                    for n in range(prev, newmax):
                        kpath = os.path.join(keys_dir, str(n))
                        with salt.utils.fopen(kpath, 'w') as fp_:
                            fp_.write(str(n))

                    opts = {
                        'max_open_files': newmax,
                        'pki_dir': tempdir
                    }

                    check_max_open_files(opts)

                    if level is None:
                        # No log message is triggered, only the DEBUG one which
                        # tells us how many minion keys were accepted.
                        self.assertEqual(
                            [logmsg_dbg.format(newmax)], handler.messages
                        )
                    else:
                        self.assertIn(
                            logmsg_dbg.format(newmax), handler.messages
                        )
                        self.assertIn(
                            logmsg_chk.format(
                                level,
                                newmax,
                                mof_test,
                                mof_h - newmax,
                            ),
                            handler.messages
                        )
                    handler.clear()
                    prev = newmax

                newmax = mof_test
                for n in range(prev, newmax):
                    kpath = os.path.join(keys_dir, str(n))
                    with salt.utils.fopen(kpath, 'w') as fp_:
                        fp_.write(str(n))

                opts = {
                    'max_open_files': newmax,
                    'pki_dir': tempdir
                }

                check_max_open_files(opts)
                self.assertIn(logmsg_dbg.format(newmax), handler.messages)
                self.assertIn(
                    logmsg_crash.format(
                        'CRITICAL',
                        newmax,
                        mof_test,
                        mof_h - newmax,
                    ),
                    handler.messages
                )
                handler.clear()
            except IOError as err:
                if err.errno == 24:
                    # Too many open files
                    self.skipTest('We\'ve hit the max open files setting')
                raise
            finally:
                shutil.rmtree(tempdir)
                resource.setrlimit(resource.RLIMIT_NOFILE, (mof_s, mof_h))

if __name__ == '__main__':
    from integration import run_tests
    run_tests(TestVerify, needs_daemon=False)
<commit_msg>Updated `unit.utils.verify_test` test suite to use `salttesting.helpers.ensure_in_syspath()`<commit_after>'''
Test the verification routines
'''

# Import Python libs
import getpass
import os
import sys
import stat
import shutil
import resource
import tempfile
import socket

# Import Salt Testing libs
from salttesting import skipIf, TestCase
from salttesting.helpers import ensure_in_syspath, TestsLoggingHandler
ensure_in_syspath('../../')

# Import salt libs
import salt.utils
from salt.utils.verify import (
    check_user,
    verify_env,
    verify_socket,
    zmq_version,
    check_max_open_files
)

class TestVerify(TestCase):
    '''
    Verify module tests
    '''

    def test_zmq_verify(self):
        self.assertTrue(zmq_version())

    def test_zmq_verify_insufficient(self):
        import zmq
        zmq.__version__ = '2.1.0'
        self.assertFalse(zmq_version())

    def test_user(self):
        self.assertTrue(check_user(getpass.getuser()))

    def test_no_user(self):
        # Catch sys.stderr here since no logging is configured and
        # check_user WILL write to sys.stderr
        class FakeWriter(object):
            def __init__(self):
                self.output = ""

            def write(self, data):
                self.output += data
        stderr = sys.stderr
        writer = FakeWriter()
        sys.stderr = writer
        # Now run the test
        self.assertFalse(check_user('nouser'))
        # Restore sys.stderr
        sys.stderr = stderr
        if writer.output != 'CRITICAL: User not found: "nouser"\n':
            # If there's a different error catch, write it to sys.stderr
            sys.stderr.write(writer.output)

    @skipIf(sys.platform.startswith('win'), 'No verify_env Windows')
    def test_verify_env(self):
        root_dir = tempfile.mkdtemp()
        var_dir = os.path.join(root_dir, 'var', 'log', 'salt')
        verify_env([var_dir], getpass.getuser())
        self.assertTrue(os.path.exists(var_dir))
        dir_stat = os.stat(var_dir)
        self.assertEqual(dir_stat.st_uid, os.getuid())
        self.assertEqual(dir_stat.st_gid, os.getgid())
        self.assertEqual(dir_stat.st_mode & stat.S_IRWXU, stat.S_IRWXU)
        self.assertEqual(dir_stat.st_mode & stat.S_IRWXG, 40)
        self.assertEqual(dir_stat.st_mode & stat.S_IRWXO, 5)

    def test_verify_socket(self):
        self.assertTrue(verify_socket('', 18000, 18001))
        if socket.has_ipv6:
            # Only run if Python is built with IPv6 support; otherwise
            # this will just fail.
            self.assertTrue(verify_socket('::', 18000, 18001))

    @skipIf(os.environ.get('TRAVIS_PYTHON_VERSION', None) is not None,
            'Travis environment does not like too many open files')
    def test_max_open_files(self):

        with TestsLoggingHandler() as handler:
            logmsg_dbg = (
                'DEBUG:This salt-master instance has accepted {0} minion keys.'
            )
            logmsg_chk = (
                '{0}:The number of accepted minion keys({1}) should be lower '
                'than 1/4 of the max open files soft setting({2}). According '
                'to the system\'s hard limit, there\'s still a margin of {3} '
                'to raise the salt\'s max_open_files setting. Please consider '
                'raising this value.'
            )
            logmsg_crash = (
                '{0}:The number of accepted minion keys({1}) should be lower '
                'than 1/4 of the max open files soft setting({2}). '
                'salt-master will crash pretty soon! According to the '
                'system\'s hard limit, there\'s still a margin of {3} to '
                'raise the salt\'s max_open_files setting. Please consider '
                'raising this value.'
            )

            mof_s, mof_h = resource.getrlimit(resource.RLIMIT_NOFILE)
            tempdir = tempfile.mkdtemp(prefix='fake-keys')
            keys_dir = os.path.join(tempdir, 'minions')
            os.makedirs(keys_dir)

            mof_test = 256

            resource.setrlimit(resource.RLIMIT_NOFILE, (mof_test, mof_h))

            try:
                prev = 0
                for newmax, level in ((24, None), (66, 'INFO'),
                                      (127, 'WARNING'), (196, 'CRITICAL')):

                    for n in range(prev, newmax):
                        kpath = os.path.join(keys_dir, str(n))
                        with salt.utils.fopen(kpath, 'w') as fp_:
                            fp_.write(str(n))

                    opts = {
                        'max_open_files': newmax,
                        'pki_dir': tempdir
                    }

                    check_max_open_files(opts)

                    if level is None:
                        # No log message is triggered, only the DEBUG one which
                        # tells us how many minion keys were accepted.
                        self.assertEqual(
                            [logmsg_dbg.format(newmax)], handler.messages
                        )
                    else:
                        self.assertIn(
                            logmsg_dbg.format(newmax), handler.messages
                        )
                        self.assertIn(
                            logmsg_chk.format(
                                level,
                                newmax,
                                mof_test,
                                mof_h - newmax,
                            ),
                            handler.messages
                        )
                    handler.clear()
                    prev = newmax

                newmax = mof_test
                for n in range(prev, newmax):
                    kpath = os.path.join(keys_dir, str(n))
                    with salt.utils.fopen(kpath, 'w') as fp_:
                        fp_.write(str(n))

                opts = {
                    'max_open_files': newmax,
                    'pki_dir': tempdir
                }

                check_max_open_files(opts)
                self.assertIn(logmsg_dbg.format(newmax), handler.messages)
                self.assertIn(
                    logmsg_crash.format(
                        'CRITICAL',
                        newmax,
                        mof_test,
                        mof_h - newmax,
                    ),
                    handler.messages
                )
                handler.clear()
            except IOError as err:
                if err.errno == 24:
                    # Too many open files
                    self.skipTest('We\'ve hit the max open files setting')
                raise
            finally:
                shutil.rmtree(tempdir)
                resource.setrlimit(resource.RLIMIT_NOFILE, (mof_s, mof_h))

if __name__ == '__main__':
    from integration import run_tests
    run_tests(TestVerify, needs_daemon=False)
<|endoftext|>


<commit_before>"""
I/O for AVS-UCD format, cf.
<https://lanl.github.io/LaGriT/pages/docs/read_avs.html>.
"""
import numpy

from .._exceptions import ReadError, WriteError
from .._files import open_file
from .._helpers import register
from .._mesh import Mesh

meshio_to_avsucd_type = {
    "vertex": "pt",
    "line": "line",
    "triangle": "tri",
    "quad": "quad",
    "tetra": "tet",
    "pyramid": "pyr",
    "wedge": "prism",
    "hexahedron": "hex",
}
avsucd_to_meshio_type = {v: k for k, v in meshio_to_avsucd_type.items()}

avsucd_to_meshio_order = {
    "vertex": [0],
    "line": [0, 1],
    "triangle": [0, 1, 2],
    "quad": [0, 1, 2, 3],
    "tetra": [0, 1, 3, 2],
    "pyramid": [4, 0, 1, 2, 3],
    "wedge": [3, 4, 5, 0, 1, 2],
    "hexahedron": [4, 5, 6, 7, 0, 1, 2, 3],
}

def read(filename):
    with open_file(filename, "r") as f:
        out = read_buffer(f)
    return out

def read_buffer(f):
    # Skip comments and unpack first line
    num_nodes, num_cells, num_node_data, num_cell_data, _ = numpy.genfromtxt(
        f, max_rows=1, dtype=int, comments="#"
    )

    # Read nodes
    point_ids, points = _read_nodes(f, num_nodes)

    # Read cells
    cell_ids, cells, cell_data = _read_cells(f, num_cells, point_ids)

    # Read node data
    if num_node_data:
        point_data = _read_node_data(f, num_nodes, num_node_data, point_ids)
    else:
        point_data = {}

    # Read cell data
    if num_cell_data:
        cell_data.update(_read_cell_data(f, num_cells, num_cell_data, cells, cell_ids))

    return Mesh(
        points,
        cells,
        point_data=point_data,
        cell_data={
            k: {kk: numpy.array(vv) for kk, vv in v.items()}
            for k, v in cell_data.items()
        },
    )

def _read_nodes(f, num_nodes):
    data = numpy.genfromtxt(f, max_rows=num_nodes)
    points_ids = {int(pid): i for i, pid in enumerate(data[:, 0])}
    return points_ids, data[:, 1:]

def _read_cells(f, num_cells, point_ids):
    cells = {}
    cell_ids = {}
    cell_data = {"avsucd:mat": {}}
    count = {k: 0 for k in meshio_to_avsucd_type.keys()}
    for _ in range(num_cells):
        line = f.readline().strip().split()
        cell_id, cell_mat = int(line[0]), int(line[1])
        cell_type = avsucd_to_meshio_type[line[2]]
        corner = [point_ids[int(pid)] for pid in line[3:]]

        if cell_type not in cells:
            cells[cell_type] = [corner]
            cell_data["avsucd:mat"][cell_type] = [cell_mat]
        else:
            cells[cell_type].append(corner)
            cell_data["avsucd:mat"][cell_type].append(cell_mat)

        cell_ids[cell_id] = (cell_type, count[cell_type])
        count[cell_type] += 1

    for k, v in cells.items():
        cells[k] = numpy.array(v)[:, avsucd_to_meshio_order[k]]
    return cell_ids, cells, cell_data

def _read_node_data(f, num_nodes, num_node_data, point_ids):
    line = f.readline()  # Not quite sure what to do with this line...

    labels = {}
    point_data = {}
    for i in range(num_node_data):
        line = f.readline().strip().split(",")
        labels[i] = line[0].strip()
        point_data[labels[i]] = numpy.empty(num_nodes)

    for _ in range(num_nodes):
        line = f.readline().strip().split()
        pid = point_ids[int(line[0])]
        for i, val in enumerate(line[1:]):
            point_data[labels[i]][pid] = float(val)
    return point_data

def _read_cell_data(f, num_cells, num_cell_data, cells, cell_ids):
    line = f.readline()  # Not quite sure what to do with this line...

    labels = {}
    cell_data = {}
    for i in range(num_cell_data):
        line = f.readline().strip().split(",")
        labels[i] = line[0].strip()
        cell_data[labels[i]] = {k: numpy.empty(len(v)) for k, v in cells.items()}

    for _ in range(num_cells):
        line = f.readline().strip().split()
        cell_type, cid = cell_ids[int(line[0])]
        for i, val in enumerate(line[1:]):
            cell_data[labels[i]][cell_type][cid] = float(val)
    return cell_data

def write(filename, mesh):
    pass

register("avsucd", [".inp"], read, {"avsucd": write})
<commit_msg>handle multicomponent data<commit_after>"""
I/O for AVS-UCD format, cf.
<https://lanl.github.io/LaGriT/pages/docs/read_avs.html>.
"""
import numpy

from .._exceptions import ReadError, WriteError
from .._files import open_file
from .._helpers import register
from .._mesh import Mesh

meshio_to_avsucd_type = {
    "vertex": "pt",
    "line": "line",
    "triangle": "tri",
    "quad": "quad",
    "tetra": "tet",
    "pyramid": "pyr",
    "wedge": "prism",
    "hexahedron": "hex",
}
avsucd_to_meshio_type = {v: k for k, v in meshio_to_avsucd_type.items()}

avsucd_to_meshio_order = {
    "vertex": [0],
    "line": [0, 1],
    "triangle": [0, 1, 2],
    "quad": [0, 1, 2, 3],
    "tetra": [0, 1, 3, 2],
    "pyramid": [4, 0, 1, 2, 3],
    "wedge": [3, 4, 5, 0, 1, 2],
    "hexahedron": [4, 5, 6, 7, 0, 1, 2, 3],
}

def read(filename):
    with open_file(filename, "r") as f:
        out = read_buffer(f)
    return out

def read_buffer(f):
    # Skip comments and unpack first line
    num_nodes, num_cells, num_node_data, num_cell_data, _ = numpy.genfromtxt(
        f, max_rows=1, dtype=int, comments="#"
    )

    # Read nodes
    point_ids, points = _read_nodes(f, num_nodes)

    # Read cells
    cell_ids, cells, cell_data = _read_cells(f, num_cells, point_ids)

    # Read node data
    if num_node_data:
        point_data = _read_node_data(f, num_nodes, point_ids)
    else:
        point_data = {}

    # Read cell data
    if num_cell_data:
        cell_data.update(_read_cell_data(f, num_cells, cells, cell_ids))

    return Mesh(
        points,
        cells,
        point_data=point_data,
        cell_data={
            k: {kk: numpy.array(vv) for kk, vv in v.items()}
            for k, v in cell_data.items()
        },
    )

def _read_nodes(f, num_nodes):
    data = numpy.genfromtxt(f, max_rows=num_nodes)
    points_ids = {int(pid): i for i, pid in enumerate(data[:, 0])}
    return points_ids, data[:, 1:]

def _read_cells(f, num_cells, point_ids):
    cells = {}
    cell_ids = {}
    cell_data = {"avsucd:mat": {}}
    count = {k: 0 for k in meshio_to_avsucd_type.keys()}
    for _ in range(num_cells):
        line = f.readline().strip().split()
        cell_id, cell_mat = int(line[0]), int(line[1])
        cell_type = avsucd_to_meshio_type[line[2]]
        corner = [point_ids[int(pid)] for pid in line[3:]]

        if cell_type not in cells:
            cells[cell_type] = [corner]
            cell_data["avsucd:mat"][cell_type] = [cell_mat]
        else:
            cells[cell_type].append(corner)
            cell_data["avsucd:mat"][cell_type].append(cell_mat)

        cell_ids[cell_id] = (cell_type, count[cell_type])
        count[cell_type] += 1

    for k, v in cells.items():
        cells[k] = numpy.array(v)[:, avsucd_to_meshio_order[k]]
    return cell_ids, cells, cell_data

def _read_node_data(f, num_nodes, point_ids):
    line = f.readline().strip().split()
    node_data_size = [int(i) for i in line[1:]]

    labels = {}
    point_data = {}
    for i, dsize in enumerate(node_data_size):
        line = f.readline().strip().split(",")
        labels[i] = line[0].strip()
        point_data[labels[i]] = numpy.empty(num_nodes) if dsize == 1 else numpy.empty((num_nodes, dsize))

    for _ in range(num_nodes):
        line = f.readline().strip().split()
        pid = point_ids[int(line[0])]
        j = 0
        for i, dsize in enumerate(node_data_size):
            if dsize == 1:
                point_data[labels[i]][pid] = float(line[j+1])
            else:
                point_data[labels[i]][pid] = [float(val) for val in line[j+1:j+1+dsize]]
            j += dsize
    return point_data

def _read_cell_data(f, num_cells, cells, cell_ids):
    line = f.readline().strip().split()
    cell_data_size = [int(i) for i in line[1:]]

    labels = {}
    cell_data = {}
    for i, dsize in enumerate(cell_data_size):
        line = f.readline().strip().split(",")
        labels[i] = line[0].strip()
        cell_data[labels[i]] = {k: numpy.empty(len(v)) if dsize == 1 else numpy.empty((len(v), dsize)) for k, v in cells.items()}

    for _ in range(num_cells):
        line = f.readline().strip().split()
        cell_type, cid = cell_ids[int(line[0])]
        j = 0
        for i, dsize in enumerate(cell_data_size):
            if dsize == 1:
                cell_data[labels[i]][cell_type][cid] = float(line[j+1])
            else:
                cell_data[labels[i]][cell_type][cid] = [float(val) for val in line[j+1:j+1+dsize]]
            j += dsize
    return cell_data

def write(filename, mesh):
    pass

register("avsucd", [".inp"], read, {"avsucd": write})
<|endoftext|>