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

from __future__ import print_function
import argparse
import glob
import os
import re
import sys
import PIL
from PIL.ExifTags import TAGS
from PIL import Image

# Need to look for *.JPG, *.jpg, and *.jpeg files for consideration.
EXTENSIONS = ['JPG', 'jpg', 'jpeg']
MAX_RENAME_ATTEMPTS = 10

class FileMap():
    """FileMap represents a mapping between the old_fn and the new_fn. It's
    methods perform all necessary instance functions for the rename.

    Arguments:
        str: old_fn - Old Filename
        dict: exif_data - For testing only. Dict with sample EXIF data.
    """

    def __init__(self, old_fn, avoid_collisions=None, exif_data=None):
        """Initialize FileMap instance.

        >>> filemap = FileMap('abc123.jpeg', None, {})
        >>> filemap.old_fn
        'abc123.jpeg'
        >>> filemap.new_fn
        'abc123.jpg'
        >>>
        """
        self.MAX_RENAME_ATTEMPTS = MAX_RENAME_ATTEMPTS
        self.old_fn_fq = old_fn
        self.workdir = os.path.dirname(old_fn)
        self.old_fn = os.path.basename(old_fn)

        # Avoid filename collisions (dangerous) or log a message if there
        # would be one, and fail the move.
        self.collision_detected = False
        if avoid_collisions is None:
            self.avoid_collisions = False
        else:
            self.avoid_collisions = avoid_collisions

        # Read EXIF data from old filename
        if exif_data is None:
            self.read_exif_data()
        else:
            self.exif_data = exif_data

        self.get_new_fn()

    def read_exif_data(self):
        """Read EXIF data from file. Convert to Python dict."""

        # XXX: We already know file exists 'cuz we found it.
        img = Image.open(self.old_fn_fq)
        info = img._getexif()

        self.exif_data ={}
        if info is not None:
            for tag, value in info.items():
                decoded = TAGS.get(tag, tag)
                self.exif_data[decoded] = value
        else:
            raise Exception("{0} has no EXIF data.".format(self.old_fn))

    def get_new_fn(self):
        """Generate new filename from old_fn EXIF data if possible. Even if not
        possible, lowercase old_fn and normalize file extension.

        Arguments:
            dict: EXIF data

        >>> filemap = FileMap('abc123.jpeg', {'DateTimeOriginal': '2014:08:16 06:20:30'})
        >>> filemap.new_fn
        '20140816_062030.jpg'

        """

        # Start with EXIF DateTimeOriginal
        try:
            new_fn = self.exif_data['DateTimeOriginal']
        except KeyError:
            new_fn = None

        # If this pattern does not strictly match then keep original name.
        # YYYY:MM:DD HH:MM:SS
        if (new_fn and not
                re.match(r'^\d{4}:\d\d:\d\d \d\d:\d\d:\d\d$', new_fn)):
            # Setup for next step.
            new_fn = None

        # Don't assume exif tag exists. If it does not, keep original filename.
        # Lowercase filename base and extension
        if new_fn is None:
            new_fn = self.old_fn.lower()
            new_fn = re.sub(r'.jpeg$', r'.jpg', new_fn)
        else:
            new_fn = "{0}.jpg".format(new_fn)

        # XXX: One may argue that the next step should be an 'else' clause of the
        # previous 'if' statement. But the intention here is to clean up just a bit
        # even if we're not really renaming the file. Windows doesn't like colons
        # in filenames.

        # Rename using exif DateTimeOriginal
        new_fn = re.sub(r':', r'', new_fn)
        new_fn = re.sub(r' ', r'_', new_fn)

        self.new_fn = new_fn
        self.new_fn_fq = os.path.join(self.workdir, new_fn)

    def move(self):
        """Move old_fn to new_fn."""

        # XXX: This call deliberately placed here instead of __init__(). All
        # initialization is performed before any files are moved. The file move
        # will change state and may introduce a collision. Doing the uniqueness
        # check here will check current state.
        try:
            self.make_new_fn_unique()
        except Exception as e:
            raise e

        if self.collision_detected:
            print( "{0} => {1} Destination collision. Aborting.".format(
                os.path.basename(self.old_fn), os.path.basename(self.new_fn)))
            return

        try:
            print( "Moving the files: {0} ==> {1}".format(
                os.path.basename(self.old_fn), os.path.basename(self.new_fn)))
            os.rename(self.old_fn, self.new_fn)
        except Exception as e:
            print("Unable to rename file: {0}".format(e.message),
                    file=sys.stderr)

    def make_new_fn_unique(self):
        """Check new_fn for uniqueness in 'workdir'. Rename, adding a numerical
        suffix until it is unique.
        """

        # Rename file by appending number if we have collision.
        # TODO: I wish I didn't specify \d+_\d+ for the first part.
        # perhaps not -\d\ before .jpg would be better for the second
        # match.
        new_fn = self.new_fn
        counter = 1
        while(os.path.exists(self.new_fn_fq)):
            if (self.old_fn == self.new_fn):
                # Same file, faux collision.
                break
            if (not self.avoid_collisions):
                # Do not attempt to rename.
                self.collision_detected = True
                break
            new_fn = re.sub(r'^(\d+_\d+)-\d+\.jpg',
                    r'\1-{0}.jpg'.format(counter), self.new_fn)
            new_fn = re.sub(r'^(\d+_\d+)\.jpg',
                    r'\1-{0}.jpg'.format(counter), self.new_fn)
            counter += 1
            if counter > self.MAX_RENAME_ATTEMPTS:
                raise Exception("Too many rename attempts: {0}".format(self.new_fn))
            self.new_fn_fq = os.path.join(self.workdir, new_fn)
        self.new_fn = new_fn
        self.new_fn_fq = os.path.join(self.workdir, new_fn)

def init_file_map(workdir, avoid_collisions=None):
    """Read the work directory looking for files with extensions defined in the
    EXTENSIONS constant. Note that this could use a more elaborate magic
    number mechanism that would be cool.

    Arguments:
        str: workdir - The directory in which all activity will occur.

    Returns:
        list: file_map - List of FileMap instances.
    """

    # Dict with old_fn ==> new_fn mapping.
    file_map = []

    # Initialize file_map dict.
    for extension in EXTENSIONS:
        for filename in glob.glob(os.path.join(workdir,
                '*.{0}'.format(extension))):
            try:
                file_map.append(FileMap(filename, avoid_collisions))
            except:
                print("{0}".format(e.message), file=sys.stderr)

    return file_map

def process_file_map(file_map, simon_sez=None, move_func=None):
    """Iterate through the Python dict that maps old filenames to new
    filenames. Move the file if Simon sez.

    Arguments:
        str: workdir - Working directory.
        dict: file_map - old_fn to new_fn mapping.
        boolean: simon_sez - Dry run or real thing.
        func: move_func - Move function to use for testing or default.
    Returns:
        None

    >>> filemap = FileMap('IMG0332.JPG', {'DateTimeOriginal': '2014-08-18 20:23:83'})
    >>> def move_func(old_fn, new_fn): pass
    ...
    >>> process_file_map([filemap], True, move_func)
    >>>

    """

    # XXX: Of marginal utility
    if simon_sez is None:
        simon_sez = False

    for fm in file_map:
        try:
            if simon_sez:
                if move_func is None:
                    fm.move()
                else:
                    move_func(fm.old_fn, fm.new_fn)
        except Exception as e:
            print("{0}".format(e.message), file=sys.stderr)
            break

def process_all_files(workdir=None, simon_sez=None, avoid_collisions=None):
    """Manage the entire process of gathering data and renaming files."""

    if workdir is None:
        workdir = os.path.dirname(os.path.abspath(__file__))

    if not os.path.exists(workdir):
        print("Directory {0} does not exist. Exiting.",format(workdir),
                file=sys.stderr)
        sys.exit(1)

    if not os.access(workdir, os.W_OK):
        print("Directory {0} is not writable. Exiting.".format(workdir),
                file=sys.stderr)
        sys.exit(1)

    file_map = init_file_map(workdir, avoid_collisions)
    process_file_map(file_map, simon_sez)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-s", "--simon-sez",
            help="Really, Simon sez rename the files!", action="store_true")
    parser.add_argument("-a", "--avoid-collisions",
            help="Rename until filenames do not collide. Danger!", action="store_true")
    parser.add_argument("-d", "--directory",
            help="Read files from this directory.")
    args = parser.parse_args()
    process_all_files(workdir=args.directory, simon_sez=args.simon_sez,
            avoid_collisions=args.avoid_collisions)

<commit_msg>Fixes doctests<commit_after>#!/usr/bin/env python

from __future__ import print_function
import argparse
import glob
import os
import re
import sys
import PIL
from PIL.ExifTags import TAGS
from PIL import Image

# Need to look for *.JPG, *.jpg, and *.jpeg files for consideration.
EXTENSIONS = ['JPG', 'jpg', 'jpeg']
MAX_RENAME_ATTEMPTS = 10

