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

from bs4 import Tag
from parser.general import GeneralParser
from .helper import get_heading_level, get_html_tree_from_url, remove_parenthesis, remove_comma_period

COMMA_SEMI_PERIOD = re.compile('[,;.]')
COMMA_OR_SEMICOLON = re.compile('[,;]')

class RuParser(GeneralParser):
    def __init__(self):
        super(RuParser, self).__init__()
        self.edition = 'ru'

    tested_url = [
        "https://ru.wiktionary.org/wiki/-ациj"#,
       # "https://ru.wiktionary.org/wiki/blood",
        #"https://ru.wiktionary.org/wiki/flower",
        #"https://ru.wiktionary.org/wiki/house",
        #"https://ru.wiktionary.org/wiki/speak",
        #"https://ru.wiktionary.org/wiki/%D0%B4%D0%BE%D0%BC",
        #"https://ru.wiktionary.org/wiki/говорить"
    ]

    # parse through a translation table in the Russian edition of Wiktionary
    def parse_translation_table_russian(self, table):

        for li in table.find_all('li'):
            if not isinstance(li, Tag):
                continue
            text = li.get_text().split(':')

            # language name is before ":"
            lang_name = text[0]

            lang_code = ''
            if li.find("sub"):
                lang_code = li.find("sub").get_text()

            # remove the lang code from the lang name
            lang_name = lang_name[:-len(lang_code)]

            if len(text) > 1:
                t = remove_parenthesis(text[1])
            else:
                t = remove_parenthesis(text[0])

            trans_list = re.split(COMMA_OR_SEMICOLON, t)

            for trans in trans_list:
                translation = trans.split('(')[0].strip()
                if not translation == '':
                    yield (translation, lang_name, lang_code)

    # parse through an ordered list in the Russian edition of wiktionary
    def parse_ordered_list_russian2(self, olist):
        for li in olist.find_all('li'):
            for href in li.find_all('a'):
                if href.has_attr('href') and href.has_attr('title'):
                    if not href['title'] == 'Викисловарь:Условные сокращения' and not href['title'] == 'ВС:ЦИТ':
                        translation = href['title']
                        yield translation

    # parse through an ordered list in the Russian edition of wiktionary
    def parse_ordered_list_russian(self, olist):
        for li in olist.find_all('li'):
            if not li.get_text() == '':
                # format the text (specific to RU)
                text = li.get_text().split(u'◆')[0]
                text = remove_parenthesis(text)
                text = re.split(COMMA_SEMI_PERIOD, text)
                for translations in text:
                    yield translations

    # grab and format the part of speech
    def get_pos_russian(self, target):
        text = target.get_text().split(',')[0]
        text = text.split(u'—')
        if len(text) > 1:
            text = text[1].strip()
        else:
            text = text[0].strip()
        return remove_comma_period(text)

    # parse through an unordered list in the Russian edition of wiktionary for part of speech
    def parse_unordered_list_russian(self, ulist):
        for li in ulist.find_all('li'):
            if not li.get_text() == '':
                # format the text (specific to RU)
                text = li.get_text().split(u'—')
                if len(text) > 1:
                    text = text[1].strip()
                else:
                    text = text[0].strip()
                return remove_comma_period(text)

    def generate_translation_tuples(self, soup):
        """
        A generator of translation tuples
        :param soup: BeautifulSoup object
        :return: tuple of the form (edition, headword, head_lang, translation, trans_lang, trans_lang_code, part_of_speech)
        """

        # Find the headword of the page
        title = soup.find('title')
        if title:
            title = title.text
        else:
            title = ''

        # Find the table of contents
        toc = soup.find('div', id='mw-content-text')

        # Variable to keep track of important information
        page_state = {'headword': title,
                      'headword_lang': '',
                      'part_of_speech': '',
                      'pronunciation': '',
                      'pos_region': False, # pos_region is specific to Russian
                      'translation_region': False,
                      'pro_region': False}

        # iterate through the childen of the table of contents to find translations
        for element in toc.children:
            if isinstance(element, Tag):
                level = get_heading_level(element.name)

                # in the Russian edition, h1s always contain the language
                if level == 1:
                    page_state['headword_lang'] = self.get_heading_text(element).strip()
                    page_state['translation_region'] = False
                    page_state['pos_region'] = True

                # Grab the part of speech, always contained in a level 3
                # sometimes the part of speech is preceded by headword
                elif level == 3 or level == 4:
                    text = element.get_text()
                    text = remove_parenthesis(text).strip()

                    if text == u'Морфологические и синтаксические свойства':
                        page_state['pos_region'] = True
                    elif text == u'Перевод':
                        page_state['translation_region'] = True
                        page_state['pos_region'] = False
                    elif text == u'Произношение':
                        page_state['translation_region'] = False
                        page_state['pro_region'] = True
                        page_state['pos_region'] = False
                    elif text == u'Значение':
                        if not page_state['headword_lang'] == u'Русский':
                            page_state['translation_region'] = True
                            page_state['pos_region'] = False
                        else:
                            page_state['translation_region'] = False
                            page_state['pos_region'] = False
                    else:
                        page_state['translation_region'] = False
                        page_state['pos_region'] = False

                # grab the part of speech
                elif element.name == 'p' and page_state['pos_region']:
                    bold_word = element.b
                    if not bold_word: # if the word is not bold, it's the pos
                        page_state['part_of_speech'] = self.get_pos_russian(element)
                        page_state['pos_region'] = False

                # parse through an unordered list to grab part of speech
                elif element.name == 'ul' and page_state['pos_region']:
                    page_state['part_of_speech'] = self.parse_unordered_list_russian(element)
                    page_state['pos_region'] = False

                # parse through a paragraph to grab translations
                elif element.name == 'p' and page_state['translation_region']:
                    translation = element.get_text()
                    yield (self.edition, page_state['headword'], page_state['headword_lang'], translation.strip(),
                           u'Русский', 'ru', page_state['part_of_speech'])

                # parse through an ordered list to grab translations
                elif element.name == 'ol' and page_state['translation_region']:
                    for translation in self.parse_ordered_list_russian(element):
                        yield (self.edition, page_state['headword'], page_state['headword_lang'], translation.strip(),
                               u'Русский', 'ru', page_state['part_of_speech'])

                # parse through a table to grab translations
                elif element.name == 'table' and page_state['translation_region']:
                    for translation, lang, lang_code in self.parse_translation_table_russian(element):
                        yield (self.edition, page_state['headword'], page_state['headword_lang'],
                               translation.strip(), lang, lang_code, page_state['part_of_speech'])
                else:
                    page_state['translation_region'] = False

        #yield('', '', '', '', '','', '')

def main():

    parser = RuParser()
    text_file = open("data/html_ru.txt", "r")
    url_list = text_file.read().split('\n')
    for url in url_list:
        soup = get_html_tree_from_url(url)
        for tup in parser.generate_translation_tuples(soup):
            if not ''.join(str(el) for el in tup) == '':
               print(','.join(str(el) for el in tup))
'''
def main():
    parser = RuParser()
    for url in parser.tested_url:
        soup = get_html_tree_from_url(url)
        for tup in parser.generate_translation_tuples(soup):
            print(','.join(tup))

'''
if __name__ == '__main__':
    main()
<commit_msg>Russian parser complete<commit_after># -*- coding: utf-8 -*-
import re

from bs4 import Tag
from parser.general import GeneralParser
from .helper import get_heading_level, get_html_tree_from_url, remove_parenthesis, remove_comma_period

COMMA_SEMI_PERIOD = re.compile('[,;.]')
COMMA_OR_SEMICOLON = re.compile('[,;]')

