<commit_before>""" A set of stream oriented parsers for http requests and responses,
inline with the current draft recommendations from the http working group.

http://tools.ietf.org/html/draft-ietf-httpbis-p1-messaging-17

Unlike other libraries, this is for clients, servers and proxies.

Missing:
    comma parsing/header folding

"""

import re
import zlib

class ParseError(StandardError):
    pass

from .semantics import Codes, Methods

CRLF = '\r\n'

class HTTPMessage(object):
    """A stream based parser for http like messages"""
    CONTENT_TYPE="application/http"
    def __init__(self, header):
        self.buffer = bytearray()
        self.offset = 0
        self.header = header

        self.body_chunks = []
        self.mode = 'start'
        self.body_reader = None

    @property
    def url(self):
        return self.header.url

    @property
    def scheme(self):
        return self.header.scheme

    @property
    def method(self):
        return self.header.method

    @property
    def host(self):
        return self.header.host

    @property
    def port(self):
        return self.header.port

    def feed_fd(self, fd):
        while True:
            length, terminator = self.feed_predict()
            if length == 0:
                return  ''
            elif terminator == '\r\n':
                text = fd.readLine()
            elif length <0:
                text = fd.read()
            elif length >0:
                text = fd.read(length)
            unread = self.feed(text)
            if unread:
                return unread

    def feed_predict(self):
        """returns size, terminator request for input. size is 0 means end. """
        if self.mode == 'start':
            return None, CRLF
        elif self.mode == 'headers':
            return None, CRLF
        elif self.mode == 'body':
            if self.body_reader is not None:
                return self.body_reader.feed_predict()
            else:
                # connection close
                return -1, None
        if self.mode == 'end':
            return 0, None
        if self.mode == 'incomplete':
            return 0, None

    def feed(self, text):
        if text and self.mode == 'start':
            text = self.feed_start(text)

        if text and self.mode == 'headers':
            text = self.feed_headers(text)
            if self.mode == 'body':
                if not self.header.has_body():
                    self.mode = 'end'
                else:
                    if self.header.body_is_chunked():
                        self.body_reader = ChunkReader()
                    else:
                        length = self.header.body_length()
                        if length >= 0:
                            self.body_reader = LengthReader(length)
                            self.body_chunks = [(self.offset, length)]
                            if length == 0:
                                self.mode = 'end'
                        else:
                            self.body_chunks = [(self.offset, 0)]
                            self.body_reader = None

        if text and self.mode == 'body':
            if self.body_reader is not None:
                text = self.body_reader.feed(self, text)
            else:
                ( (offset, length), ) = self.body_chunks
                self.buffer.extend(text)
                self.offset = len(self.buffer)
                self.body_chunks = ( (offset, length+len(text)), )
                text = ''

        return  text

    def close(self):
        if self.mode =='start' or (self.body_reader is None and self.mode == 'body'):
            self.mode = 'end'

        elif self.mode != 'end':
            if self.body_chunks:
                # check for incomplete in body_chunks
                offset, length = self.body_chunks.pop()
                position = len(self.buffer)
                length = min(length, position-offset)
                self.body_chunks.append((offset, length))
            self.mode = 'incomplete'

    def get_unread(self):
        return self.unread.getvalue()

    def headers_complete(self):
        return self.mode in ('end', 'body')

    def complete(self):
        return self.mode =='end'

    def feed_line(self, text):
        """ feed text into the buffer, returning the first line found (if found yet)"""
        self.buffer.extend(text)
        pos = self.buffer.find("\r\n", self.offset)
        if pos > -1:
            pos+=2
            text =str(self.buffer[pos:])
            del self.buffer[pos:]
            line = str(self.buffer[self.offset:])
            self.offset = len(self.buffer)
        else:
            line = None
            text = ''
        return line, text

    def feed_length(self, text, remaining):
        """ feed (at most remaining bytes) text to buffer, returning leftovers """
        body, text = text[:remaining], text[remaining:]
        remaining -= len(body)
        self.buffer.extend(body)
        self.offset = len(self.buffer)
        return remaining, text

    def feed_start(self, text):
        line, text = self.feed_line(text)
        if line is not None:
            if line != CRLF: # skip leading newlines
                self.header.set_start_line(line)
                self.mode = 'headers'

        return text

    def feed_headers(self, text):
        while text:
            line, text = self.feed_line(text)
            if line is not None:
                self.header.add_header_line(line)
                if line == CRLF:
                    self.mode = 'body'
                    break

        return text

    def get_message(self):
        return str(self.buffer)

    def get_decoded_message(self):
        buf = bytearray()
        self.write_decoded_message(buf)
        return str(buf)

    def write_message(self, buf):
        self.header.write(buf)
        buf.extend('\r\n')
        self.write_body(buf)

    def write_decoded_message(self, buf):
        self.header.write_decoded(buf)
        if self.header.has_body():
            length = sum(l for o,l in self.body_chunks)
            buf.extend('Content-Length: %d\r\n'%length)
        body = self.get_body()
        if self.header.encoding and body:
            try:
                body = zlib.decompress(body)
            except zlib.error:
                try:
                    body = zlib.decompress(body, 16 + zlib.MAX_WBITS)
                except zlib.error:
                    buf.extend('Content-Encoding: %s\r\n' % self.header.encoding)
        buf.extend('\r\n')
        buf.extend(body)

    def get_body(self):
        buf = bytearray()
        self.write_body(buf)
        return str(buf)

    def write_body(self, buf):
        for offset, length in self.body_chunks:
            buf.extend(self.buffer[offset:offset+length])

class ChunkReader(object):
    def __init__(self):
        self.mode = "start"
        self.remaining = 0

    def feed_predict(self):
        if self.mode =='start':
            return None, CRLF
        elif self.mode == 'chunk':
            if self.remaining == 0:
                return None, CRLF
            else:
                return self.remaining, None
        elif self.mode == 'trailer':
            return None, CRLF
        elif self.mode == 'end':
            return 0, None

    def feed(self, parser, text):
        while text:
            if self.mode == 'start':

                line, text = parser.feed_line(text)
                offset = len(parser.buffer)

                if line is not None:
                    chunk = int(line.split(';',1)[0], 16)
                    parser.body_chunks.append((offset, chunk))
                    self.remaining = chunk
                    if chunk == 0:
                        self.mode = 'trailer'
                    else:
                        self.mode = 'chunk'
                #print self.mode, repr(text)

            if text and self.mode == 'chunk':
                #print self.mode, repr(text), self.remaining
                if self.remaining > 0:
                    self.remaining, text = parser.feed_length(text, self.remaining)
                if self.remaining == 0:
                    end_of_chunk, text = parser.feed_line(text)
                    #print 'end',end_of_chunk
                    if end_of_chunk:
                        #print 'ended'
                        self.mode = 'start'
                #print self.mode, repr(text)

            if text and self.mode == 'trailer':
                line, text = parser.feed_line(text)
                if line is not None:
                    parser.header.add_trailer_line(line)
                    if line == CRLF:
                        self.mode = 'end'

            if self.mode == 'end':
                parser.mode ='end'
                break

        return text

class LengthReader(object):
    def __init__(self, length):
        self.remaining = length

    def feed_predict(self):
        return self.remaining, None

    def feed(self, parser, text):
        if self.remaining > 0:
            self.remaining, text = parser.feed_length(text, self.remaining)
        if self.remaining <= 0:
            parser.mode ='end'
        return text

class HTTPHeader(object):
    STRIP_HEADERS = ('Content-Length', 'Transfer-Encoding', 'Content-Encoding', 'TE', 'Expect', 'Trailer')
    def __init__(self):
        self.headers = []
        self.keep_alive = False
        self.mode = 'close'
        self.content_length = None
        self.encoding = None
        self.trailers = []
        self.expect_continue=False

    def has_body(self):
        pass

    def set_start_line(self, line):
        pass

    def write_decoded(self, buf):
        self.write_decoded_start(buf)
        strip_headers = self.STRIP_HEADERS if self.has_body() else ()
        self.write_headers(buf, strip_headers)

    def write_decoded_start(self, buf):
        pass

    def write_headers(self, buf, strip_headers=()):
        for k,v in self.headers:
            if k not in strip_headers:
                buf.extend('%s: %s\r\n'%(k,v))
        for k,v in self.trailers:
            if k not in strip_headers:
                buf.extend('%s: %s\r\n'%(k,v))

    def add_trailer_line(self, line):
        if line.startswith(' ') or line.startswith('\t'):
            k,v = self.trailers.pop()
            line = line.strip()
            v = "%s %s"%(v, line)
            self.trailers.append((k,v))
        elif line == '\r\n':
            pass
        else:
            name, value = line.split(':',1)
            name = name.strip()
            value = value.strip()
            self.trailers.append((name, value))

    def add_header(self, name, value):
        self.headers.append((name, value))

    def add_header_line(self, line):
        if line.startswith(' ') or line.startswith('\t'):
            k,v = self.headers.pop()
            line = line.strip()
            v = "%s %s"%(v, line)
            self.add_header(k, v)

        elif line == '\r\n':
            for name, value in self.headers:
                name = name.lower()
                value = value.lower()

                # todo handle multiple instances
                # of these headers
                if name == 'expect':
                    if '100-continue' in value:
                        self.expect_continue = True
                if name == 'content-length':
                    if self.mode == 'close':
                        self.content_length = int(value)
                        self.mode = 'length'

                if name == 'transfer-encoding':
                    if 'chunked' in value:
                        self.mode = 'chunked'

                if name == 'content-encoding':
                    self.encoding = value

                if name == 'connection':
                    if 'keep-alive' in value:
                        self.keep_alive = True
                    elif 'close' in value:
                        self.keep_alive = False

        else:
            #print line
            name, value = line.split(':',1)
            name = name.strip()
            value = value.strip()
            self.add_header(name, value)

    def body_is_chunked(self):
        return self.mode == 'chunked'

    def body_length(self):
        if self.mode == 'length':
            return self.content_length

