<commit_before>import csv
import errno
import hashlib
import logging
import os
import requests
from StringIO import StringIO

NCBI_URI = 'http://ftp.ncbi.nih.gov/genomes'
supported_domains = ['archaea', 'bacteria', 'fungi', 'invertebrate', 'plant',
                     'protozoa', 'unknown', 'vertebrate_mammalian',
                     'vertebrate_other', 'viral']

def download(args):
    '''Download data from NCBI'''
    if args.domain == 'all':
        for domain in supported_domains:
            _download(args.section, domain, args.uri, args.output)
    else:
        _download(args.section, args.domain, args.uri, args.output)

def _download(section, domain, uri, output):
    '''Download a specified domain form a section'''
    summary = get_summary(section, domain, uri)
    entries = parse_summary(summary)
    for entry in entries:
        download_entry(entry, section, domain, uri, output)

def get_summary(section, domain, uri):
    '''Get the assembly_summary.txt file from NCBI and return a StringIO object for it'''
    logging.debug('Downloading summary for %r/%r uri: %r', section, domain, uri)
    url = '{uri}/{section}/{domain}/assembly_summary.txt'.format(
        section=section, domain=domain, uri=uri)
    r = requests.get(url)
    return StringIO(r.text)

def parse_summary(summary):
    '''Parse the summary file from TSV format to a csv DictReader'''
    # skip the leading 2 comment characters in the header to get nicer names
    comment = summary.read(2)
    if comment != '# ':
        # Huh, unexpected header, just go back to whereever we were before
        idx = summary.tell()
        summary.seek(max(0, idx - 2))

    reader = csv.DictReader(summary, dialect='excel-tab')
    return reader

def download_entry(entry, section, domain, uri, output):
    '''Download an entry from the summary file'''
    logging.info('Downloading record %r', entry['assembly_accession'])
    full_output_dir = create_dir(entry, section, domain, output)
    checksums = grab_checksums_file(entry)

    # TODO: Only write this when the checksums file changed
    with open(os.path.join(full_output_dir, 'MD5SUMS'), 'w') as fh:
        fh.write(checksums)

    parsed_checksums = parse_checksums(checksums)
    if has_gbk_file_changed(full_output_dir, parsed_checksums):
        download_gbk_file(entry, full_output_dir, parsed_checksums)

def create_dir(entry, section, domain, output):
    '''Create the output directory for the entry if needed'''
    full_output_dir = os.path.join(output, section, domain, entry['assembly_accession'])
    try:
        os.makedirs(full_output_dir)
    except OSError as e:
        if e.errno == errno.EEXIST and os.path.isdir(full_output_dir):
            pass
        else:
            raise

    return full_output_dir

def grab_checksums_file(entry):
    '''Grab the checksum file for a given entry'''
    http_url = convert_ftp_url(entry['ftp_path'])
    full_url = '{}/md5checksums.txt'.format(http_url)
    r = requests.get(full_url)
    return r.text

def convert_ftp_url(url):
    '''Convert FTP to HTTP URLs'''
    return url.replace('ftp://', 'http://', 1)

def parse_checksums(checksums_string):
    '''Parse a file containing checksums and filenames'''
    checksums_list = []
    for line in checksums_string.split('\n'):
        try:
            checksum, filename = line.split()
            # strip leading ./
            if filename.startswith('./'):
                filename = filename[2:]
            checksums_list.append({'checksum': checksum, 'file': filename})
        except ValueError:
            logging.debug('Skipping over unexpected checksum line %r', line)
            continue

    return checksums_list

def has_gbk_file_changed(directory, checksums):
    '''Check if the checksum of a given gbk file has changed'''
    filename, expected_checksum = get_name_and_checksum(checksums, '_genomic.gbff.gz')
    full_filename = os.path.join(directory, filename)
    # if file doesn't exist, it has changed
    if not os.path.isfile(full_filename):
        return True

    actual_checksum = md5sum(full_filename)
    return expected_checksum != actual_checksum

def get_name_and_checksum(checksums, end):
    '''Extract a full filename and checksum from the checksums list for a file ending in given end'''
    for entry in checksums:
        if not entry['file'].endswith(end):
            # wrong file
            continue
        filename = entry['file']
        expected_checksum = entry['checksum']
        return filename, expected_checksum

def md5sum(filename):
    '''Calculate the md5sum of a file and return the hexdigest'''
    hash_md5 = hashlib.md5()
    with open(filename, 'rb') as fh:
        for chunk in iter(lambda: fh.read(4096), b''):
            hash_md5.update(chunk)
    return hash_md5.hexdigest()

def download_gbk_file(entry, directory, checksums):
    '''Download and verirfy a given gbk file'''
    filename, expected_checksum = get_name_and_checksum(checksums, '_genomic.gbff.gz')
    base_url = convert_ftp_url(entry['ftp_path'])
    full_url = '{}/{}'.format(base_url, filename)
    local_file = os.path.join(directory, filename)

    r = requests.get(full_url, stream=True)

    with open(local_file, 'wb') as fh:
        for chunk in r.iter_content(4096):
            fh.write(chunk)

    actual_checksum = md5sum(local_file)
    if actual_checksum != expected_checksum:
        logging.error('Checksum mismatch for %r. Expected %r, got %r',
                local_file, expected_checksum, actual_checksum)
        return False

    return True
<commit_msg>core: csv can't deal with unicode, so keep summary file as str<commit_after>import csv
import errno
import hashlib
import logging
import os
import requests
from StringIO import StringIO

NCBI_URI = 'http://ftp.ncbi.nih.gov/genomes'
supported_domains = ['archaea', 'bacteria', 'fungi', 'invertebrate', 'plant',
                     'protozoa', 'unknown', 'vertebrate_mammalian',
                     'vertebrate_other', 'viral']

def download(args):
    '''Download data from NCBI'''
    if args.domain == 'all':
        for domain in supported_domains:
            _download(args.section, domain, args.uri, args.output)
    else:
        _download(args.section, args.domain, args.uri, args.output)

def _download(section, domain, uri, output):
    '''Download a specified domain form a section'''
    summary = get_summary(section, domain, uri)
    entries = parse_summary(summary)
    for entry in entries:
        download_entry(entry, section, domain, uri, output)

def get_summary(section, domain, uri):
    '''Get the assembly_summary.txt file from NCBI and return a StringIO object for it'''
    logging.debug('Downloading summary for %r/%r uri: %r', section, domain, uri)
    url = '{uri}/{section}/{domain}/assembly_summary.txt'.format(
        section=section, domain=domain, uri=uri)
    r = requests.get(url)
    return StringIO(r.content)

def parse_summary(summary):
    '''Parse the summary file from TSV format to a csv DictReader'''
    # skip the leading 2 comment characters in the header to get nicer names
    comment = summary.read(2)
    if comment != '# ':
        # Huh, unexpected header, just go back to whereever we were before
        idx = summary.tell()
        summary.seek(max(0, idx - 2))

    reader = csv.DictReader(summary, dialect='excel-tab')
    return reader

def download_entry(entry, section, domain, uri, output):
    '''Download an entry from the summary file'''
    logging.info('Downloading record %r', entry['assembly_accession'])
    full_output_dir = create_dir(entry, section, domain, output)
    checksums = grab_checksums_file(entry)

    # TODO: Only write this when the checksums file changed
    with open(os.path.join(full_output_dir, 'MD5SUMS'), 'w') as fh:
        fh.write(checksums)

    parsed_checksums = parse_checksums(checksums)
    if has_gbk_file_changed(full_output_dir, parsed_checksums):
        download_gbk_file(entry, full_output_dir, parsed_checksums)

def create_dir(entry, section, domain, output):
    '''Create the output directory for the entry if needed'''
    full_output_dir = os.path.join(output, section, domain, entry['assembly_accession'])
    try:
        os.makedirs(full_output_dir)
    except OSError as e:
        if e.errno == errno.EEXIST and os.path.isdir(full_output_dir):
            pass
        else:
            raise

    return full_output_dir

def grab_checksums_file(entry):
    '''Grab the checksum file for a given entry'''
    http_url = convert_ftp_url(entry['ftp_path'])
    full_url = '{}/md5checksums.txt'.format(http_url)
    r = requests.get(full_url)
    return r.text

def convert_ftp_url(url):
    '''Convert FTP to HTTP URLs'''
    return url.replace('ftp://', 'http://', 1)

def parse_checksums(checksums_string):
    '''Parse a file containing checksums and filenames'''
    checksums_list = []
    for line in checksums_string.split('\n'):
        try:
            checksum, filename = line.split()
            # strip leading ./
            if filename.startswith('./'):
                filename = filename[2:]
            checksums_list.append({'checksum': checksum, 'file': filename})
        except ValueError:
            logging.debug('Skipping over unexpected checksum line %r', line)
            continue

    return checksums_list

def has_gbk_file_changed(directory, checksums):
    '''Check if the checksum of a given gbk file has changed'''
    filename, expected_checksum = get_name_and_checksum(checksums, '_genomic.gbff.gz')
    full_filename = os.path.join(directory, filename)
    # if file doesn't exist, it has changed
    if not os.path.isfile(full_filename):
        return True

    actual_checksum = md5sum(full_filename)
    return expected_checksum != actual_checksum

def get_name_and_checksum(checksums, end):
    '''Extract a full filename and checksum from the checksums list for a file ending in given end'''
    for entry in checksums:
        if not entry['file'].endswith(end):
            # wrong file
            continue
        filename = entry['file']
        expected_checksum = entry['checksum']
        return filename, expected_checksum

def md5sum(filename):
    '''Calculate the md5sum of a file and return the hexdigest'''
    hash_md5 = hashlib.md5()
    with open(filename, 'rb') as fh:
        for chunk in iter(lambda: fh.read(4096), b''):
            hash_md5.update(chunk)
    return hash_md5.hexdigest()

def download_gbk_file(entry, directory, checksums):
    '''Download and verirfy a given gbk file'''
    filename, expected_checksum = get_name_and_checksum(checksums, '_genomic.gbff.gz')
    base_url = convert_ftp_url(entry['ftp_path'])
    full_url = '{}/{}'.format(base_url, filename)
    local_file = os.path.join(directory, filename)

    r = requests.get(full_url, stream=True)

    with open(local_file, 'wb') as fh:
        for chunk in r.iter_content(4096):
            fh.write(chunk)

    actual_checksum = md5sum(local_file)
    if actual_checksum != expected_checksum:
        logging.error('Checksum mismatch for %r. Expected %r, got %r',
                local_file, expected_checksum, actual_checksum)
        return False

    return True
