<commit_before>import string
import model

class Token:
    def __init__(self, value, lineno, offset):
        self.value = value
        self.line = lineno
        self.offset = offset

    def __bool__(self):
        return bool(self.value)

    def __str__(self):
        return self.value

    def __repr__(self):
        return 'Token("{}", {}, {})'.format(self.value, self.line, self.offset)

class Generator:
    OPERATORS = ("==", "!=", "<=", ">=", "||", "&&",
                 "+",  "-",  "*",  "/",  "%",  "<", ">", "!")
    PUNCTUATORS = (",", "{", "}", "(", ")", ";")

    def __init__(self, source):
        self.source = source
        self.position = 0
        self.line = 0
        self.offset = 0

    def peek_char(self, offset = 0):
        at = self.position + offset
        if at >= len(self.source):
            return ''
        return self.source[at]

    def peek_chars(self, count = 1):
        return ''.join(self.peek_char(i) for i in range(count))

    def read_char(self):
        char = self.peek_char()
        if not char:
            return char

        self.position += 1
        if char == '\n':
            self.line, self.offset = self.line + 1, 0
        else:
            self.line, self.offset = self.line, self.offset + 1

        return char

    def read_chars(self, count = 1):
        return ''.join(self.read_char() for _ in range(count))

    def extract_token(self, chars):
        line, offset = self.line, self.offset
        word = ""
        char = self.peek_char()
        while char and char in chars:
            word += self.read_char()
            char = self.peek_char()
        return Token(word, line, offset)

    def read_spaces(self):
        return self.extract_token(string.whitespace)

    def read_name(self):
        name_chars = string.ascii_letters + string.digits + '_'
        return self.extract_token(name_chars)

    def read_number(self):
        return self.extract_token(string.digits)

    def read_operator(self):
        token = self.extract_token('+-*/<>!=%&|')
        if token.value not in Generator.OPERATORS:
            raise Exception('Undefined operator "{}"'.format(token.value))
        return token

    def read_punctuator(self):
        token = Token(self.peek_char(), self.line, self.offset)
        self.read_char()
        return token

    def __next__(self):
        self.read_spaces()
        if not self.peek_char():
            raise StopIteration()

        char = self.peek_char()
        if char.isalpha() or char == '_':
            return self.read_name()

        if char.isdigit():
            return self.read_number()

        if char in Generator.PUNCTUATORS:
            return self.read_punctuator()

        return self.read_operator()

class Scanner:
    def __init__(self, filename):
        with open(filename, "r") as source:
            self.source = source.read()

    def __iter__(self):
        return Generator(self.source)

class Parser:
    KEYWORDS = ('if', 'def', 'read', 'write')
    PRECEDENCE = {
        '*': 6, '/': 6, '%': 6,
        '-': 5, '+': 5,
        '<': 4, '>': 4, '<=': 4, '>=': 4,
        '==': 3, '!=': 3,
        '&&': 2,
        '||': 1
    }
    LOWEST_PRECEDENCE = 1
    HIGHEST_PRECEDENCE = 6

    def peek_token(self, offset = 0):
        at = self.position + offset
        if at >= len(self.tokens):
            return Token("", -1, -1)
        return self.tokens[at]

    def drop_tokens(self, count = 1):
        self.position += count

    def ensure_token(self, value):
        token = self.peek_token()
        if token.value != value:
            raise Exception('Expected "{}" at {}:{} instead of "{}"'.format(
                value, token.line, token.offset, token.value))
        self.drop_tokens()

    def is_name(self, value):
        name_chars = string.ascii_letters + string.digits + '_'
        if value in self.KEYWORDS:
            return False
        if not value[0].isalpha():
            return False
        for char in value:
            if char not in name_chars:
                return False
        return True

    def is_unary(self, value):
        return value in ('-', '!')

    def is_number(self, value):
        return value.isdigit()

    def parse_name(self):
        token = self.peek_token()
        if not self.is_name(token.value):
            raise Exception('"{}" at {}:{} is not allowed name'.format(
                token.value, token.line, token.offset))
        self.drop_tokens()
        return token.value

    def parse_number(self):
        token = self.peek_token()
        if not self.is_number(token.value):
            raise Exception('"{}" at {}:{} is not a number'.format(
                token.value, token.line, token.offset))
        self.drop_tokens()
        return int(token.value)

    def parse_arguments(self):
        self.ensure_token('(');
        token = self.peek_token()
        names = []
        while token and token.value != ')':
            names.append(self.parse_name())
            token = self.peek_token()
            if token.value == ',':
                self.drop_tokens()
                token = self.peek_token()
        self.ensure_token(')')
        return names

    def parse_def(self):
        self.ensure_token('def')
        function_name = self.parse_name()
        arguments = self.parse_arguments()
        body = self.parse_block()
        function = model.Function(arguments, body)
        return model.FunctionDefinition(function_name, function)

    def parse_if(self):
        self.ensure_token('if')
        self.ensure_token('(')
        condition = self.parse_expression()
        self.ensure_token(')')
        if_body = self.parse_block()
        else_body = None
        if self.peek_token().value == 'else':
            self.drop_tokens()
            else_body = self.parse_block()
        return model.Conditional(condition, if_body, else_body)

    def parse_read(self):
        self.ensure_token('read')
        name = self.parse_name()
        return model.Read(name)

    def parse_print(self):
        self.ensure_token('print')
        expression = self.parse_expression()
        return model.Print(expression)

    def parse_primary(self):
        token = self.peek_token()
        if self.is_number(token.value):
            return model.Number(self.parse_number())
        if self.is_name(token.value):
            return model.Reference(self.parse_name())

        if token.value == '(':
            self.drop_tokens()
            expr = self.parse_expression()
            self.ensure_token(')')
            return expr
        raise Exception('Unexpected token "{}" at {}:{}'.format(
            token.value, token.line, token.offset))

    def parse_parameters(self):
        self.ensure_token('(')
        expressions = []
        token = self.peek_token()
        while token and token.value != ')':
            expressions.append(self.parse_expression())
            token = self.peek_token()
            if token.value == ',':
                self.drop_tokens()
                token = self.peek_token()
        self.ensure_token(')')
        return expressions

    def parse_unary(self):
        token = self.peek_token()
        if self.is_unary(token.value):
            self.drop_tokens()
            return model.UnaryOperation(token.value, self.parse_unary())

        primary = self.parse_primary()
        while self.peek_token().value == '(':
            parameters = self.parse_parameters()
            primary = model.FunctionCall(primary, parameters)
        return primary

    def parse_binary(self, min_precedence = LOWEST_PRECEDENCE):
        left = self.parse_unary()
        precedence = self.PRECEDENCE.get(self.peek_token().value, 0)
        while precedence >= min_precedence:
            operator = self.peek_token()
            while self.PRECEDENCE.get(operator.value, 0) == precedence:
                self.drop_tokens()
                right = self.parse_binary(precedence + 1)
                left = model.BinaryOperation(left, operator.value, right)
                operator = self.peek_token()
            precedence -= 1
        return left

    def parse_expression(self):
        return self.parse_binary()

    def parse_block(self, parentheses = True):
        if parentheses:
            self.ensure_token('{')
        statements = []
        token = self.peek_token()
        while token and token.value != '}':
            if token.value == 'def':
                statements.append(self.parse_def())
            elif token.value == 'if':
                statements.append(self.parse_if())
            else:
                if token.value == 'read':
                    statements.append(self.parse_read())
                elif token.value == 'print':
                    statements.append(self.parse_print())
                else:
                    statements.append(self.parse_expression())
                self.ensure_token(';')
            token = self.peek_token()
        if parentheses:
            self.ensure_token('}')
        return statements

    def parse(self, tokens):
        self.tokens = [token for token in tokens]
        self.position = 0
        program = self.parse_block(False)
        return model.Function([], program)

if __name__ == "__main__":
    import sys
    for filename in sys.argv[1:]:
        print(Parser().parse(Scanner(filename)))
<commit_msg>parse empty statement<commit_after>import string
import model

class Token:
    def __init__(self, value, lineno, offset):
        self.value = value
        self.line = lineno
        self.offset = offset

    def __bool__(self):
        return bool(self.value)

    def __str__(self):
        return self.value

    def __repr__(self):
        return 'Token("{}", {}, {})'.format(self.value, self.line, self.offset)

class Generator:
    OPERATORS = ("==", "!=", "<=", ">=", "||", "&&",
                 "+",  "-",  "*",  "/",  "%",  "<", ">", "!")
    PUNCTUATORS = (",", "{", "}", "(", ")", ";")

    def __init__(self, source):
        self.source = source
        self.position = 0
        self.line = 0
        self.offset = 0

    def peek_char(self, offset = 0):
        at = self.position + offset
        if at >= len(self.source):
            return ''
        return self.source[at]

    def peek_chars(self, count = 1):
        return ''.join(self.peek_char(i) for i in range(count))

    def read_char(self):
        char = self.peek_char()
        if not char:
            return char

        self.position += 1
        if char == '\n':
            self.line, self.offset = self.line + 1, 0
        else:
            self.line, self.offset = self.line, self.offset + 1

        return char

    def read_chars(self, count = 1):
        return ''.join(self.read_char() for _ in range(count))

    def extract_token(self, chars):
        line, offset = self.line, self.offset
        word = ""
        char = self.peek_char()
        while char and char in chars:
            word += self.read_char()
            char = self.peek_char()
        return Token(word, line, offset)

    def read_spaces(self):
        return self.extract_token(string.whitespace)

    def read_name(self):
        name_chars = string.ascii_letters + string.digits + '_'
        return self.extract_token(name_chars)

    def read_number(self):
        return self.extract_token(string.digits)

    def read_operator(self):
        token = self.extract_token('+-*/<>!=%&|')
        if token.value not in Generator.OPERATORS:
            raise Exception('Undefined operator "{}"'.format(token.value))
        return token

    def read_punctuator(self):
        token = Token(self.peek_char(), self.line, self.offset)
        self.read_char()
        return token

    def __next__(self):
        self.read_spaces()
        if not self.peek_char():
            raise StopIteration()

        char = self.peek_char()
        if char.isalpha() or char == '_':
            return self.read_name()

        if char.isdigit():
            return self.read_number()

        if char in Generator.PUNCTUATORS:
            return self.read_punctuator()

        return self.read_operator()

class Scanner:
    def __init__(self, filename):
        with open(filename, "r") as source:
            self.source = source.read()

    def __iter__(self):
        return Generator(self.source)