url_rx = re.compile('(?P<scheme>https?)://(?P<authority>(?P<host>[^:/]+)(?::(?P<port>\d+))?)(?P<path>.*)', re.I)
class RequestHeader(HTTPHeader):
    def __init__(self):
        HTTPHeader.__init__(self)
        self.method = ''
        self.target_uri = ''
        self.version = ''
        self.host = ''
        self.scheme = 'http'
        self.port = 80
        self.host = ''

    def set_start_line(self, line):
        self.method, self.target_uri, self.version = line.rstrip().split(' ',2)
        if self.method.upper() == "CONNECT":
            # target_uri = host:port
            self.host, self.port = self.target_uri.split(':')
        else:
            match = url_rx.match(self.target_uri)
            if match:
                #self.add_header('Host', match.group('authority'))
                self.target_uri = match.group('path')
                self.host = match.group('host')
                port = match.group('port')
                self.port = int(port) if port else 80

                self.scheme = match.group('scheme')
                if not self.target_uri:
                    if self.method.upper() == 'OPTIONS':
                        self.target_uri = '*'
                    else:
                        self.target_uri = '/'

        if self.version =='HTTP/1.0':
            self.keep_alive = False

    def has_body(self):
        return self.mode in ('chunked', 'length')

    def write_decoded_start(self, buf):
        buf.extend('%s %s %s\r\n'%(self.method, self.target_uri, self.version))

class ResponseHeader(HTTPHeader):
    def __init__(self, request):
        HTTPHeader.__init__(self)
        self.request = request
        self.version = None
        self.code = None
        self.phrase = None

    @property
    def method(self):
        return self.request.method

    @property
    def url(self):
        return self.request.url

    @property
    def host(self):
        return self.request.host

    @property
    def port(self):
        return self.request.port

    @property
    def scheme(self):
        return self.request.scheme

    def set_start_line(self, line):
        self.version, self.code, self.phrase = line.rstrip().split(' ',2)
        self.code = int(self.code)
        if self.version =='HTTP/1.0':
            self.keep_alive = False

    def has_body(self):
        if self.request.method in Methods.no_body:
            return False
        elif self.code in Codes.no_body:
            return False

        return True

    def write_decoded_start(self, buf):
        buf.extend('%s %d %s\r\n'%(self.version, self.code, self.phrase))

class RequestMessage(HTTPMessage):
    CONTENT_TYPE="%s;msgtype=request"%HTTPMessage.CONTENT_TYPE
    def __init__(self):
        HTTPMessage.__init__(self, RequestHeader())

class ResponseMessage(HTTPMessage):
    CONTENT_TYPE="%s;msgtype=response"%HTTPMessage.CONTENT_TYPE
    def __init__(self, request):
        self.interim = []
        HTTPMessage.__init__(self, ResponseHeader(request.header))

    def got_continue(self):
        return bool(self.interim)

    @property
    def code(self):
        return self.header.code

    def feed(self, text):
        text = HTTPMessage.feed(self, text)
        if self.complete() and self.header.code == Codes.Continue:
            self.interim.append(self.header)
            self.header = ResponseHeader(self.header.request)
            self.body_chunks = []
            self.mode = 'start'
            self.body_reader = None
            text = HTTPMessage.feed(self, text)
        return text

<commit_msg>useless method<commit_after>""" A set of stream oriented parsers for http requests and responses,
inline with the current draft recommendations from the http working group.

http://tools.ietf.org/html/draft-ietf-httpbis-p1-messaging-17

Unlike other libraries, this is for clients, servers and proxies.

Missing:
    comma parsing/header folding

"""

import re
import zlib

class ParseError(StandardError):
    pass

from .semantics import Codes, Methods

CRLF = '\r\n'

class HTTPMessage(object):
    """A stream based parser for http like messages"""
    CONTENT_TYPE="application/http"
    def __init__(self, header):
        self.buffer = bytearray()
        self.offset = 0
        self.header = header

        self.body_chunks = []
        self.mode = 'start'
        self.body_reader = None

    @property
    def url(self):
        return self.header.url

    @property
    def scheme(self):
        return self.header.scheme

    @property
    def method(self):
        return self.header.method

    @property
    def host(self):
        return self.header.host

    @property
    def port(self):
        return self.header.port

    def feed_fd(self, fd):
        while True:
            length, terminator = self.feed_predict()
            if length == 0:
                return  ''
            elif terminator == '\r\n':
                text = fd.readLine()
            elif length <0:
                text = fd.read()
            elif length >0:
                text = fd.read(length)
            unread = self.feed(text)
            if unread:
                return unread

    def feed_predict(self):
        """returns size, terminator request for input. size is 0 means end. """
        if self.mode == 'start':
            return None, CRLF
        elif self.mode == 'headers':
            return None, CRLF
        elif self.mode == 'body':
            if self.body_reader is not None:
                return self.body_reader.feed_predict()
            else:
                # connection close
                return -1, None
        if self.mode == 'end':
            return 0, None
        if self.mode == 'incomplete':
            return 0, None

    def feed(self, text):
        if text and self.mode == 'start':
            text = self.feed_start(text)

        if text and self.mode == 'headers':
            text = self.feed_headers(text)
            if self.mode == 'body':
                if not self.header.has_body():
                    self.mode = 'end'
                else:
                    if self.header.body_is_chunked():
                        self.body_reader = ChunkReader()
                    else:
                        length = self.header.body_length()
                        if length >= 0:
                            self.body_reader = LengthReader(length)
                            self.body_chunks = [(self.offset, length)]
                            if length == 0:
                                self.mode = 'end'
                        else:
                            self.body_chunks = [(self.offset, 0)]
                            self.body_reader = None

        if text and self.mode == 'body':
            if self.body_reader is not None:
                text = self.body_reader.feed(self, text)
            else:
                ( (offset, length), ) = self.body_chunks
                self.buffer.extend(text)
                self.offset = len(self.buffer)
                self.body_chunks = ( (offset, length+len(text)), )
                text = ''

        return  text

    def close(self):
        if self.mode =='start' or (self.body_reader is None and self.mode == 'body'):
            self.mode = 'end'

        elif self.mode != 'end':
            if self.body_chunks:
                # check for incomplete in body_chunks
                offset, length = self.body_chunks.pop()
                position = len(self.buffer)
                length = min(length, position-offset)
                self.body_chunks.append((offset, length))
            self.mode = 'incomplete'

    def headers_complete(self):
        return self.mode in ('end', 'body')

    def complete(self):
        return self.mode =='end'

    def feed_line(self, text):
        """ feed text into the buffer, returning the first line found (if found yet)"""
        self.buffer.extend(text)
        pos = self.buffer.find("\r\n", self.offset)
        if pos > -1:
            pos+=2
            text =str(self.buffer[pos:])
            del self.buffer[pos:]
            line = str(self.buffer[self.offset:])
            self.offset = len(self.buffer)
        else:
            line = None
            text = ''
        return line, text

    def feed_length(self, text, remaining):
        """ feed (at most remaining bytes) text to buffer, returning leftovers """
        body, text = text[:remaining], text[remaining:]
        remaining -= len(body)
        self.buffer.extend(body)
        self.offset = len(self.buffer)
        return remaining, text

    def feed_start(self, text):
        line, text = self.feed_line(text)
        if line is not None:
            if line != CRLF: # skip leading newlines
                self.header.set_start_line(line)
                self.mode = 'headers'

        return text

    def feed_headers(self, text):
        while text:
            line, text = self.feed_line(text)
            if line is not None:
                self.header.add_header_line(line)
                if line == CRLF:
                    self.mode = 'body'
                    break

        return text

    def get_message(self):
        return str(self.buffer)

    def get_decoded_message(self):
        buf = bytearray()
        self.write_decoded_message(buf)
        return str(buf)

    def write_message(self, buf):
        self.header.write(buf)
        buf.extend('\r\n')
        self.write_body(buf)

    def write_decoded_message(self, buf):
        self.header.write_decoded(buf)
        if self.header.has_body():
            length = sum(l for o,l in self.body_chunks)
            buf.extend('Content-Length: %d\r\n'%length)
        body = self.get_body()
        if self.header.encoding and body:
            try:
                body = zlib.decompress(body)
            except zlib.error:
                try:
                    body = zlib.decompress(body, 16 + zlib.MAX_WBITS)
                except zlib.error:
                    buf.extend('Content-Encoding: %s\r\n' % self.header.encoding)
        buf.extend('\r\n')
        buf.extend(body)

    def get_body(self):
        buf = bytearray()
        self.write_body(buf)
        return str(buf)

    def write_body(self, buf):
        for offset, length in self.body_chunks:
            buf.extend(self.buffer[offset:offset+length])

