<commit_before>"""Helper class to quickly write a loop over all standard input files.

Typical use is:

    import fileinput
    for line in fileinput.input():
        process(line)

This iterates over the lines of all files listed in sys.argv[1:],
defaulting to sys.stdin if the list is empty.  If a filename is '-' it
is also replaced by sys.stdin.  To specify an alternative list of
filenames, pass it as the argument to input().  A single file name is
also allowed.

Functions filename(), lineno() return the filename and cumulative line
number of the line that has just been read; filelineno() returns its
line number in the current file; isfirstline() returns true iff the
line just read is the first line of its file; isstdin() returns true
iff the line was read from sys.stdin.  Function nextfile() closes the
current file so that the next iteration will read the first line from
the next file (if any); lines not read from the file will not count
towards the cumulative line count; the filename is not changed until
after the first line of the next file has been read.  Function close()
closes the sequence.

Before any lines have been read, filename() returns None and both line
numbers are zero; nextfile() has no effect.  After all lines have been
read, filename() and the line number functions return the values
pertaining to the last line read; nextfile() has no effect.

All files are opened in text mode.  If an I/O error occurs during
opening or reading a file, the IOError exception is raised.

If sys.stdin is used more than once, the second and further use will
return no lines, except perhaps for interactive use, or if it has been
explicitly reset (e.g. using sys.stdin.seek(0)).

Empty files are opened and immediately closed; the only time their
presence in the list of filenames is noticeable at all is when the
last file opened is empty.

It is possible that the last line of a file doesn't end in a newline
character; otherwise lines are returned including the trailing
newline.

Class FileInput is the implementation; its methods filename(),
lineno(), fileline(), isfirstline(), isstdin(), nextfile() and close()
correspond to the functions in the module.  In addition it has a
readline() method which returns the next input line, and a
__getitem__() method which implements the sequence behavior.  The
sequence must be accessed in strictly sequential order; sequence
access and readline() cannot be mixed.

Optional in-place filtering: if the keyword argument inplace=1 is
passed to input() or to the FileInput constructor, the file is moved
to a backup file and standard output is directed to the input file.
This makes it possible to write a filter that rewrites its input file
in place.  If the keyword argument backup=".<some extension>" is also
given, it specifies the extension for the backup file, and the backup
file remains around; by default, the extension is ".bak" and it is
deleted when the output file is closed.  In-place filtering is
disabled when standard input is read.  XXX The current implementation
does not work for MS-DOS 8+3 filesystems.

XXX Possible additions:

- optional getopt argument processing
- specify open mode ('r' or 'rb')
- specify buffer size
- fileno()
- isatty()
- read(), read(size), even readlines()

"""

import sys, os

_state = None

def input(files=(), inplace=0, backup=""):
    global _state
    if _state and _state._file:
        raise RuntimeError, "input() already active"
    _state = FileInput(files, inplace, backup)
    return _state

def close():
    global _state
    state = _state
    _state = None
    if state:
        state.close()

def nextfile():
    if not _state:
        raise RuntimeError, "no active input()"
    return _state.nextfile()

def filename():
    if not _state:
        raise RuntimeError, "no active input()"
    return _state.filename()

def lineno():
    if not _state:
        raise RuntimeError, "no active input()"
    return _state.lineno()

def filelineno():
    if not _state:
        raise RuntimeError, "no active input()"
    return _state.filelineno()

def isfirstline():
    if not _state:
        raise RuntimeError, "no active input()"
    return _state.isfirstline()

def isstdin():
    if not _state:
        raise RuntimeError, "no active input()"
    return _state.isstdin()

class FileInput:

    def __init__(self, files=(), inplace=0, backup=""):
        if type(files) == type(''):
            files = (files,)
        else:
            files = tuple(files)
            if not files:
                files = tuple(sys.argv[1:])
                if not files:
                    files = ('-',)
        self._files = files
        self._inplace = inplace
        self._backup = backup
        self._savestdout = None
        self._output = None
        self._filename = None
        self._lineno = 0
        self._filelineno = 0
        self._file = None
        self._isstdin = 0
        self._backupfilename = None

    def __del__(self):
        self.close()

    def close(self):
        self.nextfile()
        self._files = ()

    def __getitem__(self, i):
        if i != self._lineno:
            raise RuntimeError, "accessing lines out of order"
        line = self.readline()
        if not line:
            raise IndexError, "end of input reached"
        return line

    def nextfile(self):
        savestdout = self._savestdout
        self._savestdout = 0
        if savestdout:
            sys.stdout = savestdout

        output = self._output
        self._output = 0
        if output:
            output.close()

        file = self._file
        self._file = 0
        if file and not self._isstdin:
            file.close()

        backupfilename = self._backupfilename
        self._backupfilename = 0
        if backupfilename and not self._backup:
            try: os.unlink(backupfilename)
            except: pass

        self._isstdin = 0

    def readline(self):
        if not self._file:
            if not self._files:
                return ""
            self._filename = self._files[0]
            self._files = self._files[1:]
            self._filelineno = 0
            self._file = None
            self._isstdin = 0
            self._backupfilename = 0
            if self._filename == '-':
                self._filename = '<stdin>'
                self._file = sys.stdin
                self._isstdin = 1
            else:
                if self._inplace:
                    self._backupfilename = (
                        self._filename + (self._backup or ".bak"))
                    try: os.unlink(self._backupfilename)
                    except os.error: pass
                    # The next three lines may raise IOError
                    os.rename(self._filename, self._backupfilename)
                    self._file = open(self._backupfilename, "r")
                    self._output = open(self._filename, "w")
                    self._savestdout = sys.stdout
                    sys.stdout = self._output
                else:
                    # This may raise IOError
                    self._file = open(self._filename, "r")
        line = self._file.readline()
        if line:
            self._lineno = self._lineno + 1
            self._filelineno = self._filelineno + 1
            return line
        self.nextfile()
        # Recursive call
        return self.readline()

    def filename(self):
        return self._filename

    def lineno(self):
        return self._lineno

    def filelineno(self):
        return self._filelineno

    def isfirstline(self):
        return self._filelineno == 1

    def isstdin(self):
        return self._isstdin

def _test():
    import getopt
    inplace = 0
    backup = 0
    opts, args = getopt.getopt(sys.argv[1:], "ib:")
    for o, a in opts:
        if o == '-i': inplace = 1
        if o == '-b': backup = a
    for line in input(args, inplace=inplace, backup=backup):
        if line[-1:] == '\n': line = line[:-1]
        if line[-1:] == '\r': line = line[:-1]
        print "%d: %s[%d]%s %s" % (lineno(), filename(), filelineno(),
                                   isfirstline() and "*" or "", line)
    print "%d: %s[%d]" % (lineno(), filename(), filelineno())

if __name__ == '__main__':
    _test()
<commit_msg>Fix for PR#111: when using the inplace option, give the new file the same permissions as the old file, plugging a security hole. (Not using exactly the suggested bugfix.)<commit_after>"""Helper class to quickly write a loop over all standard input files.

Typical use is:

    import fileinput
    for line in fileinput.input():
        process(line)

This iterates over the lines of all files listed in sys.argv[1:],
defaulting to sys.stdin if the list is empty.  If a filename is '-' it
is also replaced by sys.stdin.  To specify an alternative list of
filenames, pass it as the argument to input().  A single file name is
also allowed.

Functions filename(), lineno() return the filename and cumulative line
number of the line that has just been read; filelineno() returns its
line number in the current file; isfirstline() returns true iff the
line just read is the first line of its file; isstdin() returns true
iff the line was read from sys.stdin.  Function nextfile() closes the
current file so that the next iteration will read the first line from
the next file (if any); lines not read from the file will not count
towards the cumulative line count; the filename is not changed until
after the first line of the next file has been read.  Function close()
closes the sequence.

Before any lines have been read, filename() returns None and both line
numbers are zero; nextfile() has no effect.  After all lines have been
read, filename() and the line number functions return the values
pertaining to the last line read; nextfile() has no effect.

All files are opened in text mode.  If an I/O error occurs during
opening or reading a file, the IOError exception is raised.

If sys.stdin is used more than once, the second and further use will
return no lines, except perhaps for interactive use, or if it has been
explicitly reset (e.g. using sys.stdin.seek(0)).

Empty files are opened and immediately closed; the only time their
presence in the list of filenames is noticeable at all is when the
last file opened is empty.

It is possible that the last line of a file doesn't end in a newline
character; otherwise lines are returned including the trailing
newline.

Class FileInput is the implementation; its methods filename(),
lineno(), fileline(), isfirstline(), isstdin(), nextfile() and close()
correspond to the functions in the module.  In addition it has a
readline() method which returns the next input line, and a
__getitem__() method which implements the sequence behavior.  The
sequence must be accessed in strictly sequential order; sequence
access and readline() cannot be mixed.

Optional in-place filtering: if the keyword argument inplace=1 is
passed to input() or to the FileInput constructor, the file is moved
to a backup file and standard output is directed to the input file.
This makes it possible to write a filter that rewrites its input file
in place.  If the keyword argument backup=".<some extension>" is also
given, it specifies the extension for the backup file, and the backup
file remains around; by default, the extension is ".bak" and it is
deleted when the output file is closed.  In-place filtering is
disabled when standard input is read.  XXX The current implementation
does not work for MS-DOS 8+3 filesystems.

XXX Possible additions:

- optional getopt argument processing
- specify open mode ('r' or 'rb')
- specify buffer size
- fileno()
- isatty()
- read(), read(size), even readlines()