class Parser:
    KEYWORDS = ('if', 'def', 'read', 'write')
    PRECEDENCE = {
        '*': 6, '/': 6, '%': 6,
        '-': 5, '+': 5,
        '<': 4, '>': 4, '<=': 4, '>=': 4,
        '==': 3, '!=': 3,
        '&&': 2,
        '||': 1
    }
    LOWEST_PRECEDENCE = 1
    HIGHEST_PRECEDENCE = 6

    def peek_token(self, offset = 0):
        at = self.position + offset
        if at >= len(self.tokens):
            return Token("", -1, -1)
        return self.tokens[at]

    def drop_tokens(self, count = 1):
        self.position += count

    def ensure_token(self, value):
        token = self.peek_token()
        if token.value != value:
            raise Exception('Expected "{}" at {}:{} instead of "{}"'.format(
                value, token.line, token.offset, token.value))
        self.drop_tokens()

    def is_name(self, value):
        name_chars = string.ascii_letters + string.digits + '_'
        if value in self.KEYWORDS:
            return False
        if not value[0].isalpha():
            return False
        for char in value:
            if char not in name_chars:
                return False
        return True

    def is_unary(self, value):
        return value in ('-', '!')

    def is_number(self, value):
        return value.isdigit()

    def parse_name(self):
        token = self.peek_token()
        if not self.is_name(token.value):
            raise Exception('"{}" at {}:{} is not allowed name'.format(
                token.value, token.line, token.offset))
        self.drop_tokens()
        return token.value

    def parse_number(self):
        token = self.peek_token()
        if not self.is_number(token.value):
            raise Exception('"{}" at {}:{} is not a number'.format(
                token.value, token.line, token.offset))
        self.drop_tokens()
        return int(token.value)

    def parse_arguments(self):
        self.ensure_token('(');
        token = self.peek_token()
        names = []
        while token and token.value != ')':
            names.append(self.parse_name())
            token = self.peek_token()
            if token.value == ',':
                self.drop_tokens()
                token = self.peek_token()
        self.ensure_token(')')
        return names

    def parse_def(self):
        self.ensure_token('def')
        function_name = self.parse_name()
        arguments = self.parse_arguments()
        body = self.parse_block()
        function = model.Function(arguments, body)
        return model.FunctionDefinition(function_name, function)

    def parse_if(self):
        self.ensure_token('if')
        self.ensure_token('(')
        condition = self.parse_expression()
        self.ensure_token(')')
        if_body = self.parse_block()
        else_body = None
        if self.peek_token().value == 'else':
            self.drop_tokens()
            else_body = self.parse_block()
        return model.Conditional(condition, if_body, else_body)

    def parse_read(self):
        self.ensure_token('read')
        name = self.parse_name()
        return model.Read(name)

    def parse_print(self):
        self.ensure_token('print')
        expression = self.parse_expression()
        return model.Print(expression)

    def parse_primary(self):
        token = self.peek_token()
        if self.is_number(token.value):
            return model.Number(self.parse_number())
        if self.is_name(token.value):
            return model.Reference(self.parse_name())

        if token.value == '(':
            self.drop_tokens()
            expr = self.parse_expression()
            self.ensure_token(')')
            return expr
        raise Exception('Unexpected token "{}" at {}:{}'.format(
            token.value, token.line, token.offset))

    def parse_parameters(self):
        self.ensure_token('(')
        expressions = []
        token = self.peek_token()
        while token and token.value != ')':
            expressions.append(self.parse_expression())
            token = self.peek_token()
            if token.value == ',':
                self.drop_tokens()
                token = self.peek_token()
        self.ensure_token(')')
        return expressions

    def parse_unary(self):
        token = self.peek_token()
        if self.is_unary(token.value):
            self.drop_tokens()
            return model.UnaryOperation(token.value, self.parse_unary())

        primary = self.parse_primary()
        while self.peek_token().value == '(':
            parameters = self.parse_parameters()
            primary = model.FunctionCall(primary, parameters)
        return primary

    def parse_binary(self, min_precedence = LOWEST_PRECEDENCE):
        left = self.parse_unary()
        precedence = self.PRECEDENCE.get(self.peek_token().value, 0)
        while precedence >= min_precedence:
            operator = self.peek_token()
            while self.PRECEDENCE.get(operator.value, 0) == precedence:
                self.drop_tokens()
                right = self.parse_binary(precedence + 1)
                left = model.BinaryOperation(left, operator.value, right)
                operator = self.peek_token()
            precedence -= 1
        return left

    def parse_expression(self):
        return self.parse_binary()

    def parse_block(self, parentheses = True):
        if parentheses:
            self.ensure_token('{')
        statements = []
        token = self.peek_token()
        while token and token.value != '}':
            if token.value == 'def':
                statements.append(self.parse_def())
            elif token.value == 'if':
                statements.append(self.parse_if())
            elif token.value == ';':
                self.drop_tokens()
            else:
                if token.value == 'read':
                    statements.append(self.parse_read())
                elif token.value == 'print':
                    statements.append(self.parse_print())
                else:
                    statements.append(self.parse_expression())
                self.ensure_token(';')
            token = self.peek_token()
        if parentheses:
            self.ensure_token('}')
        return statements

    def parse(self, tokens):
        self.tokens = [token for token in tokens]
        self.position = 0
        program = self.parse_block(False)
        return model.Function([], program)

if __name__ == "__main__":
    import sys
    for filename in sys.argv[1:]:
        print(Parser().parse(Scanner(filename)))
<|endoftext|>


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

import pycouchdb
import requests
import sys
import xml.etree.ElementTree as ET

def usage():
    print("usage: ./parser.py <NETXML-FILE>")

def class mantis:
    __db = pycouchdb.client.Database()

    def __init__(self):
        # assume default config for couchdb
        server = pycouchdb.Server()
        try:
            server.info()
        except requests.exceptions.ConnectionError:
            sys.stderr.write("connecting to server failed\n")
            sys.exit(1)
        self.__db =

    def extract_ssid_data(rawdata):
        """Extract relevant SSID data from given XML-node.

        This data contains per SSID: maximum data rate, encryption modes and
        ESSID.
        """

        if not isinstance(rawdata,ET.Element):
            raise TypeError("given rawdata not an ElementTree-element")

        ssiddata = {}
        maxrate = rawdata.find('max-rate').text
        ssiddata['max-rate'] = float(maxrate)

        encryption_modes = rawdata.findall('encryption')
        enctxt = []
        for enc in encryption_modes:
            enctxt.append(enc.text)
            ssiddata['encryption'] = enctxt

        ssiddata['essid'] = rawdata.find('essid').text

        return ssiddata

    def extract_snr_info(rawdata):
        """Extract relevant radio signal data from given XML-node.

        This data contains minimum and maximum levels of the signal and noise.
        """
        if not isinstance(rawdata,ET.Element):
            raise TypeError("given rawdata not an ElementTree-element")

        snrinfo = {}
        data = rawdata.find('min_signal_dbm').text
        snrinfo['min_signal_dbm'] = int(data)
        data = rawdata.find('min_noise_dbm').text
        snrinfo['min_noise_dbm'] = int(data)
        data = rawdata.find('min_signal_rssi').text
        snrinfo['min_signal_rssi'] = int(data)
        data = rawdata.find('min_noise_rssi').text
        snrinfo['min_noise_rssi'] = int(data)

        data = rawdata.find('max_signal_dbm').text
        snrinfo['max_signal_dbm'] = int(data)
        data = rawdata.find('max_noise_dbm').text
        snrinfo['max_noise_dbm'] = int(data)
        data = rawdata.find('max_signal_rssi').text
        snrinfo['max_signal_rssi'] = int(data)
        data = rawdata.find('max_noise_rssi').text
        snrinfo['max_noise_rssi'] = int(data)

        return snrinfo

    def extract_gps_info(rawdata):
        """Extract relevant GPS data from given XML-node.

        This data contains the coordinates of the minimum, maximum and peak
        signal level.
        """

        if not isinstance(rawdata,ET.Element):
            raise TypeError("given rawdata not an ElementTree-element")

        gpsinfo = {}
        data = rawdata.find('min-lat').text
        gpsinfo['min-lat'] = float(data)
        data = rawdata.find('min-lon').text
        gpsinfo['min-lon'] = float(data)
        data = rawdata.find('max-lat').text
        gpsinfo['max-lat'] = float(data)
        data = rawdata.find('max-lon').text
        gpsinfo['max-lon'] = float(data)
        data = rawdata.find('peak-lat').text
        gpsinfo['peak-lat'] = float(data)
        data = rawdata.find('peak-lon').text
        gpsinfo['peak-lon'] = float(data)

        return gpsinfo

    def deflate(db):
        """Remove all diplicate entries from the database."""
        mapfunction = "function(doc) {\
        var bssid, essid, uid, docid;\
        if (doc.ssid && doc.bssid) {\
            uid = [doc.bssid];\
            docid = [doc['_id'], doc['_rev']];\
            names = [];\
            for (index in doc.ssid) {\
                names.push(doc.ssid[index]['essid']);\
            }\
            uid.push(names);\
            emit(uid, docid);\
        }\
        }"

        allentries = list(db.temporary_query(mapfunction))
        print(len(allentries))

# only call if executed as script
if __name__ == '__main__':
    if 2 != len(sys.argv):
        usage()
        sys.exit(1)

    # assume default config for couchdb
    server = pycouchdb.Server()
    try:
        server.info()
    except requests.exceptions.ConnectionError:
        sys.stderr.write("connecting to server failed\n")
        sys.exit(1)
    # assume databasename, if database does not exists, create it
    dbname = "wifinetworks"
    try:
        db = server.database(dbname)
    except pycouchdb.exceptions.NotFound:
        print("database '"+dbname+"' does not exist, creating one ...")
        server.create(dbname)
        db = server.database(dbname)
    newdoccounter = 0

    tree = ET.ElementTree()
    try:
        tree = ET.parse(sys.argv[1])
    except FileNotFoundError:
        sys.stderr.write("opening file failed\n")
    root = tree.getroot()
    for network in tree.findall('wireless-network'):
        networkdata = {}
        # only handle fixed networks, ignore probes etc.
        if "infrastructure" == network.attrib['type']:
            # extract the data
            ssids = network.findall('SSID')
            ssiddata = []
            for ssidnode in ssids:
                ssiddata.append( extract_ssid_data(ssidnode) )
            networkdata['ssid'] = ssiddata
            networkdata['bssid'] = network.find('BSSID').text
            networkdata['snr-info'] = extract_snr_info( network.find('snr-info') )
            networkdata['gps-info'] = extract_gps_info( network.find('gps-info') )
            # push it into couchdb
            doc = db.save(networkdata)
            newdoccounter += 1
    print(str(newdoccounter) + " new networks added")
    deflate(db)
<commit_msg>db creation moved, typo fixed, see issue #4<commit_after>#!/usr/bin/env python3

import pycouchdb
import requests
import sys
import xml.etree.ElementTree as ET

def usage():
    print("usage: ./parser.py <NETXML-FILE>")