class ChunkReader(object):
    def __init__(self):
        self.mode = "start"
        self.remaining = 0

    def feed_predict(self):
        if self.mode =='start':
            return None, CRLF
        elif self.mode == 'chunk':
            if self.remaining == 0:
                return None, CRLF
            else:
                return self.remaining, None
        elif self.mode == 'trailer':
            return None, CRLF
        elif self.mode == 'end':
            return 0, None

    def feed(self, parser, text):
        while text:
            if self.mode == 'start':

                line, text = parser.feed_line(text)
                offset = len(parser.buffer)

                if line is not None:
                    chunk = int(line.split(';',1)[0], 16)
                    parser.body_chunks.append((offset, chunk))
                    self.remaining = chunk
                    if chunk == 0:
                        self.mode = 'trailer'
                    else:
                        self.mode = 'chunk'
                #print self.mode, repr(text)

            if text and self.mode == 'chunk':
                #print self.mode, repr(text), self.remaining
                if self.remaining > 0:
                    self.remaining, text = parser.feed_length(text, self.remaining)
                if self.remaining == 0:
                    end_of_chunk, text = parser.feed_line(text)
                    #print 'end',end_of_chunk
                    if end_of_chunk:
                        #print 'ended'
                        self.mode = 'start'
                #print self.mode, repr(text)

            if text and self.mode == 'trailer':
                line, text = parser.feed_line(text)
                if line is not None:
                    parser.header.add_trailer_line(line)
                    if line == CRLF:
                        self.mode = 'end'

            if self.mode == 'end':
                parser.mode ='end'
                break

        return text

class LengthReader(object):
    def __init__(self, length):
        self.remaining = length

    def feed_predict(self):
        return self.remaining, None

    def feed(self, parser, text):
        if self.remaining > 0:
            self.remaining, text = parser.feed_length(text, self.remaining)
        if self.remaining <= 0:
            parser.mode ='end'
        return text

class HTTPHeader(object):
    STRIP_HEADERS = ('Content-Length', 'Transfer-Encoding', 'Content-Encoding', 'TE', 'Expect', 'Trailer')
    def __init__(self):
        self.headers = []
        self.keep_alive = False
        self.mode = 'close'
        self.content_length = None
        self.encoding = None
        self.trailers = []
        self.expect_continue=False

    def has_body(self):
        pass

    def set_start_line(self, line):
        pass

    def write_decoded(self, buf):
        self.write_decoded_start(buf)
        strip_headers = self.STRIP_HEADERS if self.has_body() else ()
        self.write_headers(buf, strip_headers)

    def write_decoded_start(self, buf):
        pass

    def write_headers(self, buf, strip_headers=()):
        for k,v in self.headers:
            if k not in strip_headers:
                buf.extend('%s: %s\r\n'%(k,v))
        for k,v in self.trailers:
            if k not in strip_headers:
                buf.extend('%s: %s\r\n'%(k,v))

    def add_trailer_line(self, line):
        if line.startswith(' ') or line.startswith('\t'):
            k,v = self.trailers.pop()
            line = line.strip()
            v = "%s %s"%(v, line)
            self.trailers.append((k,v))
        elif line == '\r\n':
            pass
        else:
            name, value = line.split(':',1)
            name = name.strip()
            value = value.strip()
            self.trailers.append((name, value))

    def add_header(self, name, value):
        self.headers.append((name, value))

    def add_header_line(self, line):
        if line.startswith(' ') or line.startswith('\t'):
            k,v = self.headers.pop()
            line = line.strip()
            v = "%s %s"%(v, line)
            self.add_header(k, v)

        elif line == '\r\n':
            for name, value in self.headers:
                name = name.lower()
                value = value.lower()

                # todo handle multiple instances
                # of these headers
                if name == 'expect':
                    if '100-continue' in value:
                        self.expect_continue = True
                if name == 'content-length':
                    if self.mode == 'close':
                        self.content_length = int(value)
                        self.mode = 'length'

                if name == 'transfer-encoding':
                    if 'chunked' in value:
                        self.mode = 'chunked'

                if name == 'content-encoding':
                    self.encoding = value

                if name == 'connection':
                    if 'keep-alive' in value:
                        self.keep_alive = True
                    elif 'close' in value:
                        self.keep_alive = False

        else:
            #print line
            name, value = line.split(':',1)
            name = name.strip()
            value = value.strip()
            self.add_header(name, value)

    def body_is_chunked(self):
        return self.mode == 'chunked'

    def body_length(self):
        if self.mode == 'length':
            return self.content_length

url_rx = re.compile('(?P<scheme>https?)://(?P<authority>(?P<host>[^:/]+)(?::(?P<port>\d+))?)(?P<path>.*)', re.I)
class RequestHeader(HTTPHeader):
    def __init__(self):
        HTTPHeader.__init__(self)
        self.method = ''
        self.target_uri = ''
        self.version = ''
        self.host = ''
        self.scheme = 'http'
        self.port = 80
        self.host = ''

    def set_start_line(self, line):
        self.method, self.target_uri, self.version = line.rstrip().split(' ',2)
        if self.method.upper() == "CONNECT":
            # target_uri = host:port
            self.host, self.port = self.target_uri.split(':')
        else:
            match = url_rx.match(self.target_uri)
            if match:
                #self.add_header('Host', match.group('authority'))
                self.target_uri = match.group('path')
                self.host = match.group('host')
                port = match.group('port')
                self.port = int(port) if port else 80

                self.scheme = match.group('scheme')
                if not self.target_uri:
                    if self.method.upper() == 'OPTIONS':
                        self.target_uri = '*'
                    else:
                        self.target_uri = '/'

        if self.version =='HTTP/1.0':
            self.keep_alive = False

    def has_body(self):
        return self.mode in ('chunked', 'length')

    def write_decoded_start(self, buf):
        buf.extend('%s %s %s\r\n'%(self.method, self.target_uri, self.version))

class ResponseHeader(HTTPHeader):
    def __init__(self, request):
        HTTPHeader.__init__(self)
        self.request = request
        self.version = None
        self.code = None
        self.phrase = None

    @property
    def method(self):
        return self.request.method

    @property
    def url(self):
        return self.request.url

    @property
    def host(self):
        return self.request.host

    @property
    def port(self):
        return self.request.port

    @property
    def scheme(self):
        return self.request.scheme

    def set_start_line(self, line):
        self.version, self.code, self.phrase = line.rstrip().split(' ',2)
        self.code = int(self.code)
        if self.version =='HTTP/1.0':
            self.keep_alive = False

    def has_body(self):
        if self.request.method in Methods.no_body:
            return False
        elif self.code in Codes.no_body:
            return False

        return True

    def write_decoded_start(self, buf):
        buf.extend('%s %d %s\r\n'%(self.version, self.code, self.phrase))

class RequestMessage(HTTPMessage):
    CONTENT_TYPE="%s;msgtype=request"%HTTPMessage.CONTENT_TYPE
    def __init__(self):
        HTTPMessage.__init__(self, RequestHeader())

class ResponseMessage(HTTPMessage):
    CONTENT_TYPE="%s;msgtype=response"%HTTPMessage.CONTENT_TYPE
    def __init__(self, request):
        self.interim = []
        HTTPMessage.__init__(self, ResponseHeader(request.header))

    def got_continue(self):
        return bool(self.interim)

    @property
    def code(self):
        return self.header.code

    def feed(self, text):
        text = HTTPMessage.feed(self, text)
        if self.complete() and self.header.code == Codes.Continue:
            self.interim.append(self.header)
            self.header = ResponseHeader(self.header.request)
            self.body_chunks = []
            self.mode = 'start'
            self.body_reader = None
            text = HTTPMessage.feed(self, text)
        return text

<|endoftext|>


<commit_before>import delta
import tensor
import sympy

'''
This module is ment for dim == 3
'''

class Eps(tensor.AppliedTensor):
    pass

def evalTwoEps(eps,ePS):
    for i,j,k in [(eps.index[0], eps.index[1], eps.index[2]),
                  (eps.index[1], eps.index[2], eps.index[0]),
                  (eps.index[2], eps.index[0], eps.index[1])]:
        for I,J,K in [(ePS.index[0], ePS.index[1], ePS.index[2]),
                      (ePS.index[1], ePS.index[2], ePS.index[0]),
                      (ePS.index[2], ePS.index[0], ePS.index[1])]:
            if not i==I and isinstance(i, (sympy.Symbol, sympy.Dummy)):
                continue
            return delta.contractDeltas(
                delta.Delta(j,J)*delta.Delta(k,K)
                -delta.Delta(j,K)*delta.Delta(k,J),
                dim=3)
    return None

def evalOneEps(eps):
    e1,e2,e3 = eps.index
    if e1==e2 or e2==e3 or e3==e1:
        return 0
    for i,j,k in [(e1,e2,e3),
                  (e2,e3,e1),
                  (e3,e1,e2)]:
        for I,J,K in [(1,2,3),
                      (2,3,1),
                      (3,1,2)]:
            if not i==I: continue
            return delta.contractDeltas(
                delta.Delta(j,J)*delta.Delta(k,K)
                -delta.Delta(j,K)*delta.Delta(k,J),
                dim=3)

### Here be unittest ###
import unittest
from delta import Deta

class TestEps(unittest.TestCase):

    def setUp(self):
        pass

    def test_evaOneEsps(self):
        a,b,c = sympy.symbols('a,b,c')
        self.assertEqual(evalOneEps(Eps(1,a,b)),
                Delta(a,2)*Delta(b,3)-Delta(a,3)*Delta(b,2) )
        self.assertEqual(evalOneEps(Eps(a,b,c)), None)
        self.assertEqual(evalOneEps(Eps(a,3,2)), -Delta(a,1))
        self.assertEqual(evalOneEps(Eps(a,a,b)), 0)
        self.assertEqual(evalOneEps(Eps(1,2,3)), 1)

    def test_evaTwoEsps(self):
        a,b,c,d,e = sympy.symbols('a,b,c,d,e')
        self.assertEqual(evalTwoEps(Eps(1,a,b),Eps(1,a,b)), 6)
        self.assertEqual(evalTwoEps(Eps(a,b,c),Eps(d,e,a)),
                Delta(b,d)*Delta(c,e)-Delta(b,e)*Delta(c,d) )