<|endoftext|>


<commit_before>###
# Copyright (c) 2002-2005, Jeremiah Fincher
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#   * Redistributions of source code must retain the above copyright notice,
#     this list of conditions, and the following disclaimer.
#   * Redistributions in binary form must reproduce the above copyright notice,
#     this list of conditions, and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#   * Neither the name of the author of this software nor the name of
#     contributors to this software may be used to endorse or promote products
#     derived from this software without specific prior written consent.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
###

"""
Database module, similar to dbhash.  Uses a format similar to (if not entirely
the same as) DJB's CDB <http://cr.yp.to/cdb.html>.
"""

import os
import sys
import struct
import os.path
import cPickle as pickle

import supybot.utils as utils

def hash(s):
    """DJB's hash function for CDB."""
    h = 5381
    for c in s:
        h = ((h + (h << 5)) ^ ord(c)) & 0xFFFFFFFFL
    return h

def unpack2Ints(s):
    """Returns two ints unpacked from the binary string s."""
    return struct.unpack('<LL', s)

def pack2Ints(i, j):
    """Returns a packed binary string from the two ints."""
    return struct.pack('<LL', i, j)

def dump(map, fd=sys.stdout):
    """Dumps a dictionary-structure in CDB format."""
    for (key, value) in map.iteritems():
        fd.write('+%s,%s:%s->%s\n' % (len(key), len(value), key, value))

def open_db(filename, mode='r', **kwargs):
    """Opens a database; used for compatibility with other database modules."""
    if mode == 'r':
        return Reader(filename, **kwargs)
    elif mode == 'w':
        return ReaderWriter(filename, **kwargs)
    elif mode == 'c':
        if os.path.exists(filename):
            return ReaderWriter(filename, **kwargs)
        else:
            maker = Maker(filename)
            maker.finish()
            return ReaderWriter(filename, **kwargs)
    elif mode == 'n':
        maker = Maker(filename)
        maker.finish()
        return ReaderWriter(filename, **kwargs)
    else:
        raise ValueError, 'Invalid flag: %s' % mode

def shelf(filename, *args, **kwargs):
    """Opens a new shelf database object."""
    if os.path.exists(filename):
        return Shelf(filename, *args, **kwargs)
    else:
        maker = Maker(filename)
        maker.finish()
        return Shelf(filename, *args, **kwargs)

def _readKeyValue(fd):
    klen = 0
    dlen = 0
    s = initchar = fd.read(1)
    if s == '':
        return (None, None, None)
    s = fd.read(1)
    while s != ',':
        klen = 10 * klen + int(s)
        s = fd.read(1)
    s = fd.read(1)
    while s != ':':
        dlen = 10 * dlen + int(s)
        s = fd.read(1)
    key = fd.read(klen)
    assert fd.read(2) == '->'
    value = fd.read(dlen)
    assert fd.read(1) == '\n'
    return (initchar, key, value)

def make(dbFilename, readFilename=None):
    """Makes a database from the filename, otherwise uses stdin."""
    if readFilename is None:
        readfd = sys.stdin
    else:
        readfd = open(readFilename, 'rb')
    maker = Maker(dbFilename)
    while 1:
        (initchar, key, value) = _readKeyValue(readfd)
        if initchar is None:
            break
        assert initchar == '+'
        maker.add(key, value)
    readfd.close()
    maker.finish()

class Maker(object):
    """Class for making CDB databases."""
    def __init__(self, filename):
        self.fd = utils.file.AtomicFile(filename, 'wb')
        self.filename = filename
        self.fd.seek(2048)
        self.hashPointers = [(0, 0)] * 256
        #self.hashes = [[]] * 256 # Can't use this, [] stays the same...
        self.hashes = []
        for _ in xrange(256):
            self.hashes.append([])

    def add(self, key, data):
        """Adds a key->value pair to the database."""
        h = hash(key)
        hashPointer = h % 256
        startPosition = self.fd.tell()
        self.fd.write(pack2Ints(len(key), len(data)))
        self.fd.write(key)
        self.fd.write(data)
        self.hashes[hashPointer].append((h, startPosition))

    def finish(self):
        """Finishes the current Maker object.

        Writes the remainder of the database to disk.
        """
        for i in xrange(256):
            hash = self.hashes[i]
            self.hashPointers[i] = (self.fd.tell(), self._serializeHash(hash))
        self._serializeHashPointers()
        self.fd.flush()
        self.fd.close()

    def _serializeHash(self, hash):
        hashLen = len(hash) * 2
        a = [(0, 0)] * hashLen
        for (h, pos) in hash:
            i = (h / 256) % hashLen
            while a[i] != (0, 0):
                i = (i + 1) % hashLen
            a[i] = (h, pos)
        for (h, pos) in a:
            self.fd.write(pack2Ints(h, pos))
        return hashLen

    def _serializeHashPointers(self):
        self.fd.seek(0)
        for (hashPos, hashLen) in self.hashPointers:
            self.fd.write(pack2Ints(hashPos, hashLen))

class Reader(utils.IterableMap):
    """Class for reading from a CDB database."""
    def __init__(self, filename):
        self.filename = filename
        self.fd = open(filename, 'rb')
        self.loop = 0
        self.khash = 0
        self.kpos = 0
        self.hpos = 0
        self.hslots = 0
        self.dpos = 0
        self.dlen = 0

    def close(self):
        self.fd.close()

    def _read(self, len, pos):
        self.fd.seek(pos)
        return self.fd.read(len)

    def _match(self, key, pos):
        return self._read(len(key), pos) == key

    def iteritems(self):
        # uses loop/hslots in a strange, non-re-entrant manner.
        (self.loop,) = struct.unpack('<i', self._read(4, 0))
        self.hslots = 2048
        while self.hslots < self.loop:
            (klen, dlen) = unpack2Ints(self._read(8, self.hslots))
            dpos = self.hslots + 8 + klen
            ret = (self._read(klen, self.hslots+8), self._read(dlen, dpos))
            self.hslots = dpos + dlen
            yield ret
        self.loop = 0
        self.hslots = 0

    def _findnext(self, key):
        if not self.loop:
            self.khash = hash(key)
            (self.hpos, self.hslots) = unpack2Ints(self._read(8,
                                                    (self.khash * 8) & 2047))
            if not self.hslots:
                return False
            self.kpos = self.hpos + (((self.khash / 256) % self.hslots) * 8)
        while self.loop < self.hslots:
            (h, p) = unpack2Ints(self._read(8, self.kpos))
            if p == 0:
                return False
            self.loop += 1
            self.kpos += 8
            if self.kpos == self.hpos + (self.hslots * 8):
                self.kpos = self.hpos
            if h == self.khash:
                (u, self.dlen) = unpack2Ints(self._read(8, p))
                if u == len(key):
                    if self._match(key, p+8):
                        self.dpos = p + 8 + u
                        return True
        return False

    def _find(self, key, loop=0):
        self.loop = loop
        return self._findnext(key)

    def _getCurrentData(self):
        return self._read(self.dlen, self.dpos)

    def find(self, key, loop=0):
        if self._find(key, loop=loop):
            return self._getCurrentData()
        else:
            try:
                return self.default
            except AttributeError:
                raise KeyError, key

    def findall(self, key):
        ret = []
        while self._findnext(key):
            ret.append(self._getCurrentData())
        return ret

    def get(self, key, default=None):
        try:
            return self[key]
        except KeyError:
            return default

    def __len__(self):
        (start,) = struct.unpack('<i', self._read(4, 0))
        self.fd.seek(0, 2)
        return ((self.fd.tell() - start) / 16)

    has_key = _find
    __contains__ = has_key
    __getitem__ = find

class ReaderWriter(utils.IterableMap):
    """Uses a journal to pretend that a CDB is writable database."""
    def __init__(self, filename, journalName=None, maxmods=0):
        if journalName is None:
            journalName = filename + '.journal'
        self.journalName = journalName
        self.maxmods = maxmods
        self.mods = 0
        self.filename = filename
        self._readJournal()
        self._openFiles()
        self.adds = {}
        self.removals = set()

    def _openFiles(self):
        self.cdb = Reader(self.filename)
        self.journal = open(self.journalName, 'wb')

    def _closeFiles(self):
        self.cdb.close()
        self.journal.close()

    def _journalRemoveKey(self, key):
        s = '-%s,%s:%s->%s\n' % (len(key), 0, key, '')
        self.journal.write(s)
        self.journal.flush()

    def _journalAddKey(self, key, value):
        s = '+%s,%s:%s->%s\n' % (len(key), len(value), key, value)
        self.journal.write(s)
        self.journal.flush()

    def _readJournal(self):
        removals = set()
        adds = {}
        try:
            fd = open(self.journalName, 'rb')
            while 1:
                (initchar, key, value) = _readKeyValue(fd)
                if initchar is None:
                    break
                elif initchar == '+':
                    if key in removals:
                        removals.remove(key)
                    adds[key] = value
                elif initchar == '-':
                    if key in adds:
                        del adds[key]
                    removals.add(key)
            fd.close()
        except IOError:
            pass
        if removals or adds:
            maker = Maker(self.filename)
            cdb = Reader(self.filename)
            for (key, value) in cdb.iteritems():
                if key in removals:
                    continue
                elif key in adds:
                    value = adds[key]
                    if value is not None:
                        maker.add(key, value)
                        adds[key] = None
                else:
                    maker.add(key, value)
            for (key, value) in adds.iteritems():
                if value is not None:
                    maker.add(key, value)
            cdb.close()
            maker.finish()
        if os.path.exists(self.journalName):
            os.remove(self.journalName)

    def close(self):
        self.flush()
        self._closeFiles()

    def flush(self):
        self._closeFiles()
        self._readJournal()
        self._openFiles()

    def _flushIfOverLimit(self):
        if self.maxmods:
            if isinstance(self.maxmods, int):
                if self.mods > self.maxmods:
                    self.flush()
                    self.mods = 0
            elif isinstance(self.maxmods, float):
                assert 0 <= self.maxmods
                if self.mods / max(len(self.cdb), 100) > self.maxmods:
                    self.flush()
                    self.mods = 0

    def __getitem__(self, key):
        if key in self.removals:
            raise KeyError, key
        else:
            try:
                return self.adds[key]
            except KeyError:
                return self.cdb[key] # If this raises KeyError, we lack key.

    def __delitem__(self, key):
        if key in self.removals:
            raise KeyError, key
        else:
            if key in self.adds and key in self.cdb:
                self._journalRemoveKey(key)
                del self.adds[key]
                self.removals.add(key)
            elif key in self.adds:
                self._journalRemoveKey(key)
                del self.adds[key]
            elif key in self.cdb:
                self._journalRemoveKey(key)
            else:
                raise KeyError, key
        self.mods += 1
        self._flushIfOverLimit()

    def __setitem__(self, key, value):
        if key in self.removals:
            self.removals.remove(key)
        self._journalAddKey(key, value)
        self.adds[key] = value
        self.mods += 1
        self._flushIfOverLimit()

    def __contains__(self, key):
        if key in self.removals:
            return False
        else:
            return key in self.adds or key in self.cdb

    has_key = __contains__

    def iteritems(self):
        already = set()
        for (key, value) in self.cdb.iteritems():
            if key in self.removals or key in already:
                continue
            elif key in self.adds:
                already.add(key)
                yield (key, self.adds[key])
            else:
                yield (key, value)
        for (key, value) in self.adds.iteritems():
            if key not in already:
                yield (key, value)

    def setdefault(self, key, value):
        try:
            return self[key]
        except KeyError:
            self[key] = value
            return value

    def get(self, key, default=None):
        try:
            return self[key]
        except KeyError:
            return default