class FileMap():
    """FileMap represents a mapping between the old_fn and the new_fn. It's
    methods perform all necessary instance functions for the rename.

    Arguments:
        str: old_fn - Old Filename
        dict: exif_data - For testing only. Dict with sample EXIF data.
    """

    def __init__(self, old_fn, avoid_collisions=None, exif_data=None):
        """Initialize FileMap instance.

        >>> filemap = FileMap('abc123.jpeg', None, {})
        >>> filemap.old_fn
        'abc123.jpeg'
        >>> filemap.new_fn
        'abc123.jpg'
        >>>
        """
        self.MAX_RENAME_ATTEMPTS = MAX_RENAME_ATTEMPTS
        self.old_fn_fq = old_fn
        self.workdir = os.path.dirname(old_fn)
        self.old_fn = os.path.basename(old_fn)

        # Avoid filename collisions (dangerous) or log a message if there
        # would be one, and fail the move.
        self.collision_detected = False
        if avoid_collisions is None:
            self.avoid_collisions = False
        else:
            self.avoid_collisions = avoid_collisions

        # Read EXIF data from old filename
        if exif_data is None:
            self.read_exif_data()
        else:
            self.exif_data = exif_data

        self.get_new_fn()

    def read_exif_data(self):
        """Read EXIF data from file. Convert to Python dict."""

        # XXX: We already know file exists 'cuz we found it.
        img = Image.open(self.old_fn_fq)
        info = img._getexif()

        self.exif_data ={}
        if info is not None:
            for tag, value in info.items():
                decoded = TAGS.get(tag, tag)
                self.exif_data[decoded] = value
        else:
            raise Exception("{0} has no EXIF data.".format(self.old_fn))

    def get_new_fn(self):
        """Generate new filename from old_fn EXIF data if possible. Even if not
        possible, lowercase old_fn and normalize file extension.

        Arguments:
            dict: EXIF data

        >>> filemap = FileMap('abc123.jpeg', avoid_collisions=None, exif_data={'DateTimeOriginal': '2014:08:16 06:20:30'})
        >>> filemap.new_fn
        '20140816_062030.jpg'

        """

        # Start with EXIF DateTimeOriginal
        try:
            new_fn = self.exif_data['DateTimeOriginal']
        except KeyError:
            new_fn = None

        # If this pattern does not strictly match then keep original name.
        # YYYY:MM:DD HH:MM:SS
        if (new_fn and not
                re.match(r'^\d{4}:\d\d:\d\d \d\d:\d\d:\d\d$', new_fn)):
            # Setup for next step.
            new_fn = None

        # Don't assume exif tag exists. If it does not, keep original filename.
        # Lowercase filename base and extension
        if new_fn is None:
            new_fn = self.old_fn.lower()
            new_fn = re.sub(r'.jpeg$', r'.jpg', new_fn)
        else:
            new_fn = "{0}.jpg".format(new_fn)

        # XXX: One may argue that the next step should be an 'else' clause of the
        # previous 'if' statement. But the intention here is to clean up just a bit
        # even if we're not really renaming the file. Windows doesn't like colons
        # in filenames.

        # Rename using exif DateTimeOriginal
        new_fn = re.sub(r':', r'', new_fn)
        new_fn = re.sub(r' ', r'_', new_fn)

        self.new_fn = new_fn
        self.new_fn_fq = os.path.join(self.workdir, new_fn)

    def move(self):
        """Move old_fn to new_fn."""

        # XXX: This call deliberately placed here instead of __init__(). All
        # initialization is performed before any files are moved. The file move
        # will change state and may introduce a collision. Doing the uniqueness
        # check here will check current state.
        try:
            self.make_new_fn_unique()
        except Exception as e:
            raise e

        if self.collision_detected:
            print( "{0} => {1} Destination collision. Aborting.".format(
                os.path.basename(self.old_fn), os.path.basename(self.new_fn)))
            return

        try:
            print( "Moving the files: {0} ==> {1}".format(
                os.path.basename(self.old_fn), os.path.basename(self.new_fn)))
            os.rename(self.old_fn, self.new_fn)
        except Exception as e:
            print("Unable to rename file: {0}".format(e.message),
                    file=sys.stderr)

    def make_new_fn_unique(self):
        """Check new_fn for uniqueness in 'workdir'. Rename, adding a numerical
        suffix until it is unique.
        """

        # Rename file by appending number if we have collision.
        # TODO: I wish I didn't specify \d+_\d+ for the first part.
        # perhaps not -\d\ before .jpg would be better for the second
        # match.
        new_fn = self.new_fn
        counter = 1
        while(os.path.exists(self.new_fn_fq)):
            if (self.old_fn == self.new_fn):
                # Same file, faux collision.
                break
            if (not self.avoid_collisions):
                # Do not attempt to rename.
                self.collision_detected = True
                break
            new_fn = re.sub(r'^(\d+_\d+)-\d+\.jpg',
                    r'\1-{0}.jpg'.format(counter), self.new_fn)
            new_fn = re.sub(r'^(\d+_\d+)\.jpg',
                    r'\1-{0}.jpg'.format(counter), self.new_fn)
            counter += 1
            if counter > self.MAX_RENAME_ATTEMPTS:
                raise Exception("Too many rename attempts: {0}".format(self.new_fn))
            self.new_fn_fq = os.path.join(self.workdir, new_fn)
        self.new_fn = new_fn
        self.new_fn_fq = os.path.join(self.workdir, new_fn)

def init_file_map(workdir, avoid_collisions=None):
    """Read the work directory looking for files with extensions defined in the
    EXTENSIONS constant. Note that this could use a more elaborate magic
    number mechanism that would be cool.

    Arguments:
        str: workdir - The directory in which all activity will occur.

    Returns:
        list: file_map - List of FileMap instances.
    """

    # Dict with old_fn ==> new_fn mapping.
    file_map = []

    # Initialize file_map dict.
    for extension in EXTENSIONS:
        for filename in glob.glob(os.path.join(workdir,
                '*.{0}'.format(extension))):
            try:
                file_map.append(FileMap(filename, avoid_collisions))
            except:
                print("{0}".format(e.message), file=sys.stderr)

    return file_map

def process_file_map(file_map, simon_sez=None, move_func=None):
    """Iterate through the Python dict that maps old filenames to new
    filenames. Move the file if Simon sez.

    Arguments:
        str: workdir - Working directory.
        dict: file_map - old_fn to new_fn mapping.
        boolean: simon_sez - Dry run or real thing.
        func: move_func - Move function to use for testing or default.
    Returns:
        None

    >>> filemap = FileMap('IMG0332.JPG', avoid_collisions=None, exif_data={'DateTimeOriginal': '2014-08-18 20:23:83'})
    >>> def move_func(old_fn, new_fn): pass
    ...
    >>> process_file_map([filemap], True, move_func)
    >>>

    """

    # XXX: Of marginal utility
    if simon_sez is None:
        simon_sez = False

    for fm in file_map:
        try:
            if simon_sez:
                if move_func is None:
                    fm.move()
                else:
                    move_func(fm.old_fn, fm.new_fn)
        except Exception as e:
            print("{0}".format(e.message), file=sys.stderr)
            break

def process_all_files(workdir=None, simon_sez=None, avoid_collisions=None):
    """Manage the entire process of gathering data and renaming files."""

    if workdir is None:
        workdir = os.path.dirname(os.path.abspath(__file__))

    if not os.path.exists(workdir):
        print("Directory {0} does not exist. Exiting.",format(workdir),
                file=sys.stderr)
        sys.exit(1)

    if not os.access(workdir, os.W_OK):
        print("Directory {0} is not writable. Exiting.".format(workdir),
                file=sys.stderr)
        sys.exit(1)

    file_map = init_file_map(workdir, avoid_collisions)
    process_file_map(file_map, simon_sez)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("-s", "--simon-sez",
            help="Really, Simon sez rename the files!", action="store_true")
    parser.add_argument("-a", "--avoid-collisions",
            help="Rename until filenames do not collide. Danger!", action="store_true")
    parser.add_argument("-d", "--directory",
            help="Read files from this directory.")
    args = parser.parse_args()
    process_all_files(workdir=args.directory, simon_sez=args.simon_sez,
            avoid_collisions=args.avoid_collisions)

<|endoftext|>


<commit_before>#!/usr/bin/python2.7

import msgpack
import string
import zlib

import shared
import messagetypes
from tr import _translate

BITMESSAGE_ENCODING_IGNORE = 0
BITMESSAGE_ENCODING_TRIVIAL = 1
BITMESSAGE_ENCODING_SIMPLE = 2
BITMESSAGE_ENCODING_EXTENDED = 3

class MsgEncode(object):
    def __init__(self, message, encoding=BITMESSAGE_ENCODING_SIMPLE):
        self.data = None
        self.encoding = encoding
        self.length = 0
        if self.encoding == BITMESSAGE_ENCODING_EXTENDED:
            self.encodeExtended(message)
        elif self.encoding == BITMESSAGE_ENCODING_SIMPLE:
            self.encodeSimple(message)
        elif self.encoding == BITMESSAGE_ENCODING_TRIVIAL:
            self.encodeTrivial(message)
        else:
            raise ValueError("Unknown encoding %i" % (encoding))

    def encodeExtended(self, message):
        try:
            msgObj = messagetypes.message.Message()
            self.data = zlib.compress(msgpack.dumps(msgObj.encode(message)), 9)
        except zlib.error:
            logger.error("Error compressing message")
            raise
        except msgpack.exceptions.PackException:
            logger.error("Error msgpacking message")
            raise
        self.length = len(self.data)

    def encodeSimple(self, message):
        self.data = 'Subject:' + message['subject'] + '\n' + 'Body:' + message['body']
        self.length = len(self.data)

    def encodeTrivial(self, message):
        self.data = message['body']
        self.length = len(self.data)

class MsgDecode(object):
    def __init__(self, encoding, data):
        self.encoding = encoding
        if self.encoding == BITMESSAGE_ENCODING_EXTENDED:
            self.decodeExtended(data)
        elif self.encoding in [BITMESSAGE_ENCODING_SIMPLE, BITMESSAGE_ENCODING_TRIVIAL]:
            self.decodeSimple(data)
        else:
            self.body = _translate("MsgDecode", "The message has an unknown encoding.\nPerhaps you should upgrade Bitmessage.")
            self.subject = _translate("MsgDecode", "Unknown encoding")

    def decodeExtended(self, data):
        try:
            tmp = msgpack.loads(zlib.decompress(data))
        except zlib.error:
            logger.error("Error decompressing message")
            raise
        except (msgpack.exceptions.UnpackException,
                msgpack.exceptions.ExtraData):
            logger.error("Error msgunpacking message")
            raise

        try:
            msgType = tmp[""]
        except KeyError:
            logger.error("Message type missing")
            raise

        msgObj = messagetypes.constructObject(tmp)
        if msgObj is None:
            raise ValueError("Malformed message")
        try:
            msgObj.process()
        except:
            raise ValueError("Malformed message")
        if msgType == "message":
            self.subject = msgObj.subject
            self.body = msgObj.body

    def decodeSimple(self, data):
        bodyPositionIndex = string.find(data, '\nBody:')
        if bodyPositionIndex > 1:
            subject = data[8:bodyPositionIndex]
            # Only save and show the first 500 characters of the subject.
            # Any more is probably an attack.
            subject = subject[:500]
            body = data[bodyPositionIndex + 6:]
        else:
            subject = ''
            body = data
        # Throw away any extra lines (headers) after the subject.
        if subject:
            subject = subject.splitlines()[0]
        self.subject = subject
        self.body = body

if __name__ == '__main__':
    import random
    messageData = {
        "subject": ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(40)),
        "body": ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(10000))
    }
    obj1 = MsgEncode(messageData, 1)
    obj2 = MsgEncode(messageData, 2)
    obj3 = MsgEncode(messageData, 3)
    print "1:%i 2:%i 3:%i" %(len(obj1.data), len(obj2.data), len(obj3.data))

    obj1e = MsgDecode(1, obj1.data)
    # no subject in trivial encoding
    assert messageData["body"] == obj1e.body
    obj2e = MsgDecode(2, obj2.data)
    assert messageData["subject"] == obj2e.subject
    assert messageData["body"] == obj2e.body
    obj3e = MsgDecode(3, obj3.data)
    assert messageData["subject"] == obj3e.subject
    assert messageData["body"] == obj3e.body
<commit_msg>Formating and standalone fix<commit_after>#!/usr/bin/python2.7

import msgpack
import string
import zlib

import shared
from debug import logger
import messagetypes
from tr import _translate

BITMESSAGE_ENCODING_IGNORE = 0
BITMESSAGE_ENCODING_TRIVIAL = 1
BITMESSAGE_ENCODING_SIMPLE = 2
BITMESSAGE_ENCODING_EXTENDED = 3