if __name__ == '__main__':
    unittest.main()

<commit_msg>bug fix<commit_after>import delta
import tensor
import sympy

'''
This module is ment for dim == 3
'''

class Eps(tensor.AppliedTensor):
    pass

def evalTwoEps(eps,ePS):
    for i,j,k in [(eps.index[0], eps.index[1], eps.index[2]),
                  (eps.index[1], eps.index[2], eps.index[0]),
                  (eps.index[2], eps.index[0], eps.index[1])]:
        for I,J,K in [(ePS.index[0], ePS.index[1], ePS.index[2]),
                      (ePS.index[1], ePS.index[2], ePS.index[0]),
                      (ePS.index[2], ePS.index[0], ePS.index[1])]:
            if not i==I and isinstance(i, (sympy.Symbol, sympy.Dummy)):
                continue
            return delta.contractDeltas(
                delta.Delta(j,J)*delta.Delta(k,K)
                -delta.Delta(j,K)*delta.Delta(k,J),
                dim=3)
    return None

def evalOneEps(eps):
    e1,e2,e3 = eps.index
    if e1==e2 or e2==e3 or e3==e1:
        return 0
    for i,j,k in [(e1,e2,e3),
                  (e2,e3,e1),
                  (e3,e1,e2)]:
        for I,J,K in [(1,2,3),
                      (2,3,1),
                      (3,1,2)]:
            if not i==I: continue
            return delta.contractDeltas(
                delta.Delta(j,J)*delta.Delta(k,K)
                -delta.Delta(j,K)*delta.Delta(k,J),
                dim=3)

### Here be unittest ###
import unittest
from delta import Delta, dim, setDim

class TestEps(unittest.TestCase):

    def setUp(self):
        pass

    def test_dim(self):
        a,b = sympy.symbols('a,b')
        tempDim=dim
        setDim(4711)
        self.assertEqual(evalTwoEps(Eps(1,a,b),Eps(1,a,b)), 6)
        setDim(tempDim)

    def test_evaOneEsps(self):
        a,b,c = sympy.symbols('a,b,c')
        self.assertEqual(evalOneEps(Eps(1,a,b)),
                Delta(a,2)*Delta(b,3)-Delta(a,3)*Delta(b,2) )
        self.assertEqual(evalOneEps(Eps(a,b,c)), None)
        self.assertEqual(evalOneEps(Eps(a,3,2)), -Delta(a,1))
        self.assertEqual(evalOneEps(Eps(a,a,b)), 0)
        self.assertEqual(evalOneEps(Eps(1,2,3)), 1)

    def test_evaTwoEsps(self):
        a,b,c,d,e = sympy.symbols('a,b,c,d,e')
        self.assertEqual(evalTwoEps(Eps(1,a,b),Eps(1,a,b)), 6)
        self.assertEqual(evalTwoEps(Eps(a,b,c),Eps(d,e,a)),
                Delta(b,d)*Delta(c,e)-Delta(b,e)*Delta(c,d) )

if __name__ == '__main__':
    unittest.main()

<|endoftext|>


<commit_before>#!/usr/bin/env python
#alphabeta.py

#class to essentially memoize the search function of ABPruning_Tree.
#basically list of tuples of (Transposition_State, depth, alpha, beta, value, best_child_depth)
#implement __contains__ function.

class Transposition_Key(object):
	def __init__(self, _t_state, _depth, _alpha, _beta, _max):
		self.t_state = _t_state
		self.depth = _depth
		self.alpha = _alpha
		self.beta = _beta
		self.max = _max

class Transposition_Manager(object):
	def __init__(self, strictness = 1):
		self.searches = {}	# self.searches : (t_state, max) -> [depth, alpha, beta]
		self.results = {}	# self.results : (t_state, max) -> (value, best_depth)
		self.search_strictness = 1
		if strictness in [0,1,2,3]:
			self.search_strictness = strictness
			# search_stricness	discretion
			# 0			t_state, max
			# 1			t_state, max, depth
			# 2			t_state, max, depth, alpha or beta
			# 3			t_state, max, depth, alpha, beta

	#returns False if:
	#	search has never been done
	#	or, the new search parameters are more likely to give a better result.
	#returns True if:
	#	search has been done before
	#	and, the previously recorded search is more likely to be accurate.
	#
	#in general,
	#	bigger depth means better search.
	#	smaller alphas means better search.
	#	bigger betas means better search.
	def __contains__(self, key):
		if not isinstance(key, Transposition_Key):
			raise TypeError
		k = (key.t_state, key.max)
		is_in = (k in searches)
		if is_in:
			search_params = searches[k]
		is_valid = False
		if is_in and self.search_strictness == 0:
			is_valid = True
		if is_in and self.search_strictness == 1:
			if search_params[0] >= key.depth:	#if the search we have on record is from higher up.
				is_valid = True
		if is_in and self.search_strictness == 2:
			if search_params[0] >= key.depth:
				if search_params[1] <= key.alpha or search_params[2] >= key.beta:
					is_valid = True
		if is_in and self.search_strictness == 3:
			if search_params[0] >= key.depth:
				if search_params[1] <= key.alpha and search_params[2] >= key.beta:
					is_valid = True
		return is_valid

	def __getitem__(self, key):
		if key in self:
			return results[(key.t_state, key.max)]
		else:
			raise KeyError

	def get_search_params(self, key):
		if key in self:
			return searches[(key.t_state, key.max)]
		else:
			raise KeyError

	def __setitem__(self, key, value):
		if not isinstance(key, Transposition_Key):
			raise TypeError
		if len(value) != 2:
			raise ValueError
		else:
			if not isinstance(value[0], int) or not isinstance(value[1], int):
				raise ValueError
		searches[(key.t_state, key.max)] = [key.depth, key.alpha, key.beta]
		results[(key.t_state, key.max)] = value

	def __len__(self):
		return len(results)

	def get_strictness(self):
		return self.search_strictness

<commit_msg>fixed small errors in TM.<commit_after>#!/usr/bin/env python
#alphabeta.py

#class to essentially memoize the search function of ABPruning_Tree.
#basically list of tuples of (Transposition_State, depth, alpha, beta, value, best_child_depth)
#implement __contains__ function.

class Transposition_Key(object):
	def __init__(self, _t_state, _depth, _alpha, _beta, _max):
		self.t_state = _t_state
		self.depth = _depth
		self.alpha = _alpha
		self.beta = _beta
		self.max = _max

class Transposition_Manager(object):
	def __init__(self, strictness = 1):
		self.searches = {}	# self.searches : (t_state, max) -> [depth, alpha, beta]
		self.results = {}	# self.results : (t_state, max) -> (value, best_depth)
		self.search_strictness = 1
		if strictness in [0,1,2,3]:
			self.search_strictness = strictness
			# search_stricness	discretion
			# 0			t_state, max
			# 1			t_state, max, depth
			# 2			t_state, max, depth, alpha or beta
			# 3			t_state, max, depth, alpha, beta

	#returns False if:
	#	search has never been done
	#	or, the new search parameters are more likely to give a better result.
	#returns True if:
	#	search has been done before
	#	and, the previously recorded search is more likely to be accurate.
	#
	#in general,
	#	bigger depth means better search.
	#	smaller alphas means better search.
	#	bigger betas means better search.
	def __contains__(self, key):
		if not isinstance(key, Transposition_Key):
			raise TypeError
		k = (key.t_state, key.max)
		is_in = (k in self.searches)
		if is_in:
			search_params = self.searches[k]
		is_valid = False
		if is_in and self.search_strictness == 0:
			is_valid = True
		if is_in and self.search_strictness == 1:
			if search_params[0] >= key.depth:	#if the search we have on record is from higher up.
				is_valid = True
		if is_in and self.search_strictness == 2:
			if search_params[0] >= key.depth:
				if search_params[1] <= key.alpha or search_params[2] >= key.beta:
					is_valid = True
		if is_in and self.search_strictness == 3:
			if search_params[0] >= key.depth:
				if search_params[1] <= key.alpha and search_params[2] >= key.beta:
					is_valid = True
		return is_valid

	def __getitem__(self, key):
		if key in self:
			return self.results[(key.t_state, key.max)]
		else:
			raise KeyError

	def get_search_params(self, key):
		if key in self:
			return self.searches[(key.t_state, key.max)]
		else:
			raise KeyError

	def __setitem__(self, key, value):
		if not isinstance(key, Transposition_Key):
			raise TypeError
		if len(value) != 2:
			raise ValueError
		else:
			if not isinstance(value[0], int) or not isinstance(value[1], int):
				raise ValueError
		self.searches[(key.t_state, key.max)] = [key.depth, key.alpha, key.beta]
		self.results[(key.t_state, key.max)] = value

	def __len__(self):
		return len(self.results)

	def get_strictness(self):
		return self.search_strictness

<|endoftext|>


<commit_before>from __future__ import print_function

import os
import math
import warnings

import numpy as np

from matplotlib.collections import LineCollection
from matplotlib.transforms import Bbox
from matplotlib.patches import Polygon

from .geometry.path import Path, get_endpoints
from . import extract_pv_slice

def distance(x1, y1, x2, y2, x3, y3):
    """
    Find the shortest distance between a point (x3, y3) and the line passing
    through the points (x1, y1) and (x2, y2).
    """

    px = x2-x1
    py = y2-y1

    something = px * px + py * py

    u =  ((x3 - x1) * px + (y3 - y1) * py) / float(something)

    x = x1 + u * px
    y = y1 + u * py

    dx = x - x3
    dy = y - y3

    dist = math.sqrt(dx*dx + dy*dy)

    return dist