class RuParser(GeneralParser):
    def __init__(self):
        super(RuParser, self).__init__()
        self.edition = 'ru'

    tested_url = [
        "https://ru.wiktionary.org/wiki/-ациj"#,
       # "https://ru.wiktionary.org/wiki/blood",
        #"https://ru.wiktionary.org/wiki/flower",
        #"https://ru.wiktionary.org/wiki/house",
        #"https://ru.wiktionary.org/wiki/speak",
        #"https://ru.wiktionary.org/wiki/%D0%B4%D0%BE%D0%BC",
        #"https://ru.wiktionary.org/wiki/говорить"
    ]

    # parse through a translation table in the Russian edition of Wiktionary
    def parse_translation_table_russian(self, table):

        for li in table.find_all('li'):
            if not isinstance(li, Tag):
                continue
            text = li.get_text().split(':')

            # language name is before ":"
            lang_name = text[0]

            lang_code = ''
            if li.find("sub"):
                lang_code = li.find("sub").get_text()

            # remove the lang code from the lang name
            lang_name = lang_name[:-len(lang_code)]

            if len(text) > 1:
                t = remove_parenthesis(text[1])
            else:
                t = remove_parenthesis(text[0])

            trans_list = re.split(COMMA_OR_SEMICOLON, t)

            for trans in trans_list:
                translation = trans.split('(')[0].strip()
                if not translation == '':
                    yield (translation, lang_name, lang_code)

    def parse_translation_element_russian(self, tre):
        MAX_LEN = 1
        text = remove_parenthesis(tre.get_text()).split(u'◆')[0]
        text = re.split(COMMA_OR_SEMICOLON, text)
        for trans in text:
            trans = trans.split(' ')
            if not len(trans) > MAX_LEN:
                for tran in trans:
                    yield tran

    # grab and format the part of speech
    def get_pos_russian(self, target):
        text = target.get_text().split(',')[0]
        text = text.split(u'—')
        if len(text) > 1:
            text = text[1].strip()
        else:
            text = text[0].strip()
        return remove_comma_period(text)

    # parse through an unordered list in the Russian edition of wiktionary for part of speech
    def parse_unordered_list_russian(self, ulist):
        for li in ulist.find_all('li'):
            if not li.get_text() == '':
                # format the text (specific to RU)
                text = li.get_text().split(u'—')
                if len(text) > 1:
                    text = text[1].strip()
                else:
                    text = text[0].strip()
                return remove_comma_period(text)

    def parse_pronunciation_list_russian(self, ulist):

        text = ulist.get_text()
        text = text[text.find('['): text.find(']')+1]
        return text

    def generate_translation_tuples(self, soup):
        """
        A generator of translation tuples
        :param soup: BeautifulSoup object
        :return: tuple of the form (edition, headword, head_lang, translation, trans_lang, trans_lang_code, part_of_speech)
        """

        # Find the headword of the page
        title = soup.find('title')
        if title:
            title = title.text
        else:
            title = ''

        # Find the table of contents
        toc = soup.find('div', id='mw-content-text')

        # Variable to keep track of important information
        page_state = {'headword': title,
                      'headword_lang': '',
                      'part_of_speech': '',
                      'pronunciation': '',
                      'pos_region': False, # pos_region is specific to Russian
                      'translation_region': False,
                      'pro_region': False}

        # iterate through the childen of the table of contents to find translations
        for element in toc.children:
            if isinstance(element, Tag):
                level = get_heading_level(element.name)

                # in the Russian edition, h1s always contain the language
                if level == 1:
                    page_state['headword_lang'] = self.get_heading_text(element).strip()
                    page_state['translation_region'] = False
                    page_state['pos_region'] = True

                # Grab the part of speech, always contained in a level 3
                # sometimes the part of speech is preceded by headword
                elif level == 3 or level == 4:
                    text = element.get_text()
                    text = remove_parenthesis(text).strip()

                    if text == u'Морфологические и синтаксические свойства':
                        page_state['pos_region'] = True
                        page_state['pro_region'] = False
                    elif text == u'Перевод':
                        page_state['translation_region'] = True
                        page_state['pos_region'] = False
                        page_state['pro_region'] = False
                    elif text == u'Произношение':
                        page_state['pronunciation'] = '' # reset previous pronunciation
                        page_state['translation_region'] = False
                        page_state['pro_region'] = True
                        page_state['pos_region'] = False
                    elif text == u'Значение':
                        if not page_state['headword_lang'] == u'Русский':
                            page_state['translation_region'] = True
                            page_state['pos_region'] = False
                            page_state['pro_region'] = False
                        else:
                            page_state['translation_region'] = False
                            page_state['pos_region'] = False
                            page_state['pro_region'] = False
                    else:
                        page_state['translation_region'] = False
                        page_state['pos_region'] = False
                        page_state['pro_region'] = False

                # grab the part of speech
                elif element.name == 'p' and page_state['pos_region']:
                    bold_word = element.b
                    if not bold_word: # if the word is not bold, it's the pos
                        page_state['part_of_speech'] = self.get_pos_russian(element)
                        page_state['pos_region'] = False

                # parse through an unordered list to grab part of speech
                elif element.name == 'ul' and page_state['pos_region']:
                    page_state['part_of_speech'] = self.parse_unordered_list_russian(element)
                    page_state['pos_region'] = False

                # parse through a paragraph to grab translations
                elif element.name == 'p' and page_state['translation_region']:
                    for translation in self.parse_translation_element_russian(element):
                        if translation:
                            yield (self.edition, page_state['headword'], page_state['headword_lang'], translation.strip(),
                                   u'Русский', 'ru', page_state['part_of_speech'], page_state['pronunciation'])

                # parse through an ordered list to grab translations
                elif element.name == 'ol' and page_state['translation_region']:
                    for li in element.find_all('li'):
                        if not li.get_text() == '':
                            for translation in self.parse_translation_element_russian(li):
                                if translation:
                                    yield (self.edition, page_state['headword'], page_state['headword_lang'], translation.strip(),
                                           u'Русский', 'ru', page_state['part_of_speech'], page_state['pronunciation'])

                # parse through a table to grab translations
                elif element.name == 'table' and page_state['translation_region']:
                    for translation, lang, lang_code in self.parse_translation_table_russian(element):
                        yield (self.edition, page_state['headword'], page_state['headword_lang'],
                               translation.strip(), lang, lang_code, page_state['part_of_speech'], page_state['pronunciation'])

                # parse through an unordered list to grab pronunciations
                elif element.name == 'ul' and page_state['pro_region']:
                    for pro in self.parse_pronunciation_list_russian(element):
                        page_state['pronunciation'] = page_state['pronunciation'] + pro
                else:
                    page_state['translation_region'] = False

def main():

    parser = RuParser()
    text_file = open("data/html_ru.txt", "r")
    url_list = text_file.read().split('\n')
    for url in url_list:
        soup = get_html_tree_from_url(url)
        for tup in parser.generate_translation_tuples(soup):
            if not ''.join(str(el) for el in tup) == '':
               print(','.join(str(el) for el in tup))