class Mantis:
    __db = pycouchdb.client.Database()

    def __init__(self,dbname = "wifinetworks"):
        # assume default config for couchdb
        server = pycouchdb.Server()
        try:
            server.info()
        except requests.exceptions.ConnectionError:
            sys.stderr.write("connecting to server failed\n")
            sys.exit(1)

        # assume databasename, if database does not exists, create it
        dbname = "wifinetworks"
        try:
            db = server.database(dbname)
        except pycouchdb.exceptions.NotFound:
            server.create(dbname)
        self.__db = server.database(dbname)

    def extract_ssid_data(rawdata):
        """Extract relevant SSID data from given XML-node.

        This data contains per SSID: maximum data rate, encryption modes and
        ESSID.
        """

        if not isinstance(rawdata,ET.Element):
            raise TypeError("given rawdata not an ElementTree-element")

        ssiddata = {}
        maxrate = rawdata.find('max-rate').text
        ssiddata['max-rate'] = float(maxrate)

        encryption_modes = rawdata.findall('encryption')
        enctxt = []
        for enc in encryption_modes:
            enctxt.append(enc.text)
            ssiddata['encryption'] = enctxt

        ssiddata['essid'] = rawdata.find('essid').text

        return ssiddata

    def extract_snr_info(rawdata):
        """Extract relevant radio signal data from given XML-node.

        This data contains minimum and maximum levels of the signal and noise.
        """
        if not isinstance(rawdata,ET.Element):
            raise TypeError("given rawdata not an ElementTree-element")

        snrinfo = {}
        data = rawdata.find('min_signal_dbm').text
        snrinfo['min_signal_dbm'] = int(data)
        data = rawdata.find('min_noise_dbm').text
        snrinfo['min_noise_dbm'] = int(data)
        data = rawdata.find('min_signal_rssi').text
        snrinfo['min_signal_rssi'] = int(data)
        data = rawdata.find('min_noise_rssi').text
        snrinfo['min_noise_rssi'] = int(data)

        data = rawdata.find('max_signal_dbm').text
        snrinfo['max_signal_dbm'] = int(data)
        data = rawdata.find('max_noise_dbm').text
        snrinfo['max_noise_dbm'] = int(data)
        data = rawdata.find('max_signal_rssi').text
        snrinfo['max_signal_rssi'] = int(data)
        data = rawdata.find('max_noise_rssi').text
        snrinfo['max_noise_rssi'] = int(data)

        return snrinfo

    def extract_gps_info(rawdata):
        """Extract relevant GPS data from given XML-node.

        This data contains the coordinates of the minimum, maximum and peak
        signal level.
        """

        if not isinstance(rawdata,ET.Element):
            raise TypeError("given rawdata not an ElementTree-element")

        gpsinfo = {}
        data = rawdata.find('min-lat').text
        gpsinfo['min-lat'] = float(data)
        data = rawdata.find('min-lon').text
        gpsinfo['min-lon'] = float(data)
        data = rawdata.find('max-lat').text
        gpsinfo['max-lat'] = float(data)
        data = rawdata.find('max-lon').text
        gpsinfo['max-lon'] = float(data)
        data = rawdata.find('peak-lat').text
        gpsinfo['peak-lat'] = float(data)
        data = rawdata.find('peak-lon').text
        gpsinfo['peak-lon'] = float(data)

        return gpsinfo

    def deflate(db):
        """Remove all diplicate entries from the database."""
        mapfunction = "function(doc) {\
        var bssid, essid, uid, docid;\
        if (doc.ssid && doc.bssid) {\
            uid = [doc.bssid];\
            docid = [doc['_id'], doc['_rev']];\
            names = [];\
            for (index in doc.ssid) {\
                names.push(doc.ssid[index]['essid']);\
            }\
            uid.push(names);\
            emit(uid, docid);\
        }\
        }"

        allentries = list(db.temporary_query(mapfunction))
        print(len(allentries))

# only call if executed as script
if __name__ == '__main__':
    if 2 != len(sys.argv):
        usage()
        sys.exit(1)

    tree = ET.ElementTree()
    try:
        tree = ET.parse(sys.argv[1])
    except FileNotFoundError:
        sys.stderr.write("opening file failed\n")
    root = tree.getroot()
    for network in tree.findall('wireless-network'):
        networkdata = {}
        # only handle fixed networks, ignore probes etc.
        if "infrastructure" == network.attrib['type']:
            # extract the data
            ssids = network.findall('SSID')
            ssiddata = []
            for ssidnode in ssids:
                ssiddata.append( extract_ssid_data(ssidnode) )
            networkdata['ssid'] = ssiddata
            networkdata['bssid'] = network.find('BSSID').text
            networkdata['snr-info'] = extract_snr_info( network.find('snr-info') )
            networkdata['gps-info'] = extract_gps_info( network.find('gps-info') )
            # push it into couchdb
            doc = db.save(networkdata)
            newdoccounter += 1
    print(str(newdoccounter) + " new networks added")
    deflate(db)
<|endoftext|>


<commit_before>import sys
import numpy as np
import cv2
import os
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QPixmap, QImage
from PyQt5.QtWidgets import (QWidget, QLabel, QHBoxLayout, QVBoxLayout,
                             QApplication, QPushButton, QSlider,
                             QFileDialog)

class Image(QWidget):
    """Common base for the images"""

    def __init__(self, name):
        super(Image, self).__init__()

        self.setMinimumSize(500,500)

        # Label for the image
        self.name_lbl = QLabel(name)

        # Label (frame) where the original image will be located, with true scaling and maximum size
        self.frame_lbl = QLabel(self)
        self.frame_lbl.setMinimumSize(400,400)
        self.frame_lbl.setScaledContents(False)

        self.createLayout()

    def createLayout(self):
        self.v_lay = QVBoxLayout()

        self.v_lay.addWidget(self.name_lbl)
        self.v_lay.addWidget(self.frame_lbl)

    def updateImage(self, opencv_rgb_image):

        self.cv_img_rgb = opencv_rgb_image

        height, width, channel = self.cv_img_rgb.shape
        bytesPerLine = 3 * width
        self.q_image = QImage(self.cv_img_rgb.data, width, height, bytesPerLine, QImage.Format_RGB888)

        self.frame_lbl.setPixmap(QPixmap.fromImage(self.q_image))

class Filter(QWidget):
    """Common base class for all filters"""
    defaultK = 3
    filterCount = 0

    def __init__(self, name, minValue, maxValue, init, parent=None):
        super(Filter, self).__init__(parent=parent)

        self.setMaximumHeight(65)

        self.k = init

        # Increase the number of filters created
        Filter.filterCount += 1

        # Variable for the slider/label layout
        self.lay = QHBoxLayout(self)

        # Variable for the constant of the OpenCV filter
        self.k = self.defaultK

        # Label for the slider
        self.k_lbl = QLabel(str(self.k))

        # Name for the slider
        self.name_lbl = QLabel(name)

        # Set default parameters
        self.setParameters(minValue, maxValue)

        # Create delete button
        self.delete_btn = QPushButton('X')

        # Adds the slider and it's label to the layout
        self.createLayout()

        # Function sending the slider signal to the processing function
        self.thresh_sld.valueChanged.connect(self.changeValue)

    def setParameters(self, minValue, maxValue):
        # Creates the slider for the OpenCV filter, with min, max, default and step values
        self.thresh_sld = QSlider(Qt.Horizontal, self)
        self.thresh_sld.setFocusPolicy(Qt.NoFocus)
        self.thresh_sld.setMinimum(minValue)
        self.thresh_sld.setMaximum(maxValue)
        self.thresh_sld.setValue(self.k)
        self.thresh_sld.setSingleStep(2)

    def createLayout(self):
        # Adds the slider and its label to the bottom of the main layout

        self.lay.addWidget(self.name_lbl)
        self.lay.addWidget(self.k_lbl)
        self.lay.addWidget(self.thresh_sld)
        self.lay.addWidget(self.delete_btn)

    def changeValue(self, value):
        # Function for setting the value of k1

        if value % 2 == 1:
            self.k = value
        else:
            self.k = value + 1

        self.thresh_sld.setValue(self.k)
        self.k_lbl.setText(str(self.k))
        self.parent().updateImages()

    def resetValue(self):
        # Resets the K value to it's default
        self.changeValue(self.defaultK)