"""

import sys, os, stat

_state = None

def input(files=(), inplace=0, backup=""):
    global _state
    if _state and _state._file:
        raise RuntimeError, "input() already active"
    _state = FileInput(files, inplace, backup)
    return _state

def close():
    global _state
    state = _state
    _state = None
    if state:
        state.close()

def nextfile():
    if not _state:
        raise RuntimeError, "no active input()"
    return _state.nextfile()

def filename():
    if not _state:
        raise RuntimeError, "no active input()"
    return _state.filename()

def lineno():
    if not _state:
        raise RuntimeError, "no active input()"
    return _state.lineno()

def filelineno():
    if not _state:
        raise RuntimeError, "no active input()"
    return _state.filelineno()

def isfirstline():
    if not _state:
        raise RuntimeError, "no active input()"
    return _state.isfirstline()

def isstdin():
    if not _state:
        raise RuntimeError, "no active input()"
    return _state.isstdin()

class FileInput:

    def __init__(self, files=(), inplace=0, backup=""):
        if type(files) == type(''):
            files = (files,)
        else:
            files = tuple(files)
            if not files:
                files = tuple(sys.argv[1:])
                if not files:
                    files = ('-',)
        self._files = files
        self._inplace = inplace
        self._backup = backup
        self._savestdout = None
        self._output = None
        self._filename = None
        self._lineno = 0
        self._filelineno = 0
        self._file = None
        self._isstdin = 0
        self._backupfilename = None

    def __del__(self):
        self.close()

    def close(self):
        self.nextfile()
        self._files = ()

    def __getitem__(self, i):
        if i != self._lineno:
            raise RuntimeError, "accessing lines out of order"
        line = self.readline()
        if not line:
            raise IndexError, "end of input reached"
        return line

    def nextfile(self):
        savestdout = self._savestdout
        self._savestdout = 0
        if savestdout:
            sys.stdout = savestdout

        output = self._output
        self._output = 0
        if output:
            output.close()

        file = self._file
        self._file = 0
        if file and not self._isstdin:
            file.close()

        backupfilename = self._backupfilename
        self._backupfilename = 0
        if backupfilename and not self._backup:
            try: os.unlink(backupfilename)
            except: pass

        self._isstdin = 0

    def readline(self):
        if not self._file:
            if not self._files:
                return ""
            self._filename = self._files[0]
            self._files = self._files[1:]
            self._filelineno = 0
            self._file = None
            self._isstdin = 0
            self._backupfilename = 0
            if self._filename == '-':
                self._filename = '<stdin>'
                self._file = sys.stdin
                self._isstdin = 1
            else:
                if self._inplace:
                    self._backupfilename = (
                        self._filename + (self._backup or ".bak"))
                    try: os.unlink(self._backupfilename)
                    except os.error: pass
                    # The next few lines may raise IOError
                    os.rename(self._filename, self._backupfilename)
                    self._file = open(self._backupfilename, "r")
                    try:
                        perm = os.fstat(self._file.fileno())[stat.ST_MODE]
                    except:
                        self._output = open(self._filename, "w")
                    else:
                        fd = os.open(self._filename,
                                     os.O_CREAT | os.O_WRONLY | os.O_TRUNC,
                                     perm)
                        self._output = os.fdopen(fd, "w")
                        try:
                            os.chmod(self._filename, perm)
                        except:
                            pass
                    self._savestdout = sys.stdout
                    sys.stdout = self._output
                else:
                    # This may raise IOError
                    self._file = open(self._filename, "r")
        line = self._file.readline()
        if line:
            self._lineno = self._lineno + 1
            self._filelineno = self._filelineno + 1
            return line
        self.nextfile()
        # Recursive call
        return self.readline()

    def filename(self):
        return self._filename

    def lineno(self):
        return self._lineno

    def filelineno(self):
        return self._filelineno

    def isfirstline(self):
        return self._filelineno == 1

    def isstdin(self):
        return self._isstdin

def _test():
    import getopt
    inplace = 0
    backup = 0
    opts, args = getopt.getopt(sys.argv[1:], "ib:")
    for o, a in opts:
        if o == '-i': inplace = 1
        if o == '-b': backup = a
    for line in input(args, inplace=inplace, backup=backup):
        if line[-1:] == '\n': line = line[:-1]
        if line[-1:] == '\r': line = line[:-1]
        print "%d: %s[%d]%s %s" % (lineno(), filename(), filelineno(),
                                   isfirstline() and "*" or "", line)
    print "%d: %s[%d]" % (lineno(), filename(), filelineno())

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


<commit_before>"""Cache lines from files.

This is intended to read lines from modules imported -- hence if a filename
is not found, it will look down the module search path for a file by
that name.
"""

import sys
import os
import re

__all__ = ["getline", "clearcache", "checkcache"]

def getline(filename, lineno, module_globals=None):
    lines = getlines(filename, module_globals)
    if 1 <= lineno <= len(lines):
        return lines[lineno-1]
    else:
        return ''

# The cache

cache = {} # The cache

def clearcache():
    """Clear the cache entirely."""

    global cache
    cache = {}

def getlines(filename, module_globals=None):
    """Get the lines for a file from the cache.
    Update the cache if it doesn't contain an entry for this file already."""

    if filename in cache:
        return cache[filename][2]
    else:
        return updatecache(filename, module_globals)

def checkcache(filename=None):
    """Discard cache entries that are out of date.
    (This is not checked upon each call!)"""

    if filename is None:
        filenames = list(cache.keys())
    else:
        if filename in cache:
            filenames = [filename]
        else:
            return

    for filename in filenames:
        size, mtime, lines, fullname = cache[filename]
        if mtime is None:
            continue   # no-op for files loaded via a __loader__
        try:
            stat = os.stat(fullname)
        except os.error:
            del cache[filename]
            continue
        if size != stat.st_size or mtime != stat.st_mtime:
            del cache[filename]

def updatecache(filename, module_globals=None):
    """Update a cache entry and return its list of lines.
    If something's wrong, print a message, discard the cache entry,
    and return an empty list."""

    if filename in cache:
        del cache[filename]
    if not filename or filename[0] + filename[-1] == '<>':
        return []

    fullname = filename
    try:
        stat = os.stat(fullname)
    except os.error as msg:
        basename = os.path.split(filename)[1]

        # Try for a __loader__, if available
        if module_globals and '__loader__' in module_globals:
            name = module_globals.get('__name__')
            loader = module_globals['__loader__']
            get_source = getattr(loader, 'get_source', None)

            if name and get_source:
                if basename.startswith(name.split('.')[-1]+'.'):
                    try:
                        data = get_source(name)
                    except (ImportError, IOError):
                        pass
                    else:
                        if data is None:
                            # No luck, the PEP302 loader cannot find the source
                            # for this module.
                            return []
                        cache[filename] = (
                            len(data), None,
                            [line+'\n' for line in data.splitlines()], fullname
                        )
                        return cache[filename][2]

        # Try looking through the module search path.

        for dirname in sys.path:
            # When using imputil, sys.path may contain things other than
            # strings; ignore them when it happens.
            try:
                fullname = os.path.join(dirname, basename)
            except (TypeError, AttributeError):
                # Not sufficiently string-like to do anything useful with.
                pass
            else:
                try:
                    stat = os.stat(fullname)
                    break
                except os.error:
                    pass
        else:
            # No luck
##          print '*** Cannot stat', filename, ':', msg
            return []
    try:
        fp = open(fullname, 'rU')
        lines = fp.readlines()
        fp.close()
    except IOError as msg:
##      print '*** Cannot open', fullname, ':', msg
        return []
    coding = "utf-8"
    for line in lines[:2]:
        m = re.search(r"coding[:=]\s*([-\w.]+)", line)
        if m:
            coding = m.group(1)
            break
    try:
        lines = [line if isinstance(line, str) else str(line, coding)
                 for line in lines]
    except:
        pass  # Hope for the best
    size, mtime = stat.st_size, stat.st_mtime
    cache[filename] = size, mtime, lines, fullname
    return lines
<commit_msg>Be even more lenient when catching errors from readlines().<commit_after>"""Cache lines from files.

This is intended to read lines from modules imported -- hence if a filename
is not found, it will look down the module search path for a file by
that name.
"""

import sys
import os
import re

__all__ = ["getline", "clearcache", "checkcache"]

def getline(filename, lineno, module_globals=None):
    lines = getlines(filename, module_globals)
    if 1 <= lineno <= len(lines):
        return lines[lineno-1]
    else:
        return ''

# The cache

cache = {} # The cache

def clearcache():
    """Clear the cache entirely."""

    global cache
    cache = {}

def getlines(filename, module_globals=None):
    """Get the lines for a file from the cache.
    Update the cache if it doesn't contain an entry for this file already."""

    if filename in cache:
        return cache[filename][2]
    else:
        return updatecache(filename, module_globals)

def checkcache(filename=None):
    """Discard cache entries that are out of date.
    (This is not checked upon each call!)"""

    if filename is None:
        filenames = list(cache.keys())
    else:
        if filename in cache:
            filenames = [filename]
        else:
            return

    for filename in filenames:
        size, mtime, lines, fullname = cache[filename]
        if mtime is None:
            continue   # no-op for files loaded via a __loader__
        try:
            stat = os.stat(fullname)
        except os.error:
            del cache[filename]
            continue
        if size != stat.st_size or mtime != stat.st_mtime:
            del cache[filename]

def updatecache(filename, module_globals=None):
    """Update a cache entry and return its list of lines.
    If something's wrong, print a message, discard the cache entry,
    and return an empty list."""

    if filename in cache:
        del cache[filename]
    if not filename or filename[0] + filename[-1] == '<>':
        return []

    fullname = filename
    try:
        stat = os.stat(fullname)
    except os.error as msg:
        basename = os.path.split(filename)[1]

        # Try for a __loader__, if available
        if module_globals and '__loader__' in module_globals:
            name = module_globals.get('__name__')
            loader = module_globals['__loader__']
            get_source = getattr(loader, 'get_source', None)

            if name and get_source:
                if basename.startswith(name.split('.')[-1]+'.'):
                    try:
                        data = get_source(name)
                    except (ImportError, IOError):
                        pass
                    else:
                        if data is None:
                            # No luck, the PEP302 loader cannot find the source
                            # for this module.
                            return []
                        cache[filename] = (
                            len(data), None,
                            [line+'\n' for line in data.splitlines()], fullname
                        )
                        return cache[filename][2]

        # Try looking through the module search path.

        for dirname in sys.path:
            # When using imputil, sys.path may contain things other than
            # strings; ignore them when it happens.
            try:
                fullname = os.path.join(dirname, basename)
            except (TypeError, AttributeError):
                # Not sufficiently string-like to do anything useful with.
                pass
            else:
                try:
                    stat = os.stat(fullname)
                    break
                except os.error:
                    pass
        else:
            # No luck
##          print '*** Cannot stat', filename, ':', msg
            return []
##  print("Refreshing cache for %s..." % fullname)
    try:
        fp = open(fullname, 'rU')
        lines = fp.readlines()
        fp.close()
    except Exception as msg:
##      print '*** Cannot open', fullname, ':', msg
        return []
    coding = "utf-8"
    for line in lines[:2]:
        m = re.search(r"coding[:=]\s*([-\w.]+)", line)
        if m:
            coding = m.group(1)
            break
    try:
        lines = [line if isinstance(line, str) else str(line, coding)
                 for line in lines]
    except:
        pass  # Hope for the best
    size, mtime = stat.st_size, stat.st_mtime
    cache[filename] = size, mtime, lines, fullname
    return lines
<|endoftext|>


<commit_before>"""Various tools used by MIME-reading or MIME-writing programs."""

import os
import rfc822
import tempfile

__all__ = ["Message","choose_boundary","encode","decode","copyliteral",
           "copybinary"]

class Message(rfc822.Message):
    """A derived class of rfc822.Message that knows about MIME headers and
    contains some hooks for decoding encoded and multipart messages."""

    def __init__(self, fp, seekable = 1):
        rfc822.Message.__init__(self, fp, seekable)
        self.encodingheader = \
                self.getheader('content-transfer-encoding')
        self.typeheader = \
                self.getheader('content-type')
        self.parsetype()
        self.parseplist()

    def parsetype(self):
        str = self.typeheader
        if str is None:
            str = 'text/plain'
        if ';' in str:
            i = str.index(';')
            self.plisttext = str[i:]
            str = str[:i]
        else:
            self.plisttext = ''
        fields = str.split('/')
        for i in range(len(fields)):
            fields[i] = fields[i].strip().lower()
        self.type = '/'.join(fields)
        self.maintype = fields[0]
        self.subtype = '/'.join(fields[1:])

    def parseplist(self):
        str = self.plisttext
        self.plist = []
        while str[:1] == ';':
            str = str[1:]
            if ';' in str:
                # XXX Should parse quotes!
                end = str.index(';')
            else:
                end = len(str)
            f = str[:end]
            if '=' in f:
                i = f.index('=')
                f = f[:i].strip().lower() + \
                        '=' + f[i+1:].strip()
            self.plist.append(f.strip())
            str = str[end:]

    def getplist(self):
        return self.plist

    def getparam(self, name):
        name = name.lower() + '='
        n = len(name)
        for p in self.plist:
            if p[:n] == name:
                return rfc822.unquote(p[n:])
        return None

    def getparamnames(self):
        result = []
        for p in self.plist:
            i = p.find('=')
            if i >= 0:
                result.append(p[:i].lower())
        return result

    def getencoding(self):
        if self.encodingheader is None:
            return '7bit'
        return self.encodingheader.lower()

    def gettype(self):
        return self.type

    def getmaintype(self):
        return self.maintype

    def getsubtype(self):
        return self.subtype