'''
def main():
    parser = RuParser()
    for url in parser.tested_url:
        soup = get_html_tree_from_url(url)
        for tup in parser.generate_translation_tuples(soup):
            print(','.join(tup))

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


<commit_before>#!/usr/bin/python

import unittest

import sys; sys.path.append("../")
import Selector as SELECTORMODULE
from Selector import Selector

import Axon
from Axon.Ipc import shutdown
from Kamaelia.KamaeliaIPC import newReader, removeReader, newWriter, removeWriter, newExceptional, removeExceptional

class SmokeTests_Selector(unittest.TestCase):
    def test_SmokeTest(self):
        """__init__ - Called with no arguments succeeds"""
        S = Selector()
        self.assert_(isinstance(S, Axon.Component.component))

    def test_RunsForever(self):
        """main - Run with no messages, keeps running"""
        S = Selector()
        S.activate()
        for i in xrange(1,100):
            try:
                S.next()
            except StopIteration:
                self.fail("Component should run until told to stop. Failed on iteration: "+str(i))

    def test_PausesUntilFirstMessage(self):
        """main - Before we recieve any messages telling us what to watch for, the system should pause and yield"""
        S = Selector()
        S.activate()
        V = S.next()
        self.assert_(S._isRunnable() is not True)

    def test_shutdownMessageCausesShutdown(self):
        """main - If the component recieves a shutdown() message, the component shuts down"""
        S = Selector()
        S.activate()

        S._deliver(shutdown(),"control")

        componentExit = False
        for i in xrange(2000):
            try:
                S.next()
            except StopIteration:
                componentExit = True
                break
        if not componentExit:
            self.fail("When sent a shutdown message, the component should shutdown")

class MockSelect:
   """This is needed because we need to test that select is being used correctly"""
   def __init__(self, results=None):
       self.log = []
       self.results = results
   # We're using this simply as a namespace.
   def select(self,*args):
       self.log.append(("select", args))
       readers,writers, excepts, timeout = args
       if self.results is not None:
           try:
               result = self.results[0]
               del self.results[0]
               return result
           except IndexError:
               return [],[],[]
       else:
           return readers,writers, excepts
   def addResults(self, results):
       self.results.extend(results)

class Readables_Selector(unittest.TestCase):
    def test_SelectIsMockable(self):
        "main - The module uses select, and that is externally mockable"
        try:
            SELECTORMODULE.select
        except AttributeError:
            self.fail("Should import the select module")

    def test_SendingAReadableMessageResultsInItBeingSelectedAgainst(self):
        "main - If we send a newReader message to the notify inbox, it results in the selectable reader being selected on in the readers set"
        MOCKSELECTORMODULE = MockSelect()
        SELECTORMODULE.select = MOCKSELECTORMODULE
        S = Selector()
        S.activate()
        for i in xrange(100): S.next()
        dummyservice = (Axon.Component.component(), "inbox")
        S._deliver(newReader(S,( dummyservice, "LOOKINGFORTHIS")),"notify")
        for i in xrange(100): S.next()
        func, args = MOCKSELECTORMODULE.log[0]
        self.assertEqual("select", func, "select was called in the main loop")
        self.assertEqual(["LOOKINGFORTHIS"], args[0], "The selectable was added to the list of readables")
        self.assertEqual([], args[1], "Writable set should be empty")
        self.assertEqual([], args[2], "Exception set should be empty")
        self.assertEqual(0, args[3], "The select should be non-blocking")

    def test_WeSendTheSelectorAServiceAndSelectableOnlySelectsTheSelectable(self):
        "main - When we send the newReader message, it also includes a service upon which the selector can talk back to us. The selector only selects on the selectable part of the newReader message"

        MOCKSELECTORMODULE = MockSelect()
        SELECTORMODULE.select = MOCKSELECTORMODULE
        S = Selector()
        S.activate()
        for i in xrange(100): S.next()
        dummyservice = (Axon.Component.component(), "inbox")
        S._deliver(newReader(S,( dummyservice, "LOOKINGFORTHIS") ),"notify")
        for i in xrange(100): S.next()
        func, args = MOCKSELECTORMODULE.log[0]
        self.assertEqual("select", func, "select was called in the main loop")
        self.assertEqual(["LOOKINGFORTHIS"], args[0])#, "The selectable was added to the list of readables")
        self.assertEqual([], args[1], "Writable set should be empty")
        self.assertEqual([], args[2], "Exception set should be empty")
        self.assertEqual(0, args[3], "The select should be non-blocking")

    def test_SendingMultipleReadersResultsInAllSelected(self):
        "main - Sending multiple newReader messages results in all being select()ed"

        MOCKSELECTORMODULE = MockSelect()
        SELECTORMODULE.select = MOCKSELECTORMODULE
        S = Selector()
        S.activate()
        for i in xrange(100): S.next()
        dummyservice = (Axon.Component.component(), "inbox")
        S._deliver(newReader(S,( dummyservice, "LOOKINGFORTHIS") ),"notify")
        S._deliver(newReader(S,( dummyservice, "LOOKINGFORTHISTOO") ),"notify")
        S._deliver(newReader(S,( dummyservice, "LOOKINGFORANDTHIS") ),"notify")
        for i in xrange(100):
            S.next()
        lastfunc, lastargs = None, None
        i = 0
        func, args = MOCKSELECTORMODULE.log[i]
        while not( (lastfunc, lastargs) == (func, args)): # Search for quiescent state
            i = i + 1
            lastfunc, lastargs = func, args
            func, args = MOCKSELECTORMODULE.log[i]

        self.assertEqual("select", func, "select was called in the main loop")
        self.assertEqual(["LOOKINGFORTHIS","LOOKINGFORTHISTOO","LOOKINGFORANDTHIS"], args[0])#, "The selectable was added to the list of readables")

    def test_ActivityOnReaderResultsInMessageOnReadersService(self):
        "main - Activity on the selectable results in a message appearing in the service provided to the selector"

        MOCKSELECTORMODULE = MockSelect()
        SELECTORMODULE.select = MOCKSELECTORMODULE
        S = Selector()
        S.activate()
        for i in xrange(100): S.next()
        D = Axon.Component.component()
        dummyservice = (D, "inbox")
        S._deliver(newReader(S,( dummyservice, "LOOKINGFORTHIS") ),"notify")

        for i in xrange(100):
            S.next();
            try:
               S.postoffice.next()
            except:
               pass
        self.assert_(not ( D.inboxes["inbox"] == [] ) )
        selectable = D.recv("inbox")
        self.assertEqual(selectable,"LOOKINGFORTHIS")#, "The value returned should be the selectable we originally asked for")

    def test_ActivityOnAnyReaderResultsInMessageOnThatReadersService(self):
        "main - Activity on a selectable results in a message appearing in the service provided to the selector for that selectable"

        MOCKSELECTORMODULE = MockSelect(results=[ (["LOOKINGFORTHIS"],[],[]),
                                                  (["THENFORTHIS"],[],[]),
                                                  (["ANDTHENFORTHIS"],[],[]) ])
        SELECTORMODULE.select = MOCKSELECTORMODULE
        S = Selector()
        S.activate()
        for i in xrange(100): S.next()
        D = Axon.Component.component()
        E = Axon.Component.component()
        F = Axon.Component.component()
        dummyservice1 = (D, "inbox")
        S._deliver(newReader(S,( dummyservice1, "LOOKINGFORTHIS") ),"notify")
        dummyservice2 = (E, "inbox")
        S._deliver(newReader(S,( dummyservice2, "THENFORTHIS") ),"notify")
        dummyservice3 = (F, "inbox")
        S._deliver(newReader(S,( dummyservice3, "ANDTHENFORTHIS") ),"notify")

        for i in xrange(100):
            S.next();
            try:
               S.postoffice.next()
            except:
               pass
        selectable = D.recv("inbox")
        self.assertEqual(selectable,"LOOKINGFORTHIS")#, "The value returned should be the selectable we originally asked for")
        selectable = E.recv("inbox")
        self.assertEqual(selectable,"THENFORTHIS")#, "The value returned should be the selectable we originally asked for")
        selectable = F.recv("inbox")
        self.assertEqual(selectable,"ANDTHENFORTHIS")#, "The value returned should be the selectable we originally asked for")

    def test_RemoveReader_ResultsInReaderNoLongerBeingSelectedOrWiredIn(self):
        "main - Sending a remove reader message unwires/links a component, and also removes it's selectable from the readers list"
        MOCKSELECTORMODULE = MockSelect(results=[ ([], [], [] )])
        SELECTORMODULE.select = MOCKSELECTORMODULE
        S = Selector()
        S.activate()
        for i in xrange(100): S.next()
        D = Axon.Component.component()
        dummyservice = (D, "inbox")
        S._deliver(newReader(S,( dummyservice, "LOOKINGFORTHIS") ),"notify")
        S._deliver(removeReader(S,"LOOKINGFORTHIS"),"notify")

        for i in xrange(100):
            S.next();
            try:
               S.postoffice.next()
            except:
               pass
        self.assert_( D.inboxes["inbox"] == [] )

    def test_ActivityOnAnyWriterResultsInMessageOnThatWritersService(self):
        "main - Activity on a selectable results in a message appearing in the service provided to the selector for that selectable"

        MOCKSELECTORMODULE = MockSelect(results=[ ([],["LOOKINGFORTHIS"],[]),
                                                  ([],["THENFORTHIS"],[]),
                                                  ([],["ANDTHENFORTHIS"],[]) ])
        SELECTORMODULE.select = MOCKSELECTORMODULE
        S = Selector()
        S.activate()
        for i in xrange(100): S.next()
        D = Axon.Component.component()
        E = Axon.Component.component()
        F = Axon.Component.component()
        dummyservice1 = (D, "inbox")
        S._deliver(newWriter(S,( dummyservice1, "LOOKINGFORTHIS") ),"notify")
        dummyservice2 = (E, "inbox")
        S._deliver(newWriter(S,( dummyservice2, "THENFORTHIS") ),"notify")
        dummyservice3 = (F, "inbox")
        S._deliver(newWriter(S,( dummyservice3, "ANDTHENFORTHIS") ),"notify")

        for i in xrange(100):
            S.next();
            try:
               S.postoffice.next()
            except:
               pass
        selectable = D.recv("inbox")
        self.assertEqual(selectable,"LOOKINGFORTHIS")#, "The value returned should be the selectable we originally asked for")
        selectable = E.recv("inbox")
        self.assertEqual(selectable,"THENFORTHIS")#, "The value returned should be the selectable we originally asked for")
        selectable = F.recv("inbox")
        self.assertEqual(selectable,"ANDTHENFORTHIS")#, "The value returned should be the selectable we originally asked for")

    def test_RemoveWriter_ResultsInWriterNoLongerBeingSelectedOrWiredIn(self):
        "main - Sending a remove writer message unwires/links a component, and also removes it's selectable from the writers list"
        MOCKSELECTORMODULE = MockSelect(results=[ ([], [], [] )])
        SELECTORMODULE.select = MOCKSELECTORMODULE
        S = Selector()
        S.activate()
        for i in xrange(100): S.next()
        D = Axon.Component.component()
        dummyservice = (D, "inbox")
        S._deliver(newWriter(S,( dummyservice, "LOOKINGFORTHIS") ),"notify")
        S._deliver(removeWriter(S,"LOOKINGFORTHIS"),"notify")

        for i in xrange(100):
            S.next();
            try:
               S.postoffice.next()
            except:
               pass
        print
        self.assert_( D.inboxes["inbox"] == [] )

    def test_ActivityOnAnyExceptionalResultsInMessageOnThatExceptionalsService(self):
        "main - Activity on a selectable results in a message appearing in the service provided to the selector for that selectable"

        MOCKSELECTORMODULE = MockSelect(results=[ ([],[],["LOOKINGFORTHIS"]),
                                                  ([],[],["THENFORTHIS"]),
                                                  ([],[],["ANDTHENFORTHIS"]) ])
        SELECTORMODULE.select = MOCKSELECTORMODULE
        S = Selector()
        S.activate()
        for i in xrange(5): S.next()
        D = Axon.Component.component()
        E = Axon.Component.component()
        F = Axon.Component.component()
        dummyservice1 = (D, "inbox")
        S._deliver(newExceptional(S,( dummyservice1, "LOOKINGFORTHIS") ),"notify")
        dummyservice2 = (E, "inbox")
        S._deliver(newExceptional(S,( dummyservice2, "THENFORTHIS") ),"notify")
        dummyservice3 = (F, "inbox")
        S._deliver(newExceptional(S,( dummyservice3, "ANDTHENFORTHIS") ),"notify")

        for i in xrange(5):
            S.next();
            try:
               S.postoffice.next()
            except:
               pass
        selectable = D.recv("inbox")
        self.assertEqual(selectable,"LOOKINGFORTHIS")#, "The value returned should be the selectable we originally asked for")
        selectable = E.recv("inbox")
        self.assertEqual(selectable,"THENFORTHIS")#, "The value returned should be the selectable we originally asked for")
        selectable = F.recv("inbox")
        self.assertEqual(selectable,"ANDTHENFORTHIS")#, "The value returned should be the selectable we originally asked for")

    def test_RemoveExceptional_ResultsInExceptionalNoLongerBeingSelectedOrWiredIn(self):
        "main - Sending a remove exceptional message unwires/links a component, and also removes it's selectable from the exceptionals list"
        MOCKSELECTORMODULE = MockSelect(results=[ ([], [], [] )])
        SELECTORMODULE.select = MOCKSELECTORMODULE
        S = Selector()
        S.activate()
        for i in xrange(100): S.next()
        D = Axon.Component.component()
        dummyservice = (D, "inbox")
        S._deliver(newExceptional(S,( dummyservice, "LOOKINGFORTHIS") ),"notify")
        S._deliver(removeExceptional(S,"LOOKINGFORTHIS"),"notify")

        for i in xrange(100):
            S.next();
            try:
               S.postoffice.next()
            except:
               pass
        self.assert_( D.inboxes["inbox"] == [] )

if __name__=="__main__":
    unittest.main()
<commit_msg>Change license to Apache 2<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

# Copyright 2010 British Broadcasting Corporation and Kamaelia Contributors(1)
#
# (1) Kamaelia Contributors are listed in the AUTHORS file and at
#     http://www.kamaelia.org/AUTHORS - please extend this file,
#     not this notice.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest

import sys; sys.path.append("../")
import Selector as SELECTORMODULE
from Selector import Selector

import Axon
from Axon.Ipc import shutdown
from Kamaelia.KamaeliaIPC import newReader, removeReader, newWriter, removeWriter, newExceptional, removeExceptional

class SmokeTests_Selector(unittest.TestCase):
    def test_SmokeTest(self):
        """__init__ - Called with no arguments succeeds"""
        S = Selector()
        self.assert_(isinstance(S, Axon.Component.component))

    def test_RunsForever(self):
        """main - Run with no messages, keeps running"""
        S = Selector()
        S.activate()
        for i in xrange(1,100):
            try:
                S.next()
            except StopIteration:
                self.fail("Component should run until told to stop. Failed on iteration: "+str(i))

    def test_PausesUntilFirstMessage(self):
        """main - Before we recieve any messages telling us what to watch for, the system should pause and yield"""
        S = Selector()
        S.activate()
        V = S.next()
        self.assert_(S._isRunnable() is not True)

    def test_shutdownMessageCausesShutdown(self):
        """main - If the component recieves a shutdown() message, the component shuts down"""
        S = Selector()
        S.activate()

        S._deliver(shutdown(),"control")

        componentExit = False
        for i in xrange(2000):
            try:
                S.next()
            except StopIteration:
                componentExit = True
                break
        if not componentExit:
            self.fail("When sent a shutdown message, the component should shutdown")

class MockSelect:
   """This is needed because we need to test that select is being used correctly"""
   def __init__(self, results=None):
       self.log = []
       self.results = results
   # We're using this simply as a namespace.
   def select(self,*args):
       self.log.append(("select", args))
       readers,writers, excepts, timeout = args
       if self.results is not None:
           try:
               result = self.results[0]
               del self.results[0]
               return result
           except IndexError:
               return [],[],[]
       else:
           return readers,writers, excepts
   def addResults(self, results):
       self.results.extend(results)

class Readables_Selector(unittest.TestCase):
    def test_SelectIsMockable(self):
        "main - The module uses select, and that is externally mockable"
        try:
            SELECTORMODULE.select
        except AttributeError:
            self.fail("Should import the select module")

    def test_SendingAReadableMessageResultsInItBeingSelectedAgainst(self):
        "main - If we send a newReader message to the notify inbox, it results in the selectable reader being selected on in the readers set"
        MOCKSELECTORMODULE = MockSelect()
        SELECTORMODULE.select = MOCKSELECTORMODULE
        S = Selector()
        S.activate()
        for i in xrange(100): S.next()
        dummyservice = (Axon.Component.component(), "inbox")
        S._deliver(newReader(S,( dummyservice, "LOOKINGFORTHIS")),"notify")
        for i in xrange(100): S.next()
        func, args = MOCKSELECTORMODULE.log[0]
        self.assertEqual("select", func, "select was called in the main loop")
        self.assertEqual(["LOOKINGFORTHIS"], args[0], "The selectable was added to the list of readables")
        self.assertEqual([], args[1], "Writable set should be empty")
        self.assertEqual([], args[2], "Exception set should be empty")
        self.assertEqual(0, args[3], "The select should be non-blocking")

    def test_WeSendTheSelectorAServiceAndSelectableOnlySelectsTheSelectable(self):
        "main - When we send the newReader message, it also includes a service upon which the selector can talk back to us. The selector only selects on the selectable part of the newReader message"

        MOCKSELECTORMODULE = MockSelect()
        SELECTORMODULE.select = MOCKSELECTORMODULE
        S = Selector()
        S.activate()
        for i in xrange(100): S.next()
        dummyservice = (Axon.Component.component(), "inbox")
        S._deliver(newReader(S,( dummyservice, "LOOKINGFORTHIS") ),"notify")
        for i in xrange(100): S.next()
        func, args = MOCKSELECTORMODULE.log[0]
        self.assertEqual("select", func, "select was called in the main loop")
        self.assertEqual(["LOOKINGFORTHIS"], args[0])#, "The selectable was added to the list of readables")
        self.assertEqual([], args[1], "Writable set should be empty")
        self.assertEqual([], args[2], "Exception set should be empty")
        self.assertEqual(0, args[3], "The select should be non-blocking")

    def test_SendingMultipleReadersResultsInAllSelected(self):
        "main - Sending multiple newReader messages results in all being select()ed"

        MOCKSELECTORMODULE = MockSelect()
        SELECTORMODULE.select = MOCKSELECTORMODULE
        S = Selector()
        S.activate()
        for i in xrange(100): S.next()
        dummyservice = (Axon.Component.component(), "inbox")
        S._deliver(newReader(S,( dummyservice, "LOOKINGFORTHIS") ),"notify")
        S._deliver(newReader(S,( dummyservice, "LOOKINGFORTHISTOO") ),"notify")
        S._deliver(newReader(S,( dummyservice, "LOOKINGFORANDTHIS") ),"notify")
        for i in xrange(100):
            S.next()
        lastfunc, lastargs = None, None
        i = 0
        func, args = MOCKSELECTORMODULE.log[i]
        while not( (lastfunc, lastargs) == (func, args)): # Search for quiescent state
            i = i + 1
            lastfunc, lastargs = func, args
            func, args = MOCKSELECTORMODULE.log[i]

        self.assertEqual("select", func, "select was called in the main loop")
        self.assertEqual(["LOOKINGFORTHIS","LOOKINGFORTHISTOO","LOOKINGFORANDTHIS"], args[0])#, "The selectable was added to the list of readables")

    def test_ActivityOnReaderResultsInMessageOnReadersService(self):
        "main - Activity on the selectable results in a message appearing in the service provided to the selector"

        MOCKSELECTORMODULE = MockSelect()
        SELECTORMODULE.select = MOCKSELECTORMODULE
        S = Selector()
        S.activate()
        for i in xrange(100): S.next()
        D = Axon.Component.component()
        dummyservice = (D, "inbox")
        S._deliver(newReader(S,( dummyservice, "LOOKINGFORTHIS") ),"notify")

        for i in xrange(100):
            S.next();
            try:
               S.postoffice.next()
            except:
               pass
        self.assert_(not ( D.inboxes["inbox"] == [] ) )
        selectable = D.recv("inbox")
        self.assertEqual(selectable,"LOOKINGFORTHIS")#, "The value returned should be the selectable we originally asked for")

    def test_ActivityOnAnyReaderResultsInMessageOnThatReadersService(self):
        "main - Activity on a selectable results in a message appearing in the service provided to the selector for that selectable"

        MOCKSELECTORMODULE = MockSelect(results=[ (["LOOKINGFORTHIS"],[],[]),
                                                  (["THENFORTHIS"],[],[]),
                                                  (["ANDTHENFORTHIS"],[],[]) ])
        SELECTORMODULE.select = MOCKSELECTORMODULE
        S = Selector()
        S.activate()
        for i in xrange(100): S.next()
        D = Axon.Component.component()
        E = Axon.Component.component()
        F = Axon.Component.component()
        dummyservice1 = (D, "inbox")
        S._deliver(newReader(S,( dummyservice1, "LOOKINGFORTHIS") ),"notify")
        dummyservice2 = (E, "inbox")
        S._deliver(newReader(S,( dummyservice2, "THENFORTHIS") ),"notify")
        dummyservice3 = (F, "inbox")
        S._deliver(newReader(S,( dummyservice3, "ANDTHENFORTHIS") ),"notify")

        for i in xrange(100):
            S.next();
            try:
               S.postoffice.next()
            except:
               pass
        selectable = D.recv("inbox")
        self.assertEqual(selectable,"LOOKINGFORTHIS")#, "The value returned should be the selectable we originally asked for")
        selectable = E.recv("inbox")
        self.assertEqual(selectable,"THENFORTHIS")#, "The value returned should be the selectable we originally asked for")
        selectable = F.recv("inbox")
        self.assertEqual(selectable,"ANDTHENFORTHIS")#, "The value returned should be the selectable we originally asked for")

    def test_RemoveReader_ResultsInReaderNoLongerBeingSelectedOrWiredIn(self):
        "main - Sending a remove reader message unwires/links a component, and also removes it's selectable from the readers list"
        MOCKSELECTORMODULE = MockSelect(results=[ ([], [], [] )])
        SELECTORMODULE.select = MOCKSELECTORMODULE
        S = Selector()
        S.activate()
        for i in xrange(100): S.next()
        D = Axon.Component.component()
        dummyservice = (D, "inbox")
        S._deliver(newReader(S,( dummyservice, "LOOKINGFORTHIS") ),"notify")
        S._deliver(removeReader(S,"LOOKINGFORTHIS"),"notify")

        for i in xrange(100):
            S.next();
            try:
               S.postoffice.next()
            except:
               pass
        self.assert_( D.inboxes["inbox"] == [] )

    def test_ActivityOnAnyWriterResultsInMessageOnThatWritersService(self):
        "main - Activity on a selectable results in a message appearing in the service provided to the selector for that selectable"

        MOCKSELECTORMODULE = MockSelect(results=[ ([],["LOOKINGFORTHIS"],[]),
                                                  ([],["THENFORTHIS"],[]),
                                                  ([],["ANDTHENFORTHIS"],[]) ])
        SELECTORMODULE.select = MOCKSELECTORMODULE
        S = Selector()
        S.activate()
        for i in xrange(100): S.next()
        D = Axon.Component.component()
        E = Axon.Component.component()
        F = Axon.Component.component()
        dummyservice1 = (D, "inbox")
        S._deliver(newWriter(S,( dummyservice1, "LOOKINGFORTHIS") ),"notify")
        dummyservice2 = (E, "inbox")
        S._deliver(newWriter(S,( dummyservice2, "THENFORTHIS") ),"notify")
        dummyservice3 = (F, "inbox")
        S._deliver(newWriter(S,( dummyservice3, "ANDTHENFORTHIS") ),"notify")

        for i in xrange(100):
            S.next();
            try:
               S.postoffice.next()
            except:
               pass
        selectable = D.recv("inbox")
        self.assertEqual(selectable,"LOOKINGFORTHIS")#, "The value returned should be the selectable we originally asked for")
        selectable = E.recv("inbox")
        self.assertEqual(selectable,"THENFORTHIS")#, "The value returned should be the selectable we originally asked for")
        selectable = F.recv("inbox")
        self.assertEqual(selectable,"ANDTHENFORTHIS")#, "The value returned should be the selectable we originally asked for")

    def test_RemoveWriter_ResultsInWriterNoLongerBeingSelectedOrWiredIn(self):
        "main - Sending a remove writer message unwires/links a component, and also removes it's selectable from the writers list"
        MOCKSELECTORMODULE = MockSelect(results=[ ([], [], [] )])
        SELECTORMODULE.select = MOCKSELECTORMODULE
        S = Selector()
        S.activate()
        for i in xrange(100): S.next()
        D = Axon.Component.component()
        dummyservice = (D, "inbox")
        S._deliver(newWriter(S,( dummyservice, "LOOKINGFORTHIS") ),"notify")
        S._deliver(removeWriter(S,"LOOKINGFORTHIS"),"notify")

        for i in xrange(100):
            S.next();
            try:
               S.postoffice.next()
            except:
               pass
        print
        self.assert_( D.inboxes["inbox"] == [] )

    def test_ActivityOnAnyExceptionalResultsInMessageOnThatExceptionalsService(self):
        "main - Activity on a selectable results in a message appearing in the service provided to the selector for that selectable"

        MOCKSELECTORMODULE = MockSelect(results=[ ([],[],["LOOKINGFORTHIS"]),
                                                  ([],[],["THENFORTHIS"]),
                                                  ([],[],["ANDTHENFORTHIS"]) ])
        SELECTORMODULE.select = MOCKSELECTORMODULE
        S = Selector()
        S.activate()
        for i in xrange(5): S.next()
        D = Axon.Component.component()
        E = Axon.Component.component()
        F = Axon.Component.component()
        dummyservice1 = (D, "inbox")
        S._deliver(newExceptional(S,( dummyservice1, "LOOKINGFORTHIS") ),"notify")
        dummyservice2 = (E, "inbox")
        S._deliver(newExceptional(S,( dummyservice2, "THENFORTHIS") ),"notify")
        dummyservice3 = (F, "inbox")
        S._deliver(newExceptional(S,( dummyservice3, "ANDTHENFORTHIS") ),"notify")

        for i in xrange(5):
            S.next();
            try:
               S.postoffice.next()
            except:
               pass
        selectable = D.recv("inbox")
        self.assertEqual(selectable,"LOOKINGFORTHIS")#, "The value returned should be the selectable we originally asked for")
        selectable = E.recv("inbox")
        self.assertEqual(selectable,"THENFORTHIS")#, "The value returned should be the selectable we originally asked for")
        selectable = F.recv("inbox")
        self.assertEqual(selectable,"ANDTHENFORTHIS")#, "The value returned should be the selectable we originally asked for")

    def test_RemoveExceptional_ResultsInExceptionalNoLongerBeingSelectedOrWiredIn(self):
        "main - Sending a remove exceptional message unwires/links a component, and also removes it's selectable from the exceptionals list"
        MOCKSELECTORMODULE = MockSelect(results=[ ([], [], [] )])
        SELECTORMODULE.select = MOCKSELECTORMODULE
        S = Selector()
        S.activate()
        for i in xrange(100): S.next()
        D = Axon.Component.component()
        dummyservice = (D, "inbox")
        S._deliver(newExceptional(S,( dummyservice, "LOOKINGFORTHIS") ),"notify")
        S._deliver(removeExceptional(S,"LOOKINGFORTHIS"),"notify")

        for i in xrange(100):
            S.next();
            try:
               S.postoffice.next()
            except:
               pass
        self.assert_( D.inboxes["inbox"] == [] )

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


<commit_before>"""This module contains several helper functions which can be used to
find an address of the submitting system, for example to use as the
address parameter for HighThroughputExecutor.

The helper to use depends on the network environment around the submitter,
so some experimentation will probably be needed to choose the correct one.
"""

import logging
import os
import platform
import requests
import socket
import fcntl
import struct
import psutil

from typing import Set

logger = logging.getLogger(__name__)

def address_by_route() -> str:
    """Finds an address for the local host by querying the local routing table
       for the route to Google DNS.

       This will return an unusable value when the internet-facing address is
       not reachable from workers.
    """
    logger.debug("Finding address by querying local routing table")
    addr = os.popen("/sbin/ip route get 8.8.8.8 | awk '{print $NF;exit}'").read().strip()
    logger.debug("Address found: {}".format(addr))
    return addr

def address_by_query() -> str:
    """Finds an address for the local host by querying ipify. This may
       return an unusable value when the host is behind NAT, or when the
       internet-facing address is not reachable from workers.
    """
    logger.debug("Finding address by querying remote service")
    response = requests.get('https://api.ipify.org')

    if response.status_code == 200:
        addr = response.text
        logger.debug("Address found: {}".format(addr))
        return addr
    else:
        raise RuntimeError("Remote service returned unexpected HTTP status code {}".format(response.status_code))

def address_by_hostname() -> str:
    """Returns the hostname of the local host.

       This will return an unusable value when the hostname cannot be
       resolved from workers.
    """
    logger.debug("Finding address by using local hostname")
    addr = platform.node()
    logger.debug("Address found: {}".format(addr))
    return addr

def address_by_interface(ifname: str) -> str:
    """Returns the IP address of the given interface name, e.g. 'eth0'

    This is taken from a Stack Overflow answer: https://stackoverflow.com/questions/24196932/how-can-i-get-the-ip-address-of-eth0-in-python#24196955

    Parameters
    ----------
    ifname : str
        Name of the interface whose address is to be returned. Required.

    """
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    return socket.inet_ntoa(fcntl.ioctl(
        s.fileno(),
        0x8915,  # SIOCGIFADDR
        struct.pack('256s', bytes(ifname[:15], 'utf-8'))
    )[20:24])

def get_all_addresses() -> Set[str]:
    """ Uses a combination of methods to determine possible addresses.

    Returns:
         list of addresses as strings
    """
    net_interfaces = psutil.net_if_addrs()

    s_addresses = set()
    for interface in net_interfaces:
        try:
            s_addresses.add(address_by_interface(interface))
        except Exception:
            logger.exception("Ignoring failure to fetch address from interface {}".format(interface))
            pass

    resolution_functions = [address_by_hostname, address_by_route, address_by_query]
    for f in resolution_functions:
        try:
            s_addresses.add(f())
        except Exception:
            logger.exception("Ignoring an address finder exception")

    return s_addresses
<commit_msg>Adding a timeout for the call to get address (#1715)<commit_after>"""This module contains several helper functions which can be used to
find an address of the submitting system, for example to use as the
address parameter for HighThroughputExecutor.

The helper to use depends on the network environment around the submitter,
so some experimentation will probably be needed to choose the correct one.
"""

import logging
import os
import platform
import requests
import socket
import fcntl
import struct
import psutil

from typing import Set, List, Callable

logger = logging.getLogger(__name__)

def address_by_route() -> str:
    """Finds an address for the local host by querying the local routing table
       for the route to Google DNS.

       This will return an unusable value when the internet-facing address is
       not reachable from workers.
    """
    logger.debug("Finding address by querying local routing table")
    addr = os.popen("/sbin/ip route get 8.8.8.8 | awk '{print $NF;exit}'").read().strip()
    logger.debug("Address found: {}".format(addr))
    return addr

def address_by_query(timeout: float = 30) -> str:
    """Finds an address for the local host by querying ipify. This may
       return an unusable value when the host is behind NAT, or when the
       internet-facing address is not reachable from workers.
       Parameters:
       -----------

       timeout : float
          Timeout for the request in seconds. Default: 30s
    """
    logger.debug("Finding address by querying remote service")
    response = requests.get('https://api.ipify.org', timeout=timeout)

    if response.status_code == 200:
        addr = response.text
        logger.debug("Address found: {}".format(addr))
        return addr
    else:
        raise RuntimeError("Remote service returned unexpected HTTP status code {}".format(response.status_code))

def address_by_hostname() -> str:
    """Returns the hostname of the local host.

       This will return an unusable value when the hostname cannot be
       resolved from workers.
    """
    logger.debug("Finding address by using local hostname")
    addr = platform.node()
    logger.debug("Address found: {}".format(addr))
    return addr

def address_by_interface(ifname: str) -> str:
    """Returns the IP address of the given interface name, e.g. 'eth0'

    This is taken from a Stack Overflow answer: https://stackoverflow.com/questions/24196932/how-can-i-get-the-ip-address-of-eth0-in-python#24196955

    Parameters
    ----------
    ifname : str
        Name of the interface whose address is to be returned. Required.

    """
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    return socket.inet_ntoa(fcntl.ioctl(
        s.fileno(),
        0x8915,  # SIOCGIFADDR
        struct.pack('256s', bytes(ifname[:15], 'utf-8'))
    )[20:24])

def get_all_addresses() -> Set[str]:
    """ Uses a combination of methods to determine possible addresses.

    Returns:
         list of addresses as strings
    """
    net_interfaces = psutil.net_if_addrs()

    s_addresses = set()
    for interface in net_interfaces:
        try:
            s_addresses.add(address_by_interface(interface))
        except Exception:
            logger.exception("Ignoring failure to fetch address from interface {}".format(interface))
            pass

    resolution_functions = [address_by_hostname, address_by_route, address_by_query]  # type: List[Callable[[], str]]
    for f in resolution_functions:
        try:
            s_addresses.add(f())
        except Exception:
            logger.exception("Ignoring an address finder exception")

    return s_addresses
<|endoftext|>


<commit_before>"""Following the general logging philosophy of python libraries, by default
`Parsl <https://github.com/swift-lang/swift-e-lab/>`_ doesn't log anything.
However the following helper functions are provided for logging:

1. set_stream_logger
    This sets the logger to the StreamHandler. This is quite useful when working from
    a Jupyter notebook.

2. set_file_logger
    This sets the logging to a file. This is ideal for reporting issues to the dev team.

"""
import logging
import typeguard

from typing import Optional

@typeguard.typechecked
def set_stream_logger(name: str = 'parsl', level: int = logging.DEBUG, format_string: Optional[str] = None):
    """Add a stream log handler.

    Args:
         - name (string) : Set the logger name.
         - level (logging.LEVEL) : Set to logging.DEBUG by default.
         - format_string (string) : Set to None by default.

    Returns:
         - None
    """
    if format_string is None:
        # format_string = "%(asctime)s %(name)s [%(levelname)s] Thread:%(thread)d %(message)s"
        format_string = "%(asctime)s %(name)s:%(lineno)d [%(levelname)s]  %(message)s"

    logger = logging.getLogger(name)
    logger.setLevel(logging.DEBUG)
    handler = logging.StreamHandler()
    handler.setLevel(level)
    formatter = logging.Formatter(format_string, datefmt='%Y-%m-%d %H:%M:%S')
    handler.setFormatter(formatter)
    logger.addHandler(handler)

    # Concurrent.futures errors are also of interest, as exceptions
    # which propagate out of the top of a callback are logged this way
    # and then discarded. (see #240)
    futures_logger = logging.getLogger("concurrent.futures")
    futures_logger.addHandler(handler)

@typeguard.typechecked
def set_file_logger(filename: str, name: str = 'parsl', level: int = logging.DEBUG, format_string: Optional[str] = None):
    """Add a stream log handler.

    Args:
        - filename (string): Name of the file to write logs to
        - name (string): Logger name
        - level (logging.LEVEL): Set the logging level.
        - format_string (string): Set the format string

    Returns:
       -  None
    """
    if format_string is None:
        format_string = "%(asctime)s.%(msecs)03d %(name)s:%(lineno)d [%(levelname)s]  %(message)s"

    logger = logging.getLogger(name)
    logger.setLevel(logging.DEBUG)
    handler = logging.FileHandler(filename)
    handler.setLevel(level)
    formatter = logging.Formatter(format_string, datefmt='%Y-%m-%d %H:%M:%S')
    handler.setFormatter(formatter)
    logger.addHandler(handler)

    # see note in set_stream_logger for notes about logging
    # concurrent.futures
    futures_logger = logging.getLogger("concurrent.futures")
    futures_logger.addHandler(handler)
<commit_msg>Correct docstring for set_file_logger (#2156)<commit_after>"""Following the general logging philosophy of python libraries, by default
`Parsl <https://github.com/swift-lang/swift-e-lab/>`_ doesn't log anything.
However the following helper functions are provided for logging:

1. set_stream_logger
    This sets the logger to the StreamHandler. This is quite useful when working from
    a Jupyter notebook.

2. set_file_logger
    This sets the logging to a file. This is ideal for reporting issues to the dev team.

"""
import logging
import typeguard

from typing import Optional

@typeguard.typechecked
def set_stream_logger(name: str = 'parsl', level: int = logging.DEBUG, format_string: Optional[str] = None):
    """Add a stream log handler.

    Args:
         - name (string) : Set the logger name.
         - level (logging.LEVEL) : Set to logging.DEBUG by default.
         - format_string (string) : Set to None by default.

    Returns:
         - None
    """
    if format_string is None:
        # format_string = "%(asctime)s %(name)s [%(levelname)s] Thread:%(thread)d %(message)s"
        format_string = "%(asctime)s %(name)s:%(lineno)d [%(levelname)s]  %(message)s"

    logger = logging.getLogger(name)
    logger.setLevel(logging.DEBUG)
    handler = logging.StreamHandler()
    handler.setLevel(level)
    formatter = logging.Formatter(format_string, datefmt='%Y-%m-%d %H:%M:%S')
    handler.setFormatter(formatter)
    logger.addHandler(handler)

    # Concurrent.futures errors are also of interest, as exceptions
    # which propagate out of the top of a callback are logged this way
    # and then discarded. (see #240)
    futures_logger = logging.getLogger("concurrent.futures")
    futures_logger.addHandler(handler)

@typeguard.typechecked
def set_file_logger(filename: str, name: str = 'parsl', level: int = logging.DEBUG, format_string: Optional[str] = None):
    """Add a file log handler.

    Args:
        - filename (string): Name of the file to write logs to
        - name (string): Logger name
        - level (logging.LEVEL): Set the logging level.
        - format_string (string): Set the format string

    Returns:
       -  None
    """
    if format_string is None:
        format_string = "%(asctime)s.%(msecs)03d %(name)s:%(lineno)d [%(levelname)s]  %(message)s"

    logger = logging.getLogger(name)
    logger.setLevel(logging.DEBUG)
    handler = logging.FileHandler(filename)
    handler.setLevel(level)
    formatter = logging.Formatter(format_string, datefmt='%Y-%m-%d %H:%M:%S')
    handler.setFormatter(formatter)
    logger.addHandler(handler)

    # see note in set_stream_logger for notes about logging
    # concurrent.futures
    futures_logger = logging.getLogger("concurrent.futures")
    futures_logger.addHandler(handler)
<|endoftext|>


<commit_before>"""hug/types.py

Defines hugs built-in supported types / validators

Copyright (C) 2015  Timothy Edmund Crosley

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
documentation files (the "Software"), to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and
to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or
substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF
CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
OTHER DEALINGS IN THE SOFTWARE.

"""
from decimal import Decimal
from json import loads as load_json

def accept(formatter, doc=None, error_text=None, cli_behaviour=None, exception_handlers=None):
    '''Allows quick wrapping any Python type converter for use with Hug type annotations'''
    defined_exception_handlers = exception_handlers or {}
    if defined_exception_handlers or error_text:
        def hug_formatter(data):
            try:
                return formatter(data)
            except Exception as exception:
                for take_exception, rewrite in defined_exception_handlers.items():
                    if isinstance(exception, take_exception):
                        if isinstance(rewrite, str):
                            raise ValueError(rewrite)
                        else:
                            raise rewrite(data)
                if error_text:
                    raise ValueError(error_text)
                raise exception
    else:
        def hug_formatter(data):
            return formatter(data)

    new_cli_behaviour = getattr(formatter, 'cli_behaviour', {})
    if cli_behaviour:
        new_cli_behaviour.update(cli_behaviour)
    if new_cli_behaviour:
        hug_formatter.cli_behaviour = new_cli_behaviour
    hug_formatter.__doc__ = doc or formatter.__doc__
    return hug_formatter

def multiple(value):
    '''Multiple Values'''
    return value if isinstance(value, list) else [value]
multiple.cli_behaviour = {'action': 'append'}

def comma_separated_list(value):
    '''Multiple values, separated by a comma'''
    return value.split(",")

number = accept(int, 'A whole number', 'Invalid whole number provided')
float_number = accept(float, 'A float number', 'Invalid float number provided')
decimal = accept(Decimal, 'A decimal number', 'Invalid decimal number provided')
text = accept(str, 'Basic text / string value', 'Invalid text value provided')
boolean = accept(bool, 'Providing any value will set this to true',
                 'Invalid boolean value provided', cli_behaviour={'action': 'store_true'})

def smart_boolean(input_value):
    '''Accepts a true or false value'''
    if type(input_value) == bool or input_value in (None, 1, 0):
        return bool(input_value)

    value = input_value.lower()
    if value in ('true', 't', '1'):
        return True
    elif value in ('false', 'f', '0', ''):
        return False

    raise KeyError('Invalid value passed in for true/false field')

def inline_dictionary(input_value):
    '''A single line dictionary, where items are separted by commas and key:value are separated by a pipe'''
    return {key.strip(): value.strip() for key, value in (item.split(":") for item in input_value.split("|"))}

def one_of(values):
    '''Ensures the value is within a set of acceptable values'''
    def matches(value):
        if not value in values:
            raise KeyError('Invalid value passed. The accepted values are: ({0})'.format("|".join(values)))
        return value

    matches.__doc__ = 'Accepts one of the following values: ({0})'.format("|".join(values))
    matches.cli_behaviour = {'choices': values}
    return matches

def mapping(value_map):
    '''Ensures the value is one of an acceptable set of values mapping those values to a Python equivelent'''
    values = value_map.keys()
    def matches(value):
        if not value in value_map.keys():
            raise KeyError('Invalid value passed. The accepted values are: ({0})'.format("|".join(values)))
        return value_map[value]

    matches.__doc__ = 'Accepts one of the following values: ({0})'.format("|".join(values))
    matches.cli_behaviour = {'choices': values}
    return matches

def json(value):
    '''Accepts a JSON formatted data structure'''
    if type(value) in (str, bytes):
        try:
            return load_json(value)
        except Exception:
            raise ValueError('Incorrectly formatted JSON provided')
    else:
        return value
<commit_msg>Implement multi-type support in accordance with changelog and test definition<commit_after>"""hug/types.py

Defines hugs built-in supported types / validators

Copyright (C) 2015  Timothy Edmund Crosley

Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated
documentation files (the "Software"), to deal in the Software without restriction, including without limitation
the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and
to permit persons to whom the Software is furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all copies or
substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF
CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
OTHER DEALINGS IN THE SOFTWARE.

"""
from decimal import Decimal
from json import loads as load_json

def accept(formatter, doc=None, error_text=None, cli_behaviour=None, exception_handlers=None):
    '''Allows quick wrapping any Python type converter for use with Hug type annotations'''
    defined_exception_handlers = exception_handlers or {}
    if defined_exception_handlers or error_text:
        def hug_formatter(data):
            try:
                return formatter(data)
            except Exception as exception:
                for take_exception, rewrite in defined_exception_handlers.items():
                    if isinstance(exception, take_exception):
                        if isinstance(rewrite, str):
                            raise ValueError(rewrite)
                        else:
                            raise rewrite(data)
                if error_text:
                    raise ValueError(error_text)
                raise exception
    else:
        def hug_formatter(data):
            return formatter(data)

    new_cli_behaviour = getattr(formatter, 'cli_behaviour', {})
    if cli_behaviour:
        new_cli_behaviour.update(cli_behaviour)
    if new_cli_behaviour:
        hug_formatter.cli_behaviour = new_cli_behaviour
    hug_formatter.__doc__ = doc or formatter.__doc__
    return hug_formatter

def multiple(value):
    '''Multiple Values'''
    return value if isinstance(value, list) else [value]
multiple.cli_behaviour = {'action': 'append'}

def comma_separated_list(value):
    '''Multiple values, separated by a comma'''
    return value.split(",")

number = accept(int, 'A whole number', 'Invalid whole number provided')
float_number = accept(float, 'A float number', 'Invalid float number provided')
decimal = accept(Decimal, 'A decimal number', 'Invalid decimal number provided')
text = accept(str, 'Basic text / string value', 'Invalid text value provided')
boolean = accept(bool, 'Providing any value will set this to true',
                 'Invalid boolean value provided', cli_behaviour={'action': 'store_true'})

def smart_boolean(input_value):
    '''Accepts a true or false value'''
    if type(input_value) == bool or input_value in (None, 1, 0):
        return bool(input_value)

    value = input_value.lower()
    if value in ('true', 't', '1'):
        return True
    elif value in ('false', 'f', '0', ''):
        return False

    raise KeyError('Invalid value passed in for true/false field')

def inline_dictionary(input_value):
    '''A single line dictionary, where items are separted by commas and key:value are separated by a pipe'''
    return {key.strip(): value.strip() for key, value in (item.split(":") for item in input_value.split("|"))}

def one_of(values):
    '''Ensures the value is within a set of acceptable values'''
    def matches(value):
        if not value in values:
            raise KeyError('Invalid value passed. The accepted values are: ({0})'.format("|".join(values)))
        return value

    matches.__doc__ = 'Accepts one of the following values: ({0})'.format("|".join(values))
    matches.cli_behaviour = {'choices': values}
    return matches

def mapping(value_map):
    '''Ensures the value is one of an acceptable set of values mapping those values to a Python equivelent'''
    values = value_map.keys()
    def matches(value):
        if not value in value_map.keys():
            raise KeyError('Invalid value passed. The accepted values are: ({0})'.format("|".join(values)))
        return value_map[value]

    matches.__doc__ = 'Accepts one of the following values: ({0})'.format("|".join(values))
    matches.cli_behaviour = {'choices': values}
    return matches

def json(value):
    '''Accepts a JSON formatted data structure'''
    if type(value) in (str, bytes):
        try:
            return load_json(value)
        except Exception:
            raise ValueError('Incorrectly formatted JSON provided')
    else:
        return value

def multi(*types):
    '''Enables accepting one of multiple type methods'''
    type_strings = (type_method.__doc__ for type_method in types)
    doc_string = 'Accepts any of the following value types:{0}\n'.format('\n  - '.join(type_strings))
    def multi_type(value):
        for type_method in types:
            try:
                return type_method(value)
            except:
                pass
        raise ValueError(doc_string)
    multi_type.__doc__ = doc_string
    return multi_type
<|endoftext|>


<commit_before>from __future__ import division
import numpy as np
from pymc import deterministic, Normal
from pymc.Model import Model
from .array_utils import *
from .ModelComponents import Sky
from .model_parser import component_list_from_file
# import matplotlib.pyplot as pp

# TODO: Is there a way to use masked arrays to skip bad pixels instead?
# Pixels that have zero weight will be replaced with a very small weight
_zero_weight = 1e-20

def multicomponent_model(obs_data, obs_ivm, psf_data, psf_ivm,
                         components=None, mag_zp=0):
    """
    Multi-component model for MCMC psf fitting. Components is a list of
    ComponentBase subclasses, or a model definition python file parsable by
    model_parser
    """
    if isinstance(components, basestring):
        try:
            components = component_list_from_file(components)
        except IOError, err:
            message = 'Unable to open components file {}. Does it exist?'
            err.message = message.format(components)
            raise err

    np.seterr(divide='ignore')

    # Normalize the PSF kernel
    psf_data, psf_ivm = normed_psf(psf_data, psf_ivm)
    # psf_rms = np.where(psf_ivm == 0, 0, 1 / np.sqrt(psf_ivm))
    psf_var = np.where(psf_ivm == 0, 0, 1 / psf_ivm)

    # pad the psf arrays to the same size as the data, for fft
    f_psf = pad_and_rfft_image(psf_data, obs_data.shape)
    # f_psf_rms = pad_and_rfft_image(psf_rms, obs_data.shape)
    f_psf_var = pad_and_rfft_image(psf_var, obs_data.shape)

    obs_mask = obs_ivm <= 0
    obs_var = np.where(obs_mask, 0, 1 / obs_ivm)

    # pre-compute data x,y coordinates
    data_coords = array_coords(obs_data)

    model_comps = []
    stochastics = []
    sky = Sky(adu=0.0)

    for count, component in enumerate(components):
        component.update_trace_names(count=count)

        # Sky is added after convolution
        if isinstance(component, Sky):
            sky = component
            stochastics += [sky]
        else:
            model_comps += [component]

    @deterministic(plot=False, trace=False)
    def raw_model(model_comps=model_comps):
        modelpx = np.zeros_like(obs_data)
        for comp in model_comps:
            comp.add_to_array(modelpx, mag_zp, coords=data_coords)
        return modelpx

    @deterministic(plot=False, trace=False)
    def convolved_model(f_psf=f_psf, raw_model=raw_model):
        cmodel = convolve(raw_model, f_psf)
        return cmodel

    @deterministic(plot=False, trace=False)
    def composite_ivm(obs_var=obs_var, obs_mask=obs_mask,
                      f_psf_var=f_psf_var, #f_psf_rms=f_psf_rms,
                      raw_model=raw_model):
        # f * (g + h) = (f * g) + (f * h), so convolve PSF RMS map with model to
        # get model RMS map
        # TODO: should this be modelIVM = convolve(rawmodel**2, f_psf_var)?
        # modelRMS = convolve(raw_model, f_psf_rms)
        model_var = convolve(raw_model**2, f_psf_var)
        # Set zero-weight pixels to very small number instead
        # badpx = (modelRMS <= 0) | (obs_ivm <= 0)
        badpx = (model_var <= 0) | obs_mask
        # compIVM = np.where(badpx, _zero_weight, 1 / (modelRMS**2 + 1 / obs_ivm))
        compIVM = np.where(badpx, _zero_weight, 1 / (model_var + obs_var))
        return compIVM

    @deterministic(plot=False, trace=False)
    def residual(obs_data=obs_data, convolved_model=convolved_model, sky=sky):
        return obs_data - convolved_model - sky.adu

    data = Normal('data', value=np.zeros_like(obs_data), mu=residual,
                  tau=composite_ivm, observed=True, trace=False)

    stochastics += [raw_model, convolved_model, composite_ivm, residual, data]
    stochastics += model_comps

    return Model(stochastics)
<commit_msg>Cleanup in models<commit_after>from __future__ import division
import numpy as np
from pymc import deterministic, Normal
from pymc.Model import Model
from .array_utils import *
from .ModelComponents import Sky
from .model_parser import component_list_from_file
# import matplotlib.pyplot as pp

# TODO: Is there a way to use masked arrays to skip bad pixels instead?
# Pixels that have zero weight will be replaced with a very small weight
_zero_weight = 1e-20

def multicomponent_model(obs_data, obs_ivm, psf_data, psf_ivm,
                         components=None, mag_zp=0):
    """
    Multi-component model for MCMC psf fitting. Components is a list of
    ComponentBase subclasses, or a model definition python file parsable by
    model_parser
    """
    if isinstance(components, basestring):
        try:
            components = component_list_from_file(components)
        except IOError, err:
            message = 'Unable to open components file {}. Does it exist?'
            err.message = message.format(components)
            raise err

    np.seterr(divide='ignore')

    # Normalize the PSF kernel
    psf_data, psf_ivm = normed_psf(psf_data, psf_ivm)
    psf_var = np.where(psf_ivm == 0, 0, 1 / psf_ivm)

    # pad the psf arrays to the same size as the data, for fft
    f_psf = pad_and_rfft_image(psf_data, obs_data.shape)
    f_psf_var = pad_and_rfft_image(psf_var, obs_data.shape)

    # pre-compute variance map and bad pixel mask for observation
    obs_mask = obs_ivm <= 0
    obs_var = np.where(obs_mask, 0, 1 / obs_ivm)

    # pre-compute data x,y coordinates
    data_coords = array_coords(obs_data)

    model_comps = []
    stochastics = []
    sky = Sky(adu=0.0)

    for count, component in enumerate(components):
        component.update_trace_names(count=count)

        # Sky is added after convolution
        if isinstance(component, Sky):
            sky = component
            stochastics += [sky]
        else:
            model_comps += [component]

    @deterministic(plot=False, trace=False)
    def raw_model(model_comps=model_comps):
        modelpx = np.zeros_like(obs_data)
        for comp in model_comps:
            comp.add_to_array(modelpx, mag_zp, coords=data_coords)
        return modelpx

    @deterministic(plot=False, trace=False)
    def convolved_model(f_psf=f_psf, raw_model=raw_model, sky=sky):
        cmodel = convolve(raw_model, f_psf)
        return cmodel + sky.adu

    @deterministic(plot=False, trace=False)
    def composite_ivm(obs_var=obs_var, obs_mask=obs_mask, f_psf_var=f_psf_var,
                      raw_model=raw_model):
        # compute model variance
        model_var = convolve(raw_model**2, f_psf_var)
        # Set zero-weight pixels to very small number instead
        badpx = (model_var <= 0) | obs_mask
        compIVM = np.where(badpx, _zero_weight, 1 / (model_var + obs_var))
        return compIVM

    @deterministic(plot=False, trace=False)
    def residual(obs_data=obs_data, convolved_model=convolved_model):
        return obs_data - convolved_model

    data = Normal('data', value=obs_data, mu=convolved_model,
                  tau=composite_ivm, observed=True, trace=False)

    stochastics += [raw_model, convolved_model, composite_ivm, residual, data]
    stochastics += model_comps

    return Model(stochastics)
<|endoftext|>


<commit_before>from errno import EINTR
from subprocess import Popen, PIPE
import askpass
import color
import os
import signal
import sys
import time
import traceback

BUFFER_SIZE = 1 << 16

class Task(object):
    """Starts a process and manages its input and output."""
    def __init__(self, host, port, cmd, opts, stdin=None):
        self.host = host
        self.port = port
        self.cmd = cmd

        self.proc = None
        self.writer = None
        self.timestamp = None
        self.failures = []
        self.killed = False
        self.inputbuffer = stdin
        self.outputbuffer = ''
        self.errorbuffer = ''

        self.stdin = None
        self.stdout = None
        self.stderr = None
        self.outfile = None
        self.errfile = None

        # Set options.
        self.verbose = opts.verbose
        try:
            self.print_out = bool(opts.print_out)
        except AttributeError:
            self.print_out = False
        try:
            self.inline = bool(opts.inline)
        except AttributeError:
            self.inline = False

    def start(self, nodenum, iomap, writer, askpass_socket=None):
        """Starts the process and registers files with the IOMap."""
        self.writer = writer

        if writer:
            self.outfile, self.errfile = writer.open_files(self.host)

        # Set up the environment.
        environ = dict(os.environ)
        environ['PSSH_NODENUM'] = str(nodenum)
        # Disable the GNOME pop-up password dialog and allow ssh to use
        # askpass.py to get a provided password.  If the module file is
        # askpass.pyc, we replace the extension.
        root, ext = os.path.splitext(os.path.abspath(askpass.__file__))
        environ['SSH_ASKPASS'] = '%s.py' % root
        if askpass_socket:
            environ['PSSH_ASKPASS_SOCKET'] = askpass_socket

        # Create the subprocess.
        self.proc = Popen(self.cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE,
                close_fds=True, preexec_fn=os.setsid, env=environ)
        self.timestamp = time.time()
        if self.inputbuffer:
            self.stdin = self.proc.stdin
            iomap.register(self.stdin.fileno(), self.handle_stdin, write=True)
        else:
            self.proc.stdin.close()
        self.stdout = self.proc.stdout
        iomap.register(self.stdout.fileno(), self.handle_stdout, read=True)
        self.stderr = self.proc.stderr
        iomap.register(self.stderr.fileno(), self.handle_stderr, read=True)

    def _kill(self):
        """Signals the process to terminate."""
        if self.proc:
            os.kill(-self.proc.pid, signal.SIGKILL)
            self.killed = True

    def timedout(self):
        """Kills the process and registers a timeout error."""
        if not self.killed:
            self._kill()
            self.failures.append('Timed out')

    def interrupted(self):
        """Kills the process and registers an keyboard interrupt error."""
        if not self.killed:
            self._kill()
            self.failures.append('Interrupted')

    def cancel(self):
        """Stops a task that has not started."""
        self.failures.append('Cancelled')

    def elapsed(self):
        """Finds the time in seconds since the process was started."""
        return time.time() - self.timestamp

    def running(self):
        """Finds if the process has terminated and saves the return code."""
        if self.stdin or self.stdout or self.stderr:
            return True
        if self.proc:
            self.returncode = self.proc.poll()
            if self.returncode is None:
                if self.killed:
                    return False
                else:
                    return True
            else:
                if self.returncode < 0:
                    message = 'Killed by signal %s' % (-self.returncode)
                    self.failures.append(message)
                elif self.returncode > 0:
                    message = 'Exited with error code %s' % self.returncode
                    self.failures.append(message)
                self.proc = None
                return False

    def handle_stdin(self, fd, event, iomap):
        try:
            if self.inputbuffer:
                bytes_written = os.write(fd, self.inputbuffer)
                self.inputbuffer = self.inputbuffer[bytes_written:]
            else:
                self.close_stdin(iomap)
        except (OSError, IOError), e:
            if e.errno != EINTR:
                self.close_stdin(iomap)
                self.log_exception(e)

    def close_stdin(self, iomap):
        if self.stdin:
            iomap.unregister(self.stdin.fileno())
            self.stdin.close()
            self.stdin = None

    def handle_stdout(self, fd, event, iomap):
        try:
            buf = os.read(fd, BUFFER_SIZE)
            if buf:
                if self.inline:
                    self.outputbuffer += buf
                if self.outfile:
                    self.writer.write(self.outfile, buf)
                if self.print_out:
                    print '%s: %s' % (self.host, buf),
            else:
                self.close_stdout(iomap)
        except (OSError, IOError), e:
            if e.errno != EINTR:
                self.close_stdout(iomap)
                self.log_exception(e)

    def close_stdout(self, iomap):
        if self.stdout:
            iomap.unregister(self.stdout.fileno())
            self.stdout.close()
            self.stdout = None
        if self.outfile:
            self.writer.close(self.outfile)
            self.outfile = None

    def handle_stderr(self, fd, event, iomap):
        try:
            buf = os.read(fd, BUFFER_SIZE)
            if buf:
                if self.inline:
                    self.errorbuffer += buf
                if self.errfile:
                    self.writer.write(self.errfile, buf)
            else:
                self.close_stderr(iomap)
        except (OSError, IOError), e:
            if e.errno != EINTR:
                self.close_stderr(iomap)
                self.log_exception(e)

    def close_stderr(self, iomap):
        if self.stderr:
            iomap.unregister(self.stderr.fileno())
            self.stderr.close()
            self.stderr = None
        if self.errfile:
            self.writer.close(self.errfile)
            self.errfile = None

    def log_exception(self, e):
        """Saves a record of the most recent exception for error reporting."""
        if self.verbose:
            exc_type, exc_value, exc_traceback = sys.exc_info()
            exc = ("Exception: %s, %s, %s" %
                    (exc_type, exc_value, traceback.format_tb(exc_traceback)))
        else:
            exc = str(e)
        self.failures.append(exc)

    def report(self, n):
        """Pretty prints a status report after the Task completes."""
        error = ', '.join(self.failures)
        tstamp = time.asctime().split()[3] # Current time
        if color.has_colors(sys.stdout):
            progress = color.c("[%s]" % color.B(n))
            success = color.g("[%s]" % color.B("SUCCESS"))
            failure = color.r("[%s]" % color.B("FAILURE"))
            stderr = color.r("Stderr:")
            error = color.r(color.B(error))
        else:
            progress = "[%s]" % n
            success = "[SUCCESS]"
            failure = "[FAILURE]"
            stderr = "Stderr:"
        if self.port:
            host = '%s:%s' % (self.host, self.port)
        else:
            host = self.host
        if self.failures:
            print progress, tstamp, failure, host, error
        else:
            print progress, tstamp, success, host
        if self.outputbuffer:
            print self.outputbuffer,
        if self.errorbuffer:
            print stderr, self.errorbuffer,
        sys.stdout.flush()

<commit_msg>worked around an ssh misfeature with SSH_ASKPASS when DISPLAY is unset<commit_after>from errno import EINTR
from subprocess import Popen, PIPE
import askpass
import color
import os
import signal
import sys
import time
import traceback

BUFFER_SIZE = 1 << 16

class Task(object):
    """Starts a process and manages its input and output."""
    def __init__(self, host, port, cmd, opts, stdin=None):
        self.host = host
        self.port = port
        self.cmd = cmd

        self.proc = None
        self.writer = None
        self.timestamp = None
        self.failures = []
        self.killed = False
        self.inputbuffer = stdin
        self.outputbuffer = ''
        self.errorbuffer = ''

        self.stdin = None
        self.stdout = None
        self.stderr = None
        self.outfile = None
        self.errfile = None

        # Set options.
        self.verbose = opts.verbose
        try:
            self.print_out = bool(opts.print_out)
        except AttributeError:
            self.print_out = False
        try:
            self.inline = bool(opts.inline)
        except AttributeError:
            self.inline = False

    def start(self, nodenum, iomap, writer, askpass_socket=None):
        """Starts the process and registers files with the IOMap."""
        self.writer = writer

        if writer:
            self.outfile, self.errfile = writer.open_files(self.host)

        # Set up the environment.
        environ = dict(os.environ)
        environ['PSSH_NODENUM'] = str(nodenum)
        # Disable the GNOME pop-up password dialog and allow ssh to use
        # askpass.py to get a provided password.  If the module file is
        # askpass.pyc, we replace the extension.
        root, ext = os.path.splitext(os.path.abspath(askpass.__file__))
        environ['SSH_ASKPASS'] = '%s.py' % root
        if askpass_socket:
            environ['PSSH_ASKPASS_SOCKET'] = askpass_socket
        # Work around a mis-feature in ssh where it won't call SSH_ASKPASS
        # if DISPLAY is unset.
        if 'DISPLAY' not in environ:
            environ['DISPLAY'] = 'pssh-gibberish'

        # Create the subprocess.
        self.proc = Popen(self.cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE,
                close_fds=True, preexec_fn=os.setsid, env=environ)
        self.timestamp = time.time()
        if self.inputbuffer:
            self.stdin = self.proc.stdin
            iomap.register(self.stdin.fileno(), self.handle_stdin, write=True)
        else:
            self.proc.stdin.close()
        self.stdout = self.proc.stdout
        iomap.register(self.stdout.fileno(), self.handle_stdout, read=True)
        self.stderr = self.proc.stderr
        iomap.register(self.stderr.fileno(), self.handle_stderr, read=True)

    def _kill(self):
        """Signals the process to terminate."""
        if self.proc:
            os.kill(-self.proc.pid, signal.SIGKILL)
            self.killed = True

    def timedout(self):
        """Kills the process and registers a timeout error."""
        if not self.killed:
            self._kill()
            self.failures.append('Timed out')

    def interrupted(self):
        """Kills the process and registers an keyboard interrupt error."""
        if not self.killed:
            self._kill()
            self.failures.append('Interrupted')

    def cancel(self):
        """Stops a task that has not started."""
        self.failures.append('Cancelled')

    def elapsed(self):
        """Finds the time in seconds since the process was started."""
        return time.time() - self.timestamp

    def running(self):
        """Finds if the process has terminated and saves the return code."""
        if self.stdin or self.stdout or self.stderr:
            return True
        if self.proc:
            self.returncode = self.proc.poll()
            if self.returncode is None:
                if self.killed:
                    return False
                else:
                    return True
            else:
                if self.returncode < 0:
                    message = 'Killed by signal %s' % (-self.returncode)
                    self.failures.append(message)
                elif self.returncode > 0:
                    message = 'Exited with error code %s' % self.returncode
                    self.failures.append(message)
                self.proc = None
                return False

    def handle_stdin(self, fd, event, iomap):
        try:
            if self.inputbuffer:
                bytes_written = os.write(fd, self.inputbuffer)
                self.inputbuffer = self.inputbuffer[bytes_written:]
            else:
                self.close_stdin(iomap)
        except (OSError, IOError), e:
            if e.errno != EINTR:
                self.close_stdin(iomap)
                self.log_exception(e)

    def close_stdin(self, iomap):
        if self.stdin:
            iomap.unregister(self.stdin.fileno())
            self.stdin.close()
            self.stdin = None

    def handle_stdout(self, fd, event, iomap):
        try:
            buf = os.read(fd, BUFFER_SIZE)
            if buf:
                if self.inline:
                    self.outputbuffer += buf
                if self.outfile:
                    self.writer.write(self.outfile, buf)
                if self.print_out:
                    print '%s: %s' % (self.host, buf),
            else:
                self.close_stdout(iomap)
        except (OSError, IOError), e:
            if e.errno != EINTR:
                self.close_stdout(iomap)
                self.log_exception(e)

    def close_stdout(self, iomap):
        if self.stdout:
            iomap.unregister(self.stdout.fileno())
            self.stdout.close()
            self.stdout = None
        if self.outfile:
            self.writer.close(self.outfile)
            self.outfile = None

    def handle_stderr(self, fd, event, iomap):
        try:
            buf = os.read(fd, BUFFER_SIZE)
            if buf:
                if self.inline:
                    self.errorbuffer += buf
                if self.errfile:
                    self.writer.write(self.errfile, buf)
            else:
                self.close_stderr(iomap)
        except (OSError, IOError), e:
            if e.errno != EINTR:
                self.close_stderr(iomap)
                self.log_exception(e)

    def close_stderr(self, iomap):
        if self.stderr:
            iomap.unregister(self.stderr.fileno())
            self.stderr.close()
            self.stderr = None
        if self.errfile:
            self.writer.close(self.errfile)
            self.errfile = None

    def log_exception(self, e):
        """Saves a record of the most recent exception for error reporting."""
        if self.verbose:
            exc_type, exc_value, exc_traceback = sys.exc_info()
            exc = ("Exception: %s, %s, %s" %
                    (exc_type, exc_value, traceback.format_tb(exc_traceback)))
        else:
            exc = str(e)
        self.failures.append(exc)

    def report(self, n):
        """Pretty prints a status report after the Task completes."""
        error = ', '.join(self.failures)
        tstamp = time.asctime().split()[3] # Current time
        if color.has_colors(sys.stdout):
            progress = color.c("[%s]" % color.B(n))
            success = color.g("[%s]" % color.B("SUCCESS"))
            failure = color.r("[%s]" % color.B("FAILURE"))
            stderr = color.r("Stderr:")
            error = color.r(color.B(error))
        else:
            progress = "[%s]" % n
            success = "[SUCCESS]"
            failure = "[FAILURE]"
            stderr = "Stderr:"
        if self.port:
            host = '%s:%s' % (self.host, self.port)
        else:
            host = self.host
        if self.failures:
            print progress, tstamp, failure, host, error
        else:
            print progress, tstamp, success, host
        if self.outputbuffer:
            print self.outputbuffer,
        if self.errorbuffer:
            print stderr, self.errorbuffer,
        sys.stdout.flush()

<|endoftext|>