class MsgEncode(object):
    def __init__(self, message, encoding=BITMESSAGE_ENCODING_SIMPLE):
        self.data = None
        self.encoding = encoding
        self.length = 0
        if self.encoding == BITMESSAGE_ENCODING_EXTENDED:
            self.encodeExtended(message)
        elif self.encoding == BITMESSAGE_ENCODING_SIMPLE:
            self.encodeSimple(message)
        elif self.encoding == BITMESSAGE_ENCODING_TRIVIAL:
            self.encodeTrivial(message)
        else:
            raise ValueError("Unknown encoding %i" % (encoding))

    def encodeExtended(self, message):
        try:
            msgObj = messagetypes.message.Message()
            self.data = zlib.compress(msgpack.dumps(msgObj.encode(message)), 9)
        except zlib.error:
            logger.error("Error compressing message")
            raise
        except msgpack.exceptions.PackException:
            logger.error("Error msgpacking message")
            raise
        self.length = len(self.data)

    def encodeSimple(self, message):
        self.data = 'Subject:' + message['subject'] + '\n' + 'Body:' + message['body']
        self.length = len(self.data)

    def encodeTrivial(self, message):
        self.data = message['body']
        self.length = len(self.data)

class MsgDecode(object):
    def __init__(self, encoding, data):
        self.encoding = encoding
        if self.encoding == BITMESSAGE_ENCODING_EXTENDED:
            self.decodeExtended(data)
        elif self.encoding in [BITMESSAGE_ENCODING_SIMPLE, BITMESSAGE_ENCODING_TRIVIAL]:
            self.decodeSimple(data)
        else:
            self.body = _translate("MsgDecode", "The message has an unknown encoding.\nPerhaps you should upgrade Bitmessage.")
            self.subject = _translate("MsgDecode", "Unknown encoding")

    def decodeExtended(self, data):
        try:
            tmp = msgpack.loads(zlib.decompress(data))
        except zlib.error:
            logger.error("Error decompressing message")
            raise
        except (msgpack.exceptions.UnpackException,
                msgpack.exceptions.ExtraData):
            logger.error("Error msgunpacking message")
            raise

        try:
            msgType = tmp[""]
        except KeyError:
            logger.error("Message type missing")
            raise

        msgObj = messagetypes.constructObject(tmp)
        if msgObj is None:
            raise ValueError("Malformed message")
        try:
            msgObj.process()
        except:
            raise ValueError("Malformed message")
        if msgType == "message":
            self.subject = msgObj.subject
            self.body = msgObj.body

    def decodeSimple(self, data):
        bodyPositionIndex = string.find(data, '\nBody:')
        if bodyPositionIndex > 1:
            subject = data[8:bodyPositionIndex]
            # Only save and show the first 500 characters of the subject.
            # Any more is probably an attack.
            subject = subject[:500]
            body = data[bodyPositionIndex + 6:]
        else:
            subject = ''
            body = data
        # Throw away any extra lines (headers) after the subject.
        if subject:
            subject = subject.splitlines()[0]
        self.subject = subject
        self.body = body

if __name__ == '__main__':
    import random
    messageData = {
        "subject": ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(40)),
        "body": ''.join(random.choice(string.ascii_lowercase + string.digits) for _ in range(10000))
    }
    obj1 = MsgEncode(messageData, 1)
    obj2 = MsgEncode(messageData, 2)
    obj3 = MsgEncode(messageData, 3)
    print "1:%i 2:%i 3:%i" %(len(obj1.data), len(obj2.data), len(obj3.data))

    obj1e = MsgDecode(1, obj1.data)
    # no subject in trivial encoding
    assert messageData["body"] == obj1e.body
    obj2e = MsgDecode(2, obj2.data)
    assert messageData["subject"] == obj2e.subject
    assert messageData["body"] == obj2e.body
    obj3e = MsgDecode(3, obj3.data)
    assert messageData["subject"] == obj3e.subject
    assert messageData["body"] == obj3e.body
<|endoftext|>


<commit_before>import os
from functools import partial

from astropy import units as u
from astropy.io import fits
from astropy.coordinates import SkyCoord

from ....utils import images

def on_enter(event_data):
    """ Adjust pointing.

    * Take 60 second exposure
    * Call `sync_coordinates`
        * Plate-solve
        * Get pointing error
        * If within `_pointing_threshold`
            * goto tracking
        * Else
            * set set mount target coords to center RA/Dec
            * sync mount coords
            * slew to target
    """
    pan = event_data.model

    try:
        pan.say("Taking guide picture.")

        guide_camera = pan.observatory.get_guide_camera()

        filename = pan.observatory.construct_filename(guide=True)
        filename = filename.replace('guide.cr2', 'guide_{:03.0f}.cr2'.format(pan._pointing_iteration))
        pan.logger.debug("Path for guide: {}".format(filename))

        guide_image = guide_camera.take_exposure(seconds=pan._pointing_exptime, filename=filename)
        pan.logger.debug("Waiting for guide image: {}".format(guide_image))

        try:
            future = pan.wait_until_files_exist(guide_image)

            pan.logger.debug("Adding callback for guide image")
            future.add_done_callback(partial(sync_coordinates, pan))
        except Exception as e:
            pan.logger.error("Problem waiting for images: {}".format(e))
            pan.goto('park')

    except Exception as e:
        pan.say("Hmm, I had a problem checking the pointing error. Sending to parking. {}".format(e))
        pan.goto('park')

def sync_coordinates(pan, future):
    """ Adjusts pointing error from the most recent image.

    Receives a future from an asyncio call (e.g.,`wait_until_files_exist`) that contains
    filename of recent image. Uses utility function to return pointing error. If the error
    is off by some threshold, sync the coordinates to the center and reacquire the target.
    Iterate on process until threshold is met then start tracking.

    Parameters
    ----------
    future : {asyncio.Future}
        Future from returned from asyncio call, `.get_result` contains filename of image.

    Returns
    -------
    u.Quantity
        The separation between the center of the solved image and the target.
    """
    pan.logger.debug("Getting pointing error")
    pan.say("Ok, I've got the guide picture, let's see how close we are")

    separation = 0 * u.deg
    pan.logger.debug("Default separation: {}".format(separation))

    if future.done() and not future.cancelled():
        pan.logger.debug("Task completed successfully, getting image name")

        fname = future.result()[0]

        pan.logger.debug("Processing image: {}".format(fname))

        target = pan.observatory.current_target
        pan.logger.debug("Target: {}".format(target))

        fits_headers = pan.observatory._get_standard_headers(target=target)
        pan.logger.debug("Guide headers: {}".format(fits_headers))

        kwargs = {}
        if 'ra_center' in target.guide_wcsinfo:
            kwargs['ra'] = target.guide_wcsinfo['ra_center'].value
        else:
            kwargs['ra'] = target.ra.value
        if 'dec_center' in target.guide_wcsinfo:
            kwargs['dec'] = target.guide_wcsinfo['dec_center'].value
        else:
            kwargs['dec'] = target.dec.value
        if 'fieldw' in target.guide_wcsinfo:
            kwargs['radius'] = target.guide_wcsinfo['radius'].value
        else:
            kwargs['radius'] = 15.0

        pan.logger.debug("Processing CR2 files with kwargs: {}".format(kwargs))
        processed_info = images.process_cr2(fname, fits_headers=fits_headers, timeout=45, **kwargs)
        # pan.logger.debug("Processed info: {}".format(processed_info))

        # Use the solve file
        fits_fname = processed_info.get('solved_fits_file', None)

        if os.path.exists(fits_fname):
            pan.logger.debug("Solved guide file: {}".format(fits_fname))
            # Get the WCS info and the HEADER info
            pan.logger.debug("Getting WCS and FITS headers for: {}".format(fits_fname))

            wcs_info = images.get_wcsinfo(fits_fname)

            # Save guide wcsinfo to use for future solves
            target.guide_wcsinfo = wcs_info
            pan.logger.debug("WCS Info: {}".format(target.guide_wcsinfo))

            target = None
            with fits.open(fits_fname) as hdulist:
                hdu = hdulist[0]
                # pan.logger.debug("FITS Headers: {}".format(hdu.header))

                target = SkyCoord(ra=float(hdu.header['RA']) * u.degree, dec=float(hdu.header['Dec']) * u.degree)
                pan.logger.debug("Target coords: {}".format(target))

            # Create two coordinates
            center = SkyCoord(ra=wcs_info['ra_center'], dec=wcs_info['dec_center'])
            pan.logger.debug("Center coords: {}".format(center))

            if target is not None:
                separation = center.separation(target)
        else:
            pan.logger.warning("Could not solve guide image")
    else:
        pan.logger.debug("Future cancelled. Result from callback: {}".format(future.result()))

    pan.logger.debug("Separation: {}".format(separation))
    if separation < pan._pointing_threshold:
        pan.say("I'm pretty close to the target, starting track.")
        pan.goto('track')
    elif pan._pointing_iteration >= pan._max_iterations:
        pan.say("I've tried to get closer to the target but can't. I'll just observe where I am.")
        pan.goto('track')
    else:
        pan.say("I'm still a bit away from the target so I'm going to try and get a bit closer.")

        pan._pointing_iteration = pan._pointing_iteration + 1

        # Set the target to center
        has_target = pan.observatory.mount.set_target_coordinates(center)

        if has_target:
            # Tell the mount we are at the target, which is the center
            pan.observatory.mount.serial_query('calibrate_mount')
            pan.say("Syncing with the latest image...")

            # Now set back to target
            if target is not None:
                pan.observatory.mount.set_target_coordinates(target)

        pan.goto('slew_to_target')
<commit_msg>Always use target for centering solve<commit_after>import os
from functools import partial

from astropy import units as u
from astropy.io import fits
from astropy.coordinates import SkyCoord

from ....utils import images

def on_enter(event_data):
    """ Adjust pointing.

    * Take 60 second exposure
    * Call `sync_coordinates`
        * Plate-solve
        * Get pointing error
        * If within `_pointing_threshold`
            * goto tracking
        * Else
            * set set mount target coords to center RA/Dec
            * sync mount coords
            * slew to target
    """
    pan = event_data.model

    try:
        pan.say("Taking guide picture.")

        guide_camera = pan.observatory.get_guide_camera()

        filename = pan.observatory.construct_filename(guide=True)
        filename = filename.replace('guide.cr2', 'guide_{:03.0f}.cr2'.format(pan._pointing_iteration))
        pan.logger.debug("Path for guide: {}".format(filename))

        guide_image = guide_camera.take_exposure(seconds=pan._pointing_exptime, filename=filename)
        pan.logger.debug("Waiting for guide image: {}".format(guide_image))

        try:
            future = pan.wait_until_files_exist(guide_image)

            pan.logger.debug("Adding callback for guide image")
            future.add_done_callback(partial(sync_coordinates, pan))
        except Exception as e:
            pan.logger.error("Problem waiting for images: {}".format(e))
            pan.goto('park')

    except Exception as e:
        pan.say("Hmm, I had a problem checking the pointing error. Sending to parking. {}".format(e))
        pan.goto('park')