class MainWindow(QWidget):

    def __init__(self):
        super(MainWindow, self).__init__()

        self.initUI()
        self.filters = [10]

    def initUI(self):

        self.filter1 = Filter("Hue", 31, 225, 31)
        self.filter2 = Filter("Gaussian", 3, 51, 3)
        self.filter3 = Filter("TopHat", 3, 51, 3)
        self.filter4 = Filter("Erosion", 3, 51, 3)
        self.filter5 = Filter("ValueMin", 0, 255, 0)
        self.filter6 = Filter("ValueMax", 0, 255, 0)

        self.original_image = Image('Original')
        self.processed_image = Image('Processed')
        self.setBackground()

        self.createImagesLayout()
        self.createButtons()
        self.createMainLayout()

    def createNewFilter(self):
        self.filters[0] = Filter('Filtro' + str(Filter.filterCount), 3, 51, 3)
        print Filter.filterCount

    def createImagesLayout(self):
        # Horizontal layout for the two images
        self.h_img_lay = QHBoxLayout()
        self.h_img_lay.addStretch()
        self.h_img_lay.addWidget(self.original_image)
        self.h_img_lay.addWidget(self.processed_image)
        self.h_img_lay.addStretch()

    def createMainLayout(self):
        # Creates the main layout (vertical)
        self.v_main_lay = QVBoxLayout()
        # Adds the images horizontal layout to the main layout
        self.v_main_lay.addLayout(self.h_img_lay)
        # Adds the sliders and their labels to the bottom of the main layout
        self.v_main_lay.addWidget(self.filter1)
        self.v_main_lay.addWidget(self.filter2)
        self.v_main_lay.addWidget(self.filter3)
        self.v_main_lay.addWidget(self.filter4)
        self.v_main_lay.addWidget(self.filter5)
        self.v_main_lay.addWidget(self.filter6)
        # Adds the buttons horizontal layout to the bottom of the main layout
        self.v_main_lay.addLayout(self.h_btn_lay)
        # Sets the main layout
        self.setLayout(self.v_main_lay)

        # Sets the geometry, position, window title and window default mode
        self.setGeometry(500, 100, 0, 0)
        self.setWindowTitle('Image processing')
        self.show()

    def createButtons(self):

        # Button for adding filters
        self.add_filter_btn = QPushButton('Add filter')
        self.add_filter_btn.clicked.connect(self.createNewFilter)

        # Button for selecting image
        self.select_image_btn = QPushButton('Select image')
        self.select_image_btn.clicked.connect(self.openImage)

        # Button for cleaning image
        self.clean_image_btn = QPushButton('Save image')
        self.clean_image_btn.clicked.connect(self.saveImages)

        # Horizontal layout for the buttons
        self.h_btn_lay = QHBoxLayout()
        self.h_btn_lay.addStretch(1)
        self.h_btn_lay.addWidget(self.add_filter_btn)
        self.h_btn_lay.addWidget(self.select_image_btn)
        self.h_btn_lay.addWidget(self.clean_image_btn)
        self.h_btn_lay.addStretch(1)

    def openImage(self):
        # Function for selecting the original image

        filter = "Images (*.png *.jpg)"
        image_path, _ = QFileDialog.getOpenFileName(self, 'Open image', 'Desktop', filter)
        self.path = image_path

        cv_img_bgr = cv2.imread(self.path)
        cv_img_rgb = cv2.cvtColor(cv_img_bgr, cv2.COLOR_BGR2RGB)

        self.original_image.updateImage(cv_img_rgb)
        self.updateImages()

    def saveImages(self):
        # Function for saving the processed image

        filter = "Images (*.png *.jpg)"

        image_path_orig, _ = QFileDialog.getSaveFileName(self, filter=filter)
        self.cv_img_processed_bgr = cv2.cvtColor(self.cv_img_processed_rgb, cv2.COLOR_RGB2BGR)
        cv2.imwrite(image_path_orig, self.cv_img_processed_bgr)

        image_path_proc, _ = QFileDialog.getSaveFileName(self, filter=filter)
        self.orig_img_calculated_bgr = cv2.cvtColor(self.orig_img_calculated_rgb, cv2.COLOR_RGB2BGR)
        cv2.imwrite(image_path_proc, self.orig_img_calculated_bgr)

    def updateImages(self):
        self.processImage()
        self.calculateOriginal()

    def processImage(self):
        # Function that processes the image using the OpenCV methods

        # Open the image
        self.cv_img_bgr = cv2.imread(self.path)
        self.cv_img_rgb = cv2.cvtColor(self.cv_img_bgr, cv2.COLOR_BGR2RGB)

        # Convert image to HSV
        cv_img_hsv = cv2.cvtColor(self.cv_img_bgr, cv2.COLOR_BGR2HSV)

        # Define range of desired color in HSV
        lower_color = np.array([self.filter1.k - 30, 0, self.filter5.k])
        upper_color = np.array([self.filter1.k + 30, 255, self.filter6.k])

        hsv_mask = cv2.inRange(cv_img_hsv, lower_color, upper_color)

        # Convert image to grayscale
        cv_img_gray = cv2.cvtColor(self.cv_img_bgr, cv2.COLOR_BGR2GRAY)

        # Apply blur
        cv_img_blur = cv2.GaussianBlur(hsv_mask, (self.filter2.k, self.filter2.k), 0)

        # Apply Top Hat
        opening_kernel = np.ones((self.filter3.k, self.filter3.k), np.uint8)
        cv_img_open = cv2.morphologyEx(cv_img_blur, cv2.MORPH_TOPHAT, opening_kernel)

        # Apply erosion
        erosion_kernel = np.ones((self.filter4.k, self.filter4.k), np.uint8)
        cv_img_eros = cv2.morphologyEx(cv_img_open, cv2.MORPH_OPEN, erosion_kernel)

        # Last image phase
        self.cv_img_processed_gray = cv_img_eros

        # Convert image to RGB
        self.cv_img_processed_rgb = cv2.cvtColor(self.cv_img_processed_gray, cv2.COLOR_GRAY2RGB)

        # Updates the processed image frame
        self.processed_image.updateImage(self.cv_img_processed_rgb)

    def calculateOriginal(self):
        # Update original image with the contours

        self.orig_img_calculated_rgb = self.cv_img_rgb

        # Find contours
        img, contours, hierarchy = cv2.findContours(self.cv_img_processed_gray, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)

        # Apply contours
        cv2.drawContours(self.orig_img_calculated_rgb, contours, -1, (0, 0, 255), 3)

        self.original_image.updateImage(self.orig_img_calculated_rgb)

    def setBackground(self):
        script_path = os.path.basename(os.path.dirname(os.path.realpath(__file__)))
        MainWindow.path = os.path.dirname(script_path) + 'default_BG.jpg'

        cv_img_bgr = cv2.imread(self.path)
        cv_img_rgb = cv2.cvtColor(cv_img_bgr, cv2.COLOR_BGR2RGB)

        self.original_image.updateImage(cv_img_rgb)
        self.processed_image.updateImage(cv_img_rgb)

if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    sys.exit(app.exec_())<commit_msg>Function for new filters implemented<commit_after>import sys
import numpy as np
import cv2
import os
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QPixmap, QImage
from PyQt5.QtWidgets import (QWidget, QLabel, QHBoxLayout, QVBoxLayout,
                             QApplication, QPushButton, QSlider,
                             QFileDialog)

class Image(QWidget):
    """Common base for the images"""

    def __init__(self, name):
        super(Image, self).__init__()

        self.setMinimumSize(500,500)

        # Label for the image
        self.name_lbl = QLabel(name)

        # Label (frame) where the original image will be located, with true scaling and maximum size
        self.frame_lbl = QLabel(self)
        self.frame_lbl.setMinimumSize(400,400)
        self.frame_lbl.setScaledContents(False)

        self.createLayout()

    def createLayout(self):
        self.v_lay = QVBoxLayout()

        self.v_lay.addWidget(self.name_lbl)
        self.v_lay.addWidget(self.frame_lbl)

    def updateImage(self, opencv_rgb_image):

        self.cv_img_rgb = opencv_rgb_image

        height, width, channel = self.cv_img_rgb.shape
        bytesPerLine = 3 * width
        self.q_image = QImage(self.cv_img_rgb.data, width, height, bytesPerLine, QImage.Format_RGB888)

        self.frame_lbl.setPixmap(QPixmap.fromImage(self.q_image))

class Filter(QWidget):
    """Common base class for all filters"""
    defaultK = 3
    filterCount = 0

    def __init__(self, name, minValue, maxValue, init, parent=None):
        super(Filter, self).__init__(parent=parent)

        self.setMaximumHeight(65)

        self.k = init

        # Increase the number of filters created
        Filter.filterCount += 1

        # Variable for the slider/label layout
        self.lay = QHBoxLayout(self)

        # Variable for the constant of the OpenCV filter
        self.k = self.defaultK

        # Label for the slider
        self.k_lbl = QLabel(str(self.k))

        # Name for the slider
        self.name_lbl = QLabel(name)

        # Set default parameters
        self.setParameters(minValue, maxValue)

        # Create delete button
        self.delete_btn = QPushButton('X')

        # Adds the slider and it's label to the layout
        self.createLayout()

        # Function sending the slider signal to the processing function
        self.thresh_sld.valueChanged.connect(self.changeValue)

    def setParameters(self, minValue, maxValue):
        # Creates the slider for the OpenCV filter, with min, max, default and step values
        self.thresh_sld = QSlider(Qt.Horizontal, self)
        self.thresh_sld.setFocusPolicy(Qt.NoFocus)
        self.thresh_sld.setMinimum(minValue)
        self.thresh_sld.setMaximum(maxValue)
        self.thresh_sld.setValue(self.k)
        self.thresh_sld.setSingleStep(2)

    def createLayout(self):
        # Adds the slider and its label to the bottom of the main layout

        self.lay.addWidget(self.name_lbl)
        self.lay.addWidget(self.k_lbl)
        self.lay.addWidget(self.thresh_sld)
        self.lay.addWidget(self.delete_btn)

    def changeValue(self, value):
        # Function for setting the value of k1

        if value % 2 == 1:
            self.k = value
        else:
            self.k = value + 1

        self.thresh_sld.setValue(self.k)
        self.k_lbl.setText(str(self.k))
        self.parent().updateImages()

    def resetValue(self):
        # Resets the K value to it's default
        self.changeValue(self.defaultK)

class MainWindow(QWidget):

    def __init__(self):
        super(MainWindow, self).__init__()

        self.initUI()
        self.filters = [10]

    def initUI(self):

        self.filter1 = Filter("Hue", 31, 225, 31)
        self.filter2 = Filter("Gaussian", 3, 51, 3)
        self.filter3 = Filter("TopHat", 3, 51, 3)
        self.filter4 = Filter("Erosion", 3, 51, 3)
        self.filter5 = Filter("ValueMin", 0, 255, 0)
        self.filter6 = Filter("ValueMax", 0, 255, 0)

        self.original_image = Image('Original')
        self.processed_image = Image('Processed')
        self.setBackground()

        self.createImagesLayout()
        self.createFiltersLayout()
        self.createButtons()
        self.createMainLayout()

    def createNewFilter(self):
        self.filters[0] = Filter('Filtro' + str(Filter.filterCount), 3, 51, 3)
        self.v_filters_lay.addWidget(self.filters[0])
        print Filter.filterCount

    def createImagesLayout(self):
        # Horizontal layout for the two images
        self.h_img_lay = QHBoxLayout()
        self.h_img_lay.addStretch()
        self.h_img_lay.addWidget(self.original_image)
        self.h_img_lay.addWidget(self.processed_image)
        self.h_img_lay.addStretch()

    def createFiltersLayout(self):
        self.v_filters_lay = QVBoxLayout()
        self.v_filters_lay.addWidget(self.filter1)
        self.v_filters_lay.addWidget(self.filter2)
        self.v_filters_lay.addWidget(self.filter3)
        self.v_filters_lay.addWidget(self.filter4)
        self.v_filters_lay.addWidget(self.filter5)
        self.v_filters_lay.addWidget(self.filter6)

    def createMainLayout(self):
        # Creates the main layout (vertical)
        self.v_main_lay = QVBoxLayout()
        # Adds the images horizontal layout to the main layout
        self.v_main_lay.addLayout(self.h_img_lay)
        # Adds the sliders and their labels to the bottom of the main layout
        self.v_main_lay.addLayout(self.v_filters_lay)
        # Adds the buttons horizontal layout to the bottom of the main layout
        self.v_main_lay.addLayout(self.h_btn_lay)
        # Sets the main layout
        self.setLayout(self.v_main_lay)

        # Sets the geometry, position, window title and window default mode
        self.setGeometry(500, 100, 0, 0)
        self.setWindowTitle('Image processing')
        self.show()

    def createButtons(self):

        # Button for adding filters
        self.add_filter_btn = QPushButton('Add filter')
        self.add_filter_btn.clicked.connect(self.createNewFilter)

        # Button for selecting image
        self.select_image_btn = QPushButton('Select image')
        self.select_image_btn.clicked.connect(self.openImage)

        # Button for cleaning image
        self.clean_image_btn = QPushButton('Save image')
        self.clean_image_btn.clicked.connect(self.saveImages)

        # Horizontal layout for the buttons
        self.h_btn_lay = QHBoxLayout()
        self.h_btn_lay.addStretch(1)
        self.h_btn_lay.addWidget(self.add_filter_btn)
        self.h_btn_lay.addWidget(self.select_image_btn)
        self.h_btn_lay.addWidget(self.clean_image_btn)
        self.h_btn_lay.addStretch(1)

    def openImage(self):
        # Function for selecting the original image

        filter = "Images (*.png *.jpg)"
        image_path, _ = QFileDialog.getOpenFileName(self, 'Open image', 'Desktop', filter)
        self.path = image_path

        cv_img_bgr = cv2.imread(self.path)
        cv_img_rgb = cv2.cvtColor(cv_img_bgr, cv2.COLOR_BGR2RGB)

        self.original_image.updateImage(cv_img_rgb)
        self.updateImages()

    def saveImages(self):
        # Function for saving the processed image

        filter = "Images (*.png *.jpg)"

        image_path_orig, _ = QFileDialog.getSaveFileName(self, filter=filter)
        self.cv_img_processed_bgr = cv2.cvtColor(self.cv_img_processed_rgb, cv2.COLOR_RGB2BGR)
        cv2.imwrite(image_path_orig, self.cv_img_processed_bgr)

        image_path_proc, _ = QFileDialog.getSaveFileName(self, filter=filter)
        self.orig_img_calculated_bgr = cv2.cvtColor(self.orig_img_calculated_rgb, cv2.COLOR_RGB2BGR)
        cv2.imwrite(image_path_proc, self.orig_img_calculated_bgr)

    def updateImages(self):
        self.processImage()
        self.calculateOriginal()

    def processImage(self):
        # Function that processes the image using the OpenCV methods

        # Open the image
        self.cv_img_bgr = cv2.imread(self.path)
        self.cv_img_rgb = cv2.cvtColor(self.cv_img_bgr, cv2.COLOR_BGR2RGB)

        # Convert image to HSV
        cv_img_hsv = cv2.cvtColor(self.cv_img_bgr, cv2.COLOR_BGR2HSV)

        # Define range of desired color in HSV
        lower_color = np.array([self.filter1.k - 30, 0, self.filter5.k])
        upper_color = np.array([self.filter1.k + 30, 255, self.filter6.k])

        hsv_mask = cv2.inRange(cv_img_hsv, lower_color, upper_color)

        # Convert image to grayscale
        cv_img_gray = cv2.cvtColor(self.cv_img_bgr, cv2.COLOR_BGR2GRAY)

        # Apply blur
        cv_img_blur = cv2.GaussianBlur(hsv_mask, (self.filter2.k, self.filter2.k), 0)

        # Apply Top Hat
        opening_kernel = np.ones((self.filter3.k, self.filter3.k), np.uint8)
        cv_img_open = cv2.morphologyEx(cv_img_blur, cv2.MORPH_TOPHAT, opening_kernel)

        # Apply erosion
        erosion_kernel = np.ones((self.filter4.k, self.filter4.k), np.uint8)
        cv_img_eros = cv2.morphologyEx(cv_img_open, cv2.MORPH_OPEN, erosion_kernel)

        # Last image phase
        self.cv_img_processed_gray = cv_img_eros

        # Convert image to RGB
        self.cv_img_processed_rgb = cv2.cvtColor(self.cv_img_processed_gray, cv2.COLOR_GRAY2RGB)

        # Updates the processed image frame
        self.processed_image.updateImage(self.cv_img_processed_rgb)

    def calculateOriginal(self):
        # Update original image with the contours

        self.orig_img_calculated_rgb = self.cv_img_rgb

        # Find contours
        img, contours, hierarchy = cv2.findContours(self.cv_img_processed_gray, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)

        # Apply contours
        cv2.drawContours(self.orig_img_calculated_rgb, contours, -1, (0, 0, 255), 3)

        self.original_image.updateImage(self.orig_img_calculated_rgb)

    def setBackground(self):
        script_path = os.path.basename(os.path.dirname(os.path.realpath(__file__)))
        MainWindow.path = os.path.dirname(script_path) + 'default_BG.jpg'

        cv_img_bgr = cv2.imread(self.path)
        cv_img_rgb = cv2.cvtColor(cv_img_bgr, cv2.COLOR_BGR2RGB)

        self.original_image.updateImage(cv_img_rgb)
        self.processed_image.updateImage(cv_img_rgb)