class MovableSliceBox(object):

    def __init__(self, box, callback):
        self.box = box
        self.press = None
        self.background = None
        self.point_counter = 0
        self.callback = callback
        self.mode = 0
        self.show_poly = False
        self.cidpress = self.box.figure.canvas.mpl_connect('draw_event', self.draw_slicer)

    def connect(self):
        self.cidpress = self.box.figure.canvas.mpl_connect('key_press_event', self.key_press)
        self.cidpress = self.box.figure.canvas.mpl_connect('button_press_event', self.on_press)
        self.cidmotion = self.box.figure.canvas.mpl_connect('motion_notify_event', self.on_motion)

    def draw_slicer(self, event):

        axes = self.box.axes
        canvas = self.box.figure.canvas

        self.box.axes.draw_artist(self.box)

        if self.show_poly:

            path = Path(zip(self.box.x, self.box.y))
            path.width = self.box.width

            for poly in path.sample_polygons(1):
                self.box.axes.draw_artist(Polygon(zip(poly.x, poly.y),
                                          ec='green', fc='none',
                                          transform=self.box.axes.transData,
                                          clip_on=True, clip_box=self.box.axes.bbox))

    def on_press(self, event):

        if self.box.figure.canvas.toolbar.mode != '':
            return

        if event.inaxes != self.box.axes:
            return

        if self.mode == 1:
            self.callback(self.box)
            self.mode += 1
            return

        if self.mode == 2:
            self.box.x = []
            self.box.y = []
            self.mode = 0
            self.point_counter = 0

        self.press = event.xdata, event.ydata

        self.point_counter += 1

        axes = self.box.axes
        canvas = self.box.figure.canvas

        if self.point_counter == 1:  # first point

            self.box.x.append(event.xdata)
            self.box.x.append(event.xdata)
            self.box.y.append(event.ydata)
            self.box.y.append(event.ydata)

            self.box.width = 0.

            self.box.set_animated(True)
            canvas.draw()
            self.background = canvas.copy_from_bbox(self.box.axes.bbox)

        elif self.mode == 0:

            self.box.x.append(event.xdata)
            self.box.y.append(event.ydata)

        self.box._update_segments()

        # now redraw just the lineangle
        axes.draw_artist(self.box)

        canvas.blit(axes.bbox)

    def key_press(self, event):

        if self.box.figure.canvas.toolbar.mode != '':
            return

        if event.key == 'enter' and self.mode == 0:
            self.mode += 1
            self.box.x = self.box.x[:-1]
            self.box.y = self.box.y[:-1]

        if event.key == 'y' and self.mode == 2:
            self.show_poly = not self.show_poly
            self.draw_slicer(event)
            self.box.figure.canvas.draw()

    def on_motion(self, event):

        if self.box.figure.canvas.toolbar.mode != '':
            return

        if self.point_counter == 0:
            return

        if self.mode == 2:
            return

        canvas = self.box.figure.canvas
        axes = self.box.axes
        canvas.restore_region(self.background)

        if event.inaxes != self.box.axes:
            return

        if self.mode == 0:
            self.box.x[-1] = event.xdata
            self.box.y[-1] = event.ydata
        elif self.mode == 1:
            self.box.width = distance(self.box.x[-2], self.box.y[-2], self.box.x[-1], self.box.y[-1], event.xdata, event.ydata) * 2

        self.box._update_segments()

        # redraw just the current lineangle
        axes.draw_artist(self.box)

        # blit just the redrawn area
        canvas.blit(axes.bbox)

    def disconnect(self):
        self.box.figure.canvas.mpl_disconnect(self.cidpress)
        self.box.figure.canvas.mpl_disconnect(self.cidmotion)

class SliceBox(LineCollection):

    def __init__(self, x0=None, y0=None, x1=None, y1=None, width=None, **kwargs):

        super(SliceBox, self).__init__([], **kwargs)

        self.x0 = x0
        self.y0 = y0
        self.x1 = x1
        self.y1 = y1
        self.width = width

        self._update_segments()

    def _update_segments(self):

        if self.x0 is None:
            return

        # Find angle of normal to line
        theta = np.arctan2(self.y1 - self.y0, self.x1 - self.x0) + np.pi / 2.

        # Find displacement vectors
        dx = np.cos(theta) * self.width / 2.
        dy = np.sin(theta) * self.width / 2.

        # Find central line
        line = [(self.x0, self.y0), (self.x1, self.y1)]

        # Find bounding rectangle
        rect = [(self.x0 + dx, self.y0 + dy), (self.x0 - dx, self.y0 - dy),
                (self.x1 - dx, self.y1 - dy), (self.x1 + dx, self.y1 + dy),
                (self.x0 + dx, self.y0 + dy)]

        self.set_segments((line, rect))
        self.set_linestyles(('solid', 'dashed'))
        self.set_linewidths((2, 1))

class SliceCurve(LineCollection):

    def __init__(self, x=[], y=[], width=None, **kwargs):

        super(SliceCurve, self).__init__([], **kwargs)

        self.x = x
        self.y = y
        self.width = width

        self._update_segments()

    def _update_segments(self):

        if not self.x:
            return

        x1, y1, x2, y2 = get_endpoints(self.x, self.y, self.width)

        # Find central line
        line = zip(self.x, self.y)

        # Find bounding rectangle
        rect = zip(np.hstack([x1,x2[::-1], x1[0]]),
                   np.hstack([y1,y2[::-1], y1[0]]))

        self.set_segments((line, rect))
        self.set_linestyles(('solid', 'dashed'))
        self.set_linewidths((2, 1))

class PVSlicer(object):

    def __init__(self, filename, backend="Qt4Agg", clim=None):

        self.filename = filename

        try:
            from spectral_cube import SpectralCube
            cube = SpectralCube.read(filename, format='fits')
            self.array = cube._data
        except:
            warnings.warn("spectral_cube package is not available - using astropy.io.fits directly")
            from astropy.io import fits
            self.array = fits.getdata(filename)
            if self.array.ndim != 3:
                raise ValueError("dataset does not have 3 dimensions (install the spectral_cube package to avoid this error)")

        self.backend = backend

        import matplotlib as mpl
        mpl.use(self.backend)
        import matplotlib.pyplot as plt

        self.fig = plt.figure(figsize=(14, 8))

        self.ax1 = self.fig.add_axes([0.1, 0.1, 0.4, 0.7])

        if clim is None:
            warnings.warn("clim not defined and will be determined from the data")
            # To work with large arrays, sub-sample the data
            n1 = self.array.shape[0] / 10
            n2 = self.array.shape[1] / 10
            n3 = self.array.shape[2] / 10
            sub_array = self.array[::n1,::n2,::n3]
            cmin = np.min(sub_array[~np.isnan(sub_array) & ~np.isinf(sub_array)])
            cmax = np.max(sub_array[~np.isnan(sub_array) & ~np.isinf(sub_array)])
            crange = cmax - cmin
            self._clim = (cmin - crange, cmax + crange)
        else:
            self._clim = clim

        self.slice = int(round(self.array.shape[0] / 2.))

        from matplotlib.widgets import Slider

        self.slice_slider_ax = self.fig.add_axes([0.1, 0.95, 0.4, 0.03])
        self.slice_slider_ax.set_xticklabels("")
        self.slice_slider_ax.set_yticklabels("")
        self.slice_slider = Slider(self.slice_slider_ax, "3-d slice", 0, self.array.shape[0], valinit=self.slice, valfmt="%i")
        self.slice_slider.on_changed(self.update_slice)
        self.slice_slider.drawon = False

        self.image = self.ax1.imshow(self.array[self.slice, :,:], origin='lower', interpolation='nearest', vmin=self._clim[0], vmax=self._clim[1], cmap=plt.cm.gray)

        self.vmin_slider_ax = self.fig.add_axes([0.1, 0.90, 0.4, 0.03])
        self.vmin_slider_ax.set_xticklabels("")
        self.vmin_slider_ax.set_yticklabels("")
        self.vmin_slider = Slider(self.vmin_slider_ax, "vmin", self._clim[0], self._clim[1], valinit=self._clim[0])
        self.vmin_slider.on_changed(self.update_vmin)
        self.vmin_slider.drawon = False

        self.vmax_slider_ax = self.fig.add_axes([0.1, 0.85, 0.4, 0.03])
        self.vmax_slider_ax.set_xticklabels("")
        self.vmax_slider_ax.set_yticklabels("")
        self.vmax_slider = Slider(self.vmax_slider_ax, "vmax", self._clim[0], self._clim[1], valinit=self._clim[1])
        self.vmax_slider.on_changed(self.update_vmax)
        self.vmax_slider.drawon = False

        self.grid1 = None
        self.grid2 = None
        self.grid3 = None

        self.ax2 = self.fig.add_axes([0.55, 0.1, 0.4, 0.7])

        # Add slicing box
        self.box = SliceCurve(colors=(0.8, 0.0, 0.0))
        self.ax1.add_collection(self.box)
        self.movable = MovableSliceBox(self.box, callback=self.update_pv_slice)
        self.movable.connect()

        # Add save button
        from matplotlib.widgets import Button
        self.save_button_ax = self.fig.add_axes([0.65, 0.90, 0.20, 0.05])
        self.save_button = Button(self.save_button_ax, 'Save slice to FITS')
        self.save_button.on_clicked(self.save_fits)

        self.pv_slice = None

        self.cidpress = self.fig.canvas.mpl_connect('button_press_event', self.click)

    def click(self, event):

        if event.inaxes != self.ax2:
            return

        self.slice_slider.set_val(event.ydata)

    def save_fits(self, *args, **kwargs):

        if self.backend == 'Qt4Agg':
            from matplotlib.backends.backend_qt4 import _getSaveFileName
            plot_name = _getSaveFileName(self.fig.canvas.manager.window, "Choose filename",
                                         os.path.dirname(os.path.abspath(self.filename)), "", None)
            plot_name = str(plot_name)
        else:
            print("Enter filename: ", end='')
            plot_name = raw_input()

        if self.pv_slice is None:
            return

        from astropy.io import fits
        fits.writeto(plot_name, self.pv_slice, clobber=True)
        print("Saved file to: ", plot_name)

    def update_pv_slice(self, box):

        path = Path(zip(box.x, box.y))
        path.width = box.width

        self.pv_slice = extract_pv_slice(self.array, path)

        self.ax2.cla()
        self.ax2.imshow(self.pv_slice, origin='lower', aspect='auto', interpolation='nearest')

        self.fig.canvas.draw()

    def show(self):
        import matplotlib.pyplot as plt
        plt.show()

    def update_slice(self, pos=None):

        if self.array.ndim == 2:
            self.image.set_array(self.array)
        else:
            self.slice = int(round(pos))
            self.image.set_array(self.array[self.slice, :, :])

        self.fig.canvas.draw()

    def update_vmin(self, vmin):
        if vmin > self._clim[1]:
            self._clim = (self._clim[1], self._clim[1])
        else:
            self._clim = (vmin, self._clim[1])
        self.image.set_clim(*self._clim)
        self.fig.canvas.draw()

    def update_vmax(self, vmax):
        if vmax < self._clim[0]:
            self._clim = (self._clim[0], self._clim[0])
        else:
            self._clim = (self._clim[0], vmax)
        self.image.set_clim(*self._clim)
        self.fig.canvas.draw()