# Utility functions
# -----------------

_prefix = None

def choose_boundary():
    """Return a random string usable as a multipart boundary.
    The method used is so that it is *very* unlikely that the same
    string of characters will every occur again in the Universe,
    so the caller needn't check the data it is packing for the
    occurrence of the boundary.

    The boundary contains dots so you have to quote it in the header."""

    global _prefix
    import time
    import random
    if _prefix is None:
        import socket
        hostid = socket.gethostbyname(socket.gethostname())
        try:
            uid = `os.getuid()`
        except AttributeError:
            uid = '1'
        try:
            pid = `os.getpid()`
        except AttributeError:
            pid = '1'
        _prefix = hostid + '.' + uid + '.' + pid
    timestamp = '%.3f' % time.time()
    seed = `random.randint(0, 32767)`
    return _prefix + '.' + timestamp + '.' + seed

# Subroutines for decoding some common content-transfer-types

def decode(input, output, encoding):
    """Decode common content-transfer-encodings (base64, quopri, uuencode)."""
    if encoding == 'base64':
        import base64
        return base64.decode(input, output)
    if encoding == 'quoted-printable':
        import quopri
        return quopri.decode(input, output)
    if encoding in ('uuencode', 'x-uuencode', 'uue', 'x-uue'):
        import uu
        return uu.decode(input, output)
    if encoding in ('7bit', '8bit'):
        return output.write(input.read())
    if encoding in decodetab:
        pipethrough(input, decodetab[encoding], output)
    else:
        raise ValueError, \
              'unknown Content-Transfer-Encoding: %s' % encoding

def encode(input, output, encoding):
    """Encode common content-transfer-encodings (base64, quopri, uuencode)."""
    if encoding == 'base64':
        import base64
        return base64.encode(input, output)
    if encoding == 'quoted-printable':
        import quopri
        return quopri.encode(input, output, 0)
    if encoding in ('uuencode', 'x-uuencode', 'uue', 'x-uue'):
        import uu
        return uu.encode(input, output)
    if encoding in ('7bit', '8bit'):
        return output.write(input.read())
    if encoding in encodetab:
        pipethrough(input, encodetab[encoding], output)
    else:
        raise ValueError, \
              'unknown Content-Transfer-Encoding: %s' % encoding

# The following is no longer used for standard encodings

# XXX This requires that uudecode and mmencode are in $PATH

uudecode_pipe = '''(
TEMP=/tmp/@uu.$$
sed "s%^begin [0-7][0-7]* .*%begin 600 $TEMP%" | uudecode
cat $TEMP
rm $TEMP
)'''

decodetab = {
        'uuencode':             uudecode_pipe,
        'x-uuencode':           uudecode_pipe,
        'uue':                  uudecode_pipe,
        'x-uue':                uudecode_pipe,
        'quoted-printable':     'mmencode -u -q',
        'base64':               'mmencode -u -b',
}

encodetab = {
        'x-uuencode':           'uuencode tempfile',
        'uuencode':             'uuencode tempfile',
        'x-uue':                'uuencode tempfile',
        'uue':                  'uuencode tempfile',
        'quoted-printable':     'mmencode -q',
        'base64':               'mmencode -b',
}

def pipeto(input, command):
    pipe = os.popen(command, 'w')
    copyliteral(input, pipe)
    pipe.close()

def pipethrough(input, command, output):
    (fd, tempname) = tempfile.mkstemp()
    temp = os.fdopen(fd, 'w')
    copyliteral(input, temp)
    temp.close()
    pipe = os.popen(command + ' <' + tempname, 'r')
    copybinary(pipe, output)
    pipe.close()
    os.unlink(tempname)

def copyliteral(input, output):
    while 1:
        line = input.readline()
        if not line: break
        output.write(line)

def copybinary(input, output):
    BUFSIZE = 8192
    while 1:
        line = input.read(BUFSIZE)
        if not line: break
        output.write(line)
<commit_msg>choose_boundary():  Incorporated a threadsafe incrementing counter, so that unique boundary strings within a program run are guaranteed.  On Windows, duplicates were pretty likely, due to the coarse granularity of time.time. Toned down the absurdly optimistic claims in the docstring.<commit_after>"""Various tools used by MIME-reading or MIME-writing programs."""

import os
import rfc822
import tempfile

__all__ = ["Message","choose_boundary","encode","decode","copyliteral",
           "copybinary"]

class Message(rfc822.Message):
    """A derived class of rfc822.Message that knows about MIME headers and
    contains some hooks for decoding encoded and multipart messages."""

    def __init__(self, fp, seekable = 1):
        rfc822.Message.__init__(self, fp, seekable)
        self.encodingheader = \
                self.getheader('content-transfer-encoding')
        self.typeheader = \
                self.getheader('content-type')
        self.parsetype()
        self.parseplist()

    def parsetype(self):
        str = self.typeheader
        if str is None:
            str = 'text/plain'
        if ';' in str:
            i = str.index(';')
            self.plisttext = str[i:]
            str = str[:i]
        else:
            self.plisttext = ''
        fields = str.split('/')
        for i in range(len(fields)):
            fields[i] = fields[i].strip().lower()
        self.type = '/'.join(fields)
        self.maintype = fields[0]
        self.subtype = '/'.join(fields[1:])

    def parseplist(self):
        str = self.plisttext
        self.plist = []
        while str[:1] == ';':
            str = str[1:]
            if ';' in str:
                # XXX Should parse quotes!
                end = str.index(';')
            else:
                end = len(str)
            f = str[:end]
            if '=' in f:
                i = f.index('=')
                f = f[:i].strip().lower() + \
                        '=' + f[i+1:].strip()
            self.plist.append(f.strip())
            str = str[end:]

    def getplist(self):
        return self.plist

    def getparam(self, name):
        name = name.lower() + '='
        n = len(name)
        for p in self.plist:
            if p[:n] == name:
                return rfc822.unquote(p[n:])
        return None

    def getparamnames(self):
        result = []
        for p in self.plist:
            i = p.find('=')
            if i >= 0:
                result.append(p[:i].lower())
        return result

    def getencoding(self):
        if self.encodingheader is None:
            return '7bit'
        return self.encodingheader.lower()

    def gettype(self):
        return self.type

    def getmaintype(self):
        return self.maintype

    def getsubtype(self):
        return self.subtype

# Utility functions
# -----------------

try:
    import thread
except ImportError:
    import dummy_thread as thread
_counter_lock = thread.allocate_lock()
del thread

_counter = 0
def _get_next_counter():
    global _counter
    _counter_lock.acquire()
    _counter += 1
    result = _counter
    _counter_lock.release()
    return result

_prefix = None

def choose_boundary():
    """Return a string usable as a multipart boundary.

    The string chosen is unique within a single program run, and
    incorporates the user id (if available), process id (if available),
    and current time.  So it's very unlikely the returned string appears
    in message text, but there's no guarantee.

    The boundary contains dots so you have to quote it in the header."""

    global _prefix
    import time
    import random
    if _prefix is None:
        import socket
        hostid = socket.gethostbyname(socket.gethostname())
        try:
            uid = `os.getuid()`
        except AttributeError:
            uid = '1'
        try:
            pid = `os.getpid()`
        except AttributeError:
            pid = '1'
        _prefix = hostid + '.' + uid + '.' + pid
    return "%s.%.3f.%d" % (_prefix, time.time(), _get_next_counter())

# Subroutines for decoding some common content-transfer-types

def decode(input, output, encoding):
    """Decode common content-transfer-encodings (base64, quopri, uuencode)."""
    if encoding == 'base64':
        import base64
        return base64.decode(input, output)
    if encoding == 'quoted-printable':
        import quopri
        return quopri.decode(input, output)
    if encoding in ('uuencode', 'x-uuencode', 'uue', 'x-uue'):
        import uu
        return uu.decode(input, output)
    if encoding in ('7bit', '8bit'):
        return output.write(input.read())
    if encoding in decodetab:
        pipethrough(input, decodetab[encoding], output)
    else:
        raise ValueError, \
              'unknown Content-Transfer-Encoding: %s' % encoding

def encode(input, output, encoding):
    """Encode common content-transfer-encodings (base64, quopri, uuencode)."""
    if encoding == 'base64':
        import base64
        return base64.encode(input, output)
    if encoding == 'quoted-printable':
        import quopri
        return quopri.encode(input, output, 0)
    if encoding in ('uuencode', 'x-uuencode', 'uue', 'x-uue'):
        import uu
        return uu.encode(input, output)
    if encoding in ('7bit', '8bit'):
        return output.write(input.read())
    if encoding in encodetab:
        pipethrough(input, encodetab[encoding], output)
    else:
        raise ValueError, \
              'unknown Content-Transfer-Encoding: %s' % encoding

# The following is no longer used for standard encodings

# XXX This requires that uudecode and mmencode are in $PATH

uudecode_pipe = '''(
TEMP=/tmp/@uu.$$
sed "s%^begin [0-7][0-7]* .*%begin 600 $TEMP%" | uudecode
cat $TEMP
rm $TEMP
)'''

decodetab = {
        'uuencode':             uudecode_pipe,
        'x-uuencode':           uudecode_pipe,
        'uue':                  uudecode_pipe,
        'x-uue':                uudecode_pipe,
        'quoted-printable':     'mmencode -u -q',
        'base64':               'mmencode -u -b',
}

encodetab = {
        'x-uuencode':           'uuencode tempfile',
        'uuencode':             'uuencode tempfile',
        'x-uue':                'uuencode tempfile',
        'uue':                  'uuencode tempfile',
        'quoted-printable':     'mmencode -q',
        'base64':               'mmencode -b',
}

def pipeto(input, command):
    pipe = os.popen(command, 'w')
    copyliteral(input, pipe)
    pipe.close()

def pipethrough(input, command, output):
    (fd, tempname) = tempfile.mkstemp()
    temp = os.fdopen(fd, 'w')
    copyliteral(input, temp)
    temp.close()
    pipe = os.popen(command + ' <' + tempname, 'r')
    copybinary(pipe, output)
    pipe.close()
    os.unlink(tempname)

def copyliteral(input, output):
    while 1:
        line = input.readline()
        if not line: break
        output.write(line)

def copybinary(input, output):
    BUFSIZE = 8192
    while 1:
        line = input.read(BUFSIZE)
        if not line: break
        output.write(line)
<|endoftext|>


<commit_before># A class that makes each part of a multipart message "feel" like an
# ordinary file, as long as you use fp.readline().  Allows recursive
# use, for nested multipart messages.  Probably best used together
# with module mimetools.
#
# Suggested use:
#
# real_fp = open(...)
# fp = MultiFile(real_fp)
#
# "read some lines from fp"
# fp.push(separator)
# while 1:
#	"read lines from fp until it returns an empty string" (A)
#	if not fp.next(): break
# fp.pop()
# "read remaining lines from fp until it returns an empty string"
#
# The latter sequence may be used recursively at (A).
# It is also allowed to use multiple push()...pop() sequences.
# Note that if a nested multipart message is terminated by a separator
# for an outer message, this is not reported, even though it is really
# illegal input.