if __name__ == '__main__':
    app = QApplication(sys.argv)
    window = MainWindow()
    sys.exit(app.exec_())<|endoftext|>


<commit_before>"""
utils.py - Simple utilility funtions used in pyquante2.

 Copyright (c) 2004, Richard P. Muller. All Rights Reserved.

 PyQuante version 2.0 and later is covered by the GPL
 license. Please see the file LICENSE that is part of this
 distribution.
"""
import numpy as np
from math import factorial,gamma,lgamma
from itertools import combinations_with_replacement

def pairs(it): return combinations_with_replacement(it,2)

def isnear(a,b,tol=1e-9): return abs(a-b)<tol

def fact2(n):
    """
    fact2(n) - n!!, double factorial of n
    >>> fact2(0)
    1
    >>> fact2(1)
    1
    >>> fact2(3)
    3
    >>> fact2(8)
    384
    >>> fact2(-1)
    1
    """
    return reduce(int.__mul__,xrange(n,0,-2),1)

def norm2(a): return np.dot(a,a)

def binomial(n,k):
    """
    Binomial coefficient
    >>> binomial(5,2)
    10
    >>> binomial(10,5)
    252
    """
    if n==k: return 1
    assert n>k, "Attempting to call binomial(%d,%d)" % (n,k)
    return factorial(n)/(factorial(k)*factorial(n-k))

def Fgamma(m,x):
    """
    Incomplete gamma function
    >>> round(Fgamma(0,0),10)
    1.0
    """
    SMALL=1e-12
    x = max(x,SMALL)
    return 0.5*pow(x,-m-0.5)*gamm_inc(m+0.5,x)

def gamm_inc(a,x):
    """
    Incomple gamma function \gamma; computed from NumRec routine gammp.
    """
    gammap,gln = gammp(a,x)
    return np.exp(gln)*gammap

def gammp(a,x):
    "Returns the incomplete gamma function P(a;x). NumRec sect 6.2."
    assert (x > 0 and a >= 0), "Invalid arguments in routine gammp: %s,%s" % (x,a)

    if x < (a+1.0): #Use the series representation
        gamser,gln = _gser(a,x)
        return gamser,gln
    #Use the continued fraction representation
    gammcf,gln = _gcf(a,x)
    return 1.0-gammcf ,gln  #and take its complement.

def _gser(a,x):
    "Series representation of Gamma. NumRec sect 6.1."
    ITMAX=100
    EPS=3.e-7

    gln=lgamma(a)
    assert(x>=0),'x < 0 in gser'
    if x == 0 : return 0,gln

    ap = a
    delt = sum = 1./a
    for i in xrange(ITMAX):
        ap=ap+1.
        delt=delt*x/ap
        sum=sum+delt
        if abs(delt) < abs(sum)*EPS: break
    else:
        print 'a too large, ITMAX too small in gser'
    gamser=sum*np.exp(-x+a*np.log(x)-gln)
    return gamser,gln

def _gcf(a,x):
    "Continued fraction representation of Gamma. NumRec sect 6.1"
    ITMAX=100
    EPS=3.e-7
    FPMIN=1.e-30

    gln=lgamma(a)
    b=x+1.-a
    c=1./FPMIN
    d=1./b
    h=d
    for i in xrange(1,ITMAX+1):
        an=-i*(i-a)
        b=b+2.
        d=an*d+b
        if abs(d) < FPMIN: d=FPMIN
        c=b+an/c
        if abs(c) < FPMIN: c=FPMIN
        d=1./d
        delt=d*c
        h=h*delt
        if abs(delt-1.) < EPS: break
    else:
        print 'a too large, ITMAX too small in gcf'
    gammcf=np.exp(-x+a*np.log(x)-gln)*h
    return gammcf,gln

def trace2(A,B):
    "Return trace(AB) of matrices A and B"
    return sum(A*B)

def dmat(c,nocc):
    "Form the density matrix from the first nocc orbitals of c"
    return np.dot(c[:,:nocc],c[:,:nocc].T)

def simx(A,B):
    "Similarity transform B^TAB"
    return np.dot(B.T,dot(A,B))

def geigh(H,S):
    "Solve the generalized eigensystem Hc = ESc"
    X = np.inv(np.cholesky(S)).T
    E,U = np.eigh(simx(H,X))
    return E,dot(X,U)

if __name__ == '__main__':
    import doctest
    doctest.testmod()

<commit_msg>Fixed some bugs in the geigh and simx numpy imports<commit_after>"""
utils.py - Simple utilility funtions used in pyquante2.

 Copyright (c) 2004, Richard P. Muller. All Rights Reserved.

 PyQuante version 2.0 and later is covered by the GPL
 license. Please see the file LICENSE that is part of this
 distribution.
"""
import numpy as np
from math import factorial,gamma,lgamma
from itertools import combinations_with_replacement

def pairs(it): return combinations_with_replacement(it,2)

def isnear(a,b,tol=1e-9): return abs(a-b)<tol

def fact2(n):
    """
    fact2(n) - n!!, double factorial of n
    >>> fact2(0)
    1
    >>> fact2(1)
    1
    >>> fact2(3)
    3
    >>> fact2(8)
    384
    >>> fact2(-1)
    1
    """
    return reduce(int.__mul__,xrange(n,0,-2),1)

def norm2(a): return np.dot(a,a)

def binomial(n,k):
    """
    Binomial coefficient
    >>> binomial(5,2)
    10
    >>> binomial(10,5)
    252
    """
    if n==k: return 1
    assert n>k, "Attempting to call binomial(%d,%d)" % (n,k)
    return factorial(n)/(factorial(k)*factorial(n-k))

def Fgamma(m,x):
    """
    Incomplete gamma function
    >>> round(Fgamma(0,0),10)
    1.0
    """
    SMALL=1e-12
    x = max(x,SMALL)
    return 0.5*pow(x,-m-0.5)*gamm_inc(m+0.5,x)

def gamm_inc(a,x):
    """
    Incomple gamma function \gamma; computed from NumRec routine gammp.
    """
    gammap,gln = gammp(a,x)
    return np.exp(gln)*gammap

def gammp(a,x):
    "Returns the incomplete gamma function P(a;x). NumRec sect 6.2."
    assert (x > 0 and a >= 0), "Invalid arguments in routine gammp: %s,%s" % (x,a)

    if x < (a+1.0): #Use the series representation
        gamser,gln = _gser(a,x)
        return gamser,gln
    #Use the continued fraction representation
    gammcf,gln = _gcf(a,x)
    return 1.0-gammcf ,gln  #and take its complement.

def _gser(a,x):
    "Series representation of Gamma. NumRec sect 6.1."
    ITMAX=100
    EPS=3.e-7

    gln=lgamma(a)
    assert(x>=0),'x < 0 in gser'
    if x == 0 : return 0,gln

    ap = a
    delt = sum = 1./a
    for i in xrange(ITMAX):
        ap=ap+1.
        delt=delt*x/ap
        sum=sum+delt
        if abs(delt) < abs(sum)*EPS: break
    else:
        print 'a too large, ITMAX too small in gser'
    gamser=sum*np.exp(-x+a*np.log(x)-gln)
    return gamser,gln

def _gcf(a,x):
    "Continued fraction representation of Gamma. NumRec sect 6.1"
    ITMAX=100
    EPS=3.e-7
    FPMIN=1.e-30

    gln=lgamma(a)
    b=x+1.-a
    c=1./FPMIN
    d=1./b
    h=d
    for i in xrange(1,ITMAX+1):
        an=-i*(i-a)
        b=b+2.
        d=an*d+b
        if abs(d) < FPMIN: d=FPMIN
        c=b+an/c
        if abs(c) < FPMIN: c=FPMIN
        d=1./d
        delt=d*c
        h=h*delt
        if abs(delt-1.) < EPS: break
    else:
        print 'a too large, ITMAX too small in gcf'
    gammcf=np.exp(-x+a*np.log(x)-gln)*h
    return gammcf,gln