class Shelf(ReaderWriter):
    """Uses pickle to mimic the shelf module."""
    def __getitem__(self, key):
        return pickle.loads(ReaderWriter.__getitem__(self, key))

    def __setitem__(self, key, value):
        ReaderWriter.__setitem__(self, key, pickle.dumps(value, True))

    def iteritems(self):
        for (key, value) in ReaderWriter.iteritems(self):
            yield (key, pickle.loads(value))

if __name__ == '__main__':
    if sys.argv[0] == 'cdbdump':
        if len(sys.argv) == 2:
            fd = open(sys.argv[1], 'rb')
        else:
            fd = sys.stdin
        db = Reader(fd)
        dump(db)
    elif sys.argv[0] == 'cdbmake':
        if len(sys.argv) == 2:
            make(sys.argv[1])
        else:
            make(sys.argv[1], sys.argv[2])
# vim:set shiftwidth=4 softtabstop=4 expandtab textwidth=79:
<commit_msg>src/cdb.py: Use 'future' division and decode/encode data from/to the disk.<commit_after>###
# Copyright (c) 2002-2005, Jeremiah Fincher
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
#   * Redistributions of source code must retain the above copyright notice,
#     this list of conditions, and the following disclaimer.
#   * Redistributions in binary form must reproduce the above copyright notice,
#     this list of conditions, and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#   * Neither the name of the author of this software nor the name of
#     contributors to this software may be used to endorse or promote products
#     derived from this software without specific prior written consent.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
# ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.
###

"""
Database module, similar to dbhash.  Uses a format similar to (if not entirely
the same as) DJB's CDB <http://cr.yp.to/cdb.html>.
"""

from __future__ import division

import os
import sys
import struct
import os.path
import cPickle as pickle

import supybot.utils as utils

def hash(s):
    """DJB's hash function for CDB."""
    h = 5381
    for c in s:
        h = ((h + (h << 5)) ^ ord(c)) & 0xFFFFFFFFL
    return h

def unpack2Ints(s):
    """Returns two ints unpacked from the binary string s."""
    return struct.unpack('<LL', s)

def pack2Ints(i, j):
    """Returns a packed binary string from the two ints."""
    return struct.pack('<LL', i, j)

def dump(map, fd=sys.stdout):
    """Dumps a dictionary-structure in CDB format."""
    for (key, value) in map.iteritems():
        fd.write('+%s,%s:%s->%s\n' % (len(key), len(value), key, value))

def open_db(filename, mode='r', **kwargs):
    """Opens a database; used for compatibility with other database modules."""
    if mode == 'r':
        return Reader(filename, **kwargs)
    elif mode == 'w':
        return ReaderWriter(filename, **kwargs)
    elif mode == 'c':
        if os.path.exists(filename):
            return ReaderWriter(filename, **kwargs)
        else:
            maker = Maker(filename)
            maker.finish()
            return ReaderWriter(filename, **kwargs)
    elif mode == 'n':
        maker = Maker(filename)
        maker.finish()
        return ReaderWriter(filename, **kwargs)
    else:
        raise ValueError, 'Invalid flag: %s' % mode

def shelf(filename, *args, **kwargs):
    """Opens a new shelf database object."""
    if os.path.exists(filename):
        return Shelf(filename, *args, **kwargs)
    else:
        maker = Maker(filename)
        maker.finish()
        return Shelf(filename, *args, **kwargs)

def _readKeyValue(fd):
    klen = 0
    dlen = 0
    s = initchar = fd.read(1)
    if s == '':
        return (None, None, None)
    s = fd.read(1)
    while s != ',':
        klen = 10 * klen + int(s)
        s = fd.read(1)
    s = fd.read(1)
    while s != ':':
        dlen = 10 * dlen + int(s)
        s = fd.read(1)
    key = fd.read(klen)
    assert fd.read(2) == '->'
    value = fd.read(dlen)
    assert fd.read(1) == '\n'
    return (initchar, key, value)

def make(dbFilename, readFilename=None):
    """Makes a database from the filename, otherwise uses stdin."""
    if readFilename is None:
        readfd = sys.stdin
    else:
        readfd = open(readFilename, 'rb')
    maker = Maker(dbFilename)
    while 1:
        (initchar, key, value) = _readKeyValue(readfd)
        if initchar is None:
            break
        assert initchar == '+'
        maker.add(key, value)
    readfd.close()
    maker.finish()