import sys
import string

err = sys.stderr.write

Error = 'multifile.Error'

class MultiFile:
	#
	def __init__(self, fp):
		self.fp = fp
		self.stack = [] # Grows down
		self.level = 0
		self.last = 0
		self.start = self.fp.tell()
		self.posstack = [] # Grows down
	#
	def tell(self):
		if self.level > 0:
			return self.lastpos
		return self.fp.tell() - self.start
	#
	def seek(self, pos):
		if not 0 <= pos <= self.tell() or \
				self.level > 0 and pos > self.lastpos:
			raise Error, 'bad MultiFile.seek() call'
		self.fp.seek(pos + self.start)
		self.level = 0
		self.last = 0
	#
	def readline(self):
		if self.level > 0: return ''
		line = self.fp.readline()
		if not line:
			self.level = len(self.stack)
			self.last = (self.level > 0)
			if self.last:
				err('*** Sudden EOF in MultiFile.readline()\n')
			return ''
		if line[:2] <> '--': return line
		n = len(line)
		k = n
		while k > 0 and line[k-1] in string.whitespace: k = k-1
		mark = line[2:k]
		if mark[-2:] == '--': mark1 = mark[:-2]
		else: mark1 = None
		for i in range(len(self.stack)):
			sep = self.stack[i]
			if sep == mark:
				self.last = 0
				break
			elif mark1 <> None and sep == mark1:
				self.last = 1
				break
		else:
			return line
		# Get here after break out of loop
		self.lastpos = self.tell() - len(line)
		self.level = i+1
		if self.level > 1:
			err('*** Missing endmarker in MultiFile.readline()\n')
		return ''
	#
	def readlines(self):
		list = []
		while 1:
			line = self.readline()
			if not line: break
			list.append(line)
		return list
	#
	def read(self): # Note: no size argument -- read until EOF only!
		return string.joinfields(self.readlines(), '')
	#
	def next(self):
		while self.readline(): pass
		if self.level > 1 or self.last:
			return 0
		self.level = 0
		self.last = 0
		self.start = self.fp.tell()
		return 1
	#
	def push(self, sep):
		if self.level > 0:
			raise Error, 'bad MultiFile.push() call'
		self.stack.insert(0, sep)
		self.posstack.insert(0, self.start)
		self.start = self.fp.tell()
	#
	def pop(self):
		if self.stack == []:
			raise Error, 'bad MultiFile.pop() call'
		if self.level <= 1:
			self.last = 0
		else:
			abslastpos = self.lastpos + self.start
		self.level = max(0, self.level - 1)
		del self.stack[0]
		self.start = self.posstack[0]
		del self.posstack[0]
		if self.level > 0:
			self.lastpos = abslastpos - self.start
	#
<commit_msg>Support 'whence' parameter to seek().<commit_after># A class that makes each part of a multipart message "feel" like an
# ordinary file, as long as you use fp.readline().  Allows recursive
# use, for nested multipart messages.  Probably best used together
# with module mimetools.
#
# Suggested use:
#
# real_fp = open(...)
# fp = MultiFile(real_fp)
#
# "read some lines from fp"
# fp.push(separator)
# while 1:
#	"read lines from fp until it returns an empty string" (A)
#	if not fp.next(): break
# fp.pop()
# "read remaining lines from fp until it returns an empty string"
#
# The latter sequence may be used recursively at (A).
# It is also allowed to use multiple push()...pop() sequences.
# Note that if a nested multipart message is terminated by a separator
# for an outer message, this is not reported, even though it is really
# illegal input.

import sys
import string

err = sys.stderr.write

Error = 'multifile.Error'

class MultiFile:
	#
	def __init__(self, fp):
		self.fp = fp
		self.stack = [] # Grows down
		self.level = 0
		self.last = 0
		self.start = self.fp.tell()
		self.posstack = [] # Grows down
	#
	def tell(self):
		if self.level > 0:
			return self.lastpos
		return self.fp.tell() - self.start
	#
	def seek(self, pos, whence=0):
                here = self.tell()
                if whence:
                        if whence == 1:
                                pos = pos + here
                        elif whence == 2:
                                if self.level > 0:
                                        pos = pos + self.lastpos
                                else:
                                        raise Error, "can't use whence=2 yet"
		if not 0 <= pos <= here or \
				self.level > 0 and pos > self.lastpos:
			raise Error, 'bad MultiFile.seek() call'
		self.fp.seek(pos + self.start)
		self.level = 0
		self.last = 0
	#
	def readline(self):
		if self.level > 0: return ''
		line = self.fp.readline()
		if not line:
			self.level = len(self.stack)
			self.last = (self.level > 0)
			if self.last:
				err('*** Sudden EOF in MultiFile.readline()\n')
			return ''
		if line[:2] <> '--': return line
		n = len(line)
		k = n
		while k > 0 and line[k-1] in string.whitespace: k = k-1
		mark = line[2:k]
		if mark[-2:] == '--': mark1 = mark[:-2]
		else: mark1 = None
		for i in range(len(self.stack)):
			sep = self.stack[i]
			if sep == mark:
				self.last = 0
				break
			elif mark1 <> None and sep == mark1:
				self.last = 1
				break
		else:
			return line
		# Get here after break out of loop
		self.lastpos = self.tell() - len(line)
		self.level = i+1
		if self.level > 1:
			err('*** Missing endmarker in MultiFile.readline()\n')
		return ''
	#
	def readlines(self):
		list = []
		while 1:
			line = self.readline()
			if not line: break
			list.append(line)
		return list
	#
	def read(self): # Note: no size argument -- read until EOF only!
		return string.joinfields(self.readlines(), '')
	#
	def next(self):
		while self.readline(): pass
		if self.level > 1 or self.last:
			return 0
		self.level = 0
		self.last = 0
		self.start = self.fp.tell()
		return 1
	#
	def push(self, sep):
		if self.level > 0:
			raise Error, 'bad MultiFile.push() call'
		self.stack.insert(0, sep)
		self.posstack.insert(0, self.start)
		self.start = self.fp.tell()
	#
	def pop(self):
		if self.stack == []:
			raise Error, 'bad MultiFile.pop() call'
		if self.level <= 1:
			self.last = 0
		else:
			abslastpos = self.lastpos + self.start
		self.level = max(0, self.level - 1)
		del self.stack[0]
		self.start = self.posstack[0]
		del self.posstack[0]
		if self.level > 0:
			self.lastpos = abslastpos - self.start
	#
<|endoftext|>


<commit_before>"""
URL lookup for Company app
"""

from django.conf.urls import url, include

from . import views

company_detail_urls = [
    url(r'delete/?', views.CompanyDelete.as_view(), name='company-delete'),

    # url(r'orders/?', views.CompanyDetail.as_view(template_name='company/orders.html'), name='company-detail-orders'),

    url(r'^supplier-parts/', views.CompanyDetail.as_view(template_name='company/detail_supplier_part.html'), name='company-detail-supplier-parts'),
    url(r'^manufacturer-parts/', views.CompanyDetail.as_view(template_name='company/detail_manufacturer_part.html'), name='company-detail-manufacturer-parts'),
    url(r'^stock/', views.CompanyDetail.as_view(template_name='company/detail_stock.html'), name='company-detail-stock'),
    url(r'^purchase-orders/', views.CompanyDetail.as_view(template_name='company/purchase_orders.html'), name='company-detail-purchase-orders'),
    url(r'^assigned-stock/', views.CompanyDetail.as_view(template_name='company/assigned_stock.html'), name='company-detail-assigned-stock'),
    url(r'^sales-orders/', views.CompanyDetail.as_view(template_name='company/sales_orders.html'), name='company-detail-sales-orders'),
    url(r'^notes/', views.CompanyNotes.as_view(), name='company-notes'),

    url(r'^thumb-download/', views.CompanyImageDownloadFromURL.as_view(), name='company-image-download'),

    # Any other URL
    url(r'^.*$', views.CompanyDetail.as_view(), name='company-detail'),
]

company_urls = [

    url(r'^(?P<pk>\d+)/', include(company_detail_urls)),

    url(r'suppliers/', views.CompanyIndex.as_view(), name='supplier-index'),
    url(r'manufacturers/', views.CompanyIndex.as_view(), name='manufacturer-index'),
    url(r'customers/', views.CompanyIndex.as_view(), name='customer-index'),

    # Redirect any other patterns to the 'company' index which displays all companies
    url(r'^.*$', views.CompanyIndex.as_view(), name='company-index'),
]

manufacturer_part_urls = [
    url(r'^new/?', views.ManufacturerPartCreate.as_view(), name='manufacturer-part-create'),

    url(r'^delete/', views.ManufacturerPartDelete.as_view(), name='manufacturer-part-delete'),

    # URLs for ManufacturerPartParameter views (create / edit / delete)
    url(r'^parameter/', include([
        url(r'^new/', views.ManufacturerPartParameterCreate.as_view(), name='manufacturer-part-parameter-create'),
        url(r'^(?P<pk>\d)/', include([
            url(r'^edit/', views.ManufacturerPartParameterEdit.as_view(), name='manufacturer-part-parameter-edit'),
            url(r'^delete/', views.ManufacturerPartParameterDelete.as_view(), name='manufacturer-part-parameter-delete'),
        ])),
    ])),

    url(r'^(?P<pk>\d+)/', include([
        url(r'^edit/?', views.ManufacturerPartEdit.as_view(), name='manufacturer-part-edit'),
        url(r'^suppliers/', views.ManufacturerPartDetail.as_view(template_name='company/manufacturer_part_suppliers.html'), name='manufacturer-part-suppliers'),
        url('^.*$', views.ManufacturerPartDetail.as_view(template_name='company/manufacturer_part_suppliers.html'), name='manufacturer-part-detail'),
    ])),
]

supplier_part_detail_urls = [
    url(r'^edit/?', views.SupplierPartEdit.as_view(), name='supplier-part-edit'),

    url(r'^manufacturers/', views.SupplierPartDetail.as_view(template_name='company/supplier_part_manufacturers.html'), name='supplier-part-manufacturers'),
    url(r'^pricing/', views.SupplierPartDetail.as_view(template_name='company/supplier_part_pricing.html'), name='supplier-part-pricing'),
    url(r'^orders/', views.SupplierPartDetail.as_view(template_name='company/supplier_part_orders.html'), name='supplier-part-orders'),
    url(r'^stock/', views.SupplierPartDetail.as_view(template_name='company/supplier_part_stock.html'), name='supplier-part-stock'),

    url('^.*$', views.SupplierPartDetail.as_view(template_name='company/supplier_part_pricing.html'), name='supplier-part-detail'),
]

supplier_part_urls = [
    url(r'^new/?', views.SupplierPartCreate.as_view(), name='supplier-part-create'),

    url(r'delete/', views.SupplierPartDelete.as_view(), name='supplier-part-delete'),

    url(r'^(?P<pk>\d+)/', include(supplier_part_detail_urls)),
]
<commit_msg>Remove old URL endpoints<commit_after>"""
URL lookup for Company app
"""

from django.conf.urls import url, include

from . import views