def trace2(A,B):
    "Return trace(AB) of matrices A and B"
    return sum(A*B)

def dmat(c,nocc):
    "Form the density matrix from the first nocc orbitals of c"
    return np.dot(c[:,:nocc],c[:,:nocc].T)

def simx(A,B):
    "Similarity transform B^TAB"
    return np.dot(B.T,np.dot(A,B))

def geigh(H,S):
    "Solve the generalized eigensystem Hc = ESc"
    X = np.linalg.inv(np.linalg.cholesky(S)).T
    E,U = np.linalg.eigh(simx(H,X))
    return E,np.dot(X,U)

if __name__ == '__main__':
    import doctest
    doctest.testmod()

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
@author: Federico Cerchiari <federicocerchiari@gmail.com>
"""
import types

from exceptions import TagError

SPECIAL_ATTRS = {
    'klass': 'class',
    'typ': 'type',
}

class pyQueryFinder(object):

    def __init__(self):
        self.__all_tags = []

    def __add_new_tag(self, tag):
        self.__all_tags.append(tag)

    def __call__(self, pattern):
        find_func = lambda t: True
        if isinstance(pattern, Tag):
            find_func = lambda t: isinstance(t, pattern)
        elif isinstance(pattern, str):
            pattern = pattern.strip()
            if pattern.startswith('.'):
                find_func = lambda t: t.attr.get('klass') == pattern[1:]
            elif pattern.startswith('#'):
                find_func = lambda t: t.attr.get('id') == pattern[1:]
            else:
                tag = getattr(self, '_{}__tag'.format(self.__class__.__name__))
                find_func = lambda t: tag == pattern
        return filter(find_func, self.__all_tags)

pyQ = pyQueryFinder()

class Tag(object):
    _template = '<{tag}{attrs}>{inner}</{tag}>'
    _void = False
    _needed = None

    def __init__(self, **kwargs):
        self.attrs = {}
        if self._needed and not set(self._needed).issubset(set(kwargs)):
            raise TagError()
        self.attr(**kwargs)
        self.childs = []
        self.parent = self._own_index = None
        self._tab_count = 0
        pyQ._pyQueryFinder__add_new_tag(self)

    def __call__(self, *args):
        for arg in args:
            if type(arg) in (
                types.ListType,
                types.TupleType,
                types.GeneratorType,
            ):
                for tag in arg:
                    self.__call__(tag)
            elif isinstance(arg, Tag):
                self._add_to_childs(arg)
            else:
                self.childs.append(arg)
        return self

    def __getitem__(self, i):
        return self.childs[i]

    def __iter__(self):
        return iter(self.childs)

    def _add_to_childs(self, tag):
        self.childs.append(tag)
        tag.parent = self
        tag._tab_count = self._tab_count + 1
        tag._own_index = self.childs.index(tag)

    def append(self, *childs):
        return self(*childs)

    def append_to(self, father):
        return father.append(self)

    def remove(self):
        self.parent.childs.pop(self._own_index)
        return self

    def pop(self, i=0):
        self.childs.pop(i).parent = None
        return self

    def empty(self):
        for child in self.childs:
            self.childs.remove()
        return self

    def attr(self, **kwargs):
        for karg in kwargs:
            self.attrs[SPECIAL_ATTRS.get(karg, karg)] = kwargs[karg]
        return self

    def find(self, pattern):
        pass

    def render(self, pretty=False):
        prettying = '\t' * self._tab_count + '\n'
        tag_data = {
            'tag': getattr(self, '_{}__tag'.format(self.__class__.__name__)),
            'attrs': ''.join(' {}="{}"'.format(k, v) for k, v in self.attrs.iteritems())
        }
        if not self._void:
            tag_data['inner'] = self._render_childs(pretty)
        template = self._template + prettying if pretty else self._template
        return template.format(**tag_data)

    def _render_childs(self, pretty):
        return ''.join(child.render(pretty) if isinstance(child, Tag) else str(child) for child in self.childs)

class VoidTag(Tag):
    _void = True
    _template = '<{tag}{attrs}/>'
<commit_msg>Added several jQuery-like apis<commit_after># -*- coding: utf-8 -*-
"""
@author: Federico Cerchiari <federicocerchiari@gmail.com>
"""
import types
from copy import deepcopy
from collections import Mapping

from exceptions import TagError

class Attrs(dict):
    MULTI_VALUES_ATTRS = ('klass', 'typ')

    def __init__(self):
        self['style'] = {}

    def __setitem__(self, key, value):
        if key in self.MULTI_VALUES_ATTRS:
            if key not in self:
                super(Attrs, self).__setitem__(key, [])
            self[key].append(value)
        else:
            super(Attrs, self).__setitem__(key, value)

    def update(self, other=None, **kwargs):
        if other is not None:
            for k, v in other.iteritems() if isinstance(other, Mapping) else other:
                self[k] = v
        for k, v in kwargs.iteritems():
            self[k] = v

class TagContainer(object):
    _void = False

    def __init__(self):
        self.childs = []
        self.parent = None

    def __call__(self, *args, prepend=False):
        for arg in args:
            if type(arg) in (
                types.ListType,
                types.TupleType,
                types.GeneratorType,
            ):
                for tag in arg:
                    self.__call__(tag)
            elif isinstance(arg, Tag):
                self._add_to_childs(arg, prepend=prepend)
            else:
                if not prepend:
                    self.childs.append(arg)
                else:
                    self.childs.insert(0, arg)
        return self

    def __getitem__(self, i):
        return self.childs[i]

    def __iter__(self):
        return iter(self.childs)

    def _add_to_childs(self, tag, prepend=False):
        if not prepend:
            self.childs.append(tag)
        else:
            self.childs.insert(0, tag)
        tag.parent = self
        tag._tab_count = self._tab_count + 1
        tag._own_index = self.childs.index(tag)

    def after(self, other):
        self.parent.childs.insert(self._own_index + 1, other)
        return self

    def before(self, other):
        self.parent.childs.insert(self._own_index - 1, other)
        return self

    def prepend(self, *childs):
        return self(*childs, prepend=True)

    def prepend_to(self, father):
        return father.prepend(self)

    def append(self, *childs):
        return self(*childs)

    def append_to(self, father):
        return father.append(self)

    def remove(self):
        self.parent.childs.pop(self._own_index)
        return self

    def pop(self, i=0):
        self.childs.pop(i).parent = None
        return self

    def empty(self):
        for child in self.childs:
            self.childs.remove()
        return self

    def children(self):
        return filter(lambda x: isinstance(x, Tag), self.childs)

    def filter(self, pattern):
        # TODO
        pass

    def find(self, pattern):
        # TODO
        pass

    def first(self):
        return self.childs[0]

    def last(self):
        return self.childs[-1]

    @property
    def length(self):
        return len(self.childs)

    def has(self, pattern):
        # TODO
        # return pattern or tag in self.childs
        pass

    def next(self):
        return self.parent.child[self._own_index + 1]

    def prev(self):
        return self.parent.child[self._own_index - 1]

    def prev_all(self):
        return self.parent.child[:self._own_index - 1]

    def parent(self):
        return self.parent

    def siblings(self):
        return filter(lambda x: isinstance(x, Tag), self.parent.childs)

    def slice(self, start, end):
        return self.childs[start:end]

class Tag(TagContainer):
    _template = '<{tag}{attrs}>{inner}</{tag}>'
    _needed = None

    SPECIAL_ATTRS = {
        'klass': 'class',
        'typ': 'type'
    }

    def __init__(self, **kwargs):
        self._own_index = None
        self.attrs = Attrs()
        self.data = {}
        self.style = {}
        if self._needed and not set(self._needed).issubset(set(kwargs)):
            raise TagError()
        self.attr(**kwargs)
        self._tab_count = 0
        super(Tag, self).__init__()

    def index(self):
        return self._own_index

    def attr(self, other=None, **kwargs):
        self.attrs.update(other or kwargs)
        return self

    def prop(self, other=None, **kwargs):
        return self.attr(other, **kwargs)

    def remove_attr(self, attr):
        self.attrs.pop(attr, None)
        return self

    def remove_prop(self, attr):
        self.attrs.pop(attr, None)
        return self

    def add_class(self, cssclass):
        self.attrs['klass'].append(cssclass)
        return self

    def remove_class(self, cssclass):
        self.attrs['klass'].remove(cssclass)
        return self

    def clone(self):
        return deepcopy(self)

    def remove(self):
        del self

    def contents(self):
        return self.childs

    def css(self):
        #TODO
        pass

    def hide(self):
        self.attrs['style']['display'] = None
        return self

    def show(self):
        self.attrs['style'].pop('display')
        return self

    def toggle(self):
        return self.show() if self.attrs['style']['display'] == None else self.hide()

    def data(self, key, value=None):
        if value:
            self.data[key] = value
            return self
        else:
            return self.data[key]

    def has_class(self, cssclass):
        return cssclass in self.attrs['klass']

    def toggleClass(self, cssclass):
        return self.remove_class(cssclass) if self.has_class(cssclass) else self.add_class(cssclass)

    def html(self):
        return self._render_childs()

    def is_a(self, pattern):
        #TODO
        # return true if self == pattern
        pass

    def replace_all():
        #TODO
        pass

    def replace_with(self, other):
        if isinstance(other, Tag):
            self = other
        else:
            raise TagError()
        return self

    def wrap(self, other):
        return self.before(other.append(self)).parent.pop(self._own_index)

    def wrap_inner(self, other):
        self.childs

    def text(self):
        return ''.join(child.text() for child in self.childs if isinstance(child, Tag) else child)

    def render(self, pretty=False):
        prettying = '\t' * self._tab_count + '\n'
        tag_data = {
            'tag': getattr(self, '_{}__tag'.format(self.__class__.__name__)),
            'attrs': ''.join(' {}="{}"'.format(self.SPECIAL_ATTRS.get(k, k), v) for k, v in self.attrs.iteritems())
        }
        if not self._void:
            tag_data['inner'] = self._render_childs(pretty)
        template = self._template + prettying if pretty else self._template
        return template.format(**tag_data)

    def _render_childs(self, pretty):
        return ''.join(child.render(pretty) if isinstance(child, Tag) else str(child) for child in self.childs)

class VoidTag(Tag):
    _void = True
    _template = '<{tag}{attrs}/>'
<|endoftext|>


<commit_before>"""
This file is part of Arakoon, a distributed key-value store. Copyright
(C) 2010 Incubaid BVBA

Licensees holding a valid Incubaid license may use this file in
accordance with Incubaid's Arakoon commercial license agreement. For
more information on how to enter into this agreement, please contact
Incubaid (contact details can be found on www.arakoon.org/licensing).

Alternatively, this file may be redistributed and/or modified under
the terms of the GNU Affero General Public License version 3, as
published by the Free Software Foundation. Under this license, this
file is distributed in the hope that it will be useful, but WITHOUT
ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
FITNESS FOR A PARTICULAR PURPOSE.