class Maker(object):
    """Class for making CDB databases."""
    def __init__(self, filename):
        self.fd = utils.file.AtomicFile(filename, 'wb')
        self.filename = filename
        self.fd.seek(2048)
        self.hashPointers = [(0, 0)] * 256
        #self.hashes = [[]] * 256 # Can't use this, [] stays the same...
        self.hashes = []
        for _ in xrange(256):
            self.hashes.append([])

    def add(self, key, data):
        """Adds a key->value pair to the database."""
        h = hash(key)
        hashPointer = h % 256
        startPosition = self.fd.tell()
        self.fd.write(pack2Ints(len(key), len(data)))
        self.fd.write(key.encode())
        self.fd.write(data.encode())
        self.hashes[hashPointer].append((h, startPosition))

    def finish(self):
        """Finishes the current Maker object.

        Writes the remainder of the database to disk.
        """
        for i in xrange(256):
            hash = self.hashes[i]
            self.hashPointers[i] = (self.fd.tell(), self._serializeHash(hash))
        self._serializeHashPointers()
        self.fd.flush()
        self.fd.close()

    def _serializeHash(self, hash):
        hashLen = len(hash) * 2
        a = [(0, 0)] * hashLen
        for (h, pos) in hash:
            i = (h // 256) % hashLen
            while a[i] != (0, 0):
                i = (i + 1) % hashLen
            a[i] = (h, pos)
        for (h, pos) in a:
            self.fd.write(pack2Ints(h, pos))
        return hashLen

    def _serializeHashPointers(self):
        self.fd.seek(0)
        for (hashPos, hashLen) in self.hashPointers:
            self.fd.write(pack2Ints(hashPos, hashLen))

class Reader(utils.IterableMap):
    """Class for reading from a CDB database."""
    def __init__(self, filename):
        self.filename = filename
        self.fd = open(filename, 'rb')
        self.loop = 0
        self.khash = 0
        self.kpos = 0
        self.hpos = 0
        self.hslots = 0
        self.dpos = 0
        self.dlen = 0

    def close(self):
        self.fd.close()

    def _read(self, len, pos):
        self.fd.seek(pos)
        return self.fd.read(len)

    def _match(self, key, pos):
        return self._read(len(key), pos) == key

    def iteritems(self):
        # uses loop/hslots in a strange, non-re-entrant manner.
        (self.loop,) = struct.unpack('<i', self._read(4, 0))
        self.hslots = 2048
        while self.hslots < self.loop:
            (klen, dlen) = unpack2Ints(self._read(8, self.hslots))
            dpos = self.hslots + 8 + klen
            ret = (self._read(klen, self.hslots+8).decode(),
                    self._read(dlen, dpos).decode())
            self.hslots = dpos + dlen
            yield ret
        self.loop = 0
        self.hslots = 0

    def _findnext(self, key):
        if not self.loop:
            self.khash = hash(key)
            (self.hpos, self.hslots) = unpack2Ints(self._read(8,
                                                    (self.khash * 8) & 2047))
            if not self.hslots:
                return False
            self.kpos = self.hpos + (((self.khash // 256) % self.hslots) * 8)
        while self.loop < self.hslots:
            (h, p) = unpack2Ints(self._read(8, self.kpos))
            if p == 0:
                return False
            self.loop += 1
            self.kpos += 8
            if self.kpos == self.hpos + (self.hslots * 8):
                self.kpos = self.hpos
            if h == self.khash:
                (u, self.dlen) = unpack2Ints(self._read(8, p))
                if u == len(key):
                    if self._match(key, p+8):
                        self.dpos = p + 8 + u
                        return True
        return False

    def _find(self, key, loop=0):
        self.loop = loop
        return self._findnext(key)

    def _getCurrentData(self):
        return self._read(self.dlen, self.dpos).decode()

    def find(self, key, loop=0):
        if self._find(key, loop=loop):
            return self._getCurrentData()
        else:
            try:
                return self.default
            except AttributeError:
                raise KeyError, key

    def findall(self, key):
        ret = []
        while self._findnext(key):
            ret.append(self._getCurrentData())
        return ret

    def get(self, key, default=None):
        try:
            return self[key]
        except KeyError:
            return default

    def __len__(self):
        (start,) = struct.unpack('<i', self._read(4, 0))
        self.fd.seek(0, 2)
        return ((self.fd.tell() - start) // 16)

    has_key = _find
    __contains__ = has_key
    __getitem__ = find

class ReaderWriter(utils.IterableMap):
    """Uses a journal to pretend that a CDB is writable database."""
    def __init__(self, filename, journalName=None, maxmods=0):
        if journalName is None:
            journalName = filename + '.journal'
        self.journalName = journalName
        self.maxmods = maxmods
        self.mods = 0
        self.filename = filename
        self._readJournal()
        self._openFiles()
        self.adds = {}
        self.removals = set()

    def _openFiles(self):
        self.cdb = Reader(self.filename)
        self.journal = open(self.journalName, 'w')

    def _closeFiles(self):
        self.cdb.close()
        self.journal.close()

    def _journalRemoveKey(self, key):
        s = '-%s,%s:%s->%s\n' % (len(key), 0, key, '')
        self.journal.write(s)
        self.journal.flush()

    def _journalAddKey(self, key, value):
        s = '+%s,%s:%s->%s\n' % (len(key), len(value), key, value)
        self.journal.write(s)
        self.journal.flush()

    def _readJournal(self):
        removals = set()
        adds = {}
        try:
            fd = open(self.journalName, 'r')
            while 1:
                (initchar, key, value) = _readKeyValue(fd)
                if initchar is None:
                    break
                elif initchar == '+':
                    if key in removals:
                        removals.remove(key)
                    adds[key] = value
                elif initchar == '-':
                    if key in adds:
                        del adds[key]
                    removals.add(key)
            fd.close()
        except IOError:
            pass
        if removals or adds:
            maker = Maker(self.filename)
            cdb = Reader(self.filename)
            for (key, value) in cdb.iteritems():
                if key in removals:
                    continue
                elif key in adds:
                    value = adds[key]
                    if value is not None:
                        maker.add(key, value)
                        adds[key] = None
                else:
                    maker.add(key, value)
            for (key, value) in adds.iteritems():
                if value is not None:
                    maker.add(key, value)
            cdb.close()
            maker.finish()
        if os.path.exists(self.journalName):
            os.remove(self.journalName)

    def close(self):
        self.flush()
        self._closeFiles()

    def flush(self):
        self._closeFiles()
        self._readJournal()
        self._openFiles()

    def _flushIfOverLimit(self):
        if self.maxmods:
            if isinstance(self.maxmods, int):
                if self.mods > self.maxmods:
                    self.flush()
                    self.mods = 0
            elif isinstance(self.maxmods, float):
                assert 0 <= self.maxmods
                if self.mods / max(len(self.cdb), 100) > self.maxmods:
                    self.flush()
                    self.mods = 0

    def __getitem__(self, key):
        if key in self.removals:
            raise KeyError, key
        else:
            try:
                return self.adds[key]
            except KeyError:
                return self.cdb[key] # If this raises KeyError, we lack key.

    def __delitem__(self, key):
        if key in self.removals:
            raise KeyError, key
        else:
            if key in self.adds and key in self.cdb:
                self._journalRemoveKey(key)
                del self.adds[key]
                self.removals.add(key)
            elif key in self.adds:
                self._journalRemoveKey(key)
                del self.adds[key]
            elif key in self.cdb:
                self._journalRemoveKey(key)
            else:
                raise KeyError, key
        self.mods += 1
        self._flushIfOverLimit()

    def __setitem__(self, key, value):
        if key in self.removals:
            self.removals.remove(key)
        self._journalAddKey(key, value)
        self.adds[key] = value
        self.mods += 1
        self._flushIfOverLimit()

    def __contains__(self, key):
        if key in self.removals:
            return False
        else:
            return key in self.adds or key in self.cdb

    has_key = __contains__

    def iteritems(self):
        already = set()
        for (key, value) in self.cdb.iteritems():
            if key in self.removals or key in already:
                continue
            elif key in self.adds:
                already.add(key)
                yield (key, self.adds[key])
            else:
                yield (key, value)
        for (key, value) in self.adds.iteritems():
            if key not in already:
                yield (key, value)

    def setdefault(self, key, value):
        try:
            return self[key]
        except KeyError:
            self[key] = value
            return value

    def get(self, key, default=None):
        try:
            return self[key]
        except KeyError:
            return default

class Shelf(ReaderWriter):
    """Uses pickle to mimic the shelf module."""
    def __getitem__(self, key):
        return pickle.loads(ReaderWriter.__getitem__(self, key))

    def __setitem__(self, key, value):
        ReaderWriter.__setitem__(self, key, pickle.dumps(value, True))

    def iteritems(self):
        for (key, value) in ReaderWriter.iteritems(self):
            yield (key, pickle.loads(value))

if __name__ == '__main__':
    if sys.argv[0] == 'cdbdump':
        if len(sys.argv) == 2:
            fd = open(sys.argv[1], 'rb')
        else:
            fd = sys.stdin
        db = Reader(fd)
        dump(db)
    elif sys.argv[0] == 'cdbmake':
        if len(sys.argv) == 2:
            make(sys.argv[1])
        else:
            make(sys.argv[1], sys.argv[2])
# vim:set shiftwidth=4 softtabstop=4 expandtab textwidth=79:
<|endoftext|>


<commit_before>#   Copyright 2008 Red Hat, Inc.
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import ConfigParser
import os
import socket
import re
import tarfile
import zipfile
import logging
import time
try:
   from subprocess import *
   use_popen2 = False
except:
   from popen2 import *
   use_popen2 = True
from cStringIO import StringIO

SUCCESS = 0
FAILURE = 1

# Generic enumerated type
class wf_type:
   def __init__(self,list=[]):
      count = 0
      for item in list.split():
         vars(self)[item] = count
         count = count + 1

   def __setattr__(self, name, value):
      raise ValueError('Modifying values is not allowed')

   def __delattr__(self, name):
      raise ValueError('Deleting entries is not allowed')

# The type of requests from condor to act upon
condor_wf_types = wf_type('get_work prepare_job reply_claim_accept reply_claim_reject update_job_status exit_exit exit_remove exit_hold exit_evict')

# The information about the job
class condor_wf(object):
    def set_type(self, thetype):
           self.__type__ = thetype

    def get_type(self):
        return self.__type__

    type = property(get_type, set_type)

    def set_data(self, string):
           self.__data__ = string

    def get_data(self):
        return self.__data__

    data = property(get_data, set_data)

# General failure exception
class GeneralError(Exception):
   def __init__(self, *msgs):
      self.msgs = msgs

# Configuration exception
class config_err(Exception):
   def __init__(self, msg):
      self.msg = msg

def run_cmd(cmd, args, environ=None):
   """Runs the command specified in 'cmd' with the given 'args' using
      either the subprocess module or the open2 module, depending on which
      is supported.  The subprocess module is favored.  Returns
      (return_code, stdout, stderr)"""
   retcode = -1
   std_out = ''
   std_err = ''
   if environ != None:
      try:
         for var in environ.keys():
            os.environ[var] = environ[var]
      except:
         return ([-1, None, None])

   if use_popen2 == False:
      cmd_list = [cmd]
      cmd_list += args.split(' ')
      obj = Popen(cmd_list, stdout=PIPE, stderr=PIPE)
      (std_out, std_err) = obj.communicate()
      retcode = obj.returncode
   else:
      obj = Popen3('%s %s' % (cmd, args), True)
      retcode = obj.wait()
      try:
         std_out = obj.fromchild.readlines()[0]
         obj.fromchild.close()
      except:
         pass
      try:
         obj.tochild.close()
      except:
         pass
      try:
         std_err = obj.childerr.readlines()[0]
         obj.childerr.close()
      except:
         pass

   if environ != None:
      for var in environ.keys():
         try:
            del os.environ[var]
         except:
            pass
   return (retcode, std_out, std_err)

def read_condor_config(subsys, attr_list):
   """ Uses condor_config_val to look up values in condor's configuration.
       First looks for subsys_param, then for the newer subsys.param.
       Returns map(param, value)"""
   config = {}
   if attr_list == []:
      (rcode, value, stderr) = run_cmd('condor_config_val', '%s' % (subsys))
      if rcode == 0:
         config[subsys.lower()] = value.strip()
      else:
         # Config value not found.  Raise an exception
         raise config_err('"%s" is not defined' % subsys)
   else:
      for attr in attr_list:
         (rcode, value, stderr) = run_cmd('condor_config_val', '%s_%s' % (subsys, attr))
         if rcode == 0:
            config[attr.lower()] = value.strip()
         else:
            # Try the newer <subsys>.param form
            (rcode, value, stderr) = run_cmd('condor_config_val', '%s.%s' % (subsys, attr))
            if rcode == 0:
               config[attr.lower()] = value.strip()
            else:
               # Config value not found.  Raise an exception
               raise config_err('"%s_%s" is not defined' % (subsys, attr))
   return config

def read_config_file(config_file, section):
   """Given configuration file and section names, returns a list of all value
     	 pairs in the config file as a dictionary"""
   parser = ConfigParser.ConfigParser()

   # Parse the config file
   if os.path.exists(config_file) == False:
      raise config_err('%s configuration file does not exist.' % config_file)
   else:
      try:
         parser.readfp(open(config_file))
         items = parser.items(section)
      except Exception, error:
         raise config_err('Problems reading %s.' % config_file, str(error))

      # Take the list of lists and convert into a dictionary
      dict = {}
      for list in items:
         dict[list[0]] = list[1].strip()
      return dict

def socket_read_all(sock):
   """Read all data waiting to be read on the given socket.  The first attempt
      to read will be a blocking call, and all subsequent reads will be
      non-blocking."""
   msg = ''
   old_timeout = sock.gettimeout()
   data = sock.recv(1024)
   sock.settimeout(0.1)
   try:
      while len(data) != 0:
         msg += data
         data = sock.recv(1024)
   except socket.timeout, timeout:
      pass
   except Exception, error:
      close_socket(sock)
      if error != None:
         raise GeneralError('socket error %d: %s' % (error[0], error[1]))
   sock.settimeout(old_timeout)
   return msg

def close_socket(connection):
   """Performs a shutdown and a close on the socket.  Any errors
      are logged to the system logging service.  Raises a GeneralError
      if any errors are encountered"""
   try:
      try:
         # python 2.3 doesn't have SHUT_WR defined, so use it's value (1)
         connection.shutdown(1)
      except Exception, error:
         if error != None:
            raise GeneralError('socket error %d: %s' % (error[0], error[1]))
      try:
         data = connection.recv(4096)
         while len(data) != 0:
            data = connection.recv(4096)
      except:
         pass
   finally:
      connection.close()
      connection = None

def log_messages(level, exception, logger_name=''):
   """Logs messages in the passed GeneralError exception to the
      system logger"""
   logger = logging.getLogger(logger_name)
   for msg in exception.msgs:
      if (msg != ''):
         logger.log(level, msg)

def write_file(filename, data):
   """Writes the given data into the given filename in pieces to account for
      large files"""
   file_ptr = open(filename, 'wb')
   file_str = StringIO(data)
   buffer_length = 2 ** 20
   contents = file_str.read(buffer_length)
   while contents:
      file_ptr.write(contents)
      contents = file_str.read(buffer_length)
   file_ptr.flush()
   file_ptr.close()

def grep(pattern, data):
   """Returns the first instance found of pattern in data.  If pattern
      doesn't exist in data, None is returned.  If data contains groups,
      returns the matching subgroups instead"""
   for line in data.split('\n'):
      match = re.match(pattern, line)
      if match != None:
         break
   if match != None and match.groups() != None:
      found = match.groups()
   else:
      found = match
   return(found)

def zip_extract(filename):
   """Unzips the contents of the filename given, preserving permissions and
      ownership"""
   zip = zipfile.ZipFile(filename, 'r')
   contents = zip.namelist()

   # Loop through the archive names and extract the contents
   # Make sure to create higher level directories before creating a lower
   # level file or directory (if applicable)
   contents.sort()
   for item in contents:
      if item.endswith('/'):
         # This item is all directories, so recursively create them if
         # they don't already exist
         if not os.path.exists(item):
            os.makedirs(item)
      else:
         # This is a directory + file or a file by itself, so create the
         # directory hierarchy first (if applicable) then extract the file.
         dirs = item.split('/')
         dir_structure = ''
         for count in range(0,len(dirs)-1):
            full_path = os.path.join(dir_structure, dirs[count])
            if dirs[count] != '' and not os.path.exists(full_path):
               os.mkdir(full_path)
            dir_structure = full_path

         write_file(item, zip.read(item))

   # Set the perserved permissions and timestamp for the extracted
   # files/directories
   for name in contents:
      info = zip.getinfo(name)
      file_time = time.mktime(info.date_time + (0, 0, -1))
      os.chmod(name, info.external_attr >> 16L)
      os.utime(name, (file_time, file_time))

def tarball_extract(filename):
   """Extracts the tarball (.tar.gz) given"""
   tarball = tarfile.open(filename, 'r:gz')
   for file in tarball.getnames():
      tarball.extract(file)
   tarball.close()
<commit_msg>- Removed GeneralError and changed exception names - Reordered args to log_messages - log_messages no long takes an exception, but a list of strings to print<commit_after>#   Copyright 2008 Red Hat, Inc.
#
#   Licensed under the Apache License, Version 2.0 (the "License");
#   you may not use this file except in compliance with the License.
#   You may obtain a copy of the License at
#
#       http://www.apache.org/licenses/LICENSE-2.0
#
#   Unless required by applicable law or agreed to in writing, software
#   distributed under the License is distributed on an "AS IS" BASIS,
#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#   See the License for the specific language governing permissions and
#   limitations under the License.

import ConfigParser
import os
import socket
import re
import tarfile
import zipfile
import logging
import time
try:
   from subprocess import *
   use_popen2 = False
except:
   from popen2 import *
   use_popen2 = True
from cStringIO import StringIO

SUCCESS = 0
FAILURE = 1

# Generic enumerated type
class wf_type:
   def __init__(self,list=[]):
      count = 0
      for item in list.split():
         vars(self)[item] = count
         count = count + 1

   def __setattr__(self, name, value):
      raise ValueError('Modifying values is not allowed')

   def __delattr__(self, name):
      raise ValueError('Deleting entries is not allowed')

# The type of requests from condor to act upon
condor_wf_types = wf_type('get_work prepare_job reply_claim_accept reply_claim_reject update_job_status exit_exit exit_remove exit_hold exit_evict')

# The information about the job
class condor_wf(object):
    def set_type(self, thetype):
           self.__type__ = thetype

    def get_type(self):
        return self.__type__

    type = property(get_type, set_type)

    def set_data(self, string):
           self.__data__ = string

    def get_data(self):
        return self.__data__

    data = property(get_data, set_data)

# General failure exception
class SocketError(Exception):
   def __init__(self, msg):
      self.msgs = msg

# Configuration exception
class ConfigError(Exception):
   def __init__(self, msg):
      self.msg = msg

def run_cmd(cmd, args, environ=None):
   """Runs the command specified in 'cmd' with the given 'args' using
      either the subprocess module or the open2 module, depending on which
      is supported.  The subprocess module is favored.  Returns
      (return_code, stdout, stderr)"""
   retcode = -1
   std_out = ''
   std_err = ''
   if environ != None:
      try:
         for var in environ.keys():
            os.environ[var] = environ[var]
      except:
         return ([-1, None, None])

   if use_popen2 == False:
      cmd_list = [cmd]
      cmd_list += args.split(' ')
      obj = Popen(cmd_list, stdout=PIPE, stderr=PIPE)
      (std_out, std_err) = obj.communicate()
      retcode = obj.returncode
   else:
      obj = Popen3('%s %s' % (cmd, args), True)
      retcode = obj.wait()
      try:
         std_out = obj.fromchild.readlines()[0]
         obj.fromchild.close()
      except:
         pass
      try:
         obj.tochild.close()
      except:
         pass
      try:
         std_err = obj.childerr.readlines()[0]
         obj.childerr.close()
      except:
         pass

   if environ != None:
      for var in environ.keys():
         try:
            del os.environ[var]
         except:
            pass
   return (retcode, std_out, std_err)

def read_condor_config(subsys, attr_list):
   """ Uses condor_config_val to look up values in condor's configuration.
       First looks for subsys_param, then for the newer subsys.param.
       Returns map(param, value)"""
   config = {}
   if attr_list == []:
      (rcode, value, stderr) = run_cmd('condor_config_val', '%s' % (subsys))
      if rcode == 0:
         config[subsys.lower()] = value.strip()
      else:
         # Config value not found.  Raise an exception
         raise ConfigError('"%s" is not defined' % subsys)
   else:
      for attr in attr_list:
         (rcode, value, stderr) = run_cmd('condor_config_val', '%s_%s' % (subsys, attr))
         if rcode == 0:
            config[attr.lower()] = value.strip()
         else:
            # Try the newer <subsys>.param form
            (rcode, value, stderr) = run_cmd('condor_config_val', '%s.%s' % (subsys, attr))
            if rcode == 0:
               config[attr.lower()] = value.strip()
            else:
               # Config value not found.  Raise an exception
               raise ConfigError('"%s_%s" is not defined' % (subsys, attr))
   return config

def read_config_file(config_file, section):
   """Given configuration file and section names, returns a list of all value
     	 pairs in the config file as a dictionary"""
   parser = ConfigParser.ConfigParser()

   # Parse the config file
   if os.path.exists(config_file) == False:
      raise ConfigError('%s configuration file does not exist.' % config_file)
   else:
      try:
         parser.readfp(open(config_file))
         items = parser.items(section)
      except Exception, error:
         raise ConfigError('Problems reading %s.' % config_file, str(error))

      # Take the list of lists and convert into a dictionary
      dict = {}
      for list in items:
         dict[list[0]] = list[1].strip()
      return dict

def socket_read_all(sock):
   """Read all data waiting to be read on the given socket.  The first attempt
      to read will be a blocking call, and all subsequent reads will be
      non-blocking."""
   msg = ''
   old_timeout = sock.gettimeout()
   data = sock.recv(1024)
   sock.settimeout(0.1)
   try:
      while len(data) != 0:
         msg += data
         data = sock.recv(1024)
   except socket.timeout, timeout:
      pass
   except Exception, error:
      close_socket(sock)
      if error != None:
         raise SocketError('socket error %d: %s' % (error[0], error[1]))
   sock.settimeout(old_timeout)
   return msg

def close_socket(connection):
   """Performs a shutdown and a close on the socket.  Any errors
      are logged to the system logging service.  Raises a SocketError
      if any errors are encountered"""
   try:
      try:
         # python 2.3 doesn't have SHUT_WR defined, so use it's value (1)
         connection.shutdown(1)
      except Exception, error:
         if error != None:
            raise SocketError('socket error %d: %s' % (error[0], error[1]))
      try:
         data = connection.recv(4096)
         while len(data) != 0:
            data = connection.recv(4096)
      except:
         pass
   finally:
      connection.close()
      connection = None

def log_messages(level, logger_name, *msgs):
   """Logs messages in the passed msgs to the logger with the provided
      logger_name."""
   logger = logging.getLogger(logger_name)
   for msg in msgs:
      if (msg != ''):
         logger.log(level, msg)

def write_file(filename, data):
   """Writes the given data into the given filename in pieces to account for
      large files"""
   file_ptr = open(filename, 'wb')
   file_str = StringIO(data)
   buffer_length = 2 ** 20
   contents = file_str.read(buffer_length)
   while contents:
      file_ptr.write(contents)
      contents = file_str.read(buffer_length)
   file_ptr.flush()
   file_ptr.close()

def grep(pattern, data):
   """Returns the first instance found of pattern in data.  If pattern
      doesn't exist in data, None is returned.  If data contains groups,
      returns the matching subgroups instead"""
   for line in data.split('\n'):
      match = re.match(pattern, line)
      if match != None:
         break
   if match != None and match.groups() != None:
      found = match.groups()
   else:
      found = match
   return(found)

def zip_extract(filename):
   """Unzips the contents of the filename given, preserving permissions and
      ownership"""
   zip = zipfile.ZipFile(filename, 'r')
   contents = zip.namelist()

   # Loop through the archive names and extract the contents
   # Make sure to create higher level directories before creating a lower
   # level file or directory (if applicable)
   contents.sort()
   for item in contents:
      if item.endswith('/'):
         # This item is all directories, so recursively create them if
         # they don't already exist
         if not os.path.exists(item):
            os.makedirs(item)
      else:
         # This is a directory + file or a file by itself, so create the
         # directory hierarchy first (if applicable) then extract the file.
         dirs = item.split('/')
         dir_structure = ''
         for count in range(0,len(dirs)-1):
            full_path = os.path.join(dir_structure, dirs[count])
            if dirs[count] != '' and not os.path.exists(full_path):
               os.mkdir(full_path)
            dir_structure = full_path

         write_file(item, zip.read(item))

   # Set the perserved permissions and timestamp for the extracted
   # files/directories
   for name in contents:
      info = zip.getinfo(name)
      file_time = time.mktime(info.date_time + (0, 0, -1))
      os.chmod(name, info.external_attr >> 16L)
      os.utime(name, (file_time, file_time))

def tarball_extract(filename):
   """Extracts the tarball (.tar.gz) given"""
   tarball = tarfile.open(filename, 'r:gz')
   for file in tarball.getnames():
      tarball.extract(file)
   tarball.close()
<|endoftext|>


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

# Usage: ./split_fasta.py /path/to/fasta.fa 2000
# ./split_fasta.py /path/to/reads.fastq 100 fastq
from __future__ import print_function
from Bio import SeqIO
import sys

fasta = sys.argv[1]
n = int(sys.argv[2])
if len(sys.argv) > 3:
    fmt = sys.argv[3]
else:
    fmt = "fasta"
out_fasta = os.path.join(os.getcwd(), os.path.basename(fasta))

with open(fasta, 'r') as in_handle:
    out_handle = open("{}.1.{}".format(out_fasta, fmt), 'w') # start from 1 so PBS likes us
    i = 0
    for record in SeqIO.parse(in_handle, fmt):
        if i % n == 0 and i > 0:
            out_handle.close()
            out_handle = open("{}.{}.{}".format(out_fasta, (i // n) + 1, fmt), 'w')
        SeqIO.write(record, out_handle, fmt)
        i += 1
out_handle.close()
print("Successfully split {} into {} files.".format(fasta, i // n))
<commit_msg>import os<commit_after>#!/usr/bin/env python

# Usage: ./split_fasta.py /path/to/fasta.fa 2000
# ./split_fasta.py /path/to/reads.fastq 100 fastq
from __future__ import print_function
from Bio import SeqIO
import sys
import os

fasta = sys.argv[1]
n = int(sys.argv[2])
if len(sys.argv) > 3:
    fmt = sys.argv[3]
else:
    fmt = "fasta"
out_fasta = os.path.join(os.getcwd(), os.path.basename(fasta))

with open(fasta, 'r') as in_handle:
    out_handle = open("{}.1.{}".format(out_fasta, fmt), 'w') # start from 1 so PBS likes us
    i = 0
    for record in SeqIO.parse(in_handle, fmt):
        if i % n == 0 and i > 0:
            out_handle.close()
            out_handle = open("{}.{}.{}".format(out_fasta, (i // n) + 1, fmt), 'w')
        SeqIO.write(record, out_handle, fmt)
        i += 1
out_handle.close()
print("Successfully split {} into {} files.".format(fasta, i // n))
<|endoftext|>


<commit_before>#!/bin/python

#from PIL import Image, ImageFilter

#with Image.open("forum.jpg") as img:
#    print(img.size)
#    blurred = img.filter(ImageFilter.GaussianBlur(96))
#    blurred.save("blurred.jpg", "JPEG")

import sqlite3
import os
from flask import Flask, jsonify, send_file, g, request, redirect, abort
from werkzeug.utils import secure_filename

DATABASE = "database.db"
UPLOAD_FOLDER = "./images/"

app = Flask(__name__)
app.config.from_object(__name__)

@app.route("/")
def index():
    return "Wall Art"

images = [
    {
        "id": 1,
        "name": "Forum"
    }
]

def connect_db():
    return sqlite3.connect(app.config["DATABASE"])

def get_db():
    db = getattr(g, "_database", None)
    if db is None:
        db = g._database = connect_db()
    return db

@app.teardown_appcontext
def close_db(exception):
    db = getattr(g, "_database", None)
    if db is not None:
        db.close()

def query_db(query, args=(), one=False):
    cur = get_db().execute(query, args)
    rv = cur.fetchall()
    cur.close()
    return (rv[0] if rv else None) if one else rv

@app.route("/api/image", methods=["GET"])
def get_images():
    return jsonify({ "images": images })

@app.route("/api/image/<int:image_id>", methods=["GET"])
def get_image(image_id):
    # Note: This assumes the image id is valid
    img = query_db("select * from entries where id = ?", [image_id], one=True)
    if img is None:
        abort(404)
    return send_file(img[2], mimetype="image/jpeg")

@app.route("/upload", methods=["GET", "POST"])
def upload_image():
    if request.method == "POST":
        f = request.files["file"]
        if f:
            filename = os.path.join(app.config["UPLOAD_FOLDER"], secure_filename(f.filename))
            f.save(filename)
            if query_db("select * from entries where title = ?", [request.form["title"]], one=True) is None:
                db = get_db()
                print("Will insert new image")
                db.execute("insert into entries (title, filename) values (?, ?)", [request.form["title"], filename])
                db.commit()
        return redirect("/upload")

    return """
    <!DOCTYPE html5>
    <title>Upload image</title>

Upload File

    <form action="" method=post enctype=multipart/form-data>

           <input type=file name=file>
           Title
           <input type=text name=title>

           <input type=submit value=Upload>

    </form>
    """

if __name__ == "__main__":
    app.run(host="0.0.0.0", debug=True)

<commit_msg>Have image route now list all images in database<commit_after>#!/bin/python

#from PIL import Image, ImageFilter

#with Image.open("forum.jpg") as img:
#    print(img.size)
#    blurred = img.filter(ImageFilter.GaussianBlur(96))
#    blurred.save("blurred.jpg", "JPEG")

import sqlite3
import os
from flask import Flask, jsonify, send_file, g, request, redirect, abort
from werkzeug.utils import secure_filename

DATABASE = "database.db"
UPLOAD_FOLDER = "./images/"

app = Flask(__name__)
app.config.from_object(__name__)

@app.route("/")
def index():
    return "Wall Art"

def connect_db():
    return sqlite3.connect(app.config["DATABASE"])

def get_db():
    db = getattr(g, "_database", None)
    if db is None:
        db = g._database = connect_db()
    return db

@app.teardown_appcontext
def close_db(exception):
    db = getattr(g, "_database", None)
    if db is not None:
        db.close()

def query_db(query, args=(), one=False):
    cur = get_db().execute(query, args)
    rv = cur.fetchall()
    cur.close()
    return (rv[0] if rv else None) if one else rv

@app.route("/api/image", methods=["GET"])
def get_images():
    images = []
    for img in query_db("select * from entries"):
        print("image id = {}, title = {}, file = {}".format(img[0], img[1], img[2]))
        images.append({"id": img[0], "title": img[1], "filename": img[2]})
    return jsonify({"images": images })

@app.route("/api/image/<int:image_id>", methods=["GET"])
def get_image(image_id):
    # Note: This assumes the image id is valid
    img = query_db("select * from entries where id = ?", [image_id], one=True)
    if img is None:
        abort(404)
    return send_file(img[2], mimetype="image/jpeg")

@app.route("/upload", methods=["GET", "POST"])
def upload_image():
    if request.method == "POST":
        f = request.files["file"]
        if f:
            filename = os.path.join(app.config["UPLOAD_FOLDER"], secure_filename(f.filename))
            f.save(filename)
            if query_db("select * from entries where title = ?", [request.form["title"]], one=True) is None:
                db = get_db()
                print("Will insert new image")
                db.execute("insert into entries (title, filename) values (?, ?)", [request.form["title"], filename])
                db.commit()
        return redirect("/upload")

    return """
    <!DOCTYPE html5>
    <title>Upload image</title>

Upload File

    <form action="" method=post enctype=multipart/form-data>

           <input type=file name=file>
           Title
           <input type=text name=title>

           <input type=submit value=Upload>

    </form>
    """

if __name__ == "__main__":
    app.run(host="0.0.0.0", debug=True)

<|endoftext|>


<commit_before>import os
import posixpath

import yaml

from docutils import nodes
from docutils.frontend import OptionParser
from docutils.statemachine import ViewList

from sphinx import addnodes
from sphinx.parsers import RSTParser
from sphinx.util.fileutil import copy_asset
from sphinx.util.nodes import nested_parse_with_titles
from sphinx.util.docutils import SphinxDirective, new_document

from .md2rst import md2rst
from .nodes import (add_asdf_nodes, toc_link, schema_header_title,
                    schema_title, schema_description, schema_properties,
                    schema_property, schema_property_name,
                    schema_property_details, schema_combiner_body,
                    schema_combiner_item, section_header, asdf_tree, asdf_ref,
                    example_section, example_item, example_description)

SCHEMA_DEF_SECTION_TITLE = 'Schema Definitions'
EXAMPLE_SECTION_TITLE = 'Examples'
ORIGINAL_SCHEMA_SECTION_TITLE = 'Original Schema'

class schema_def(nodes.comment):
    pass

class AsdfSchemas(SphinxDirective):

    required_arguments = 0
    optional_arguments = 0
    has_content = True

    def _process_asdf_toctree(self):

        links = []
        for name in self.content:
            if not name:
                continue
            schema = self.env.path2doc(name.strip() + '.rst')
            link = posixpath.join('generated', schema)
            links.append((schema, link))

        tocnode = addnodes.toctree()
        tocnode['includefiles'] = [x[1] for x in links]
        tocnode['entries'] = links
        tocnode['maxdepth'] = -1
        tocnode['glob'] = None

        paragraph = nodes.paragraph(text="Here's where the schemas go")
        return [paragraph, tocnode]

    def run(self):

        # This is the case when we are actually using Sphinx to generate
        # documentation
        if not getattr(self.env, 'autoasdf_generate', False):
            return self._process_asdf_toctree()

        # This case allows us to use docutils to parse input documents during
        # the 'builder-inited' phase so that we can determine which new
        # document need to be created by 'autogenerate_schema_docs'. This seems
        # much cleaner than writing a custom parser to extract the schema
        # information.
        return [schema_def(text=c.strip().split()[0]) for c in self.content]

class AsdfSchema(SphinxDirective):

    has_content = True

    def run(self):

        self.schema_name = self.content[0]
        schema_dir = self.state.document.settings.env.config.asdf_schema_path
        srcdir = self.state.document.settings.env.srcdir

        schema_file = posixpath.join(srcdir, schema_dir, self.schema_name) + '.yaml'

        with open(schema_file) as ff:
            raw_content = ff.read()
            schema = yaml.safe_load(raw_content)

        title = self._parse_title(schema.get('title', ''), schema_file)

        docnodes = [title]

        description = schema.get('description', '')
        if description:
            docnodes.append(schema_header_title(text='Description'))
            docnodes.append(self._parse_description(description, schema_file))

        docnodes.append(schema_header_title(text='Outline'))
        docnodes.append(self._create_toc(schema))

        docnodes.append(section_header(text=SCHEMA_DEF_SECTION_TITLE))
        docnodes.append(self._process_properties(schema, top=True))

        examples = schema.get('examples', [])
        if examples:
            docnodes.append(section_header(text=EXAMPLE_SECTION_TITLE))
            docnodes.append(self._process_examples(examples, schema_file))

        docnodes.append(section_header(text=ORIGINAL_SCHEMA_SECTION_TITLE))
        docnodes.append(nodes.literal_block(text=raw_content))

        return docnodes

    def _create_toc(self, schema):

        toc = nodes.compound()
        toc.append(toc_link(text=SCHEMA_DEF_SECTION_TITLE))
        if 'examples' in schema:
            toc.append(toc_link(text=EXAMPLE_SECTION_TITLE))
        toc.append(toc_link(text=ORIGINAL_SCHEMA_SECTION_TITLE))
        return toc

    def _markdown_to_nodes(self, text, filename):
        rst = ViewList()
        for i, line in enumerate(md2rst(text).split('\n')):
            rst.append(line, filename, i+1)

        node = nodes.section()
        node.document = self.state.document

        nested_parse_with_titles(self.state, rst, node)

        return node.children

    def _parse_title(self, title, filename):
        nodes = self._markdown_to_nodes(title, filename)
        return schema_title(None, *nodes)

    def _parse_description(self, description, filename):
        nodes = self._markdown_to_nodes(description, filename)
        return schema_description(None, *nodes)

    def _create_ref_node(self, ref):
        treenodes = asdf_tree()
        treenodes.append(asdf_ref(text=ref))
        return treenodes

    def _process_validation_keywords(self, schema, typename=None):
        node_list = []
        typename = typename or schema['type']

        if typename == 'string':
            if not ('minLength' in schema or 'maxLength' in schema):
                node_list.append(nodes.emphasis(text='No length restriction'))
            if schema.get('minLength', 0):
                text = 'Minimum length: {}'.format(schema['minLength'])
                node_list.append(nodes.line(text=text))
            if 'maxLength' in schema:
                text = 'Maximum length: {}'.format(schema['maxLength'])
                node_list.append(nodes.line(text=text))
            if 'pattern' in schema:
                node_list.append(nodes.line(text='Must match the following pattern:'))
                node_list.append(nodes.literal_block(text=schema['pattern']))

        elif typename == 'array':
            if not ('minItems' in schema or 'maxItems' in schema):
                node_list.append(nodes.emphasis(text='No length restriction'))
            if schema.get('minItems', 0):
                text = 'Minimum length: {}'.format(schema['minItems'])
                node_list.append(nodes.line(text=text))
            if 'maxItems' in schema:
                text = 'Maximum length: {}'.format(schema['maxItems'])
                node_list.append(nodes.line(text=text))

        # TODO: numerical validation keywords

        return node_list

    def _process_top_type(self, schema):
        tree = nodes.compound()
        prop = nodes.compound()
        typename = schema['type']
        prop.append(schema_property_name(text=typename))
        prop.extend(self._process_validation_keywords(schema))
        tree.append(prop)
        return tree

    def _process_properties(self, schema, nodetype=schema_properties, top=False):
        if 'properties' in schema:
            treenodes = asdf_tree()
            required = schema.get('required', [])
            for key, node in schema['properties'].items():
                treenodes.append(self._create_top_property(key, node,
                                                           key in required))
            comment = nodes.line(text='This type is an object with the following properties:')
            return nodetype(None, *[comment, treenodes])
        elif 'type' in schema:
            details = self._process_top_type(schema)
            return nodetype(None, details)
        elif 'anyOf' in schema:
            children = self._create_combiner(schema['anyOf'], 'any', top=top)
            return nodetype(None, *children)
        elif 'allOf' in schema:
            children = self._create_combiner(schema['allOf'], 'all', top=top)
            return nodetype(None, *children)
        elif '$ref' in schema:
            comment = nodes.line(text='This schema node is a reference:')
            ref = self._create_ref_node(schema['$ref'])
            return nodetype(None, *[comment, ref])
        # TODO: handle case of top-level allOf
        else:
            text = nodes.emphasis(text='This node has no type definition')
            return nodetype(None, text)

    def _create_combiner(self, items, combiner, top=False):
        text = 'This node must validate against **{}** of the following'
        text_nodes = self._markdown_to_nodes(text.format(combiner), '')

        body = schema_combiner_body(top=top)
        body.extend(text_nodes)
        for i, tree in enumerate(items):
            body.append(self._process_properties(tree, nodetype=schema_combiner_item))

        return [body]

    def _create_reference(self, refname):
        return refname + '.html'

    def _create_top_property(self, name, tree, required):

        description = tree.get('description', '')

        if '$ref' in tree:
            # TODO: make the reference a link
            typ = tree.get('$ref')
            ref = self._create_reference(typ)
        else:
            typ = tree.get('type', 'object')
            ref = None

        prop = schema_property()
        prop.append(schema_property_name(text=name))
        prop.append(schema_property_details(typ, required, ref))
        prop.append(self._parse_description(description, ''))
        if typ != 'object':
            prop.extend(self._process_validation_keywords(tree, typename=typ))
        else:
            prop.append(self._process_properties(tree))
        return prop

    def _process_examples(self, tree, filename):
        examples = example_section(num=len(tree))
        for i, example in enumerate(tree):
            node = example_item()
            desc_text = self._markdown_to_nodes(example[0]+':', filename)
            description = example_description(None, *desc_text)
            node.append(description)
            node.append(nodes.literal_block(text=example[1]))
            examples.append(node)
        return examples

def find_autoasdf_directives(env, filename):

    parser = RSTParser()
    settings = OptionParser(components=(RSTParser,)).get_default_values()
    settings.env = env
    document = new_document(filename, settings)

    with open(filename) as ff:
        parser.parse(ff.read(), document)

    return [x.children[0].astext() for x in document.traverse()
            if isinstance(x, schema_def)]

def find_autoschema_references(app, genfiles):

    # We set this environment variable to indicate that the AsdfSchemas
    # directive should be parsed as a simple list of schema references
    # rather than as the toctree that will be generated when the documentation
    # is actually built.
    app.env.autoasdf_generate = True

    schemas = set()
    for fn in genfiles:
        # Create documentation files based on contents of asdf-schema directives
        path = posixpath.join(app.env.srcdir, fn)
        app.env.temp_data['docname'] = app.env.path2doc(path)
        schemas = schemas.union(find_autoasdf_directives(app.env, path))

    # Unset this variable now that we're done.
    app.env.autoasdf_generate = False

    return list(schemas)

def create_schema_docs(app, schemas):

    output_dir = posixpath.join(app.srcdir, 'generated')
    os.makedirs(output_dir, exist_ok=True)

    for schema_name in schemas:
        doc_path = posixpath.join(output_dir, schema_name + '.rst')

        if posixpath.exists(doc_path):
            continue

        os.makedirs(posixpath.dirname(doc_path), exist_ok=True)

        with open(doc_path, 'w') as ff:
            ff.write(schema_name + '\n')
            ff.write('=' * len(schema_name) + '\n\n')
            ff.write('.. asdf-schema::\n\n')
            ff.write('    {}\n'.format(schema_name))

def autogenerate_schema_docs(app):

    env = app.env

    genfiles = [env.doc2path(x, base=None) for x in env.found_docs
                if posixpath.isfile(env.doc2path(x))]

    if not genfiles:
        return

    ext = list(app.config.source_suffix)
    genfiles = [genfile + (not genfile.endswith(tuple(ext)) and ext[0] or '')
                for genfile in genfiles]

    # Read all source documentation files and parse all asdf-schema directives
    schemas = find_autoschema_references(app, genfiles)
    # Create the documentation files that correspond to the schemas listed
    create_schema_docs(app, schemas)

def on_build_finished(app, exc):
    if exc is None:
        for asset in ['asdf_schema.css', 'asdf.js']:
            src = posixpath.join(posixpath.dirname(__file__), asset)
            dst = posixpath.join(app.outdir, '_static')
            copy_asset(src, dst)

def setup(app):

    # Describes a path relative to the sphinx source directory
    app.add_config_value('asdf_schema_path', 'schemas', 'env')
    app.add_directive('asdf-autoschemas', AsdfSchemas)
    app.add_directive('asdf-schema', AsdfSchema)

    add_asdf_nodes(app)

    app.add_css_file('asdf_schema.css')
    app.add_javascript('asdf.js')

    app.connect('builder-inited', autogenerate_schema_docs)
    app.connect('build-finished', on_build_finished)

    return dict(version='0.1')
<commit_msg>Get rid of extraneous parametrization of nodetype<commit_after>import os
import posixpath

import yaml

from docutils import nodes
from docutils.frontend import OptionParser
from docutils.statemachine import ViewList

from sphinx import addnodes
from sphinx.parsers import RSTParser
from sphinx.util.fileutil import copy_asset
from sphinx.util.nodes import nested_parse_with_titles
from sphinx.util.docutils import SphinxDirective, new_document

from .md2rst import md2rst
from .nodes import (add_asdf_nodes, toc_link, schema_header_title,
                    schema_title, schema_description, schema_properties,
                    schema_property, schema_property_name,
                    schema_property_details, schema_combiner_body,
                    schema_combiner_item, section_header, asdf_tree, asdf_ref,
                    example_section, example_item, example_description)

SCHEMA_DEF_SECTION_TITLE = 'Schema Definitions'
EXAMPLE_SECTION_TITLE = 'Examples'
ORIGINAL_SCHEMA_SECTION_TITLE = 'Original Schema'

class schema_def(nodes.comment):
    pass

class AsdfSchemas(SphinxDirective):

    required_arguments = 0
    optional_arguments = 0
    has_content = True

    def _process_asdf_toctree(self):

        links = []
        for name in self.content:
            if not name:
                continue
            schema = self.env.path2doc(name.strip() + '.rst')
            link = posixpath.join('generated', schema)
            links.append((schema, link))

        tocnode = addnodes.toctree()
        tocnode['includefiles'] = [x[1] for x in links]
        tocnode['entries'] = links
        tocnode['maxdepth'] = -1
        tocnode['glob'] = None

        paragraph = nodes.paragraph(text="Here's where the schemas go")
        return [paragraph, tocnode]

    def run(self):

        # This is the case when we are actually using Sphinx to generate
        # documentation
        if not getattr(self.env, 'autoasdf_generate', False):
            return self._process_asdf_toctree()

        # This case allows us to use docutils to parse input documents during
        # the 'builder-inited' phase so that we can determine which new
        # document need to be created by 'autogenerate_schema_docs'. This seems
        # much cleaner than writing a custom parser to extract the schema
        # information.
        return [schema_def(text=c.strip().split()[0]) for c in self.content]

class AsdfSchema(SphinxDirective):

    has_content = True

    def run(self):

        self.schema_name = self.content[0]
        schema_dir = self.state.document.settings.env.config.asdf_schema_path
        srcdir = self.state.document.settings.env.srcdir

        schema_file = posixpath.join(srcdir, schema_dir, self.schema_name) + '.yaml'

        with open(schema_file) as ff:
            raw_content = ff.read()
            schema = yaml.safe_load(raw_content)

        title = self._parse_title(schema.get('title', ''), schema_file)

        docnodes = [title]

        description = schema.get('description', '')
        if description:
            docnodes.append(schema_header_title(text='Description'))
            docnodes.append(self._parse_description(description, schema_file))

        docnodes.append(schema_header_title(text='Outline'))
        docnodes.append(self._create_toc(schema))

        docnodes.append(section_header(text=SCHEMA_DEF_SECTION_TITLE))
        docnodes.append(self._process_properties(schema, top=True))

        examples = schema.get('examples', [])
        if examples:
            docnodes.append(section_header(text=EXAMPLE_SECTION_TITLE))
            docnodes.append(self._process_examples(examples, schema_file))

        docnodes.append(section_header(text=ORIGINAL_SCHEMA_SECTION_TITLE))
        docnodes.append(nodes.literal_block(text=raw_content))

        return docnodes

    def _create_toc(self, schema):

        toc = nodes.compound()
        toc.append(toc_link(text=SCHEMA_DEF_SECTION_TITLE))
        if 'examples' in schema:
            toc.append(toc_link(text=EXAMPLE_SECTION_TITLE))
        toc.append(toc_link(text=ORIGINAL_SCHEMA_SECTION_TITLE))
        return toc

    def _markdown_to_nodes(self, text, filename):
        rst = ViewList()
        for i, line in enumerate(md2rst(text).split('\n')):
            rst.append(line, filename, i+1)

        node = nodes.section()
        node.document = self.state.document

        nested_parse_with_titles(self.state, rst, node)

        return node.children

    def _parse_title(self, title, filename):
        nodes = self._markdown_to_nodes(title, filename)
        return schema_title(None, *nodes)

    def _parse_description(self, description, filename):
        nodes = self._markdown_to_nodes(description, filename)
        return schema_description(None, *nodes)

    def _create_ref_node(self, ref):
        treenodes = asdf_tree()
        treenodes.append(asdf_ref(text=ref))
        return treenodes

    def _process_validation_keywords(self, schema, typename=None):
        node_list = []
        typename = typename or schema['type']

        if typename == 'string':
            if not ('minLength' in schema or 'maxLength' in schema):
                node_list.append(nodes.emphasis(text='No length restriction'))
            if schema.get('minLength', 0):
                text = 'Minimum length: {}'.format(schema['minLength'])
                node_list.append(nodes.line(text=text))
            if 'maxLength' in schema:
                text = 'Maximum length: {}'.format(schema['maxLength'])
                node_list.append(nodes.line(text=text))
            if 'pattern' in schema:
                node_list.append(nodes.line(text='Must match the following pattern:'))
                node_list.append(nodes.literal_block(text=schema['pattern']))

        elif typename == 'array':
            if not ('minItems' in schema or 'maxItems' in schema):
                node_list.append(nodes.emphasis(text='No length restriction'))
            if schema.get('minItems', 0):
                text = 'Minimum length: {}'.format(schema['minItems'])
                node_list.append(nodes.line(text=text))
            if 'maxItems' in schema:
                text = 'Maximum length: {}'.format(schema['maxItems'])
                node_list.append(nodes.line(text=text))

        # TODO: numerical validation keywords

        return node_list

    def _process_top_type(self, schema):
        tree = nodes.compound()
        prop = nodes.compound()
        typename = schema['type']
        prop.append(schema_property_name(text=typename))
        prop.extend(self._process_validation_keywords(schema))
        tree.append(prop)
        return tree

    def _process_properties(self, schema, top=False):
        if 'properties' in schema:
            treenodes = asdf_tree()
            required = schema.get('required', [])
            for key, node in schema['properties'].items():
                treenodes.append(self._create_top_property(key, node,
                                                           key in required))
            comment = nodes.line(text='This type is an object with the following properties:')
            return schema_properties(None, *[comment, treenodes])
        elif 'type' in schema:
            details = self._process_top_type(schema)
            return schema_properties(None, details)
        elif 'anyOf' in schema:
            children = self._create_combiner(schema['anyOf'], 'any', top=top)
            return schema_properties(None, *children)
        elif 'allOf' in schema:
            children = self._create_combiner(schema['allOf'], 'all', top=top)
            return schema_properties(None, *children)
        elif '$ref' in schema:
            comment = nodes.line(text='This schema node is a reference:')
            ref = self._create_ref_node(schema['$ref'])
            return schema_properties(None, *[comment, ref])
        else:
            text = nodes.emphasis(text='This node has no type definition')
            return schema_properties(None, text)

    def _create_combiner(self, items, combiner, top=False):
        text = 'This node must validate against **{}** of the following'
        text_nodes = self._markdown_to_nodes(text.format(combiner), '')

        body = schema_combiner_body(top=top)
        body.extend(text_nodes)
        for i, tree in enumerate(items):
            body.append(self._process_properties(tree))

        return [body]

    def _create_reference(self, refname):
        return refname + '.html'

    def _create_top_property(self, name, tree, required):

        description = tree.get('description', '')

        if '$ref' in tree:
            # TODO: make the reference a link
            typ = tree.get('$ref')
            ref = self._create_reference(typ)
        else:
            typ = tree.get('type', 'object')
            ref = None

        prop = schema_property()
        prop.append(schema_property_name(text=name))
        prop.append(schema_property_details(typ, required, ref))
        prop.append(self._parse_description(description, ''))
        if typ != 'object':
            prop.extend(self._process_validation_keywords(tree, typename=typ))
        else:
            prop.append(self._process_properties(tree))
        return prop

    def _process_examples(self, tree, filename):
        examples = example_section(num=len(tree))
        for i, example in enumerate(tree):
            node = example_item()
            desc_text = self._markdown_to_nodes(example[0]+':', filename)
            description = example_description(None, *desc_text)
            node.append(description)
            node.append(nodes.literal_block(text=example[1]))
            examples.append(node)
        return examples

def find_autoasdf_directives(env, filename):

    parser = RSTParser()
    settings = OptionParser(components=(RSTParser,)).get_default_values()
    settings.env = env
    document = new_document(filename, settings)

    with open(filename) as ff:
        parser.parse(ff.read(), document)

    return [x.children[0].astext() for x in document.traverse()
            if isinstance(x, schema_def)]

def find_autoschema_references(app, genfiles):

    # We set this environment variable to indicate that the AsdfSchemas
    # directive should be parsed as a simple list of schema references
    # rather than as the toctree that will be generated when the documentation
    # is actually built.
    app.env.autoasdf_generate = True

    schemas = set()
    for fn in genfiles:
        # Create documentation files based on contents of asdf-schema directives
        path = posixpath.join(app.env.srcdir, fn)
        app.env.temp_data['docname'] = app.env.path2doc(path)
        schemas = schemas.union(find_autoasdf_directives(app.env, path))

    # Unset this variable now that we're done.
    app.env.autoasdf_generate = False

    return list(schemas)

def create_schema_docs(app, schemas):

    output_dir = posixpath.join(app.srcdir, 'generated')
    os.makedirs(output_dir, exist_ok=True)

    for schema_name in schemas:
        doc_path = posixpath.join(output_dir, schema_name + '.rst')

        if posixpath.exists(doc_path):
            continue

        os.makedirs(posixpath.dirname(doc_path), exist_ok=True)

        with open(doc_path, 'w') as ff:
            ff.write(schema_name + '\n')
            ff.write('=' * len(schema_name) + '\n\n')
            ff.write('.. asdf-schema::\n\n')
            ff.write('    {}\n'.format(schema_name))

def autogenerate_schema_docs(app):

    env = app.env

    genfiles = [env.doc2path(x, base=None) for x in env.found_docs
                if posixpath.isfile(env.doc2path(x))]

    if not genfiles:
        return

    ext = list(app.config.source_suffix)
    genfiles = [genfile + (not genfile.endswith(tuple(ext)) and ext[0] or '')
                for genfile in genfiles]

    # Read all source documentation files and parse all asdf-schema directives
    schemas = find_autoschema_references(app, genfiles)
    # Create the documentation files that correspond to the schemas listed
    create_schema_docs(app, schemas)

def on_build_finished(app, exc):
    if exc is None:
        for asset in ['asdf_schema.css', 'asdf.js']:
            src = posixpath.join(posixpath.dirname(__file__), asset)
            dst = posixpath.join(app.outdir, '_static')
            copy_asset(src, dst)

def setup(app):

    # Describes a path relative to the sphinx source directory
    app.add_config_value('asdf_schema_path', 'schemas', 'env')
    app.add_directive('asdf-autoschemas', AsdfSchemas)
    app.add_directive('asdf-schema', AsdfSchema)

    add_asdf_nodes(app)

    app.add_css_file('asdf_schema.css')
    app.add_javascript('asdf.js')

    app.connect('builder-inited', autogenerate_schema_docs)
    app.connect('build-finished', on_build_finished)

    return dict(version='0.1')
<|endoftext|>