company_detail_urls = [
    url(r'delete/?', views.CompanyDelete.as_view(), name='company-delete'),

    # url(r'orders/?', views.CompanyDetail.as_view(template_name='company/orders.html'), name='company-detail-orders'),

    url(r'^supplier-parts/', views.CompanyDetail.as_view(template_name='company/detail_supplier_part.html'), name='company-detail-supplier-parts'),
    url(r'^manufacturer-parts/', views.CompanyDetail.as_view(template_name='company/detail_manufacturer_part.html'), name='company-detail-manufacturer-parts'),
    url(r'^stock/', views.CompanyDetail.as_view(template_name='company/detail_stock.html'), name='company-detail-stock'),
    url(r'^purchase-orders/', views.CompanyDetail.as_view(template_name='company/purchase_orders.html'), name='company-detail-purchase-orders'),
    url(r'^assigned-stock/', views.CompanyDetail.as_view(template_name='company/assigned_stock.html'), name='company-detail-assigned-stock'),
    url(r'^sales-orders/', views.CompanyDetail.as_view(template_name='company/sales_orders.html'), name='company-detail-sales-orders'),
    url(r'^notes/', views.CompanyNotes.as_view(), name='company-notes'),

    url(r'^thumb-download/', views.CompanyImageDownloadFromURL.as_view(), name='company-image-download'),

    # Any other URL
    url(r'^.*$', views.CompanyDetail.as_view(), name='company-detail'),
]

company_urls = [

    url(r'^(?P<pk>\d+)/', include(company_detail_urls)),

    url(r'suppliers/', views.CompanyIndex.as_view(), name='supplier-index'),
    url(r'manufacturers/', views.CompanyIndex.as_view(), name='manufacturer-index'),
    url(r'customers/', views.CompanyIndex.as_view(), name='customer-index'),

    # Redirect any other patterns to the 'company' index which displays all companies
    url(r'^.*$', views.CompanyIndex.as_view(), name='company-index'),
]

manufacturer_part_urls = [
    url(r'^new/?', views.ManufacturerPartCreate.as_view(), name='manufacturer-part-create'),

    url(r'^delete/', views.ManufacturerPartDelete.as_view(), name='manufacturer-part-delete'),

    url(r'^(?P<pk>\d+)/', include([
        url(r'^edit/?', views.ManufacturerPartEdit.as_view(), name='manufacturer-part-edit'),
        url(r'^suppliers/', views.ManufacturerPartDetail.as_view(template_name='company/manufacturer_part_suppliers.html'), name='manufacturer-part-suppliers'),
        url('^.*$', views.ManufacturerPartDetail.as_view(template_name='company/manufacturer_part_suppliers.html'), name='manufacturer-part-detail'),
    ])),
]

supplier_part_detail_urls = [
    url(r'^edit/?', views.SupplierPartEdit.as_view(), name='supplier-part-edit'),

    url(r'^manufacturers/', views.SupplierPartDetail.as_view(template_name='company/supplier_part_manufacturers.html'), name='supplier-part-manufacturers'),
    url(r'^pricing/', views.SupplierPartDetail.as_view(template_name='company/supplier_part_pricing.html'), name='supplier-part-pricing'),
    url(r'^orders/', views.SupplierPartDetail.as_view(template_name='company/supplier_part_orders.html'), name='supplier-part-orders'),
    url(r'^stock/', views.SupplierPartDetail.as_view(template_name='company/supplier_part_stock.html'), name='supplier-part-stock'),

    url('^.*$', views.SupplierPartDetail.as_view(template_name='company/supplier_part_pricing.html'), name='supplier-part-detail'),
]

supplier_part_urls = [
    url(r'^new/?', views.SupplierPartCreate.as_view(), name='supplier-part-create'),

    url(r'delete/', views.SupplierPartDelete.as_view(), name='supplier-part-delete'),

    url(r'^(?P<pk>\d+)/', include(supplier_part_detail_urls)),
]
<|endoftext|>


<commit_before>"""
This module tests multi.py
"""

import sys
import os
sys.path.append(os.path.abspath("../"))
sys.path.append(os.path.abspath("../core"))
sys.path.append(os.path.abspath("../agent_types"))
sys.path.append(os.path.abspath("../helper_functions"))

from agent import Agent
from stream import Stream, StreamArray
from stream import _no_value, _multivalue
from check_agent_parameter_types import *
from recent_values import *
#from helper_functions.check_agent_parameter_types import *
#from helper_functions.recent_values import recent_values
from sink import sink_list, sink_list_f
from op import *
from split import split_list, split_list_f
from merge import merge_list, merge_list_f
from multi import multi_element

def test_multi():
    scheduler = Stream.scheduler
    import numpy as np

    u = Stream('u')
    v = Stream('v')
    x = Stream('x')
    y = Stream('y')

    def f(a_list):
        return max(a_list), sum(a_list)

    multi_element(f, [u,v], [x,y])

    u.extend(range(8, 18, 2))
    v.extend(range(10, 14, 1))
    scheduler.step()
    assert recent_values(u) == [8, 10, 12, 14, 16]
    assert recent_values(v) == [10, 11, 12, 13]
    assert recent_values(x) == [10, 11, 12, 14]
    assert recent_values(y) == [18, 21, 24, 27]

    # Test StreamArray
    u = StreamArray('u')
    v = StreamArray('v')
    x = StreamArray('x')
    y = StreamArray('y')
    z = StreamArray('z')

    def f(a_list):
        return max(a_list), np.mean(a_list), min(a_list)

    multi_element(f, [u,v], [x,y,z])

    u.extend(np.arange(5.0))
    v.extend(np.arange(10., 16., 1.))
    scheduler.step()
    assert np.array_equal(
        recent_values(u), np.array([0.0, 1.0, 2.0, 3.0, 4.0]))
    assert np.array_equal(
        recent_values(v), np.array([10.0, 11.0, 12.0, 13.0, 14.0, 15.0]))
    assert np.array_equal(
        recent_values(x), np.array([10.0, 11.0, 12.0, 13.0, 14.0]))
    assert np.array_equal(
        recent_values(y), np.array([5.0, 6.0, 7.0, 8.0, 9.0]))
    assert np.array_equal(
        recent_values(z), np.array([0.0, 1.0, 2.0, 3.0, 4.0]))

    # Test StreamArray, dtype=int
    u = StreamArray('u', dtype=int)
    v = StreamArray('v', dtype=int)
    x = StreamArray('x', dtype=int)
    y = StreamArray('y', dtype=int)

    def f(a_list):
        return max(a_list), sum(a_list)+min(a_list)

    multi_element(f, [u,v], [x,y])

    u.extend(np.arange(5))
    v.extend(np.arange(10, 17, 1))
    scheduler.step()

    assert np.array_equal(
        recent_values(u), np.array([0, 1, 2, 3, 4]))
    assert np.array_equal(
        recent_values(v), np.array([10, 11, 12, 13, 14, 15, 16]))
    assert np.array_equal(
        recent_values(x), np.array([10, 11, 12, 13, 14]))
    assert np.array_equal(
        recent_values(y), np.array([10, 13, 16, 19, 22]))

    print "TEST OF MULTI IS SUCCESSFUL"

if __name__ == '__main__':
    test_multi()
<commit_msg>Updated from FINAL<commit_after>"""
This module tests multi.py
"""

import sys
import os
sys.path.append(os.path.abspath("../"))
sys.path.append(os.path.abspath("../core"))
sys.path.append(os.path.abspath("../agent_types"))
sys.path.append(os.path.abspath("../helper_functions"))

from agent import Agent
from stream import Stream, StreamArray
from stream import _no_value, _multivalue
from check_agent_parameter_types import *
from recent_values import *
from sink import sink_list
from op import *
from split import split_list, split_list_f
from merge import merge_list, merge_list_f
from multi import multi_element, multi_list, multi_window

def test_multi():
    scheduler = Stream.scheduler
    import numpy as np

    u = Stream('u')
    v = Stream('v')
    x = Stream('x')
    y = Stream('y')

    def f(a_list):
        return max(a_list), sum(a_list)

    multi_element(f, [u,v], [x,y])

    u.extend(range(8, 18, 2))
    v.extend(range(10, 14, 1))
    scheduler.step()
    assert recent_values(u) == [8, 10, 12, 14, 16]
    assert recent_values(v) == [10, 11, 12, 13]
    assert recent_values(x) == [10, 11, 12, 14]
    assert recent_values(y) == [18, 21, 24, 27]

    # test of single input and single output
    u = Stream('u')
    x = Stream('x')

    def f(lst):
        return [lst[0]*2]

    multi_element(f, [u], [x])

    u.extend(range(5))
    scheduler.step()
    assert recent_values(x) == [0, 2, 4, 6, 8]

    # Test StreamArray
    u = StreamArray('u')
    v = StreamArray('v')
    x = StreamArray('x')
    y = StreamArray('y')
    z = StreamArray('z')

    def f(a_list):
        return max(a_list), np.mean(a_list), min(a_list)

    multi_element(f, [u,v], [x,y,z])

    u.extend(np.arange(5.0))
    v.extend(np.arange(10., 16., 1.))
    scheduler.step()
    assert np.array_equal(
        recent_values(u), np.array([0.0, 1.0, 2.0, 3.0, 4.0]))
    assert np.array_equal(
        recent_values(v), np.array([10.0, 11.0, 12.0, 13.0, 14.0, 15.0]))
    assert np.array_equal(
        recent_values(x), np.array([10.0, 11.0, 12.0, 13.0, 14.0]))
    assert np.array_equal(
        recent_values(y), np.array([5.0, 6.0, 7.0, 8.0, 9.0]))
    assert np.array_equal(
        recent_values(z), np.array([0.0, 1.0, 2.0, 3.0, 4.0]))

    # Test StreamArray, dtype=int
    u = StreamArray('u', dtype=int)
    v = StreamArray('v', dtype=int)
    x = StreamArray('x', dtype=int)
    y = StreamArray('y', dtype=int)

    def f(a_list):
        return max(a_list), sum(a_list)+min(a_list)

    multi_element(f, [u,v], [x,y])

    u.extend(np.arange(5))
    v.extend(np.arange(10, 17, 1))
    scheduler.step()

    assert np.array_equal(
        recent_values(u), np.array([0, 1, 2, 3, 4]))
    assert np.array_equal(
        recent_values(v), np.array([10, 11, 12, 13, 14, 15, 16]))
    assert np.array_equal(
        recent_values(x), np.array([10, 11, 12, 13, 14]))
    assert np.array_equal(
        recent_values(y), np.array([10, 13, 16, 19, 22]))

    print "TEST OF MULTI IS SUCCESSFUL"

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


<commit_before>"""A slow but simple dbm clone for the Mac.

For database spam, spam.dir contains the index (a text file),
spam.bak *may* contain a backup of the index (also a text file),
while spam.dat contains the data (a binary file).

XXX TO DO:

- reclaim free space (currently, space once occupied by deleted or expanded
items is never reused)

- support concurrent access (currently, if two processes take turns making
updates, they can mess up the index)

- support efficient access to large databases (currently, the whole index
is read when the database is opened, and some updates rewrite the whole index)
"""

_os = __import__('os')
import __builtin__

_open = __builtin__.open

_BLOCKSIZE = 512