<commit_msg>use the .data parameter when imshowing pv_slice<commit_after>from __future__ import print_function

import os
import math
import warnings

import numpy as np

from matplotlib.collections import LineCollection
from matplotlib.transforms import Bbox
from matplotlib.patches import Polygon

from .geometry.path import Path, get_endpoints
from . import extract_pv_slice

def distance(x1, y1, x2, y2, x3, y3):
    """
    Find the shortest distance between a point (x3, y3) and the line passing
    through the points (x1, y1) and (x2, y2).
    """

    px = x2-x1
    py = y2-y1

    something = px * px + py * py

    u =  ((x3 - x1) * px + (y3 - y1) * py) / float(something)

    x = x1 + u * px
    y = y1 + u * py

    dx = x - x3
    dy = y - y3

    dist = math.sqrt(dx*dx + dy*dy)

    return dist

class MovableSliceBox(object):

    def __init__(self, box, callback):
        self.box = box
        self.press = None
        self.background = None
        self.point_counter = 0
        self.callback = callback
        self.mode = 0
        self.show_poly = False
        self.cidpress = self.box.figure.canvas.mpl_connect('draw_event', self.draw_slicer)

    def connect(self):
        self.cidpress = self.box.figure.canvas.mpl_connect('key_press_event', self.key_press)
        self.cidpress = self.box.figure.canvas.mpl_connect('button_press_event', self.on_press)
        self.cidmotion = self.box.figure.canvas.mpl_connect('motion_notify_event', self.on_motion)

    def draw_slicer(self, event):

        axes = self.box.axes
        canvas = self.box.figure.canvas

        self.box.axes.draw_artist(self.box)

        if self.show_poly:

            path = Path(zip(self.box.x, self.box.y))
            path.width = self.box.width

            for poly in path.sample_polygons(1):
                self.box.axes.draw_artist(Polygon(zip(poly.x, poly.y),
                                          ec='green', fc='none',
                                          transform=self.box.axes.transData,
                                          clip_on=True, clip_box=self.box.axes.bbox))

    def on_press(self, event):

        if self.box.figure.canvas.toolbar.mode != '':
            return

        if event.inaxes != self.box.axes:
            return

        if self.mode == 1:
            self.callback(self.box)
            self.mode += 1
            return

        if self.mode == 2:
            self.box.x = []
            self.box.y = []
            self.mode = 0
            self.point_counter = 0

        self.press = event.xdata, event.ydata

        self.point_counter += 1

        axes = self.box.axes
        canvas = self.box.figure.canvas

        if self.point_counter == 1:  # first point

            self.box.x.append(event.xdata)
            self.box.x.append(event.xdata)
            self.box.y.append(event.ydata)
            self.box.y.append(event.ydata)

            self.box.width = 0.

            self.box.set_animated(True)
            canvas.draw()
            self.background = canvas.copy_from_bbox(self.box.axes.bbox)

        elif self.mode == 0:

            self.box.x.append(event.xdata)
            self.box.y.append(event.ydata)

        self.box._update_segments()

        # now redraw just the lineangle
        axes.draw_artist(self.box)

        canvas.blit(axes.bbox)

    def key_press(self, event):

        if self.box.figure.canvas.toolbar.mode != '':
            return

        if event.key == 'enter' and self.mode == 0:
            self.mode += 1
            self.box.x = self.box.x[:-1]
            self.box.y = self.box.y[:-1]

        if event.key == 'y' and self.mode == 2:
            self.show_poly = not self.show_poly
            self.draw_slicer(event)
            self.box.figure.canvas.draw()

    def on_motion(self, event):

        if self.box.figure.canvas.toolbar.mode != '':
            return

        if self.point_counter == 0:
            return

        if self.mode == 2:
            return

        canvas = self.box.figure.canvas
        axes = self.box.axes
        canvas.restore_region(self.background)

        if event.inaxes != self.box.axes:
            return

        if self.mode == 0:
            self.box.x[-1] = event.xdata
            self.box.y[-1] = event.ydata
        elif self.mode == 1:
            self.box.width = distance(self.box.x[-2], self.box.y[-2], self.box.x[-1], self.box.y[-1], event.xdata, event.ydata) * 2

        self.box._update_segments()

        # redraw just the current lineangle
        axes.draw_artist(self.box)

        # blit just the redrawn area
        canvas.blit(axes.bbox)

    def disconnect(self):
        self.box.figure.canvas.mpl_disconnect(self.cidpress)
        self.box.figure.canvas.mpl_disconnect(self.cidmotion)

class SliceBox(LineCollection):

    def __init__(self, x0=None, y0=None, x1=None, y1=None, width=None, **kwargs):

        super(SliceBox, self).__init__([], **kwargs)

        self.x0 = x0
        self.y0 = y0
        self.x1 = x1
        self.y1 = y1
        self.width = width

        self._update_segments()

    def _update_segments(self):

        if self.x0 is None:
            return

        # Find angle of normal to line
        theta = np.arctan2(self.y1 - self.y0, self.x1 - self.x0) + np.pi / 2.

        # Find displacement vectors
        dx = np.cos(theta) * self.width / 2.
        dy = np.sin(theta) * self.width / 2.

        # Find central line
        line = [(self.x0, self.y0), (self.x1, self.y1)]

        # Find bounding rectangle
        rect = [(self.x0 + dx, self.y0 + dy), (self.x0 - dx, self.y0 - dy),
                (self.x1 - dx, self.y1 - dy), (self.x1 + dx, self.y1 + dy),
                (self.x0 + dx, self.y0 + dy)]

        self.set_segments((line, rect))
        self.set_linestyles(('solid', 'dashed'))
        self.set_linewidths((2, 1))

class SliceCurve(LineCollection):

    def __init__(self, x=[], y=[], width=None, **kwargs):

        super(SliceCurve, self).__init__([], **kwargs)

        self.x = x
        self.y = y
        self.width = width

        self._update_segments()

    def _update_segments(self):

        if not self.x:
            return

        x1, y1, x2, y2 = get_endpoints(self.x, self.y, self.width)

        # Find central line
        line = zip(self.x, self.y)

        # Find bounding rectangle
        rect = zip(np.hstack([x1,x2[::-1], x1[0]]),
                   np.hstack([y1,y2[::-1], y1[0]]))

        self.set_segments((line, rect))
        self.set_linestyles(('solid', 'dashed'))
        self.set_linewidths((2, 1))