See the GNU Affero General Public License for more details.
You should have received a copy of the
GNU Affero General Public License along with this program (file "COPYING").
If not, see <http://www.gnu.org/licenses/>.
"""

from system_tests_common import *
import logging
import subprocess

@with_custom_setup( default_setup, basic_teardown )
def test_single_client_100000_sets():
    iterate_n_times( 100000, simple_set )

@with_custom_setup( setup_3_nodes_forced_master, basic_teardown )
def test_delete_non_existing_with_catchup ():
    stopOne( node_names[1] )
    key='key'
    value='value'
    cli = get_client()
    try:
        cli.delete( key )
    except:
        pass
    cli.set(key,value)
    cli.set(key,value)
    cli.set(key,value)

    slave = node_names[1]
    startOne( slave )
    time.sleep(2.0)
    log_dir = q.config.arakoon.getNodeConfig(cluster_id, slave ) ['log_dir']
    log_file = q.system.fs.joinPaths( log_dir, '%s.log' % slave )
    log = q.system.fs.fileGetContents( log_file )
    assert_equals( log.find( "don't fit" ), -1, "Store counter out of sync" )

@with_custom_setup( setup_2_nodes_forced_master, basic_teardown )
def test_expect_progress_fixed_master ():
    stopOne( node_names[1] )
    key='key'
    value='value'
    cli = get_client()
    try:
        cli.set(key,value)
    except:
        pass
    restart_all()
    time.sleep(1.0)
    assert_true( cli.expectProgressPossible(),
                 "Master store counter is ahead of slave" )

@with_custom_setup( setup_3_nodes_forced_master, basic_teardown )
def test_restart_single_slave_long ():
    restart_single_slave_scenario( 100, 10000 )

@with_custom_setup( default_setup, basic_teardown )
def test_20_clients_1000_sets() :
    arakoon.ArakoonProtocol.ARA_CFG_TIMEOUT = 60.0
    create_and_wait_for_threads ( 20, 1000, simple_set, 100.0 )

@with_custom_setup( setup_3_nodes, basic_teardown)
def test_tlog_rollover():
    iterate_n_times( 150000, simple_set )
    stop_all()
    start_all()
    iterate_n_times( 150000, simple_set )

@with_custom_setup( default_setup, basic_teardown )
def test_restart_master_long ():
    restart_iter_cnt = 10
    write_loop = lambda: iterate_n_times( 100000, retrying_set_get_and_delete, failure_max=2*restart_iter_cnt, valid_exceptions=[ArakoonSockNotReadable,ArakoonNotFound] )
    restart_loop = lambda: delayed_master_restart_loop( restart_iter_cnt , 1.5*lease_duration )
    global test_failed
    test_failed = False
    create_and_wait_for_thread_list( [restart_loop, write_loop] )

    cli = get_client()
    time.sleep(2.0)
    key = "key"
    value = "value"
    cli.set(key, value)
    set_value = cli.get(key)
    assert_equals(  value, set_value ,
        "Key '%s' does not have expected value ('%s' iso '%s')" % (key, set_value, value) )

    stop_all()
    start_all()
    stop_all()

    assert_last_i_in_sync( node_names[0], node_names[1] )
    assert_last_i_in_sync( node_names[2], node_names[1] )
    compare_stores( node_names[0], node_names[1] )
    compare_stores( node_names[2], node_names[1] )
    cli._dropConnections()
    logging.info("end of `test_restart_master_long`")

@with_custom_setup( default_setup, basic_teardown )
def test_master_reelect():
    cli = get_client()
    master_id = cli.whoMaster()
    assert_not_equals ( master_id, None, "No master to begin with. Aborting.")

    key = "k"
    value = "v"
    cli.set(key ,value )

    stopOne( master_id )

    time.sleep( 1.5 * lease_duration )

    cli._masterId = None
    new_master_id = cli.whoMaster()
    assert_not_equals ( new_master_id, None, "No new master elected, no master. Aborting.")
    assert_not_equals ( new_master_id, master_id, "No new master elected, same master. Aborting.")

    assert_equals( cli.get(key), value)
    startOne( master_id )

    # Give old master some time to catch up
    time.sleep( 5.0 )

    stopOne ( new_master_id )

    time.sleep( 2.0 * lease_duration )

    cli = get_client()
    newest_master_id = cli.whoMaster()
    assert_not_equals ( newest_master_id, None, "No new master elected, no master. Aborting.")
    assert_not_equals ( newest_master_id, new_master_id, "No new master elected, same master. Aborting.")

@with_custom_setup( setup_3_nodes, basic_teardown)
def test_large_tlog_collection_restart():

    iterate_n_times( 100002, simple_set )
    stop_all()
    start_all()
    iterate_n_times( 100, set_get_and_delete )

@with_custom_setup( setup_3_nodes, basic_teardown )
def test_3_node_stop_master_slaves_restart():

    logging.info( "starting test case")
    iterate_n_times( 1000, simple_set )
    cli = get_client()
    master = cli.whoMaster()
    slaves = filter( lambda node: node != master, node_names )
    stopOne( master )
    logging.info ( lease_duration )
    nap_time = 2 * lease_duration
    logging.info( "Stopped master. Sleeping for %0.2f secs" % nap_time )

    print nap_time

    time.sleep( nap_time )
    logging.info( "Stopping old slaves")
    for node in slaves:
        print "Stopping %s" % node
        stopOne( node )

    logging.info( "Starting old master")
    startOne( master )
    time.sleep(0.2)

    logging.info( "Starting old slaves")
    for node in slaves:
        startOne( node )

    cli._dropConnections()
    cli = get_client()

    logging.info( "Sleeping a while" )
    time.sleep( lease_duration / 2 )

    iterate_n_times( 1000, set_get_and_delete )
    cli._dropConnections()

@with_custom_setup( setup_2_nodes_forced_master , basic_teardown )
def test_missed_accept ():

    # Give the new node some time to recognize the master
    time.sleep(0.5)
    stopOne( node_names[1] )

    cli = get_client()
    try:
        cli.set("k","v")
    except Exception, ex:
        logging.info( "Caught exception (%s: '%s'" , ex.__class__.__name__, ex )

    startOne ( node_names[1] )
    # Give the node some time to catch up
    time.sleep( 1.0 )

    iterate_n_times( 1000, set_get_and_delete )
    time.sleep(1.0)
    stop_all()
    assert_last_i_in_sync( node_names[0], node_names[1] )
    compare_stores( node_names[0], node_names[1] )

@with_custom_setup( setup_2_nodes_forced_master, basic_teardown)
def test_is_progress_possible():
    time.sleep(0.2)
    write_loop = lambda: iterate_n_times( 50000, retrying_set_get_and_delete  )
    create_and_wait_for_thread_list( [write_loop] )

    logging.info( "Stored all keys" )
    stop_all()

    whipe(node_names[1])

    cli = get_client()
    start_all()
    logging.info( "nodes started" )
    assert_false( cli.expectProgressPossible() )

    counter = 0
    max_wait = 60*5
    up2date = False

    while not up2date and counter < max_wait :
        time.sleep( 1.0 )
        counter += 1
        up2date = cli.expectProgressPossible()

    if counter >= max_wait :
        raise Exception ("Node did not catchup in a timely fashion")

    cli.set('k','v')

@with_custom_setup(setup_2_nodes, basic_teardown)
def test_collapse():
    zero = node_names[0]
    one = node_names[1]
    n = 298765
    iterate_n_times(n, simple_set)
    logging.info("did %i sets, now going into collapse scenario" % n)
    collapse(zero,1)
    logging.info("collapsing done")
    stopOne(one)
    whipe(one)
    startOne(one)
    cli = get_client()
    assert_false(cli.expectProgressPossible())
    up2date = False
    counter = 0
    while not up2date and counter < 100:
        time.sleep(1.0)
        counter = counter + 1
        up2date = cli.expectProgressPossible()
    logging.info("catchup from collapsed node finished")

@with_custom_setup(setup_2_nodes_forced_master, basic_teardown)
def test_catchup_exercises():
    time.sleep(1.0) # ??

    def do_one(n, max_wait):
        iterate_n_times(n, simple_set)
        stop_all()
        logging.info("stopped all nodes")
        whipe(node_names[1])

        start_all()
        cli = get_client ()

        assert_false(cli.expectProgressPossible())

        counter = 0
        up2date = False

        while not up2date and counter < max_wait :
            time.sleep( 1.0 )
            counter += 1
            up2date = cli.expectProgressPossible()

        if counter >= max_wait :
            raise Exception ("Node did not catchup in a timely fashion")

    n = 1000
    w = 10
    for i in range(7):
        do_one(n,w)
        n = n * 2
        w = w * 2

@with_custom_setup( setup_1_node_forced_master, basic_teardown )
def test_sso_deployment():
    global test_failed
    test_failed = False

    write_loop = lambda: iterate_n_times( 10000, retrying_set_get_and_delete )
    large_write_loop = lambda: iterate_n_times( 280000, retrying_set_get_and_delete, startSuffix = 1000000 )
    write_thr1 = create_and_start_thread ( write_loop )
    non_retrying_write_loop = lambda: iterate_n_times( 10000, set_get_and_delete, startSuffix = 2000000  )

    add_node( 1 )
    config = q.config.getInifile(cluster_id)
    config.setParam(node_names[1],"log_level","debug")
    config.write()

    if cluster_id in q.config.list():
        sn = "%s_nodes" % cluster_id
        q.config.remove(sn)
    q.config.arakoonnodes.generateClientConfigFromServerConfig(cluster_id)

    restart_nodes_wf_sim( 1 )
    startOne( node_names[1] )

    create_and_wait_for_thread_list ( [ large_write_loop ] )

    add_node( 2 )
    config = q.config.getInifile(cluster_id)
    config.setParam(node_names[2],"log_level","debug")
    config.write()

    q.config.arakoon.forceMaster(cluster_id, None )

    if cluster_id in q.config.list():
        q.config.remove(cluster_id)

    q.config.arakoonnodes.generateClientConfigFromServerConfig(cluster_id)

    restart_nodes_wf_sim( 2 )
    startOne( node_names[2] )
    time.sleep( 0.3 )
    assert_running_nodes ( 3 )

    write_thr3 = create_and_start_thread ( non_retrying_write_loop )

    write_thr1.join()
    write_thr3.join()

    assert_false ( test_failed )

    assert_running_nodes( 3 )

@with_custom_setup( setup_3_nodes, basic_teardown )
def test_3_nodes_2_slaves_down ():

    cli = get_client()
    master_id = cli.whoMaster()

    slaves = filter( lambda n: n != master_id, node_names )
    for slave in slaves:
        stopOne( slave )

    assert_raises( ArakoonSockNotReadable, cli.set, 'k', 'v' )

    cli._dropConnections()

<commit_msg>catchup_exercises<commit_after>"""
This file is part of Arakoon, a distributed key-value store. Copyright
(C) 2010 Incubaid BVBA

Licensees holding a valid Incubaid license may use this file in
accordance with Incubaid's Arakoon commercial license agreement. For
more information on how to enter into this agreement, please contact
Incubaid (contact details can be found on www.arakoon.org/licensing).