def sync_coordinates(pan, future):
    """ Adjusts pointing error from the most recent image.

    Receives a future from an asyncio call (e.g.,`wait_until_files_exist`) that contains
    filename of recent image. Uses utility function to return pointing error. If the error
    is off by some threshold, sync the coordinates to the center and reacquire the target.
    Iterate on process until threshold is met then start tracking.

    Parameters
    ----------
    future : {asyncio.Future}
        Future from returned from asyncio call, `.get_result` contains filename of image.

    Returns
    -------
    u.Quantity
        The separation between the center of the solved image and the target.
    """
    pan.logger.debug("Getting pointing error")
    pan.say("Ok, I've got the guide picture, let's see how close we are")

    separation = 0 * u.deg
    pan.logger.debug("Default separation: {}".format(separation))

    if future.done() and not future.cancelled():
        pan.logger.debug("Task completed successfully, getting image name")

        fname = future.result()[0]

        pan.logger.debug("Processing image: {}".format(fname))

        target = pan.observatory.current_target
        pan.logger.debug("Target: {}".format(target))

        fits_headers = pan.observatory._get_standard_headers(target=target)
        pan.logger.debug("Guide headers: {}".format(fits_headers))

        kwargs = {}
        kwargs['ra'] = target.ra.value
        kwargs['dec'] = target.dec.value
        kwargs['radius'] = 15.0

        # if 'ra_center' in target.guide_wcsinfo:
        #     kwargs['ra'] = target.guide_wcsinfo['ra_center'].value
        # else:
        #     kwargs['ra'] = target.ra.value
        # if 'dec_center' in target.guide_wcsinfo:
        #     kwargs['dec'] = target.guide_wcsinfo['dec_center'].value
        # else:
        #     kwargs['dec'] = target.dec.value
        # if 'fieldw' in target.guide_wcsinfo:
        #     kwargs['radius'] = target.guide_wcsinfo['radius'].value
        # else:
        #     kwargs['radius'] = 15.0

        pan.logger.debug("Processing CR2 files with kwargs: {}".format(kwargs))
        processed_info = images.process_cr2(fname, fits_headers=fits_headers, timeout=45, **kwargs)
        # pan.logger.debug("Processed info: {}".format(processed_info))

        # Use the solve file
        fits_fname = processed_info.get('solved_fits_file', None)

        if os.path.exists(fits_fname):
            pan.logger.debug("Solved guide file: {}".format(fits_fname))
            # Get the WCS info and the HEADER info
            pan.logger.debug("Getting WCS and FITS headers for: {}".format(fits_fname))

            wcs_info = images.get_wcsinfo(fits_fname)

            # Save guide wcsinfo to use for future solves
            target.guide_wcsinfo = wcs_info
            pan.logger.debug("WCS Info: {}".format(target.guide_wcsinfo))

            target = None
            with fits.open(fits_fname) as hdulist:
                hdu = hdulist[0]
                # pan.logger.debug("FITS Headers: {}".format(hdu.header))

                target = SkyCoord(ra=float(hdu.header['RA']) * u.degree, dec=float(hdu.header['Dec']) * u.degree)
                pan.logger.debug("Target coords: {}".format(target))

            # Create two coordinates
            center = SkyCoord(ra=wcs_info['ra_center'], dec=wcs_info['dec_center'])
            pan.logger.debug("Center coords: {}".format(center))

            if target is not None:
                separation = center.separation(target)
        else:
            pan.logger.warning("Could not solve guide image")
    else:
        pan.logger.debug("Future cancelled. Result from callback: {}".format(future.result()))

    pan.logger.debug("Separation: {}".format(separation))
    if separation < pan._pointing_threshold:
        pan.say("I'm pretty close to the target, starting track.")
        pan.goto('track')
    elif pan._pointing_iteration >= pan._max_iterations:
        pan.say("I've tried to get closer to the target but can't. I'll just observe where I am.")
        pan.goto('track')
    else:
        pan.say("I'm still a bit away from the target so I'm going to try and get a bit closer.")

        pan._pointing_iteration = pan._pointing_iteration + 1

        # Set the target to center
        has_target = pan.observatory.mount.set_target_coordinates(center)

        if has_target:
            # Tell the mount we are at the target, which is the center
            pan.observatory.mount.serial_query('calibrate_mount')
            pan.say("Syncing with the latest image...")

            # Now set back to target
            if target is not None:
                pan.observatory.mount.set_target_coordinates(target)

        pan.goto('slew_to_target')
<|endoftext|>


<commit_before>import calendar
import logging
import json

from datetime import datetime, timedelta
from simpledoge.model_lib import base
from sqlalchemy.schema import CheckConstraint
from sqlalchemy.ext.declarative import AbstractConcreteBase
from sqlalchemy.dialects.postgresql import HSTORE
from cryptokit import bits_to_difficulty

from . import db

class Blob(base):
    """ Used to store misc single value blobs of data, such as the current
    block height and difficulty. """
    key = db.Column(db.String, primary_key=True)
    data = db.Column(HSTORE, default=dict)

class Block(base):
    """ This class stores metadata on all blocks found by the pool """
    # the hash of the block for orphan checking
    hash = db.Column(db.String, primary_key=True)
    height = db.Column(db.Integer, nullable=False)
    # User who discovered block
    user = db.Column(db.String)
    # When block was found
    found_at = db.Column(db.DateTime, default=datetime.utcnow)
    # # Time started on block
    time_started = db.Column(db.DateTime, nullable=False)
    # Is block now orphaned?
    orphan = db.Column(db.Boolean, default=False)
    # Is the block matured?
    mature = db.Column(db.Boolean, default=False)
    # Total shares that were required to solve the block
    shares_to_solve = db.Column(db.BigInteger)
    # Block value (does not include transaction fees recieved)
    total_value = db.Column(db.BigInteger)
    # Associated transaction fees
    transaction_fees = db.Column(db.BigInteger)
    # total going to pool from fees
    fees = db.Column(db.BigInteger)
    # Difficulty of block when solved
    bits = db.Column(db.String(8), nullable=False)
    # the last share id that was processed when the block was entered.
    # used as a marker for calculating last n shares
    last_share_id = db.Column(db.BigInteger, db.ForeignKey('share.id'))
    last_share = db.relationship('Share', foreign_keys=[last_share_id])
    # have payments been generated for it?
    processed = db.Column(db.Boolean, default=False)

    @property
    def status(self):
        if self.mature:
            return "Mature"
        if self.orphan:
            return "Orphan"
        return "Unconfirmed"

    @classmethod
    def create(cls, user, height, total_value, transaction_fees, bits, hash,
               time_started):
        share = Share.query.order_by(Share.id.desc()).first()
        block = cls(user=user,
                    height=height,
                    total_value=total_value,
                    transaction_fees=transaction_fees,
                    bits=bits,
                    last_share=share,
                    hash=hash,
                    time_started=time_started)
        # add and flush
        db.session.add(block)
        db.session.flush()
        return block

    @property
    def difficulty(self):
        return bits_to_difficulty(self.bits)

    @property
    def duration(self):
        seconds = round((self.found_at - self.time_started).total_seconds())
        formatted_time = str(timedelta(seconds=seconds))
        return formatted_time

def last_block_time():
    """ Retrieves the last time a block was solved using progressively less
    accurate methods. Essentially used to calculate round time. """
    last_block = Block.query.order_by(Block.height.desc()).first()
    if last_block:
        return last_block.found_at

    hour = OneHourShare.query.order_by(OneHourShare.time).first()
    if hour:
        return hour.time

    five = FiveMinuteShare.query.order_by(FiveMinuteShare.time).first()
    if five:
        return five.time

    minute = OneMinuteShare.query.order_by(OneMinuteShare.time).first()
    if minute:
        return minute.time

    return datetime.utcnow()

def last_block_share_id():
    last_block = Block.query.order_by(Block.height.desc()).first()
    if not last_block:
        return 0
    return last_block.last_share_id

class Share(base):
    """ This class generates a table containing every share accepted for a
    round """
    id = db.Column(db.BigInteger, primary_key=True)
    user = db.Column(db.String)
    shares = db.Column(db.Integer)

    @classmethod
    def create(cls, user, shares):
        share = cls(user=user, shares=shares)
        db.session.add(share)
        return share

class Transaction(base):
    txid = db.Column(db.String, primary_key=True)
    confirmed = db.Column(db.Boolean, default=False)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)

    @classmethod
    def create(cls, txid):
        trans = cls(txid=txid)
        db.session.add(trans)
        return trans

class Status(base):
    """ This class generates a table containing every share accepted for a
    round """
    user = db.Column(db.String, primary_key=True)
    worker = db.Column(db.String, primary_key=True)
    status = db.Column(db.String)
    time = db.Column(db.DateTime)

    @property
    def parsed_status(self):
        return json.loads(self.status)

    def pretty_json(self, gpu=0):
        print type(gpu)
        return json.dumps(json.loads(self.status)['gpus'][gpu], indent=4, sort_keys=True)

    @property
    def stale(self):
        ten_min_ago = datetime.utcnow() - timedelta(minutes=10)
        if ten_min_ago >= self.time:
            return True
        else:
            return False

class Payout(base):
    """ Represents a users payout for a single round """
    id = db.Column(db.Integer, primary_key=True)
    blockhash = db.Column(db.String, db.ForeignKey('block.hash'))
    block = db.relationship('Block', foreign_keys=[blockhash])
    user = db.Column(db.String)
    shares = db.Column(db.BigInteger)
    amount = db.Column(db.BigInteger, CheckConstraint('amount>0', 'min_payout_amount'))
    transaction_id = db.Column(db.String, db.ForeignKey('transaction.txid'))
    transaction = db.relationship('Transaction', foreign_keys=[transaction_id])
    __table_args__ = (
        db.UniqueConstraint("user", "blockhash"),
    )

    @classmethod
    def create(cls, user, amount, block, shares):
        payout = cls(user=user, amount=amount, block=block, shares=shares)
        db.session.add(payout)
        return payout