class _Database:

	def __init__(self, file):
		self._dirfile = file + '.dir'
		self._datfile = file + '.dat'
		self._bakfile = file + '.bak'
		# Mod by Jack: create data file if needed
		try:
			f = _open(self._datfile, 'r')
		except IOError:
			f = _open(self._datfile, 'w')
		f.close()
		self._update()

	def _update(self):
		self._index = {}
		try:
			f = _open(self._dirfile)
		except IOError:
			pass
		else:
			while 1:
				line = f.readline()
				if not line: break
				key, (pos, siz) = eval(line)
				self._index[key] = (pos, siz)
			f.close()

	def _commit(self):
		try: _os.unlink(self._bakfile)
		except _os.error: pass
		try: _os.rename(self._dirfile, self._bakfile)
		except _os.error: pass
		f = _open(self._dirfile, 'w')
		for key, (pos, siz) in self._index.items():
			f.write("%s, (%s, %s)\n" % (`key`, `pos`, `siz`))
		f.close()

	def __getitem__(self, key):
		pos, siz = self._index[key]	# may raise KeyError
		f = _open(self._datfile, 'rb')
		f.seek(pos)
		dat = f.read(siz)
		f.close()
		return dat

	def _addval(self, val):
		f = _open(self._datfile, 'rb+')
		f.seek(0, 2)
		pos = f.tell()
## Does not work under MW compiler
##		pos = ((pos + _BLOCKSIZE - 1) / _BLOCKSIZE) * _BLOCKSIZE
##		f.seek(pos)
		npos = ((pos + _BLOCKSIZE - 1) / _BLOCKSIZE) * _BLOCKSIZE
		f.write('\0'*(npos-pos))
		pos = npos

		f.write(val)
		f.close()
		return (pos, len(val))

	def _setval(self, pos, val):
		f = _open(self._datfile, 'rb+')
		f.seek(pos)
		f.write(val)
		f.close()
		return pos, (val)

	def _addkey(self, key, (pos, siz)):
		self._index[key] = (pos, siz)
		f = _open(self._dirfile, 'a')
		f.write("%s, (%s, %s)\n" % (`key`, `pos`, `siz`))
		f.close()

	def __setitem__(self, key, val):
		if not type(key) == type('') == type(val):
			raise TypeError, "dbmac keys and values must be strings"
		if not self._index.has_key(key):
			(pos, siz) = self._addval(val)
			self._addkey(key, (pos, siz))
		else:
			pos, siz = self._index[key]
			oldblocks = (siz + _BLOCKSIZE - 1) / _BLOCKSIZE
			newblocks = (len(val) + _BLOCKSIZE - 1) / _BLOCKSIZE
			if newblocks <= oldblocks:
				pos, siz = self._setval(pos, val)
				self._index[key] = pos, siz
			else:
				pos, siz = self._addval(val)
				self._index[key] = pos, siz
			self._addkey(key, (pos, siz))

	def __delitem__(self, key):
		del self._index[key]
		self._commit()

	def keys(self):
		return self._index.keys()

	def has_key(self, key):
		return self._index.has_key(key)

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

	def close(self):
		self._index = self._datfile = self._dirfile = self._bakfile = None

def open(file, mode = None):
	return _Database(file)
<commit_msg>add (dummy) mode arg to open()<commit_after>"""A slow but simple dbm clone for the Mac.

For database spam, spam.dir contains the index (a text file),
spam.bak *may* contain a backup of the index (also a text file),
while spam.dat contains the data (a binary file).

XXX TO DO:

- reclaim free space (currently, space once occupied by deleted or expanded
items is never reused)

- support concurrent access (currently, if two processes take turns making
updates, they can mess up the index)

- support efficient access to large databases (currently, the whole index
is read when the database is opened, and some updates rewrite the whole index)

- support opening for read-only (flag = 'm')

"""

_os = __import__('os')
import __builtin__

_open = __builtin__.open

_BLOCKSIZE = 512

class _Database:

	def __init__(self, file):
		self._dirfile = file + '.dir'
		self._datfile = file + '.dat'
		self._bakfile = file + '.bak'
		# Mod by Jack: create data file if needed
		try:
			f = _open(self._datfile, 'r')
		except IOError:
			f = _open(self._datfile, 'w')
		f.close()
		self._update()

	def _update(self):
		self._index = {}
		try:
			f = _open(self._dirfile)
		except IOError:
			pass
		else:
			while 1:
				line = f.readline()
				if not line: break
				key, (pos, siz) = eval(line)
				self._index[key] = (pos, siz)
			f.close()

	def _commit(self):
		try: _os.unlink(self._bakfile)
		except _os.error: pass
		try: _os.rename(self._dirfile, self._bakfile)
		except _os.error: pass
		f = _open(self._dirfile, 'w')
		for key, (pos, siz) in self._index.items():
			f.write("%s, (%s, %s)\n" % (`key`, `pos`, `siz`))
		f.close()

	def __getitem__(self, key):
		pos, siz = self._index[key]	# may raise KeyError
		f = _open(self._datfile, 'rb')
		f.seek(pos)
		dat = f.read(siz)
		f.close()
		return dat

	def _addval(self, val):
		f = _open(self._datfile, 'rb+')
		f.seek(0, 2)
		pos = f.tell()
## Does not work under MW compiler
##		pos = ((pos + _BLOCKSIZE - 1) / _BLOCKSIZE) * _BLOCKSIZE
##		f.seek(pos)
		npos = ((pos + _BLOCKSIZE - 1) / _BLOCKSIZE) * _BLOCKSIZE
		f.write('\0'*(npos-pos))
		pos = npos

		f.write(val)
		f.close()
		return (pos, len(val))

	def _setval(self, pos, val):
		f = _open(self._datfile, 'rb+')
		f.seek(pos)
		f.write(val)
		f.close()
		return pos, (val)

	def _addkey(self, key, (pos, siz)):
		self._index[key] = (pos, siz)
		f = _open(self._dirfile, 'a')
		f.write("%s, (%s, %s)\n" % (`key`, `pos`, `siz`))
		f.close()

	def __setitem__(self, key, val):
		if not type(key) == type('') == type(val):
			raise TypeError, "dbmac keys and values must be strings"
		if not self._index.has_key(key):
			(pos, siz) = self._addval(val)
			self._addkey(key, (pos, siz))
		else:
			pos, siz = self._index[key]
			oldblocks = (siz + _BLOCKSIZE - 1) / _BLOCKSIZE
			newblocks = (len(val) + _BLOCKSIZE - 1) / _BLOCKSIZE
			if newblocks <= oldblocks:
				pos, siz = self._setval(pos, val)
				self._index[key] = pos, siz
			else:
				pos, siz = self._addval(val)
				self._index[key] = pos, siz
			self._addkey(key, (pos, siz))

	def __delitem__(self, key):
		del self._index[key]
		self._commit()

	def keys(self):
		return self._index.keys()

	def has_key(self, key):
		return self._index.has_key(key)

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

	def close(self):
		self._index = self._datfile = self._dirfile = self._bakfile = None

def open(file, flag = None, mode = None):
	# flag, mode arguments are currently ignored
	return _Database(file)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Tests for the ArgsCreator module.