Alternatively, this file may be redistributed and/or modified under
the terms of the GNU Affero General Public License version 3, as
published by the Free Software Foundation. Under this license, this
file is distributed in the hope that it will be useful, but WITHOUT
ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
FITNESS FOR A PARTICULAR PURPOSE.

See the GNU Affero General Public License for more details.
You should have received a copy of the
GNU Affero General Public License along with this program (file "COPYING").
If not, see <http://www.gnu.org/licenses/>.
"""

from system_tests_common import *
import logging
import subprocess

@with_custom_setup( default_setup, basic_teardown )
def test_single_client_100000_sets():
    iterate_n_times( 100000, simple_set )

@with_custom_setup( setup_3_nodes_forced_master, basic_teardown )
def test_delete_non_existing_with_catchup ():
    stopOne( node_names[1] )
    key='key'
    value='value'
    cli = get_client()
    try:
        cli.delete( key )
    except:
        pass
    cli.set(key,value)
    cli.set(key,value)
    cli.set(key,value)

    slave = node_names[1]
    startOne( slave )
    time.sleep(2.0)
    log_dir = q.config.arakoon.getNodeConfig(cluster_id, slave ) ['log_dir']
    log_file = q.system.fs.joinPaths( log_dir, '%s.log' % slave )
    log = q.system.fs.fileGetContents( log_file )
    assert_equals( log.find( "don't fit" ), -1, "Store counter out of sync" )

@with_custom_setup( setup_2_nodes_forced_master, basic_teardown )
def test_expect_progress_fixed_master ():
    stopOne( node_names[1] )
    key='key'
    value='value'
    cli = get_client()
    try:
        cli.set(key,value)
    except:
        pass
    restart_all()
    time.sleep(1.0)
    assert_true( cli.expectProgressPossible(),
                 "Master store counter is ahead of slave" )

@with_custom_setup( setup_3_nodes_forced_master, basic_teardown )
def test_restart_single_slave_long ():
    restart_single_slave_scenario( 100, 10000 )

@with_custom_setup( default_setup, basic_teardown )
def test_20_clients_1000_sets() :
    arakoon.ArakoonProtocol.ARA_CFG_TIMEOUT = 60.0
    create_and_wait_for_threads ( 20, 1000, simple_set, 100.0 )

@with_custom_setup( setup_3_nodes, basic_teardown)
def test_tlog_rollover():
    iterate_n_times( 150000, simple_set )
    stop_all()
    start_all()
    iterate_n_times( 150000, simple_set )

@with_custom_setup( default_setup, basic_teardown )
def test_restart_master_long ():
    restart_iter_cnt = 10
    write_loop = lambda: iterate_n_times( 100000, retrying_set_get_and_delete, failure_max=2*restart_iter_cnt, valid_exceptions=[ArakoonSockNotReadable,ArakoonNotFound] )
    restart_loop = lambda: delayed_master_restart_loop( restart_iter_cnt , 1.5*lease_duration )
    global test_failed
    test_failed = False
    create_and_wait_for_thread_list( [restart_loop, write_loop] )

    cli = get_client()
    time.sleep(2.0)
    key = "key"
    value = "value"
    cli.set(key, value)
    set_value = cli.get(key)
    assert_equals(  value, set_value ,
        "Key '%s' does not have expected value ('%s' iso '%s')" % (key, set_value, value) )

    stop_all()
    start_all()
    stop_all()

    assert_last_i_in_sync( node_names[0], node_names[1] )
    assert_last_i_in_sync( node_names[2], node_names[1] )
    compare_stores( node_names[0], node_names[1] )
    compare_stores( node_names[2], node_names[1] )
    cli._dropConnections()
    logging.info("end of `test_restart_master_long`")

@with_custom_setup( default_setup, basic_teardown )
def test_master_reelect():
    cli = get_client()
    master_id = cli.whoMaster()
    assert_not_equals ( master_id, None, "No master to begin with. Aborting.")

    key = "k"
    value = "v"
    cli.set(key ,value )

    stopOne( master_id )

    time.sleep( 1.5 * lease_duration )

    cli._masterId = None
    new_master_id = cli.whoMaster()
    assert_not_equals ( new_master_id, None, "No new master elected, no master. Aborting.")
    assert_not_equals ( new_master_id, master_id, "No new master elected, same master. Aborting.")

    assert_equals( cli.get(key), value)
    startOne( master_id )

    # Give old master some time to catch up
    time.sleep( 5.0 )

    stopOne ( new_master_id )

    time.sleep( 2.0 * lease_duration )

    cli = get_client()
    newest_master_id = cli.whoMaster()
    assert_not_equals ( newest_master_id, None, "No new master elected, no master. Aborting.")
    assert_not_equals ( newest_master_id, new_master_id, "No new master elected, same master. Aborting.")

@with_custom_setup( setup_3_nodes, basic_teardown)
def test_large_tlog_collection_restart():

    iterate_n_times( 100002, simple_set )
    stop_all()
    start_all()
    iterate_n_times( 100, set_get_and_delete )

@with_custom_setup( setup_3_nodes, basic_teardown )
def test_3_node_stop_master_slaves_restart():

    logging.info( "starting test case")
    iterate_n_times( 1000, simple_set )
    cli = get_client()
    master = cli.whoMaster()
    slaves = filter( lambda node: node != master, node_names )
    stopOne( master )
    logging.info ( lease_duration )
    nap_time = 2 * lease_duration
    logging.info( "Stopped master. Sleeping for %0.2f secs" % nap_time )

    print nap_time

    time.sleep( nap_time )
    logging.info( "Stopping old slaves")
    for node in slaves:
        print "Stopping %s" % node
        stopOne( node )

    logging.info( "Starting old master")
    startOne( master )
    time.sleep(0.2)

    logging.info( "Starting old slaves")
    for node in slaves:
        startOne( node )

    cli._dropConnections()
    cli = get_client()

    logging.info( "Sleeping a while" )
    time.sleep( lease_duration / 2 )

    iterate_n_times( 1000, set_get_and_delete )
    cli._dropConnections()

@with_custom_setup( setup_2_nodes_forced_master , basic_teardown )
def test_missed_accept ():

    # Give the new node some time to recognize the master
    time.sleep(0.5)
    stopOne( node_names[1] )

    cli = get_client()
    try:
        cli.set("k","v")
    except Exception, ex:
        logging.info( "Caught exception (%s: '%s'" , ex.__class__.__name__, ex )

    startOne ( node_names[1] )
    # Give the node some time to catch up
    time.sleep( 1.0 )

    iterate_n_times( 1000, set_get_and_delete )
    time.sleep(1.0)
    stop_all()
    assert_last_i_in_sync( node_names[0], node_names[1] )
    compare_stores( node_names[0], node_names[1] )

@with_custom_setup( setup_2_nodes_forced_master, basic_teardown)
def test_is_progress_possible():
    time.sleep(0.2)
    write_loop = lambda: iterate_n_times( 50000, retrying_set_get_and_delete  )
    create_and_wait_for_thread_list( [write_loop] )

    logging.info( "Stored all keys" )
    stop_all()

    whipe(node_names[1])

    cli = get_client()
    start_all()
    logging.info( "nodes started" )
    assert_false( cli.expectProgressPossible() )

    counter = 0
    max_wait = 60*5
    up2date = False

    while not up2date and counter < max_wait :
        time.sleep( 1.0 )
        counter += 1
        up2date = cli.expectProgressPossible()

    if counter >= max_wait :
        raise Exception ("Node did not catchup in a timely fashion")

    cli.set('k','v')

@with_custom_setup(setup_2_nodes, basic_teardown)
def test_collapse():
    zero = node_names[0]
    one = node_names[1]
    n = 298765
    iterate_n_times(n, simple_set)
    logging.info("did %i sets, now going into collapse scenario" % n)
    collapse(zero,1)
    logging.info("collapsing done")
    stopOne(one)
    whipe(one)
    startOne(one)
    cli = get_client()
    assert_false(cli.expectProgressPossible())
    up2date = False
    counter = 0
    while not up2date and counter < 100:
        time.sleep(1.0)
        counter = counter + 1
        up2date = cli.expectProgressPossible()
    logging.info("catchup from collapsed node finished")

@with_custom_setup(setup_2_nodes_forced_master, basic_teardown)
def test_catchup_exercises():
    time.sleep(1.0) # ??

    def do_one(n, max_wait):
        iterate_n_times(n, simple_set)
        stop_all()
        logging.info("stopped all nodes")
        whipe(node_names[1])

        start_all()
        cli = get_client ()
        counter = 0
        up2date = False

        while not up2date and counter < max_wait :
            time.sleep( 1.0 )
            counter += 1
            up2date = cli.expectProgressPossible()
            logging.info("up2date=%s", up2date)

        if counter >= max_wait :
            raise Exception ("Node did not catchup in a timely fashion")

    n = 20000
    w = 20
    for i in range(5):
        do_one(n,w)
        n = n * 2
        w = w * 2

@with_custom_setup( setup_1_node_forced_master, basic_teardown )
def test_sso_deployment():
    global test_failed
    test_failed = False

    write_loop = lambda: iterate_n_times( 10000, retrying_set_get_and_delete )
    large_write_loop = lambda: iterate_n_times( 280000, retrying_set_get_and_delete, startSuffix = 1000000 )
    write_thr1 = create_and_start_thread ( write_loop )
    non_retrying_write_loop = lambda: iterate_n_times( 10000, set_get_and_delete, startSuffix = 2000000  )

    add_node( 1 )
    config = q.config.getInifile(cluster_id)
    config.setParam(node_names[1],"log_level","debug")
    config.write()

    if cluster_id in q.config.list():
        sn = "%s_nodes" % cluster_id
        q.config.remove(sn)
    q.config.arakoonnodes.generateClientConfigFromServerConfig(cluster_id)

    restart_nodes_wf_sim( 1 )
    startOne( node_names[1] )

    create_and_wait_for_thread_list ( [ large_write_loop ] )

    add_node( 2 )
    config = q.config.getInifile(cluster_id)
    config.setParam(node_names[2],"log_level","debug")
    config.write()

    q.config.arakoon.forceMaster(cluster_id, None )

    if cluster_id in q.config.list():
        q.config.remove(cluster_id)

    q.config.arakoonnodes.generateClientConfigFromServerConfig(cluster_id)

    restart_nodes_wf_sim( 2 )
    startOne( node_names[2] )
    time.sleep( 0.3 )
    assert_running_nodes ( 3 )

    write_thr3 = create_and_start_thread ( non_retrying_write_loop )

    write_thr1.join()
    write_thr3.join()

    assert_false ( test_failed )

    assert_running_nodes( 3 )

@with_custom_setup( setup_3_nodes, basic_teardown )
def test_3_nodes_2_slaves_down ():

    cli = get_client()
    master_id = cli.whoMaster()

    slaves = filter( lambda n: n != master_id, node_names )
    for slave in slaves:
        stopOne( slave )

    assert_raises( ArakoonSockNotReadable, cli.set, 'k', 'v' )

    cli._dropConnections()

<|endoftext|>