class PVSlicer(object):

    def __init__(self, filename, backend="Qt4Agg", clim=None):

        self.filename = filename

        try:
            from spectral_cube import SpectralCube
            cube = SpectralCube.read(filename, format='fits')
            self.array = cube._data
        except:
            warnings.warn("spectral_cube package is not available - using astropy.io.fits directly")
            from astropy.io import fits
            self.array = fits.getdata(filename)
            if self.array.ndim != 3:
                raise ValueError("dataset does not have 3 dimensions (install the spectral_cube package to avoid this error)")

        self.backend = backend

        import matplotlib as mpl
        mpl.use(self.backend)
        import matplotlib.pyplot as plt

        self.fig = plt.figure(figsize=(14, 8))

        self.ax1 = self.fig.add_axes([0.1, 0.1, 0.4, 0.7])

        if clim is None:
            warnings.warn("clim not defined and will be determined from the data")
            # To work with large arrays, sub-sample the data
            n1 = self.array.shape[0] / 10
            n2 = self.array.shape[1] / 10
            n3 = self.array.shape[2] / 10
            sub_array = self.array[::n1,::n2,::n3]
            cmin = np.min(sub_array[~np.isnan(sub_array) & ~np.isinf(sub_array)])
            cmax = np.max(sub_array[~np.isnan(sub_array) & ~np.isinf(sub_array)])
            crange = cmax - cmin
            self._clim = (cmin - crange, cmax + crange)
        else:
            self._clim = clim

        self.slice = int(round(self.array.shape[0] / 2.))

        from matplotlib.widgets import Slider

        self.slice_slider_ax = self.fig.add_axes([0.1, 0.95, 0.4, 0.03])
        self.slice_slider_ax.set_xticklabels("")
        self.slice_slider_ax.set_yticklabels("")
        self.slice_slider = Slider(self.slice_slider_ax, "3-d slice", 0, self.array.shape[0], valinit=self.slice, valfmt="%i")
        self.slice_slider.on_changed(self.update_slice)
        self.slice_slider.drawon = False

        self.image = self.ax1.imshow(self.array[self.slice, :,:], origin='lower', interpolation='nearest', vmin=self._clim[0], vmax=self._clim[1], cmap=plt.cm.gray)

        self.vmin_slider_ax = self.fig.add_axes([0.1, 0.90, 0.4, 0.03])
        self.vmin_slider_ax.set_xticklabels("")
        self.vmin_slider_ax.set_yticklabels("")
        self.vmin_slider = Slider(self.vmin_slider_ax, "vmin", self._clim[0], self._clim[1], valinit=self._clim[0])
        self.vmin_slider.on_changed(self.update_vmin)
        self.vmin_slider.drawon = False

        self.vmax_slider_ax = self.fig.add_axes([0.1, 0.85, 0.4, 0.03])
        self.vmax_slider_ax.set_xticklabels("")
        self.vmax_slider_ax.set_yticklabels("")
        self.vmax_slider = Slider(self.vmax_slider_ax, "vmax", self._clim[0], self._clim[1], valinit=self._clim[1])
        self.vmax_slider.on_changed(self.update_vmax)
        self.vmax_slider.drawon = False

        self.grid1 = None
        self.grid2 = None
        self.grid3 = None

        self.ax2 = self.fig.add_axes([0.55, 0.1, 0.4, 0.7])

        # Add slicing box
        self.box = SliceCurve(colors=(0.8, 0.0, 0.0))
        self.ax1.add_collection(self.box)
        self.movable = MovableSliceBox(self.box, callback=self.update_pv_slice)
        self.movable.connect()

        # Add save button
        from matplotlib.widgets import Button
        self.save_button_ax = self.fig.add_axes([0.65, 0.90, 0.20, 0.05])
        self.save_button = Button(self.save_button_ax, 'Save slice to FITS')
        self.save_button.on_clicked(self.save_fits)

        self.pv_slice = None

        self.cidpress = self.fig.canvas.mpl_connect('button_press_event', self.click)

    def click(self, event):

        if event.inaxes != self.ax2:
            return

        self.slice_slider.set_val(event.ydata)

    def save_fits(self, *args, **kwargs):

        if self.backend == 'Qt4Agg':
            from matplotlib.backends.backend_qt4 import _getSaveFileName
            plot_name = _getSaveFileName(self.fig.canvas.manager.window, "Choose filename",
                                         os.path.dirname(os.path.abspath(self.filename)), "", None)
            plot_name = str(plot_name)
        else:
            print("Enter filename: ", end='')
            plot_name = raw_input()

        if self.pv_slice is None:
            return

        from astropy.io import fits
        fits.writeto(plot_name, self.pv_slice, clobber=True)
        print("Saved file to: ", plot_name)

    def update_pv_slice(self, box):

        path = Path(zip(box.x, box.y))
        path.width = box.width

        self.pv_slice = extract_pv_slice(self.array, path)

        self.ax2.cla()
        self.ax2.imshow(self.pv_slice.data, origin='lower', aspect='auto', interpolation='nearest')

        self.fig.canvas.draw()

    def show(self):
        import matplotlib.pyplot as plt
        plt.show()

    def update_slice(self, pos=None):

        if self.array.ndim == 2:
            self.image.set_array(self.array)
        else:
            self.slice = int(round(pos))
            self.image.set_array(self.array[self.slice, :, :])

        self.fig.canvas.draw()

    def update_vmin(self, vmin):
        if vmin > self._clim[1]:
            self._clim = (self._clim[1], self._clim[1])
        else:
            self._clim = (vmin, self._clim[1])
        self.image.set_clim(*self._clim)
        self.fig.canvas.draw()

    def update_vmax(self, vmax):
        if vmax < self._clim[0]:
            self._clim = (self._clim[0], self._clim[0])
        else:
            self._clim = (self._clim[0], vmax)
        self.image.set_clim(*self._clim)
        self.fig.canvas.draw()
<|endoftext|>


<commit_before>import logging
import os
import sys
import tempfile
from argparse import ArgumentParser
from pathlib import Path

import daemon
import psutil

from i3configger import __version__
from i3configger.lib import I3Configger, IniConfig, IpcControl

log = logging.getLogger()

def configure_logging(verbose, logfile=None, daemon=False):
    level = logging.DEBUG if verbose else logging.INFO
    fmt = '%(asctime)s %(name)s %(levelname)s: %(message)s'
    if daemon:
        if not logfile:
            logfile = Path(tempfile.gettempdir()) / 'i3configger.log'
        logging.basicConfig(filename=logfile, format=fmt, level=level)
        return
    logging.basicConfig(stream=sys.stdout, format=fmt, level=level)
    if not logfile:
        return
    fileHandler = logging.FileHandler(Path(logfile).expanduser())
    fileHandler.setFormatter(logging.Formatter(fmt))
    fileHandler.setLevel(level)
    log.addHandler(fileHandler)
    log.debug("logging to %s", logfile)

def parse_args():
    p = ArgumentParser('i3configger')
    p.add_argument('--ini-path', action="store",
                   help="path to i3configger.ini", default=None)
    p.add_argument('--verbose', action="store_true", default=False)
    p.add_argument('--version', action='version', version=__version__)
    p.add_argument('--daemon', action="store_true",
                   help="watch and build as daemon", default=False)
    p.add_argument('--watch', action="store_true",
                   help="watch and build in foreground", default=False)
    p.add_argument('--kill', action="store_true",
                   help="kill the deamon if it is running", default=False)
    return p.parse_args()

def get_foreign_processes():
    procs = [p for p in psutil.process_iter() if p.name() == 'i3configger']
    return [p for p in procs if p.pid != os.getpid()]

def daemonize(buildDefs, maxerrors, verbose, logfile=None):
    foreignProcesses = get_foreign_processes()
    if foreignProcesses:
        sys.exit("i3configger already running (%s)" % foreignProcesses)
    context = daemon.DaemonContext(
        working_directory=Path(__file__).parent,
        # TODO check if this umask is ok
        umask=0o002)
    # todo handle signals properly
    # context.signal_map = {
    #     signal.SIGTERM: program_cleanup,
    #     signal.SIGHUP: 'terminate',
    #     signal.SIGUSR1: reload_program_config}
    if verbose:
        # spew output to terminal from where daemon was started
        context.stdout = sys.stdout
        context.stderr = sys.stderr
    with context:
        configure_logging(verbose, logfile, daemon=True)
        i3configger = I3Configger(buildDefs, maxerrors)
        i3configger.watch_guarded()

def main():
    args = parse_args()
    cnf = IniConfig(IniConfig.get_config(args.ini_path))
    log.debug("config: %s", cnf)
    if args.kill:
        # todo some error handling
        for process in get_foreign_processes():
            print("killing %s" % process.pid)
            process.kill()
        return 0
    if args.daemon:
        daemonize(cnf.buildDefs, cnf.maxerrors, args.verbose, cnf.logfile)
    configure_logging(args.verbose, cnf.logfile)
    i3Configger = I3Configger(cnf.buildDefs, cnf.maxerrors)
    if args.watch:
        i3Configger.watch()
    else:
        i3Configger.build()
        # todo is there a way to reload the status bar without restarting i3?
        IpcControl.restart_i3()

if __name__ == '__main__':
    # for dev
    sys.argv = ['dev', '--verbose', '--watch']
    sys.exit(main())
<commit_msg>catch keyboard interrupt<commit_after>import logging
import os
import sys
import tempfile
from argparse import ArgumentParser
from pathlib import Path

import daemon
import psutil

from i3configger import __version__
from i3configger.lib import I3Configger, IniConfig, IpcControl

log = logging.getLogger()

def configure_logging(verbose, logfile=None, daemon=False):
    level = logging.DEBUG if verbose else logging.INFO
    fmt = '%(asctime)s %(name)s %(levelname)s: %(message)s'
    if daemon:
        if not logfile:
            logfile = Path(tempfile.gettempdir()) / 'i3configger.log'
        logging.basicConfig(filename=logfile, format=fmt, level=level)
        return
    logging.basicConfig(stream=sys.stdout, format=fmt, level=level)
    if not logfile:
        return
    fileHandler = logging.FileHandler(Path(logfile).expanduser())
    fileHandler.setFormatter(logging.Formatter(fmt))
    fileHandler.setLevel(level)
    log.addHandler(fileHandler)
    log.debug("logging to %s", logfile)

def parse_args():
    p = ArgumentParser('i3configger')
    p.add_argument('--ini-path', action="store",
                   help="path to i3configger.ini", default=None)
    p.add_argument('--verbose', action="store_true", default=False)
    p.add_argument('--version', action='version', version=__version__)
    p.add_argument('--daemon', action="store_true",
                   help="watch and build as daemon", default=False)
    p.add_argument('--watch', action="store_true",
                   help="watch and build in foreground", default=False)
    p.add_argument('--kill', action="store_true",
                   help="kill the deamon if it is running", default=False)
    return p.parse_args()