class TimeSlice(AbstractConcreteBase, base):
    user = db.Column(db.String, primary_key=True)
    # datetime floored to the minute
    time = db.Column(db.DateTime, primary_key=True)
    worker = db.Column(db.String, primary_key=True)
    value = db.Column(db.Integer)

    @classmethod
    def create(cls, user, value, time, worker):
        dt = cls.floor_time(time)
        slc = cls(user=user, value=value, time=dt, worker=worker)
        db.session.add(slc)
        return slc

    @classmethod
    def add_value(cls, user, value, time, worker):
        dt = cls.floor_time(time)
        slc = cls.query.with_lockmode('update').filter_by(
            user=user, time=dt, worker=worker).one()
        slc.value += value

    @classmethod
    def floor_time(cls, time):
        """ Changes an integer timestamp to the minute for which it falls in.
        Allows abstraction of create and add share logic for each time slice
        object. """
        if isinstance(time, datetime):
            time = calendar.timegm(time.utctimetuple())
        return datetime.utcfromtimestamp(
            (time // cls.slice_seconds) * cls.slice_seconds)

    @classmethod
    def compress(cls):
        """ Moves statistics that are past the `window` time into the next
        time slice size, effectively compressing the data. """
        # get the minute shares that are old enough to be compressed and
        # deleted
        recent = cls.floor_time(datetime.utcnow()) - cls.window
        # the five minute slice currently being processed
        current_slice = None
        # dictionary of lists keyed by user
        users = {}

        def create_upper():
            # add a time slice for each user in this pending period
            for (user, worker), slices in users.iteritems():
                total = sum([slc.value for slc in slices])

                # put it in the database
                upper = cls.upper.query.filter_by(user=user, worker=worker, time=current_slice).with_lockmode('update').first()
                # wasn't in the db? create it
                if not upper:
                    upper = cls.upper.create(user, total, current_slice, worker)
                # was in the db? add the shares for each worker, and
                # increment total share count
                else:
                    upper.value += total

                for slc in slices:
                    db.session.delete(slc)

        # traverse minute shares that are old enough in time order
        for slc in (cls.query.filter(cls.time < recent).
                    order_by(cls.time)):
            slice_time = cls.upper.floor_time(slc.time)

            if current_slice is None:
                current_slice = slice_time

            # we've encountered the next time slice, so commit the pending one
            if slice_time != current_slice:
                logging.debug("Processing slice " + str(current_slice))
                create_upper()
                users.clear()
                current_slice = slice_time

            # add the one min shares for this user the list of pending shares
            # to be grouped together
            key = (slc.user, slc.worker)
            users.setdefault(key, [])
            users[key].append(slc)

        create_upper()

class OneHourShare(TimeSlice):
    __tablename__ = 'one_hour_share'
    window = timedelta(days=30)
    slice = timedelta(hours=1)
    slice_seconds = slice.total_seconds()
    __mapper_args__ = {
        'polymorphic_identity': 'one_hour_share',
        'concrete': True
    }

class FiveMinuteShare(TimeSlice):
    __tablename__ = 'five_minute_share'
    window = timedelta(days=1)
    slice = timedelta(minutes=5)
    slice_seconds = slice.total_seconds()
    upper = OneHourShare
    __mapper_args__ = {
        'polymorphic_identity': 'five_minute_share',
        'concrete': True
    }

class OneMinuteShare(TimeSlice):
    __tablename__ = 'one_minute_share'
    window = timedelta(hours=1)
    slice = timedelta(minutes=1)
    slice_seconds = slice.total_seconds()
    upper = FiveMinuteShare
    __mapper_args__ = {
        'polymorphic_identity': 'one_minute_share',
        'concrete': True
    }
<commit_msg>started in on genericizing compression function for different unique key values<commit_after>import calendar
import logging
import json

from collections import namedtuple
from datetime import datetime, timedelta
from simpledoge.model_lib import base
from sqlalchemy.schema import CheckConstraint
from sqlalchemy.ext.declarative import AbstractConcreteBase
from sqlalchemy.dialects.postgresql import HSTORE
from cryptokit import bits_to_difficulty

from . import db

class Blob(base):
    """ Used to store misc single value blobs of data, such as the current
    block height and difficulty. """
    key = db.Column(db.String, primary_key=True)
    data = db.Column(HSTORE, default=dict)

class Block(base):
    """ This class stores metadata on all blocks found by the pool """
    # the hash of the block for orphan checking
    hash = db.Column(db.String, primary_key=True)
    height = db.Column(db.Integer, nullable=False)
    # User who discovered block
    user = db.Column(db.String)
    # When block was found
    found_at = db.Column(db.DateTime, default=datetime.utcnow)
    # # Time started on block
    time_started = db.Column(db.DateTime, nullable=False)
    # Is block now orphaned?
    orphan = db.Column(db.Boolean, default=False)
    # Is the block matured?
    mature = db.Column(db.Boolean, default=False)
    # Total shares that were required to solve the block
    shares_to_solve = db.Column(db.BigInteger)
    # Block value (does not include transaction fees recieved)
    total_value = db.Column(db.BigInteger)
    # Associated transaction fees
    transaction_fees = db.Column(db.BigInteger)
    # total going to pool from fees
    fees = db.Column(db.BigInteger)
    # Difficulty of block when solved
    bits = db.Column(db.String(8), nullable=False)
    # the last share id that was processed when the block was entered.
    # used as a marker for calculating last n shares
    last_share_id = db.Column(db.BigInteger, db.ForeignKey('share.id'))
    last_share = db.relationship('Share', foreign_keys=[last_share_id])
    # have payments been generated for it?
    processed = db.Column(db.Boolean, default=False)

    @property
    def status(self):
        if self.mature:
            return "Mature"
        if self.orphan:
            return "Orphan"
        return "Unconfirmed"

    @classmethod
    def create(cls, user, height, total_value, transaction_fees, bits, hash,
               time_started):
        share = Share.query.order_by(Share.id.desc()).first()
        block = cls(user=user,
                    height=height,
                    total_value=total_value,
                    transaction_fees=transaction_fees,
                    bits=bits,
                    last_share=share,
                    hash=hash,
                    time_started=time_started)
        # add and flush
        db.session.add(block)
        db.session.flush()
        return block

    @property
    def difficulty(self):
        return bits_to_difficulty(self.bits)

    @property
    def duration(self):
        seconds = round((self.found_at - self.time_started).total_seconds())
        formatted_time = str(timedelta(seconds=seconds))
        return formatted_time

def last_block_time():
    """ Retrieves the last time a block was solved using progressively less
    accurate methods. Essentially used to calculate round time. """
    last_block = Block.query.order_by(Block.height.desc()).first()
    if last_block:
        return last_block.found_at

    hour = OneHourShare.query.order_by(OneHourShare.time).first()
    if hour:
        return hour.time

    five = FiveMinuteShare.query.order_by(FiveMinuteShare.time).first()
    if five:
        return five.time

    minute = OneMinuteShare.query.order_by(OneMinuteShare.time).first()
    if minute:
        return minute.time

    return datetime.utcnow()

def last_block_share_id():
    last_block = Block.query.order_by(Block.height.desc()).first()
    if not last_block:
        return 0
    return last_block.last_share_id

class Share(base):
    """ This class generates a table containing every share accepted for a
    round """
    id = db.Column(db.BigInteger, primary_key=True)
    user = db.Column(db.String)
    shares = db.Column(db.Integer)

    @classmethod
    def create(cls, user, shares):
        share = cls(user=user, shares=shares)
        db.session.add(share)
        return share

class Transaction(base):
    txid = db.Column(db.String, primary_key=True)
    confirmed = db.Column(db.Boolean, default=False)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)

    @classmethod
    def create(cls, txid):
        trans = cls(txid=txid)
        db.session.add(trans)
        return trans

class Status(base):
    """ This class generates a table containing every share accepted for a
    round """
    user = db.Column(db.String, primary_key=True)
    worker = db.Column(db.String, primary_key=True)
    status = db.Column(db.String)
    time = db.Column(db.DateTime)

    @property
    def parsed_status(self):
        return json.loads(self.status)

    def pretty_json(self, gpu=0):
        print type(gpu)
        return json.dumps(json.loads(self.status)['gpus'][gpu], indent=4, sort_keys=True)

    @property
    def stale(self):
        ten_min_ago = datetime.utcnow() - timedelta(minutes=10)
        if ten_min_ago >= self.time:
            return True
        else:
            return False

class Payout(base):
    """ Represents a users payout for a single round """
    id = db.Column(db.Integer, primary_key=True)
    blockhash = db.Column(db.String, db.ForeignKey('block.hash'))
    block = db.relationship('Block', foreign_keys=[blockhash])
    user = db.Column(db.String)
    shares = db.Column(db.BigInteger)
    amount = db.Column(db.BigInteger, CheckConstraint('amount>0', 'min_payout_amount'))
    transaction_id = db.Column(db.String, db.ForeignKey('transaction.txid'))
    transaction = db.relationship('Transaction', foreign_keys=[transaction_id])
    __table_args__ = (
        db.UniqueConstraint("user", "blockhash"),
    )

    @classmethod
    def create(cls, user, amount, block, shares):
        payout = cls(user=user, amount=amount, block=block, shares=shares)
        db.session.add(payout)
        return payout

class TimeSlice(AbstractConcreteBase, base):
    user = db.Column(db.String, primary_key=True)
    # datetime floored to the minute
    time = db.Column(db.DateTime, primary_key=True)
    worker = db.Column(db.String, primary_key=True)
    value = db.Column(db.Integer)

    @classmethod
    def create(cls, user, value, time, worker):
        dt = cls.floor_time(time)
        slc = cls(user=user, value=value, time=dt, worker=worker)
        db.session.add(slc)
        return slc

    @classmethod
    def add_value(cls, user, value, time, worker):
        dt = cls.floor_time(time)
        slc = cls.query.with_lockmode('update').filter_by(
            user=user, time=dt, worker=worker).one()
        slc.value += value

    @classmethod
    def floor_time(cls, time):
        """ Changes an integer timestamp to the minute for which it falls in.
        Allows abstraction of create and add share logic for each time slice
        object. """
        if isinstance(time, datetime):
            time = calendar.timegm(time.utctimetuple())
        return datetime.utcfromtimestamp(
            (time // cls.slice_seconds) * cls.slice_seconds)

    @classmethod
    def compress(cls):
        """ Moves statistics that are past the `window` time into the next
        time slice size, effectively compressing the data. """
        # get the minute shares that are old enough to be compressed and
        # deleted
        recent = cls.floor_time(datetime.utcnow()) - cls.window
        # the five minute slice currently being processed
        current_slice = None
        # dictionary of lists keyed by user
        users = {}

        def create_upper():
            # add a time slice for each user in this pending period
            for key, slices in users.iteritems():
                total = sum([slc.value for slc in slices])

                # put it in the database
                upper = cls.upper.query.filter_by(time=current_slice, **key._asdict()).with_lockmode('update').first()
                # wasn't in the db? create it
                if not upper:
                    dt = cls.floor_time(current_slice)
                    upper = cls.upper(time=dt, value=total, **key._asdict())
                else:
                    upper.value += total

                for slc in slices:
                    db.session.delete(slc)

        # traverse minute shares that are old enough in time order
        for slc in (cls.query.filter(cls.time < recent).
                    order_by(cls.time)):
            slice_time = cls.upper.floor_time(slc.time)

            if current_slice is None:
                current_slice = slice_time

            # we've encountered the next time slice, so commit the pending one
            if slice_time != current_slice:
                logging.debug("Processing slice " + str(current_slice))
                create_upper()
                users.clear()
                current_slice = slice_time

            # add the one min shares for this user the list of pending shares
            # to be grouped together
            key = slc.make_key()
            users.setdefault(key, [])
            users[key].append(slc)

        create_upper()

    key = namedtuple('Key', ['user', 'worker'])
    def make_key(self):
        return self.key(user=self.user, worker=self.worker)

class OneHourShare(TimeSlice):
    __tablename__ = 'one_hour_share'
    window = timedelta(days=30)
    slice = timedelta(hours=1)
    slice_seconds = slice.total_seconds()
    __mapper_args__ = {
        'polymorphic_identity': 'one_hour_share',
        'concrete': True
    }

class FiveMinuteShare(TimeSlice):
    __tablename__ = 'five_minute_share'
    window = timedelta(days=1)
    slice = timedelta(minutes=5)
    slice_seconds = slice.total_seconds()
    upper = OneHourShare
    __mapper_args__ = {
        'polymorphic_identity': 'five_minute_share',
        'concrete': True
    }

class OneMinuteShare(TimeSlice):
    __tablename__ = 'one_minute_share'
    window = timedelta(hours=1)
    slice = timedelta(minutes=1)
    slice_seconds = slice.total_seconds()
    upper = FiveMinuteShare
    __mapper_args__ = {
        'polymorphic_identity': 'one_minute_share',
        'concrete': True
    }
<|endoftext|>


<commit_before>"""Make modules of observations based on cooccurence networks and collapse table"""

import argparse
import os

import networkx as nx
import numpy as np

from biom import load_table, Table
from scipy.stats import spearmanr, rankdata, pearsonr
import matplotlib.pyplot as plt

def bh_adjust(p_vals):
    """benjamini-hochberg p-value adjustment"""
    p_vals = np.array(p_vals)
    return p_vals*len(p_vals)/rankdata(p_vals, method='dense')

def bonferroni_adjust(p_vals):
    """bonferroni p-value adjustment"""
    return [i*len(p_vals) for i in p_vals]

def print_delimited(out_fp, lines, header=None):
    """print a tab delimited file with optional header"""
    out = open(out_fp, 'w')
    if header != None:
        out.write('\t'.join([str(i) for i in header])+'\n')
    for line in lines:
        out.write('\t'.join([str(i) for i in line])+'\n')
    out.close()

def plot_networkx(graph):
    """plot networkx object in matplotlib"""
    graph_pos = nx.circular_layout(graph)

    nx.draw_networkx_nodes(graph, graph_pos, node_size=1000, node_color='blue', alpha=0.3)
    nx.draw_networkx_edges(graph, graph_pos, width=2, alpha=0.3, edge_color='green')
    nx.draw_networkx_labels(graph, graph_pos, font_size=12, font_family='sans-serif')

    plt.show()

def filter_table(table, min_samples=None, to_file=True):
    """filter relative abundance table, by default throw away things greater than 1/3 zeros"""
    table = table.copy()
    # first sample filter
    if min_samples != None:
        to_keep = [i for i in table.ids(axis='observation') \
                   if sum(table.data(i, axis='observation') != 0) >= min_samples]
    else:
        to_keep = [i for i in table.ids(axis='observation') \
                   if sum(table.data(i, axis='observation') != 0) >= table.shape[1]/3]
    table.filter(to_keep, axis='observation')

    if to_file == True:
        table.to_json('filter_table',open("filtered_tab.biom",'w'))
        # open("filtered_rel_abund.txt", 'w').write(table.to_tsv())

    return table

def paired_correlations_from_table(table, correl_method=spearmanr, p_adjust=bh_adjust):
    """Takes a biom table and finds correlations between all pairs of otus."""
    otus = table.ids(axis='observation')

    correls = list()

    # do all correlations
    for i in xrange(1, len(otus)):
        otu_i = table.data(otus[i], axis='observation')
        for j in xrange(i+1, len(otus)):
            otu_j = table.data(otus[j], axis='observation')
            correl = correl_method(otu_i, otu_j)
            correls.append([otus[i], otus[j], correl[0], correl[1]])

    # adjust p-value if desired
    if p_adjust != None:
        p_adjusted = p_adjust([i[3] for i in correls])
        for i in xrange(len(correls)):
            correls[i].append(p_adjusted[i])

    header = ['feature1', 'feature2', 'r', 'p']
    if p_adjust != None:
        header.append('adjusted_p')

    return correls, header

def make_net_from_correls(correls, min_p=.05):
    """make network from set of correlations with values less than a minimum"""
    # filter to only include significant correlations
    try:
        correls = list(i for i in correls if i[4] < min_p and i[2] > 0)
    except IndexError:
        correls = list(i for i in correls if i[3] < min_p and i[2] > 0)

    graph = nx.Graph()
    for correl in correls:
        graph.add_node(correl[0])
        graph.add_node(correl[1])
        graph.add_edge(correl[0], correl[1], r=correl[2], \
                   p=correl[3], p_adj=correl[4])
    return graph

def make_modules(graph, k=None):
    """make modules with networkx k-clique communities and annotate network"""
    if k == None:
        k = [2, 3, 4, 5, 6]
    communities = dict()
    for k_val in list(k):
        cliques = list(nx.k_clique_communities(graph, k_val))
        cliques = [list(i) for i in cliques]
        communities[k_val] = cliques
        for i, clique in enumerate(cliques):
            for node in clique:
                graph.node[node]['k_' + str(k_val)] = i
    return graph, {k: list(v) for k, v in communities.iteritems()}

def collapse_modules(table, cliques):
    """collapse created modules in a biom table"""
    in_module = set()
    modules = np.zeros((len(cliques), table.shape[1]))

    # for each clique merge values
    for i, clique in enumerate(cliques):
        in_module = in_module | set(clique)
        for feature in clique:
            modules[i] += table.data(feature, axis='observation')
    table.filter(in_module, axis='observation')

    # make new table
    new_table_matrix = np.concatenate((table.matrix_data.toarray(), modules))
    new_table_obs = list(table.ids(axis='observation')) + \
                    ["module_" + str(i) for i in range(0, len(cliques))]
    return Table(new_table_matrix, new_table_obs, table.ids())

def main():
    """main, takes argparser"""
    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--input", help="location of input biom file")
    parser.add_argument("-o", "--output", help="output file location")
    parser.add_argument("-m", "--correl_method", help="correlation method", default="spearman")
    parser.add_argument("-a", "--p_adjust", help="p-value adjustment", default="bh")
    parser.add_argument("-s", "--min_sample", help="minimum number of samples present in", type=int)
    args = parser.parse_args()

    # correlation and p-value adjustment methods
    correl_methods = {'spearman':spearmanr, 'pearson':pearsonr}
    p_methods = {'bh':bh_adjust, 'bon':bonferroni_adjust}
    correl_method = correl_methods[args.correl_method]
    if args.p_adjust != None:
        p_adjust = p_methods[args.p_adjust]

    # get features to be correlated
    table = load_table(args.input)
    print "Table loaded"

    # make new output directory and change to it
    os.makedirs(args.output)
    os.chdir(args.output)

    # convert to relative abundance and filter
    table_filt = filter_table(table, args.min_sample)
    print "Table filtered"

    # correlate feature
    correls, correl_header = paired_correlations_from_table(table_filt, correl_method, p_adjust)
    print_delimited('correls.txt', correls, correl_header)
    print "Features Correlated"

    # make correlation network
    net = make_net_from_correls(correls)
    print "Network Generated"

    # make modules
    net, cliques = make_modules(net)
    print "Modules Formed"

    # print network
    nx.write_gml(net, 'conetwork.gml')

    # collapse modules
    coll_table = collapse_modules(table, cliques[3])
    print "Table Collapsed"

    # print new table
    coll_table.to_json('make_modules.py', open('collapsed.biom', 'w'))

if __name__ == '__main__':
    main()
<commit_msg>Added --prefix option and "module_" default for collapsing biom table into modules, PEP8 fixes<commit_after>"""Make modules of observations based on cooccurence networks and collapse table"""

import argparse
import os

import networkx as nx
import numpy as np

from biom import load_table, Table
from scipy.stats import spearmanr, rankdata, pearsonr
import matplotlib.pyplot as plt

def bh_adjust(p_vals):
    """benjamini-hochberg p-value adjustment"""
    p_vals = np.array(p_vals)
    return p_vals*len(p_vals)/rankdata(p_vals, method='dense')

def bonferroni_adjust(p_vals):
    """bonferroni p-value adjustment"""
    return [i*len(p_vals) for i in p_vals]

def print_delimited(out_fp, lines, header=None):
    """print a tab delimited file with optional header"""
    out = open(out_fp, 'w')
    if header is not None:
        out.write('\t'.join([str(i) for i in header])+'\n')
    for line in lines:
        out.write('\t'.join([str(i) for i in line])+'\n')
    out.close()

def plot_networkx(graph):
    """plot networkx object in matplotlib"""
    graph_pos = nx.circular_layout(graph)

    nx.draw_networkx_nodes(graph, graph_pos, node_size=1000, node_color='blue', alpha=0.3)
    nx.draw_networkx_edges(graph, graph_pos, width=2, alpha=0.3, edge_color='green')
    nx.draw_networkx_labels(graph, graph_pos, font_size=12, font_family='sans-serif')

    plt.show()

def filter_table(table, min_samples=None, to_file=True):
    """filter relative abundance table, by default throw away things greater than 1/3 zeros"""
    table = table.copy()
    # first sample filter
    if min_samples is not None:
        to_keep = [i for i in table.ids(axis='observation')
                   if sum(table.data(i, axis='observation') != 0) >= min_samples]
    else:
        to_keep = [i for i in table.ids(axis='observation')
                   if sum(table.data(i, axis='observation') != 0) >= table.shape[1]/3]
    table.filter(to_keep, axis='observation')

    if to_file:
        table.to_json('filter_table', open("filtered_tab.biom", 'w'))
        # open("filtered_rel_abund.txt", 'w').write(table.to_tsv())

    return table

def paired_correlations_from_table(table, correl_method=spearmanr, p_adjust=bh_adjust):
    """Takes a biom table and finds correlations between all pairs of otus."""
    otus = table.ids(axis='observation')

    correls = list()

    # do all correlations
    for i in xrange(1, len(otus)):
        otu_i = table.data(otus[i], axis='observation')
        for j in xrange(i+1, len(otus)):
            otu_j = table.data(otus[j], axis='observation')
            correl = correl_method(otu_i, otu_j)
            correls.append([otus[i], otus[j], correl[0], correl[1]])

    # adjust p-value if desired
    if p_adjust is not None:
        p_adjusted = p_adjust([i[3] for i in correls])
        for i in xrange(len(correls)):
            correls[i].append(p_adjusted[i])

    header = ['feature1', 'feature2', 'r', 'p']
    if p_adjust is not None:
        header.append('adjusted_p')

    return correls, header

def make_net_from_correls(correls, min_p=.05):
    """make network from set of correlations with values less than a minimum"""
    # filter to only include significant correlations
    try:
        correls = list(i for i in correls if i[4] < min_p and i[2] > 0)
    except IndexError:
        correls = list(i for i in correls if i[3] < min_p and i[2] > 0)

    graph = nx.Graph()
    for correl in correls:
        graph.add_node(correl[0])
        graph.add_node(correl[1])
        graph.add_edge(correl[0], correl[1], r=correl[2],
                       p=correl[3], p_adj=correl[4])
    return graph

def make_modules(graph, k=None):
    """make modules with networkx k-clique communities and annotate network"""
    if k is None:
        k = [2, 3, 4, 5, 6]
    communities = dict()
    for k_val in list(k):
        cliques = list(nx.k_clique_communities(graph, k_val))
        cliques = [list(i) for i in cliques]
        communities[k_val] = cliques
        for i, clique in enumerate(cliques):
            for node in clique:
                graph.node[node]['k_' + str(k_val)] = i
    return graph, {k: list(v) for k, v in communities.iteritems()}

def collapse_modules(table, cliques, prefix="module_"):
    """collapse created modules in a biom table"""
    in_module = set()
    modules = np.zeros((len(cliques), table.shape[1]))

    # for each clique merge values
    for i, clique in enumerate(cliques):
        in_module = in_module | set(clique)
        for feature in clique:
            modules[i] += table.data(feature, axis='observation')
    table.filter(in_module, axis='observation')

    # make new table
    new_table_matrix = np.concatenate((table.matrix_data.toarray(), modules))
    new_table_obs = list(table.ids(axis='observation')) + [prefix + str(i) for i in range(0, len(cliques))]
    return Table(new_table_matrix, new_table_obs, table.ids())

def main():
    """main, takes argparser"""
    parser = argparse.ArgumentParser()
    parser.add_argument("-i", "--input", help="location of input biom file")
    parser.add_argument("-o", "--output", help="output file location")
    parser.add_argument("-m", "--correl_method", help="correlation method", default="spearman")
    parser.add_argument("-a", "--p_adjust", help="p-value adjustment", default="bh")
    parser.add_argument("-s", "--min_sample", help="minimum number of samples present in", type=int)
    parser.add_argument("--prefix", help="prefix for module names in collapsed file", default="module_")
    args = parser.parse_args()

    # correlation and p-value adjustment methods
    correl_methods = {'spearman': spearmanr, 'pearson': pearsonr}
    p_methods = {'bh': bh_adjust, 'bon': bonferroni_adjust}
    correl_method = correl_methods[args.correl_method]
    if args.p_adjust is not None:
        p_adjust = p_methods[args.p_adjust]
    else:
        p_adjust = None

    # get features to be correlated
    table = load_table(args.input)
    print "Table loaded"

    # make new output directory and change to it
    os.makedirs(args.output)
    os.chdir(args.output)

    # convert to relative abundance and filter
    table_filt = filter_table(table, args.min_sample)
    print "Table filtered"

    # correlate feature
    correls, correl_header = paired_correlations_from_table(table_filt, correl_method, p_adjust)
    print_delimited('correls.txt', correls, correl_header)
    print "Features Correlated"

    # make correlation network
    net = make_net_from_correls(correls)
    print "Network Generated"

    # make modules
    net, cliques = make_modules(net)
    print "Modules Formed"

    # print network
    nx.write_gml(net, 'conetwork.gml')

    # collapse modules
    coll_table = collapse_modules(table, cliques[3], args.prefix)
    print "Table Collapsed"

    # print new table
    coll_table.to_json('make_modules.py', open('collapsed.biom', 'w'))

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


<commit_before>'''
Base Dataset class
==================
Defines the metadata fields and populates some basic info
'''

# Copyright (c) 2009 Australian Government, Department of Environment, Heritage, Water and the Arts
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

import os,time,sys,glob,time,math
import UserDict
import utilities, geometry, uuid

#Import fieldnames
import __fields__

try:
    from osgeo import gdal
    from osgeo import gdalconst
    from osgeo import osr
    from osgeo import ogr
except ImportError:
    import gdal
    import gdalconst
    import osr
    import ogr
gdal.AllRegister()
class Dataset(object):
    '''A base Dataset class'''
    def __new__(self,f):
        ##Initialise the class object
        self=object.__new__(self)

        self._gdaldataset=None
        self._metadata={}
        self._extent=[]
        self._filelist=[]
        self._stretch=None

        ##Populate some basic info
        self.__setfileinfo__(f)

        ##Initialise the fields
        self.fields=idict(__fields__.fields)#We don't want any fields added/deleted

        return self

    # ==================== #
    # Public Class Methods
    # ==================== #
    def getoverview(self,outfile=None,width=800,format='JPG'):
        '''
        Generate overviews for generic imagery

        Override this class if you like, otherwise simply expose a GDALDataset object as self._gdaldataset_

        @type  outfile: string
        @param outfile: a filepath to the output overview image. If supplied, format is determined from the file extension
        @type  width:   integer
        @param width:   image width
        @type  format:  string
        @param format:  format to generate overview image, one of ['JPG','PNG','GIF','BMP','TIF']. Not required if outfile is supplied.

        @return:
            - B{filepath} (if outfile is supplied) B{OR}
            - B{binary image data} (if outfile is not supplied)
        '''
        import overviews

        md=self.metadata
        ds=self._gdaldataset
        if not ds:raise AttributeError, 'No GDALDataset object available, overview image can not be generated'

        #Don't rely on the metadata as self._gdaldataset might be a custom VRT
        ##nbands=md['nbands']
        ##cols=md['cols']
        ##rows=md['rows']
        ##nbits=md['nbits']
        rb=ds.GetRasterBand(1)
        nbands=ds.RasterCount
        cols=ds.RasterXSize
        rows=ds.RasterYSize
        nbits=gdal.GetDataTypeSize(rb.DataType)
        datatype=gdal.GetDataTypeName(rb.DataType)
        nodata=rb.GetNoDataValue()
        stretch_type=None
        stretch_args=None
        rgb_bands = {}

        #Check for pre-defined stretch
        if self._stretch:
            stretch_type,stretch_args=self._stretch

        #Check for pre-defined rgb bands
        for i in range(1,nbands+1):
            gci=ds.GetRasterBand(i).GetRasterColorInterpretation()
            if   gci == gdal.GCI_RedBand:
                rgb_bands[0]=i
            elif gci == gdal.GCI_GreenBand:
                rgb_bands[1]=i
            elif gci == gdal.GCI_BlueBand:
                rgb_bands[2]=i
            if len(rgb_bands)==3:
                rgb_bands=rgb_bands[0],rgb_bands[1],rgb_bands[2] #Make a list from the dict

        #Set some defaults
        if not stretch_type and not stretch_args:
            if nbands < 3:
                #Default - assume greyscale
                stretch_type='PERCENT'
                stretch_args=[2,98]
                rgb_bands=[1]
                #But check if there's an attribute table or color table
                #and change the stretch type to colour table
                if datatype in ['Byte', 'Int16', 'UInt16']:
                    ct=rb.GetColorTable()
                    at=rb.GetDefaultRAT()
                    if ct and ct.GetCount() > 0:
                        stretch_type='COLOURTABLE'
                        stretch_args=[]
                    elif at and at.GetRowCount() > 0:
                        #if at.GetRowCount() <=256:
                        #    stretch_type='RANDOM'
                        #    stretch_args=[]
                        stretch_type='RANDOM'
                        stretch_args=[]

            elif nbands == 3:
                #Assume RGB
                if nbits > 8:
                    stretch_type='PERCENT'
                    stretch_args=[2,98]
                else:
                    stretch_type='NONE'
                    stretch_args=[]
                if len(rgb_bands) < 3:rgb_bands=[1,2,3]
            elif nbands >= 4:
                stretch_type='PERCENT'
                stretch_args=[2,98]
                #stretch_type='STDDEV'
                #stretch_args=[2]
                if len(rgb_bands) < 3:rgb_bands=[3,2,1]

        return overviews.getoverview(ds,outfile,width,format,rgb_bands,stretch_type,*stretch_args)

    # ===================== #
    # Private Class Methods
    # ===================== #
    def __setfileinfo__(self,f):
        self.fileinfo=utilities.FileInfo(f)
        #self.guid=str(uuid.uuid4())
        #Make the guid reproducible based on filename
        self.guid=str(uuid.uuid3(uuid.NAMESPACE_DNS,f))
        self.fileinfo['filename']=os.path.basename(f)
        self.fileinfo['filepath']=f
        self.fileinfo['guid']=self.guid
        #self.fileinfo['metadatadate']=time.strftime(utilities.datetimeformat,time.localtime()) #Geonetwork is baulking at the yyy-mm-ddThh:mm:ss format
        self.fileinfo['metadatadate']=time.strftime(utilities.dateformat,time.localtime())  #Just use yyy-mm-dd
    def __getfilelist__(self,*args,**kwargs):
        '''Get all files that have the same name (sans .ext), or are related according to gdalinfo.
            Special cases may be handled separately in their respective format drivers.'''
        f=self.fileinfo['filepath']
        files=[]
        try:
            files=glob.glob(os.path.splitext(f)[0]+'.*')
            if os.path.exists(os.path.splitext(f)[0]):files.append(os.path.splitext(f)[0])
            hdr_dir=os.path.join(os.path.split(f)[0], 'headers') #Cause ACRES creates a 'headers' directory
            if os.path.exists(hdr_dir):
                files.extend(glob.glob(os.path.join(hdr_dir,'*')))
        except:pass # Need to handle errors when dealing with an VRT XML string better...

        if self._gdaldataset:
            #try:files.extend(self._gdaldataset.GetFileList()) #GetFileList can return files that don't exist... - e.g. the aux.xml in a read-only directory
            try:files.extend([os.path.abspath(f) for f in self._gdaldataset.GetFileList() if os.path.exists(f)])
            except:pass

        self._filelist=list(set(utilities.fixSeparators(files))) #list(set([])) filters out duplicates

    def __getmetadata__(self,*args,**kwargs):
        '''In case subclasses don't override...'''
        pass
    def __init__(self,*args,**kwargs):
        '''In case subclasses don't override...'''
        pass

    # ================ #
    # Class Properties
    # ================ #
    def __classproperty__(fcn):
        '''The class property decorator function'''
        try:return property( **fcn() )
        except:pass

    @__classproperty__
    def metadata():
        '''The metadata property.'''

        def fget(self):
            if not self._metadata:
                #Initialise the metadata idict
                for field in self.fields:
                    if field in self.fileinfo:self._metadata[field]=self.fileinfo[field]
                    else:self._metadata[field]=''
                self._metadata=idict(self._metadata) #We don't want any fields added/deleted
                self.__getmetadata__()
                #Pretty print the SRS
                srs=osr.SpatialReference(self._metadata['srs'])
                self._metadata['srs']=srs.ExportToPrettyWkt()
                if self._metadata['compressionratio'] > 10000: #Possibly a dodgy JP2 that will grash GDAL and therefore python...
                    raise IOError, 'Unable to extract metadata from %s\nFile may be corrupt' % self.fileinfo['filepath']

            return self._metadata

        def fset(self, *args, **kwargs):
            if len(args) == 1:raise AttributeError('Can\'t overwrite metadata property')
            elif len(args) == 2:self._metadata[args[0]] = args[1]

        def fdel(self):pass #raise AttributeError('Can\'t delete metadata property')???????

        return locals()

    @__classproperty__
    def extent():
        '''The extent property.'''

        def fget(self, *args, **kwargs):
            if not self._extent:self.__getmetadata__() #extent gets set during metadata extraction
            return self._extent

        def fset(self, *args, **kwargs):
            if len(args) == 1:self._extent = args[0]
            elif len(args) == 2:self._extent[args[0]] = args[1]

        def fdel(self, *args, **kwargs):pass

        return locals()

    @__classproperty__
    def filelist():
        '''The filelist property.'''

        def fget(self):
            if not self._filelist:self.__getfilelist__()
            return self._filelist

        def fset(self, *args, **kwargs):
            if len(args) == 1:self._filelist = args[0]
            elif len(args) == 2:self._filelist[args[0]] = args[1]

        def fdel(self):pass

        return locals()

class idict(UserDict.IterableUserDict):
    '''An immutable dictionary.
       Modified from http://code.activestate.com/recipes/498072/
       to inherit UserDict.IterableUserDict
    '''
    def __setitem__(self, key, val):
        if key in self.data.keys():
            self.data[key]=val
        else:raise KeyError("Can't add keys")

    def __delitem__(self, key):
        raise KeyError("Can't delete keys")

    def pop(self, key):
        raise KeyError("Can't delete keys")

    def popitem(self):
        raise KeyError("Can't delete keys")
<commit_msg>Workaround to stop crashing when trying to access rasters stored in an ESRI PGDB<commit_after>'''
Base Dataset class
==================
Defines the metadata fields and populates some basic info

@todo: implement ESRI PGDB rasters (depends on GDAL I think)
'''

# Copyright (c) 2009 Australian Government, Department of Environment, Heritage, Water and the Arts
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
# THE SOFTWARE.

import os,time,sys,glob,time,math
import UserDict
import utilities, geometry, uuid

#Import fieldnames
import __fields__

try:
    from osgeo import gdal
    from osgeo import gdalconst
    from osgeo import osr
    from osgeo import ogr
except ImportError:
    import gdal
    import gdalconst
    import osr
    import ogr
gdal.AllRegister()
class Dataset(object):
    '''A base Dataset class'''
    def __new__(self,f):
        ##Initialise the class object
        self=object.__new__(self)

        #Little kludge to avoid ESRI PGDB rasters for now
        d=os.path.abspath(os.path.dirname(f))
        if d[-4:].lower() == '.gdb': raise Exception, 'Unable to open rasters stored in an ESRI PGDB %s'%f

        self._gdaldataset=None
        self._metadata={}
        self._extent=[]
        self._filelist=[]
        self._stretch=None

        ##Populate some basic info
        self.__setfileinfo__(f)

        ##Initialise the fields
        self.fields=idict(__fields__.fields)#We don't want any fields added/deleted

        return self

    # ==================== #
    # Public Class Methods
    # ==================== #
    def getoverview(self,outfile=None,width=800,format='JPG'):
        '''
        Generate overviews for generic imagery

        Override this class if you like, otherwise simply expose a GDALDataset object as self._gdaldataset_

        @type  outfile: string
        @param outfile: a filepath to the output overview image. If supplied, format is determined from the file extension
        @type  width:   integer
        @param width:   image width
        @type  format:  string
        @param format:  format to generate overview image, one of ['JPG','PNG','GIF','BMP','TIF']. Not required if outfile is supplied.

        @return:
            - B{filepath} (if outfile is supplied) B{OR}
            - B{binary image data} (if outfile is not supplied)
        '''
        import overviews

        md=self.metadata
        ds=self._gdaldataset
        if not ds:raise AttributeError, 'No GDALDataset object available, overview image can not be generated'

        #Don't rely on the metadata as self._gdaldataset might be a custom VRT
        ##nbands=md['nbands']
        ##cols=md['cols']
        ##rows=md['rows']
        ##nbits=md['nbits']
        rb=ds.GetRasterBand(1)
        nbands=ds.RasterCount
        cols=ds.RasterXSize
        rows=ds.RasterYSize
        nbits=gdal.GetDataTypeSize(rb.DataType)
        datatype=gdal.GetDataTypeName(rb.DataType)
        nodata=rb.GetNoDataValue()
        stretch_type=None
        stretch_args=None
        rgb_bands = {}

        #Check for pre-defined stretch
        if self._stretch:
            stretch_type,stretch_args=self._stretch

        #Check for pre-defined rgb bands
        for i in range(1,nbands+1):
            gci=ds.GetRasterBand(i).GetRasterColorInterpretation()
            if   gci == gdal.GCI_RedBand:
                rgb_bands[0]=i
            elif gci == gdal.GCI_GreenBand:
                rgb_bands[1]=i
            elif gci == gdal.GCI_BlueBand:
                rgb_bands[2]=i
            if len(rgb_bands)==3:
                rgb_bands=rgb_bands[0],rgb_bands[1],rgb_bands[2] #Make a list from the dict

        #Set some defaults
        if not stretch_type and not stretch_args:
            if nbands < 3:
                #Default - assume greyscale
                stretch_type='PERCENT'
                stretch_args=[2,98]
                rgb_bands=[1]
                #But check if there's an attribute table or color table
                #and change the stretch type to colour table
                if datatype in ['Byte', 'Int16', 'UInt16']:
                    ct=rb.GetColorTable()
                    at=rb.GetDefaultRAT()
                    if ct and ct.GetCount() > 0:
                        stretch_type='COLOURTABLE'
                        stretch_args=[]
                    elif at and at.GetRowCount() > 0:
                        #if at.GetRowCount() <=256:
                        #    stretch_type='RANDOM'
                        #    stretch_args=[]
                        stretch_type='RANDOM'
                        stretch_args=[]

            elif nbands == 3:
                #Assume RGB
                if nbits > 8:
                    stretch_type='PERCENT'
                    stretch_args=[2,98]
                else:
                    stretch_type='NONE'
                    stretch_args=[]
                if len(rgb_bands) < 3:rgb_bands=[1,2,3]
            elif nbands >= 4:
                stretch_type='PERCENT'
                stretch_args=[2,98]
                #stretch_type='STDDEV'
                #stretch_args=[2]
                if len(rgb_bands) < 3:rgb_bands=[3,2,1]

        return overviews.getoverview(ds,outfile,width,format,rgb_bands,stretch_type,*stretch_args)

    # ===================== #
    # Private Class Methods
    # ===================== #
    def __setfileinfo__(self,f):
        self.fileinfo=utilities.FileInfo(f)
        #self.guid=str(uuid.uuid4())
        #Make the guid reproducible based on filename
        self.guid=str(uuid.uuid3(uuid.NAMESPACE_DNS,f))
        self.fileinfo['filename']=os.path.basename(f)
        self.fileinfo['filepath']=f
        self.fileinfo['guid']=self.guid
        #self.fileinfo['metadatadate']=time.strftime(utilities.datetimeformat,time.localtime()) #Geonetwork is baulking at the yyy-mm-ddThh:mm:ss format
        self.fileinfo['metadatadate']=time.strftime(utilities.dateformat,time.localtime())  #Just use yyy-mm-dd
    def __getfilelist__(self,*args,**kwargs):
        '''Get all files that have the same name (sans .ext), or are related according to gdalinfo.
            Special cases may be handled separately in their respective format drivers.'''
        f=self.fileinfo['filepath']
        files=[]
        try:
            files=glob.glob(os.path.splitext(f)[0]+'.*')
            if os.path.exists(os.path.splitext(f)[0]):files.append(os.path.splitext(f)[0])
            hdr_dir=os.path.join(os.path.split(f)[0], 'headers') #Cause ACRES creates a 'headers' directory
            if os.path.exists(hdr_dir):
                files.extend(glob.glob(os.path.join(hdr_dir,'*')))
        except:pass # Need to handle errors when dealing with an VRT XML string better...

        if self._gdaldataset:
            #try:files.extend(self._gdaldataset.GetFileList()) #GetFileList can return files that don't exist... - e.g. the aux.xml in a read-only directory
            try:files.extend([os.path.abspath(f) for f in self._gdaldataset.GetFileList() if os.path.exists(f)])
            except:pass

        self._filelist=list(set(utilities.fixSeparators(files))) #list(set([])) filters out duplicates

    def __getmetadata__(self,*args,**kwargs):
        '''In case subclasses don't override...'''
        pass
    def __init__(self,*args,**kwargs):
        '''In case subclasses don't override...'''
        pass

    # ================ #
    # Class Properties
    # ================ #
    def __classproperty__(fcn):
        '''The class property decorator function'''
        try:return property( **fcn() )
        except:pass

    @__classproperty__
    def metadata():
        '''The metadata property.'''

        def fget(self):
            if not self._metadata:
                #Initialise the metadata idict
                for field in self.fields:
                    if field in self.fileinfo:self._metadata[field]=self.fileinfo[field]
                    else:self._metadata[field]=''
                self._metadata=idict(self._metadata) #We don't want any fields added/deleted
                self.__getmetadata__()
                #Pretty print the SRS
                srs=osr.SpatialReference(self._metadata['srs'])
                self._metadata['srs']=srs.ExportToPrettyWkt()
                if self._metadata['compressionratio'] > 10000: #Possibly a dodgy JP2 that will grash GDAL and therefore python...
                    raise IOError, 'Unable to extract metadata from %s\nFile may be corrupt' % self.fileinfo['filepath']

            return self._metadata

        def fset(self, *args, **kwargs):
            if len(args) == 1:raise AttributeError('Can\'t overwrite metadata property')
            elif len(args) == 2:self._metadata[args[0]] = args[1]

        def fdel(self):pass #raise AttributeError('Can\'t delete metadata property')???????

        return locals()

    @__classproperty__
    def extent():
        '''The extent property.'''

        def fget(self, *args, **kwargs):
            if not self._extent:self.__getmetadata__() #extent gets set during metadata extraction
            return self._extent

        def fset(self, *args, **kwargs):
            if len(args) == 1:self._extent = args[0]
            elif len(args) == 2:self._extent[args[0]] = args[1]

        def fdel(self, *args, **kwargs):pass

        return locals()

    @__classproperty__
    def filelist():
        '''The filelist property.'''

        def fget(self):
            if not self._filelist:self.__getfilelist__()
            return self._filelist

        def fset(self, *args, **kwargs):
            if len(args) == 1:self._filelist = args[0]
            elif len(args) == 2:self._filelist[args[0]] = args[1]

        def fdel(self):pass

        return locals()

class idict(UserDict.IterableUserDict):
    '''An immutable dictionary.
       Modified from http://code.activestate.com/recipes/498072/
       to inherit UserDict.IterableUserDict
    '''
    def __setitem__(self, key, val):
        if key in self.data.keys():
            self.data[key]=val
        else:raise KeyError("Can't add keys")

    def __delitem__(self, key):
        raise KeyError("Can't delete keys")

    def pop(self, key):
        raise KeyError("Can't delete keys")

    def popitem(self):
        raise KeyError("Can't delete keys")
<|endoftext|>