"""

from __future__ import division, print_function, with_statement

import itertools

from nose.tools import eq_

from orges.core.args import ArgsCreator
from orges.tests.unit.core.util import get_intervals_from_function
from orges.tests.util.integer_functions import INTEGER_FUNCTIONS

try:
    xrange  # will work in python2, only
except NameError:
    xrange = range  # rename range to xrange in python3

def test_ArgsCreator_product_versus_itertools_product():
    for function in INTEGER_FUNCTIONS:
        # log
        print(function)

        # mock
        intervals = get_intervals_from_function(function)
        valid_ranges = []
        for interval in intervals.values():
            start, stop = interval[0], interval[1]
            interval_range = xrange(start, stop + 1)  # [b, e) vs. [b, e]
            valid_ranges.append(interval_range)
        args_dummy = itertools.product(valid_ranges)
        for _ in args_dummy:
            pass  # forces execution of the range generators
        print(args_dummy)

        # use
        args_product = [(a.value, b.value)
                        for a, b in ArgsCreator(function.param_spec).product()]
        print(args_product)

        # compare
        for arg_p, arg_d in zip(args_product, args_dummy):
            eq_(arg_p[0], arg_d[0])
            eq_(arg_p[1], arg_d[1])

def test_ArgsCreator_args_versus_nested_loop():
    for function in INTEGER_FUNCTIONS:
        # log
        print(function)

        # mock
        intervals = get_intervals_from_function(function)
        args_dummy = []
        for a in intervals['a']:
            args_dummy.append(a)
            for b in intervals['b']:
                args_dummy.append(b)

        # use
        args_args = ArgsCreator(function.param_spec).args()

        # compare
        for args, prod in zip(args_args, args_dummy):
            eq_(args.value, prod)

if __name__ == '__main__':
    import nose
    nose.runmodule()
<commit_msg>quick fix for the first test.<commit_after># -*- coding: utf-8 -*-
"""
Tests for the ArgsCreator module.
"""

from __future__ import division, print_function, with_statement

import itertools

from nose.tools import eq_

from orges.core.args import ArgsCreator
from orges.tests.unit.core.util import get_intervals_from_function
from orges.tests.util.integer_functions import INTEGER_FUNCTIONS

try:
    xrange  # will work in python2, only
except NameError:
    xrange = range  # rename range to xrange in python3

def test_ArgsCreator_product_versus_itertools_product():
    for function in INTEGER_FUNCTIONS:
        # log
        print(function)

        # mock
        intervals = get_intervals_from_function(function)
        valid_ranges = []
        for interval in intervals.values():
            start, stop = interval[0], interval[1]
            interval_range = xrange(start, stop + 1)  # [b, e) vs. [b, e]
            valid_ranges.append(interval_range)
        args_dummy = itertools.product(valid_ranges)
        for _ in args_dummy:
            pass  # forces execution of the range generators
        print(args_dummy)

        # TODO look for a better way to determine the number of params
        # than trying and failing.

        # use
        try:
            args_product = [(a.value, b.value)
                            for a, b in
                            ArgsCreator(function.param_spec).product()]
        except ValueError:
            args_product = [(a[0].value)
                            for a in
                            ArgsCreator(function.param_spec).product()]
        print(args_product)

        # compare
        for arg_p, arg_d in zip(args_product, args_dummy):
            eq_(arg_p[0], arg_d[0])
            eq_(arg_p[1], arg_d[1])

def test_ArgsCreator_args_versus_nested_loop():
    for function in INTEGER_FUNCTIONS:
        # log
        print(function)

        # mock
        intervals = get_intervals_from_function(function)
        args_dummy = []
        for a in intervals['a']:
            args_dummy.append(a)
            for b in intervals['b']:
                args_dummy.append(b)

        # use
        args_args = ArgsCreator(function.param_spec).args()

        # compare
        for args, prod in zip(args_args, args_dummy):
            eq_(args.value, prod)

if __name__ == '__main__':
    import nose
    nose.runmodule()
<|endoftext|>


<commit_before>

import sys
import os
from distutils.dir_util import mkpath
import re
import numpy as np

from tasks import tclean, tget

'''
Cleans a single channel given the channel name
'''

# Load in the SPW dict in the repo on cedar
execfile(os.path.expanduser("~/code/VLA_Lband/17B-162/spw_setup.py"))

chan_num = int(sys.argv[-2])

# Load in the imaging parameters from the given file name
parameter_file = sys.argv[-1]

# Load parameters
tget(tclean, parameter_file)

# Get the output path and create directories, if needed, based on
# the imagename
if imagename.split("/") > 1:

    output_path = "/".join(imagename.split("/")[:-1])

    # Since the imaging is not run from the parent path for the channel output,
    # use `mkpath`
    if not os.path.exists(output_path):
        mkpath(output_path)

# Now update the imagename with the channel number
imagename = "{0}_channel_{1}".format(imagename, chan_num)

# Based on the channel number, update the start velocity for this channel

# Get the value out from the string, removing the unit
split_start = filter(None, re.split(r'(\d+)', start))

init_start = float("".join(split_start[:-1]))
spec_unit = split_start[-1]

chan_width = float("".join(filter(None, re.split(r'(\d+)', width))[:-1]))

start_vel = init_start + chan_width * chan_num

# Check if the products already exist and we should avoid recomputing the PSF
# and residuals

# Check for the PSF and assume the rest of the products are there too.
if os.path.exists("{}.psf".format(imagename)):
    do_calcres = False
    do_calcpsf = False

    # Force startmodel to use the model on disk
    startmodel = None

else:
    do_calcres = True
    do_calcpsf = True

# If model or mask names are given, ensure they exist.
# These should already be split into individual channels for use here
# The naming scheme should split imagename.image to imagename_channel_{}.image
# The file MUST end in ".image"
if startmodel is not None and len(startmodel) > 0:

    print("Start model??")
    print(type(startmodel))
    print(startmodel is not None)
    print(len(startmodel))

    startmodel = "{0}_channel_{1}.image"(startmodel.split(".image")[0],
                                         chan_num)

    if not os.path.exists(startmodel):
        raise ValueError("Given startmodel does not exist")

# The naming scheme should split name.mask to name_channel_{}.mask
# The file MUST end in ".mask"
if mask is not None and len(mask) > 0 and usemask == "user":

    print("Mask??")

    mask = "{0}_channel_{1}.mask"(mask.split(".image")[0], chan_num)

    if not os.path.exists(mask):
        raise ValueError("Given mask name ({0}) does not exist".format(mask))

# Grab freq from the SPW dict
spw_num = int(spw)

# Only update a few parameters, as needed

start = "{0}{1}".format(start_vel, spec_unit)
width = "{0}{1}".format(chan_width, spec_unit)
nchan = 1
restfreq = linespw_dict[spw_num][1]
restart = True
calcres = do_calcres
calcpsf = do_calcpsf
interactive = 0  # Returns a summary dictionary

out_dict = tclean()

# Save the output dictionary. Numpy should be fine for this as the individual
# channels will get concatenated together

np.save(imagename + ".results_dict.npy", out_dict)
<commit_msg>Remove extra printing<commit_after>

import sys
import os
from distutils.dir_util import mkpath
import re
import numpy as np

from tasks import tclean, tget

'''
Cleans a single channel given the channel name
'''

# Load in the SPW dict in the repo on cedar
execfile(os.path.expanduser("~/code/VLA_Lband/17B-162/spw_setup.py"))

chan_num = int(sys.argv[-2])

# Load in the imaging parameters from the given file name
parameter_file = sys.argv[-1]

# Load parameters
tget(tclean, parameter_file)

# Get the output path and create directories, if needed, based on
# the imagename
if imagename.split("/") > 1:

    output_path = "/".join(imagename.split("/")[:-1])

    # Since the imaging is not run from the parent path for the channel output,
    # use `mkpath`
    if not os.path.exists(output_path):
        mkpath(output_path)

# Now update the imagename with the channel number
imagename = "{0}_channel_{1}".format(imagename, chan_num)

# Based on the channel number, update the start velocity for this channel

# Get the value out from the string, removing the unit
split_start = filter(None, re.split(r'(\d+)', start))

init_start = float("".join(split_start[:-1]))
spec_unit = split_start[-1]

chan_width = float("".join(filter(None, re.split(r'(\d+)', width))[:-1]))

start_vel = init_start + chan_width * chan_num

# Check if the products already exist and we should avoid recomputing the PSF
# and residuals

# Check for the PSF and assume the rest of the products are there too.
if os.path.exists("{}.psf".format(imagename)):
    do_calcres = False
    do_calcpsf = False

    # Force startmodel to use the model on disk
    startmodel = None

else:
    do_calcres = True
    do_calcpsf = True

# If model or mask names are given, ensure they exist.
# These should already be split into individual channels for use here
# The naming scheme should split imagename.image to imagename_channel_{}.image
# The file MUST end in ".image"
if startmodel is not None and len(startmodel) > 0:

    startmodel = "{0}_channel_{1}.image"(startmodel.split(".image")[0],
                                         chan_num)

    if not os.path.exists(startmodel):
        raise ValueError("Given startmodel does not exist")

# The naming scheme should split name.mask to name_channel_{}.mask
# The file MUST end in ".mask"
if mask is not None and len(mask) > 0 and usemask == "user":

    mask = "{0}_channel_{1}.mask"(mask.split(".image")[0], chan_num)

    if not os.path.exists(mask):
        raise ValueError("Given mask name ({0}) does not exist".format(mask))

# Grab freq from the SPW dict
spw_num = int(spw)

# Only update a few parameters, as needed

start = "{0}{1}".format(start_vel, spec_unit)
width = "{0}{1}".format(chan_width, spec_unit)
nchan = 1
restfreq = linespw_dict[spw_num][1]
restart = True
calcres = do_calcres
calcpsf = do_calcpsf
interactive = 0  # Returns a summary dictionary

out_dict = tclean()

# Save the output dictionary. Numpy should be fine for this as the individual
# channels will get concatenated together

np.save(imagename + ".results_dict.npy", out_dict)
<|endoftext|>


<commit_before><commit_msg>[subset] Speed up cmap format 14 closure<commit_after><|endoftext|>


<commit_before>from __future__ import division
import os
import sys
import warnings

from postgres import Postgres
from psycopg2.extras import Json as postgres_jsonify

from .base import JobsExhaustedError, JobFailed

###############################################################################

TABLE_EXISTS_QUERY = r"""
SELECT EXISTS(
    SELECT *
    FROM information_schema.tables
    WHERE table_name = 'job_{}'
);
"""

CREATE_TABLE_QUERY = r"""
CREATE TABLE job_{}(
  job_id TEXT UNIQUE NOT NULL PRIMARY KEY,
  data json,
  completed BOOLEAN NOT NULL DEFAULT FALSE,
  claimed BOOLEAN NOT NULL DEFAULT FALSE
)
"""

INSERT_JOB_QUERY = r"""
INSERT INTO job_{0} (job_id) VALUES (%(job_id)s) RETURNING job_id
"""

UPDATE_DATA_QUERY = r"""
UPDATE job_{0} SET data=%(data)s WHERE job_id=%(job_id)s
"""

RANDOM_UNCOMPLETED_UNCLAIMED_ROW_QUERY = r"""
SELECT *
FROM job_{0} WHERE completed=false AND claimed=false OFFSET floor(random() * (
    SELECT COUNT(*) FROM job_{0} WHERE completed=false AND claimed=false
))
LIMIT 1;
"""

RANDOM_UNCOMPLETED_ROW_QUERY = r"""
SELECT *
FROM job_{0} WHERE completed=false OFFSET floor(random() * (
    SELECT COUNT(*) FROM job_{0} WHERE completed=false
))
LIMIT 1;
"""

SET_ROW_COMPLETED_BY_ID_QUERY = r"""
UPDATE job_{} SET completed=true WHERE job_id=%(job_id)s
"""

SET_ROW_CLAIMED_BY_ID_QUERY = r"""
UPDATE job_{} SET claimed=true WHERE job_id=%(job_id)s
"""

TOTAL_ROWS_QUERY = r"""
SELECT COUNT(*) FROM job_{0}
"""

COMPLETED_ROWS_QUERY = r"""
SELECT COUNT(*) FROM job_{0} WHERE completed=true
"""

###############################################################################
def table_exists(db_handle, table_id):
    return db_handle.one(TABLE_EXISTS_QUERY.format(table_id))

def setup_table(db_handle, table_id, job_ids, job_dicts=None):
    does_table_exist = table_exists(db_handle, table_id)

    if job_dicts is None:
        job_dicts = {}

    if not does_table_exist:
        print('Creating table...')
        with db_handle.get_cursor() as cursor:  # Single Transaction
            # Create table
            cursor.run(CREATE_TABLE_QUERY.format(table_id))

            # Fill in list of jobs
            for job_id in job_ids:
                cursor.one(INSERT_JOB_QUERY.format(table_id),
                           parameters={'job_id': job_id})
                if job_id in job_dicts:
                    cursor.run(UPDATE_DATA_QUERY.format(table_id),
                               parameters={'data': postgres_jsonify(job_dicts[job_id]),
                                           'job_id': job_id})
    return not does_table_exist

def create_db_handle(host=None, database=None, user=None):
    if host is None:
        host = os.environ['PGHOST']
    if database is None:
        database = os.environ['PGDATABASE']
    if user is None:
        user = os.environ['PGUSER']
    return Postgres("host={} dbname={} user={}".format(host, database, user))

def get_random_uncompleted_unclaimed_job(db_handle, table_id):
    return db_handle.one(RANDOM_UNCOMPLETED_UNCLAIMED_ROW_QUERY.format(table_id))

def get_random_uncompleted_job(db_handle, table_id):
    return db_handle.one(RANDOM_UNCOMPLETED_ROW_QUERY.format(table_id))

def set_job_as_complete(db_handle, table_id, job_id):
    db_handle.run(SET_ROW_COMPLETED_BY_ID_QUERY.format(table_id),
                  parameters={'job_id': job_id})

def set_job_as_claimed(db_handle, table_id, job_id):
    db_handle.run(SET_ROW_CLAIMED_BY_ID_QUERY.format(table_id),
                  parameters={'job_id': job_id})

def get_total_job_count(db_handle, table_id):
    return db_handle.one(TOTAL_ROWS_QUERY.format(table_id))

def get_completed_job_count(db_handle, table_id):
    return db_handle.one(COMPLETED_ROWS_QUERY.format(table_id))

################################################################################

def postgres_exhaust_all(perform_job_f, table_id, job_ids=None, job_dicts=None,
                         host=None, database=None, user=None, verbose=True):
    db_handle = create_db_handle(host=host, database=database, user=user)
    if job_ids is not None:
        if not setup_table(db_handle, table_id, job_ids, job_dicts=job_dicts):
            warnings.warn('The input table name "{}" already existed and the '
                          'newly provided job IDs were not used.'.format(table_id))
    else:
        if not table_exists(db_handle, table_id):
            raise ValueError('Job IDs were not provided and the given '
                             'table ID does not exist.')
    n_inputs = get_total_job_count(db_handle, table_id)
    if verbose:
        i = 0
    try:
        while True:
            if verbose:
                sys.stdout.write(str(i) + ': ')
                sys.stdout.flush()
            a_row = get_random_uncompleted_unclaimed_job(db_handle, table_id)
            if a_row is None:
                # there is nothing left that is unclaimed, so we may as well
                # 'repeat' already claimed work - maybe we can beat another
                # bee to complete it.
                a_row = get_random_uncompleted_job(db_handle, table_id)
                if verbose:
                    sys.stdout.write("no unclaimed work remains - ")
                    sys.stdout.flush()

            if a_row is None:
                raise JobsExhaustedError()

            # let's claim the job
            set_job_as_claimed(db_handle, table_id, a_row.job_id)

            if verbose:
                sys.stdout.write("claimed '{}'".format(a_row.job_id))
                sys.stdout.flush()
            try:
                perform_job_f(a_row.job_id, data=a_row.data)
            except JobFailed as e:
                sys.stdout.write('FAILED.\n')
                sys.stdout.flush()
            else:
                set_job_as_complete(db_handle, table_id, a_row.job_id)
                if verbose:
                    sys.stdout.write('...done.\n')
                    sys.stdout.flush()
            if verbose:
                i += 1
                if i % 10 == 0:
                    n_completed_rows = get_completed_job_count(db_handle,
                                                               table_id)
                    print('{:.2f}% ({}/{}) Completed'.format(
                        (n_completed_rows / n_inputs) * 100,
                        n_completed_rows, n_inputs))
    except JobsExhaustedError:
        sys.stdout.write('\rJobs exhausted')
        sys.stdout.flush()
<commit_msg>improved verbose output, cleaner naming conventions<commit_after>from __future__ import division
import os
import sys
import warnings

from postgres import Postgres
from psycopg2.extras import Json as postgres_jsonify

from .base import JobsExhaustedError, JobFailed

from string import letters, digits
ALLOWED_CHARACTERS_IN_TABLE_NAME = set(letters.lower()) | set(digits) | set('_')

def check_name(table_name):
    if not isinstance(table_name, basestring):
        raise TypeError("Experiment ID '{}' is of type {}, not string".format(
            table_name, type(table_name)))
    invalid = set(table_name) - ALLOWED_CHARACTERS_IN_TABLE_NAME
    if len(invalid) > 0:
        raise ValueError("Invalid characters in experiment ID: {} "
                         "(allowed [a...z, 0...9, '_']".format(
            ', '.join(["'{}'".format(l) for l in sorted(list(invalid))])))

###############################################################################

TABLE_EXISTS_QUERY = r"""
SELECT EXISTS(
    SELECT *
    FROM information_schema.tables
    WHERE table_name = 'job_{}'
);
"""

CREATE_TABLE_QUERY = r"""
CREATE TABLE job_{}(
  job_id TEXT UNIQUE NOT NULL PRIMARY KEY,
  data json,
  completed BOOLEAN NOT NULL DEFAULT FALSE,
  claimed BOOLEAN NOT NULL DEFAULT FALSE
)
"""

INSERT_JOB_QUERY = r"""
INSERT INTO job_{0} (job_id) VALUES (%(job_id)s) RETURNING job_id
"""

UPDATE_DATA_QUERY = r"""
UPDATE job_{0} SET data=%(data)s WHERE job_id=%(job_id)s
"""

RANDOM_UNCOMPLETED_UNCLAIMED_ROW_QUERY = r"""
SELECT *
FROM job_{0} WHERE completed=false AND claimed=false OFFSET floor(random() * (
    SELECT COUNT(*) FROM job_{0} WHERE completed=false AND claimed=false
))
LIMIT 1;
"""

RANDOM_UNCOMPLETED_ROW_QUERY = r"""
SELECT *
FROM job_{0} WHERE completed=false OFFSET floor(random() * (
    SELECT COUNT(*) FROM job_{0} WHERE completed=false
))
LIMIT 1;
"""

SET_ROW_COMPLETED_BY_ID_QUERY = r"""
UPDATE job_{} SET completed=true WHERE job_id=%(job_id)s
"""

SET_ROW_CLAIMED_BY_ID_QUERY = r"""
UPDATE job_{} SET claimed=true WHERE job_id=%(job_id)s
"""

TOTAL_ROWS_QUERY = r"""
SELECT COUNT(*) FROM job_{0}
"""

COMPLETED_ROWS_QUERY = r"""
SELECT COUNT(*) FROM job_{0} WHERE completed=true
"""

###############################################################################
def table_exists(db_handle, table_id):
    return db_handle.one(TABLE_EXISTS_QUERY.format(table_id))

def setup_table(db_handle, experiment_id, job_ids=None, job_data=None,
                verbose=False):
    check_name(experiment_id)
    does_table_exist = table_exists(db_handle, experiment_id)

    if job_data is None:
        job_data = {}
    if job_ids is not None:
        extra_job_data = set(job_data) - set(job_ids)
        if len(extra_job_data) > 0:
            raise ValueError("job_data provided for {} job IDs that aren't in this "
                             "experiment: {}".format(len(extra_job_data),
                                                     list(extra_job_data)))
    if does_table_exist:
        if verbose:
            print("Table already exists for experiment '{}'".format(experiment_id))
            if job_ids is not None:
                print('Warning: {} provided job IDs were ignored (presumably '
                      'they are already in the table?)'.format(len(job_ids)))
    else:
        if verbose:
            print("Creating table for experiment '{}'...".format(experiment_id))
        with db_handle.get_cursor() as cursor:  # Single Transaction
            # Create table
            cursor.run(CREATE_TABLE_QUERY.format(experiment_id))
            if verbose:
                print("Adding {} jobs to experiment "
                      "'{}'...".format(len(job_ids), experiment_id))
                if verbose:
                    if len(job_data) > 0:
                        print('job_data provided for {} jobs'.format(len(job_data)))
                    else:
                        print('no job_data provided.')
            # Fill in list of jobs
            for job_id in job_ids:
                cursor.one(INSERT_JOB_QUERY.format(experiment_id),
                           parameters={'job_id': job_id})
                if job_id in job_data:
                    cursor.run(UPDATE_DATA_QUERY.format(experiment_id),
                               parameters={'data': postgres_jsonify(job_data[job_id]),
                                           'job_id': job_id})
        if verbose:
            print("Experiment '{}' set up.".format(experiment_id))
    if verbose:
        print('-' * 80)
    return not does_table_exist

def create_db_handle(host=None, port=None, user=None, database=None, verbose=False):
    if host is None:
        host = os.environ['PGHOST']
    if port is None:
        port = os.environ['PGPORT']
    if user is None:
        user = os.environ['PGUSER']
    if database is None:
        database = os.environ['PGDATABASE']
    if verbose:
        print('Connecting to database {} on {}@{}:{}...'.format(database, user, host, port))
        if 'PGPASS' in os.environ:
            print(' - Password is set via environment variable PGPASS.')
        else:
            print(' - No password is set. (If needed, set the environment '
                  'variable PGPASS.)')
    return Postgres("host={} port={} user={} dbname={} ".format(host, port, user, database))

def get_random_uncompleted_unclaimed_job(db_handle, table_id):
    return db_handle.one(RANDOM_UNCOMPLETED_UNCLAIMED_ROW_QUERY.format(table_id))

def get_random_uncompleted_job(db_handle, table_id):
    return db_handle.one(RANDOM_UNCOMPLETED_ROW_QUERY.format(table_id))

def set_job_as_complete(db_handle, table_id, job_id):
    db_handle.run(SET_ROW_COMPLETED_BY_ID_QUERY.format(table_id),
                  parameters={'job_id': job_id})

def set_job_as_claimed(db_handle, table_id, job_id):
    db_handle.run(SET_ROW_CLAIMED_BY_ID_QUERY.format(table_id),
                  parameters={'job_id': job_id})

def get_total_job_count(db_handle, table_id):
    return db_handle.one(TOTAL_ROWS_QUERY.format(table_id))

def get_completed_job_count(db_handle, table_id):
    return db_handle.one(COMPLETED_ROWS_QUERY.format(table_id))

################################################################################

def postgres_experiment(experiment_id, job_function,
                        job_ids=None, job_data=None,
                        host=None, port=None, user=None, database=None,
                        verbose=True):
    """

    Parameters
    ----------
    experiment_id : `str`
        A unique identifier for an experiment. Must be a valid Postgres Table
        name (no whitespace etc).
    job_function : `callable` taking (job_id: str, job_data : {})
        A callable that performs a unit of work in this experiment. The
        function will be provided with two arguments - the ``job_id`` for this
        job (a string) and the data payload for this job (a dictionary). This
        function then uses these inputs to perform the relevent work for the
        experiment. If the function completes without error, the job will be
        automatically marked complete in the database.
    job_ids : `list` of `str`, optional if experiment_id exists.
        A list of Job IDs to be used in this experiment. If the experiment
        already exists, these jobs will be ignored. If the experiment ID
        doesn't exist yet, the experiment will be created with these jobs.
    job_data : `dict` with `str` keys pointing to `dict`s, optional
        An optional dictionary of extra data that can be stored per-job. This
        dictionary's top-level keys should be job_id's. top-level values
        should be dictionaries of basic datatypes. This dictionary can for
        example store experiment parameters that should be used for the job
        in question.
    host : `str`, optional
        The hostname of the Postgres instance. If ``None``, the environment
        variable ``PGHOST`` will be used if set.
    port : `str`, optional
        The port which the Postgres instance is running on. If ``None``, the
        environment variable ``PGPORT`` will be used if set.
    user : `str`, optional
        The Postgres user to join as. If ``None``, the
        environment variable ``PGUSER`` will be used if set.
    database : `str`, optional
        The database of the Postgres server to use. If ``None``, the
        environment variable ``PGDATABASE`` will be used if set.
    verbose : `bool`, optional
        If ``True``, additional logging will be provided.
    """
    db_handle = create_db_handle(host=host, port=port, user=user,
                                 database=database, verbose=verbose)
    if job_ids is not None:
        setup_table(db_handle, experiment_id, job_ids=job_ids,
                    job_data=job_data, verbose=verbose)
    else:
        if not table_exists(db_handle, experiment_id):
            raise ValueError('Job IDs were not provided and the given '
                             'Experiment ID "{}" does not exist.'.format(experiment_id))
    n_inputs = get_total_job_count(db_handle, experiment_id)
    if verbose:
        i = 0
    try:
        while True:
            if verbose:
                sys.stdout.write(str(i) + ': ')
                sys.stdout.flush()
            a_row = get_random_uncompleted_unclaimed_job(db_handle, experiment_id)
            if a_row is None:
                # there is nothing left that is unclaimed, so we may as well
                # 'repeat' already claimed work - maybe we can beat another
                # bee to complete it.
                a_row = get_random_uncompleted_job(db_handle, experiment_id)
                if verbose:
                    sys.stdout.write("no unclaimed work remains - ")
                    sys.stdout.flush()

            if a_row is None:
                raise JobsExhaustedError()

            # let's claim the job
            set_job_as_claimed(db_handle, experiment_id, a_row.job_id)

            if verbose:
                sys.stdout.write("claimed '{}'".format(a_row.job_id))
                sys.stdout.flush()
            try:
                job_function(a_row.job_id, data=a_row.data if a_row.data is not None else {})
            except JobFailed as e:
                sys.stdout.write('FAILED.\n')
                sys.stdout.flush()
            else:
                set_job_as_complete(db_handle, experiment_id, a_row.job_id)
                if verbose:
                    sys.stdout.write('...done.\n')
                    sys.stdout.flush()
            if verbose:
                i += 1
                if i % 10 == 0:
                    n_completed_rows = get_completed_job_count(db_handle,
                                                               experiment_id)
                    print('{:.2f}% ({}/{}) Completed'.format(
                        (n_completed_rows / n_inputs) * 100,
                        n_completed_rows, n_inputs))
    except JobsExhaustedError:
        sys.stdout.write('\rAll jobs are exhausted, terminating.')
        sys.stdout.flush()
<|endoftext|>