def get_foreign_processes():
    procs = [p for p in psutil.process_iter() if p.name() == 'i3configger']
    return [p for p in procs if p.pid != os.getpid()]

def daemonize(buildDefs, maxerrors, verbose, logfile=None):
    foreignProcesses = get_foreign_processes()
    if foreignProcesses:
        sys.exit("i3configger already running (%s)" % foreignProcesses)
    context = daemon.DaemonContext(
        working_directory=Path(__file__).parent,
        # TODO check if this umask is ok
        umask=0o002)
    # todo handle signals properly
    # context.signal_map = {
    #     signal.SIGTERM: program_cleanup,
    #     signal.SIGHUP: 'terminate',
    #     signal.SIGUSR1: reload_program_config}
    if verbose:
        # spew output to terminal from where daemon was started
        context.stdout = sys.stdout
        context.stderr = sys.stderr
    with context:
        configure_logging(verbose, logfile, daemon=True)
        i3configger = I3Configger(buildDefs, maxerrors)
        i3configger.watch_guarded()

def main():
    args = parse_args()
    cnf = IniConfig(IniConfig.get_config(args.ini_path))
    log.debug("config: %s", cnf)
    if args.kill:
        # todo some error handling
        for process in get_foreign_processes():
            print("killing %s" % process.pid)
            process.kill()
        return 0
    if args.daemon:
        daemonize(cnf.buildDefs, cnf.maxerrors, args.verbose, cnf.logfile)
    configure_logging(args.verbose, cnf.logfile)
    i3Configger = I3Configger(cnf.buildDefs, cnf.maxerrors)
    if args.watch:
        try:
            i3Configger.watch()
        except KeyboardInterrupt:
            sys.exit("bye")
    else:
        i3Configger.build()
        # todo is there a way to reload the status bar without restarting i3?
        IpcControl.restart_i3()

if __name__ == '__main__':
    # for dev
    sys.argv = ['dev', '--verbose', '--watch']
    sys.exit(main())
<|endoftext|>


<commit_before>import json
import logging
from datetime import datetime
from requests import Request, Session

log = logging.getLogger(__name__)

def to_utf(s):
    return s.encode('utf8')

class HAProxyService(object):
    """
    Generic service object representing a proxy component
    params:
     - fields(list): Fieldnames as read from haproxy stats export header
     - values(list): Stats for corresponding fields given above for this
                     frontend, backend, or listener
    """
    def __init__(self,fields,values,proxy_name=None):
        self.proxy_name = proxy_name

        #zip field names and values
        self.__dict__ = dict(zip(fields, self._read(values)))

        if self.svname == 'FRONTEND' or self.svname == 'BACKEND':
            self.name = self.pxname
        else:
            self.name =  self.svname

    def to_json(self):
        return json.dumps(self.__dict__)

    def _read(self,values):
        """
        Read stat str, convert unicode to utf and string to int where needed
        and return as list
        """
        ret = []

        for v in values:
            if v.isdigit():
                v = int(v)
            if isinstance(v,unicode):
                v = to_utf(v)
            ret.append(v)

        return ret

    def __str__(self):
        return to_json()

class HAProxyServer(object):
    """
    HAProxyServer object is created for each haproxy server we poll along with
    corresponding frontend, backend, and listener services.
    """
    def __init__(self,base_url,auth=None):
        self._auth = auth
        self.failed = False

        self.name = base_url.split(':')[0]
        self.url = 'http://' +  base_url + '/;csv;norefresh'

    def fetch_stats(self):
        """
        Fetch and parse stats from this Haproxy instance
        """
        self.frontends = []
        self.backends = []
        self.listeners = []

        csv = [ l for l in self._get(self.url).strip(' #').split('\n') if l ]
        #read fields header to create keys
        fields = [ to_utf(f) for f in csv.pop(0).split(',') if f ]

        #add frontends and backends first
        for line in csv:
            service = HAProxyService(fields, line.split(','), proxy_name=self.name)

            if service.svname == 'FRONTEND':
                self.frontends.append(service)
            elif service.svname == 'BACKEND':
                service.listener_names = []
                self.backends.append(service)
            else:
                self.listeners.append(service)

        #now add listener  names to corresponding backends
        for listener in self.listeners:
            for backend in self.backends:
                if backend.iid == listener.iid:
                    backend.listener_names.append(listener.name)

        self.stats = { 'frontends': self.frontends,
                       'backends': self.backends,
                       'listeners': self.listeners }

        self.last_update = datetime.utcnow()

    def _get(self,url):
        s = Session()
        if None in self._auth:
            req = Request('GET',url)
        else:
            req = Request('GET',url,auth=self._auth)

        try:
            r = s.send(req.prepare(),timeout=10)
        except Exception as e:
            log.warn('Error fetching stats from %s:\n%s' % (url,e))
            self.failed = True
            return ''

        return r.text

class HaproxyStats(object):
    """
    params:
     - servers(list) - List of haproxy instances defined as
       hostname:port or ip:port
     - user(str) -  User to authenticate with via basic auth(optional)
     - user_pass(str) -  Password to authenticate with via basic auth(optional)
    """
    def __init__(self,servers,user=None,user_pass=None):
        self.servers = [ HAProxyServer(s,auth=(user,user_pass)) for s in servers ]

        self.update()

    def update(self):
        start = datetime.utcnow()

        for s in self.servers:
            s.fetch_stats()

        duration = (datetime.utcnow() - start).total_seconds()
        log.info('Polled stats from %s servers in %s seconds' % \
                (len(self.servers),duration))

        if self.get_failed():
            return False

        return True

    def all_stats(self):
        return { s.name : s.stats for s in self.servers }

    def to_json(self):
        print(self.all_stats())
        return json.dumps(self.all_stats())

    def get_failed(self):
        return [ s for s in self.servers if s.failed ]
<commit_msg>remove to_json method, make proxy stats a dict of dicts<commit_after>import json
import logging
from datetime import datetime
from requests import Request, Session

log = logging.getLogger(__name__)

def to_utf(s):
    return s.encode('utf8')

class HAProxyService(object):
    """
    Generic service object representing a proxy component
    params:
     - fields(list): Fieldnames as read from haproxy stats export header
     - values(list): Stats for corresponding fields given above for this
                     frontend, backend, or listener
    """
    def __init__(self,fields,values,proxy_name=None):
        self.proxy_name = proxy_name

        #zip field names and values
        self.__dict__ = dict(zip(fields, self._read(values)))

        if self.svname == 'FRONTEND' or self.svname == 'BACKEND':
            self.name = self.pxname
        else:
            self.name =  self.svname

    def _read(self,values):
        """
        Read stat str, convert unicode to utf and string to int where needed
        and return as list
        """
        ret = []

        for v in values:
            if v.isdigit():
                v = int(v)
            if isinstance(v,unicode):
                v = to_utf(v)
            ret.append(v)

        return ret

class HAProxyServer(object):
    """
    HAProxyServer object is created for each haproxy server we poll along with
    corresponding frontend, backend, and listener services.
    """
    def __init__(self,base_url,auth=None):
        self._auth = auth
        self.failed = False

        self.name = base_url.split(':')[0]
        self.url = 'http://' +  base_url + '/;csv;norefresh'

    def fetch_stats(self):
        """
        Fetch and parse stats from this Haproxy instance
        """
        self.frontends = []
        self.backends = []
        self.listeners = []

        csv = [ l for l in self._get(self.url).strip(' #').split('\n') if l ]
        if not csv:
            self.failed = True
            return

        #read fields header to create keys
        fields = [ to_utf(f) for f in csv.pop(0).split(',') if f ]

        #add frontends and backends first
        for line in csv:
            service = HAProxyService(fields, line.split(','), proxy_name=self.name)

            if service.svname == 'FRONTEND':
                self.frontends.append(service)
            elif service.svname == 'BACKEND':
                service.listener_names = []
                self.backends.append(service)
            else:
                self.listeners.append(service)

        #now add listener  names to corresponding backends
        for listener in self.listeners:
            for backend in self.backends:
                if backend.iid == listener.iid:
                    backend.listener_names.append(listener.name)

        self.stats = { 'frontends': [ s.__dict__ for s in self.frontends ],
                       'backends': [ s.__dict__ for s in self.backends ],
                       'listeners': [ s.__dict__ for s in self.listeners ] }

        self.last_update = datetime.utcnow()

    def _get(self,url):
        s = Session()
        if None in self._auth:
            req = Request('GET',url)
        else:
            req = Request('GET',url,auth=self._auth)

        try:
            r = s.send(req.prepare(),timeout=10)
        except Exception as e:
            log.warn('Error fetching stats from %s:\n%s' % (url,e))
            return ''

        return r.text

class HaproxyStats(object):
    """
    params:
     - servers(list) - List of haproxy instances defined as
       hostname:port or ip:port
     - user(str) -  User to authenticate with via basic auth(optional)
     - user_pass(str) -  Password to authenticate with via basic auth(optional)
    """
    def __init__(self,servers,user=None,user_pass=None):
        self.servers = [ HAProxyServer(s,auth=(user,user_pass)) for s in servers ]

        self.update()

    def update(self):
        start = datetime.utcnow()

        for s in self.servers:
            s.fetch_stats()

        duration = (datetime.utcnow() - start).total_seconds()
        log.info('Polled stats from %s servers in %s seconds' % \
                (len(self.servers),duration))

        if self.get_failed():
            return False

        return True

    def all_stats(self):
        return { s.name : s.stats for s in self.servers }

    def to_json(self):
        return json.dumps(self.all_stats())

    def get_failed(self):
        return [ s for s in self.servers if s.failed ]
<|endoftext|>