<commit_before>import argparse
import pkg_resources
import re
import sys
import abc
from junitparser import JUnitXml, Failure, Error
import glob

DOXYGEN_WARNING_REGEX = r"(?:(?:((?:[/.]|[A-Za-z]:).+?):(-?\d+):\s*([Ww]arning|[Ee]rror)|<.+>:-?\d+(?::\s*([Ww]arning|[Ee]rror))?): (.+(?:\n(?!\s*(?:[Nn]otice|[Ww]arning|[Ee]rror): )[^/<\n][^:\n][^/\n].+)*)|\s*([Nn]otice|[Ww]arning|[Ee]rror): (.+))$"
doxy_pattern = re.compile(DOXYGEN_WARNING_REGEX)

SPHINX_WARNING_REGEX = r"(.+?:(?:\d+|None)):\s*(DEBUG|INFO|WARNING|ERROR|SEVERE):\s*(.+)\n?"
sphinx_pattern = re.compile(SPHINX_WARNING_REGEX)

class WarningsChecker(object):

    def __init__(self, name, verbose=False):
        ''' Constructor

        Args:
            name (str): Name of the checker
            verbose (bool): Enable/disable verbose logging
        '''
        self.name = name
        self.verbose = verbose
        self.reset()

    def reset(self):
        ''' Reset function (resets min, max and counter values) '''
        self.count = 0
        self.warn_min = 0
        self.warn_max = 0

    @abc.abstractmethod
    def check(self, content):
        '''
        Function for counting the number of warnings in a specific text

        Args:
            content (str): The content to parse
        '''
        return

    def set_maximum(self, maximum):
        ''' Setter function for the maximum amount of warnings

        Args:
            maximum (int): maximum amount of warnings allowed

        Raises:
            ValueError: Invalid argument (min limit higher than max limit)
        '''
        if self.warn_min == 0:
            self.warn_max = maximum
        elif self.warn_min > maximum:
            raise ValueError("Invalid argument: mininum limit ({min}) is higher than maximum limit ({max}). Cannot enter {value}". format(min=self.warn_min, max=self.warn_max, value=maximum))
        else:
            self.warn_max = maximum

    def get_maximum(self):
        ''' Getter function for the maximum amount of warnings

        Returns:
            int: Maximum amount of warnings
        '''
        return self.warn_max

    def set_minimum(self, minimum):
        ''' Setter function for the minimum amount of warnings

        Args:
            minimum (int): minimum amount of warnings allowed

        Raises:
            ValueError: Invalid argument (min limit higher than max limit)
        '''
        if minimum > self.warn_max:
            raise ValueError("Invalid argument: mininum limit ({min}) is higher than maximum limit ({max}). Cannot enter {value}". format(min=self.warn_min, max=self.warn_max, value=minimum))
        else:
            self.warn_min = minimum

    def get_minimum(self):
        ''' Getter function for the minimum amount of warnings

        Returns:
            int: Minimum amount of warnings
        '''
        return self.warn_min

    def return_count(self):
        ''' Getter function for the amount of warnings found

        Returns:
            int: Number of warnings found
        '''
        print("{count} {name} warnings found".format(count=self.count, name=self.name))
        return self.count

    def return_check_limits(self):
        ''' Function for checking whether the warning count is within the configured limits

        Returns:
            int: 0 if the amount of warnings is within limits. the count of warnings otherwise
        '''
        if self.count > self.warn_max:
            print("Number of warnings ({count}) is higher than the maximum limit ({max}). Returning error code 1.".format(count=self.count, max=self.warn_max))
            return self.count
        elif self.count < self.warn_min:
            print("Number of warnings ({count}) is lower than the minimum limit ({min}). Returning error code 1.".format(count=self.count, min=self.warn_min))
            return self.count
        else:
            print("Number of warnings ({count}) is between limits {min} and {max}. Well done.".format(count=self.count, min=self.warn_min, max=self.warn_max))
            return 0

class RegexChecker(WarningsChecker):

    def __init__(self, name, pattern, verbose=False):
        ''' Constructor

        Args:
            name (str): Name of the checker
            pattern (str): Regular expression used by the checker in order to find warnings
        '''
        super(RegexChecker, self).__init__(name=name, verbose=verbose)
        self.pattern = pattern

    def check(self, content):
        '''
        Function for counting the number of warnings in a specific text

        Args:
            content (str): The content to parse
        '''
        matches = re.finditer(self.pattern, content)
        for match in matches:
            self.count += 1
            if self.verbose:
                print(match.group(0).strip())

class SphinxChecker(RegexChecker):
    name = 'sphinx'

    def __init__(self, verbose=False):
        ''' Constructor

        Args:
            verbose (bool): Enable/disable verbose logging
        '''
        super(SphinxChecker, self).__init__(name=SphinxChecker.name, pattern=sphinx_pattern, verbose=verbose)

class DoxyChecker(RegexChecker):
    name = 'doxygen'

    def __init__(self, verbose=False):
        ''' Constructor

        Args:
            verbose (bool): Enable/disable verbose logging
        '''
        super(DoxyChecker, self).__init__(name=DoxyChecker.name, pattern=doxy_pattern, verbose=verbose)

class JUnitChecker(WarningsChecker):
    name = 'junit'

    def __init__(self, verbose=False):
        ''' Constructor

        Args:
            verbose (bool): Enable/disable verbose logging
        '''
        super(JUnitChecker, self).__init__(name=JUnitChecker.name, verbose=verbose)

    def check(self, content):
        '''
        Function for counting the number of JUnit failures in a specific text

        Args:
            content (str): The content to parse
        '''
        result = JUnitXml.fromstring(content)
        if self.verbose:
            for suite in result:
                for testcase in suite:
                    if isinstance(testcase.result, Failure) or isinstance(testcase.result, Error):
                        print('{classname}.{testname}'.format(classname=testcase.classname,
                                                              testname=testcase.name))
        result.update_statistics()
        self.count += result.errors + result.failures

class WarningsPlugin:

    def __init__(self, sphinx = False, doxygen = False, junit = False, verbose = False):
        '''
        Function for initializing the parsers

        Args:
            sphinx (bool, optional):    enable sphinx parser
            doxygen (bool, optional):   enable doxygen parser
            junit (bool, optional):     enable junit parser
            verbose (bool, optional):   enable verbose logging
        '''
        self.checkerList = {}
        self.verbose = False
        if verbose:
            self.verbose = True
        if sphinx:
            self.activate_checker(SphinxChecker(self.verbose))
        if doxygen:
            self.activate_checker(DoxyChecker(self.verbose))
        if junit:
            self.activate_checker(JUnitChecker(self.verbose))

        self.warn_min = 0
        self.warn_max = 0
        self.count = 0

    def activate_checker(self, checker):
        '''
        Activate additional checkers after initialization

        Args:
            checker (WarningsChecker):         checker object
        '''
        checker.reset()
        self.checkerList[checker.name] = checker

    def get_checker(self, name):
        ''' Get checker by name

        Args:
            name (str): checker name
        Return:
            checker object (WarningsChecker)
        '''
        return self.checkerList[name]

    def check(self, content):
        '''
        Function for counting the number of warnings in a specific text

        Args:
            content (str): The text to parse
        '''
        if len(self.checkerList) == 0:
            print("No checkers activated. Please use activate_checker function")
        else:
            for name, checker in self.checkerList.items():
                checker.check(content)

    def set_maximum(self, maximum):
        ''' Setter function for the maximum amount of warnings

        Args:
            maximum (int): maximum amount of warnings allowed
        '''
        for name, checker in self.checkerList.items():
            checker.set_maximum(maximum)

    def set_minimum(self, minimum):
        ''' Setter function for the minimum amount of warnings

        Args:
            minimum (int): minimum amount of warnings allowed
        '''
        for name, checker in self.checkerList.items():
            checker.set_minimum(minimum)

    def return_count(self, name = None):
        ''' Getter function for the amount of found warnings

        If the name parameter is set, this function will return the amount of
        warnings found by that checker. If not, the function will return the sum
        of the warnings found by all registered checkers.

        Args:
            name (WarningsChecker): The checker for which to return the amount of warnings (if set)

        Returns:
            int: Amount of found warnings
        '''
        self.count = 0
        if name is None:
            for name, checker in self.checkerList.items():
                self.count += checker.return_count()
        else:
            self.count = self.checkerList[name].return_count()
        return self.count

    def return_check_limits(self, name = None):
        ''' Function for determining the return value of the script

        If the name parameter is set, this function will check (and return) the
        return value of that checker. If not, this function checks whether the
        warnings for each registred checker are within the configured limits.

        Args:
            name (WarningsChecker): The checker for which to check the return value

        Return:
            int: 0 if the amount warnings are within limits otherwise 1
        '''
        if name is None:
            for name, checker in self.checkerList.items():
                retval = checker.return_check_limits()
                if retval != 0:
                    return retval
        else:
            return self.checkerList[name].return_check_limits()

        return 0

def main():
    parser = argparse.ArgumentParser(prog='mlx-warnings')
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('-d', '--doxygen', dest='doxygen', action='store_true')
    group.add_argument('-s', '--sphinx', dest='sphinx', action='store_true')
    group.add_argument('-j', '--junit', dest='junit', action='store_true')
    parser.add_argument('-v', '--verbose', dest='verbose', action='store_true')
    parser.add_argument('-m', '--maxwarnings', type=int, required=False, default=0,
                        help='Maximum amount of warnings accepted')
    parser.add_argument('--minwarnings', type=int, required=False, default=0,
                        help='Minimum amount of warnings accepted')
    parser.add_argument('--version', action='version', version='%(prog)s {version}'.format(version=pkg_resources.require('mlx.warnings')[0].version))

    parser.add_argument('logfile', nargs='+', help='Logfile that might contain warnings')
    args = parser.parse_args()

    warnings = WarningsPlugin(sphinx=args.sphinx, doxygen=args.doxygen, junit=args.junit, verbose=args.verbose)
    warnings.set_maximum(args.maxwarnings)
    warnings.set_minimum(args.minwarnings)

    # args.logfile doesn't necessarily contain wildcards, but just to be safe, we
    # assume it does, and try to expand them.
    # This mechanism is put in place to allow wildcards to be passed on even when
    # executing the script on windows (in that case there is no shell expansion of wildcards)
    # so that the script can be used in the exact same way even when moving from one
    # OS to another.
    for file_wildcard in args.logfile:
        for logfile in glob.glob(file_wildcard):
            with open(logfile, 'r') as loghandle:
                warnings.check(loghandle.read().encode('utf-8'))

    warnings.return_count()
    sys.exit(warnings.return_check_limits())

if __name__ == '__main__':
    main()

<commit_msg>doxygen regex: don't match till end of given content<commit_after>import argparse
import pkg_resources
import re
import sys
import abc
from junitparser import JUnitXml, Failure, Error
import glob

DOXYGEN_WARNING_REGEX = r"(?:((?:[/.]|[A-Za-z]:).+?):(-?\d+):\s*([Ww]arning|[Ee]rror)|<.+>:-?\d+(?::\s*([Ww]arning|[Ee]rror))?): (.+(?:\n(?!\s*(?:[Nn]otice|[Ww]arning|[Ee]rror): )[^/<\n][^:\n][^/\n].+)*)|\s*([Nn]otice|[Ww]arning|[Ee]rror): (.+)\n?"
doxy_pattern = re.compile(DOXYGEN_WARNING_REGEX)

SPHINX_WARNING_REGEX = r"(.+?:(?:\d+|None)):\s*(DEBUG|INFO|WARNING|ERROR|SEVERE):\s*(.+)\n?"
sphinx_pattern = re.compile(SPHINX_WARNING_REGEX)

class WarningsChecker(object):

    def __init__(self, name, verbose=False):
        ''' Constructor

        Args:
            name (str): Name of the checker
            verbose (bool): Enable/disable verbose logging
        '''
        self.name = name
        self.verbose = verbose
        self.reset()

    def reset(self):
        ''' Reset function (resets min, max and counter values) '''
        self.count = 0
        self.warn_min = 0
        self.warn_max = 0

    @abc.abstractmethod
    def check(self, content):
        '''
        Function for counting the number of warnings in a specific text

        Args:
            content (str): The content to parse
        '''
        return

    def set_maximum(self, maximum):
        ''' Setter function for the maximum amount of warnings

        Args:
            maximum (int): maximum amount of warnings allowed

        Raises:
            ValueError: Invalid argument (min limit higher than max limit)
        '''
        if self.warn_min == 0:
            self.warn_max = maximum
        elif self.warn_min > maximum:
            raise ValueError("Invalid argument: mininum limit ({min}) is higher than maximum limit ({max}). Cannot enter {value}". format(min=self.warn_min, max=self.warn_max, value=maximum))
        else:
            self.warn_max = maximum

    def get_maximum(self):
        ''' Getter function for the maximum amount of warnings

        Returns:
            int: Maximum amount of warnings
        '''
        return self.warn_max

    def set_minimum(self, minimum):
        ''' Setter function for the minimum amount of warnings

        Args:
            minimum (int): minimum amount of warnings allowed

        Raises:
            ValueError: Invalid argument (min limit higher than max limit)
        '''
        if minimum > self.warn_max:
            raise ValueError("Invalid argument: mininum limit ({min}) is higher than maximum limit ({max}). Cannot enter {value}". format(min=self.warn_min, max=self.warn_max, value=minimum))
        else:
            self.warn_min = minimum

    def get_minimum(self):
        ''' Getter function for the minimum amount of warnings

        Returns:
            int: Minimum amount of warnings
        '''
        return self.warn_min

    def return_count(self):
        ''' Getter function for the amount of warnings found

        Returns:
            int: Number of warnings found
        '''
        print("{count} {name} warnings found".format(count=self.count, name=self.name))
        return self.count

    def return_check_limits(self):
        ''' Function for checking whether the warning count is within the configured limits

        Returns:
            int: 0 if the amount of warnings is within limits. the count of warnings otherwise
        '''
        if self.count > self.warn_max:
            print("Number of warnings ({count}) is higher than the maximum limit ({max}). Returning error code 1.".format(count=self.count, max=self.warn_max))
            return self.count
        elif self.count < self.warn_min:
            print("Number of warnings ({count}) is lower than the minimum limit ({min}). Returning error code 1.".format(count=self.count, min=self.warn_min))
            return self.count
        else:
            print("Number of warnings ({count}) is between limits {min} and {max}. Well done.".format(count=self.count, min=self.warn_min, max=self.warn_max))
            return 0

class RegexChecker(WarningsChecker):

    def __init__(self, name, pattern, verbose=False):
        ''' Constructor

        Args:
            name (str): Name of the checker
            pattern (str): Regular expression used by the checker in order to find warnings
        '''
        super(RegexChecker, self).__init__(name=name, verbose=verbose)
        self.pattern = pattern

    def check(self, content):
        '''
        Function for counting the number of warnings in a specific text

        Args:
            content (str): The content to parse
        '''
        matches = re.finditer(self.pattern, content)
        for match in matches:
            self.count += 1
            if self.verbose:
                print(match.group(0).strip())

class SphinxChecker(RegexChecker):
    name = 'sphinx'

    def __init__(self, verbose=False):
        ''' Constructor

        Args:
            verbose (bool): Enable/disable verbose logging
        '''
        super(SphinxChecker, self).__init__(name=SphinxChecker.name, pattern=sphinx_pattern, verbose=verbose)

class DoxyChecker(RegexChecker):
    name = 'doxygen'

    def __init__(self, verbose=False):
        ''' Constructor

        Args:
            verbose (bool): Enable/disable verbose logging
        '''
        super(DoxyChecker, self).__init__(name=DoxyChecker.name, pattern=doxy_pattern, verbose=verbose)

class JUnitChecker(WarningsChecker):
    name = 'junit'

    def __init__(self, verbose=False):
        ''' Constructor

        Args:
            verbose (bool): Enable/disable verbose logging
        '''
        super(JUnitChecker, self).__init__(name=JUnitChecker.name, verbose=verbose)

    def check(self, content):
        '''
        Function for counting the number of JUnit failures in a specific text

        Args:
            content (str): The content to parse
        '''
        result = JUnitXml.fromstring(content)
        if self.verbose:
            for suite in result:
                for testcase in suite:
                    if isinstance(testcase.result, Failure) or isinstance(testcase.result, Error):
                        print('{classname}.{testname}'.format(classname=testcase.classname,
                                                              testname=testcase.name))
        result.update_statistics()
        self.count += result.errors + result.failures

class WarningsPlugin:

    def __init__(self, sphinx = False, doxygen = False, junit = False, verbose = False):
        '''
        Function for initializing the parsers

        Args:
            sphinx (bool, optional):    enable sphinx parser
            doxygen (bool, optional):   enable doxygen parser
            junit (bool, optional):     enable junit parser
            verbose (bool, optional):   enable verbose logging
        '''
        self.checkerList = {}
        self.verbose = False
        if verbose:
            self.verbose = True
        if sphinx:
            self.activate_checker(SphinxChecker(self.verbose))
        if doxygen:
            self.activate_checker(DoxyChecker(self.verbose))
        if junit:
            self.activate_checker(JUnitChecker(self.verbose))

        self.warn_min = 0
        self.warn_max = 0
        self.count = 0

    def activate_checker(self, checker):
        '''
        Activate additional checkers after initialization

        Args:
            checker (WarningsChecker):         checker object
        '''
        checker.reset()
        self.checkerList[checker.name] = checker

    def get_checker(self, name):
        ''' Get checker by name

        Args:
            name (str): checker name
        Return:
            checker object (WarningsChecker)
        '''
        return self.checkerList[name]

    def check(self, content):
        '''
        Function for counting the number of warnings in a specific text

        Args:
            content (str): The text to parse
        '''
        if len(self.checkerList) == 0:
            print("No checkers activated. Please use activate_checker function")
        else:
            for name, checker in self.checkerList.items():
                checker.check(content)

    def set_maximum(self, maximum):
        ''' Setter function for the maximum amount of warnings

        Args:
            maximum (int): maximum amount of warnings allowed
        '''
        for name, checker in self.checkerList.items():
            checker.set_maximum(maximum)

    def set_minimum(self, minimum):
        ''' Setter function for the minimum amount of warnings

        Args:
            minimum (int): minimum amount of warnings allowed
        '''
        for name, checker in self.checkerList.items():
            checker.set_minimum(minimum)

    def return_count(self, name = None):
        ''' Getter function for the amount of found warnings

        If the name parameter is set, this function will return the amount of
        warnings found by that checker. If not, the function will return the sum
        of the warnings found by all registered checkers.

        Args:
            name (WarningsChecker): The checker for which to return the amount of warnings (if set)

        Returns:
            int: Amount of found warnings
        '''
        self.count = 0
        if name is None:
            for name, checker in self.checkerList.items():
                self.count += checker.return_count()
        else:
            self.count = self.checkerList[name].return_count()
        return self.count

    def return_check_limits(self, name = None):
        ''' Function for determining the return value of the script

        If the name parameter is set, this function will check (and return) the
        return value of that checker. If not, this function checks whether the
        warnings for each registred checker are within the configured limits.

        Args:
            name (WarningsChecker): The checker for which to check the return value

        Return:
            int: 0 if the amount warnings are within limits otherwise 1
        '''
        if name is None:
            for name, checker in self.checkerList.items():
                retval = checker.return_check_limits()
                if retval != 0:
                    return retval
        else:
            return self.checkerList[name].return_check_limits()

        return 0

def main():
    parser = argparse.ArgumentParser(prog='mlx-warnings')
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('-d', '--doxygen', dest='doxygen', action='store_true')
    group.add_argument('-s', '--sphinx', dest='sphinx', action='store_true')
    group.add_argument('-j', '--junit', dest='junit', action='store_true')
    parser.add_argument('-v', '--verbose', dest='verbose', action='store_true')
    parser.add_argument('-m', '--maxwarnings', type=int, required=False, default=0,
                        help='Maximum amount of warnings accepted')
    parser.add_argument('--minwarnings', type=int, required=False, default=0,
                        help='Minimum amount of warnings accepted')
    parser.add_argument('--version', action='version', version='%(prog)s {version}'.format(version=pkg_resources.require('mlx.warnings')[0].version))

    parser.add_argument('logfile', nargs='+', help='Logfile that might contain warnings')
    args = parser.parse_args()

    warnings = WarningsPlugin(sphinx=args.sphinx, doxygen=args.doxygen, junit=args.junit, verbose=args.verbose)
    warnings.set_maximum(args.maxwarnings)
    warnings.set_minimum(args.minwarnings)

    # args.logfile doesn't necessarily contain wildcards, but just to be safe, we
    # assume it does, and try to expand them.
    # This mechanism is put in place to allow wildcards to be passed on even when
    # executing the script on windows (in that case there is no shell expansion of wildcards)
    # so that the script can be used in the exact same way even when moving from one
    # OS to another.
    for file_wildcard in args.logfile:
        for logfile in glob.glob(file_wildcard):
            with open(logfile, 'r') as loghandle:
                warnings.check(loghandle.read().encode('utf-8'))

    warnings.return_count()
    sys.exit(warnings.return_check_limits())

if __name__ == '__main__':
    main()

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

"""
GistAPI.py -- A Python wrapper for GitHub's Gist API
(c) 2011 Kenneth Reitz. MIT License.

Example usage:

>>> Gist('d4507e882a07ac6f9f92').repo
u'd4507e882a07ac6f9f92'

>>> Gist('d4507e882a07ac6f9f92').owner
u'kennethreitz'

>>> Gist('d4507e882a07ac6f9f92').description
u'Example Gist for gist.py'

>>> Gist('d4507e882a07ac6f9f92').created_at.isoformat()
'2010-05-16T10:51:15-07:00'

>>> Gist('d4507e882a07ac6f9f92').public
False

>>> Gist('d4507e882a07ac6f9f92').filenames
['exampleEmptyFile', 'exampleFile']

>>> Gists.fetch_by_user('kennethreitz')[-1].description
u'My .bashrc configuration'

>>> Gist(885658).comments[0].body
u'Great Stuff.'
"""

import cStringIO
import os.path

import requests
from datetime import datetime

try:
    import simplejson as json
except ImportError:
    import json

__all__ = ['Gist', 'Gists']

# Set your own credentials
USERNAME = 'YOUR GITHUB USERNAME'
TOKEN = 'YOUR GITHUB TOKEN'

GIST_BASE = 'http://gist.github.com/%s'
GIST_JSON = GIST_BASE % 'api/v1/json/%s'

class Gist(object):
    """Gist Object"""

    def __init__(self, id=None, json=None, username=None, token=None):
        self.id = id
        self._json = json
        self._username = username
        self._token = token

        # Map given repo id to gist id if none exists
        if self._json:
            self.id = json['repo']

        self.url = url = GIST_BASE % self.id
        self.embed_url = url + '.js'
        self.epic_embed_url = url + '.pibb'
        self.json_url = url + '.json'
        self.post_url = GIST_BASE % 'gists/%s' % self.id
        self.comments  = []

    def __repr__(self):
        return '<gist %s>' % self.id

    def __getattribute__(self, name):
        """Get attributes, but only if needed."""

        # Only make external API calls if needed
        if name in ('owner', 'description', 'created_at', 'public',
                    'files', 'filenames', 'repo', 'comments'):
            if not hasattr(self, '_meta'):
                self._meta = self._get_meta()

        return object.__getattribute__(self, name)

    def _get_meta(self):
        """Fetch Gist metadata."""

        # Use json data provided if available
        if self._json:
            _meta = self._json
            setattr(self, 'id', _meta['repo'])
        else:
            # Fetch Gist metadata
            _meta_url = GIST_JSON % self.id
            _meta = json.loads(requests.get(_meta_url).content)['gists'][0]

        for key, value in _meta.iteritems():

            if key == 'files':
                # Remap file key from API.
                setattr(self, 'filenames', value)
                # Store the {current_name: original_name} mapping.
                setattr(self, '_renames', dict((fn, fn) for fn in value))
            elif key == 'public':
                # Attach booleans
                setattr(self, key, value)
            elif key == 'created_at':
                # Attach datetime
                setattr(self, 'created_at', datetime.strptime(value[:-6], '%Y/%m/%d %H:%M:%S'))

            elif key == 'comments':
                _comments = []
                for comment in value:
                    c = GistComment().from_api(comment)
                    _comments.append(c)
                    setattr(self, 'comments', _comments)
            else:
                # Attach properties to object
                setattr(self, key, unicode(value))

        return _meta

    def _post(self, params, headers={}):
        """POST to the web form (internal method)."""
        r = requests.post(self.post_url, params, headers=headers)

        return r.status_code, r.content

    def reset(self):
        """Clear the local cache."""
        if hasattr(self, '_files'):
            del self._files
        if hasattr(self, '_meta'):
            del self._meta

    def auth(self, username, token):
        """Set credentials."""
        self._username = username
        self._token = token

    def rename(self, from_name, to_name):
        """Rename a file."""
        if from_name not in self.files:
            raise KeyError('File %r does not exist' % from_name)
        if to_name in self.files:
            raise KeyError('File %r already exist' % to_name)
        self.files[to_name] = self.files.pop(from_name)
        try:
            self._renames[to_name] = self._renames.pop(from_name)
        except KeyError:
            # New file
            pass

    def save(self):
        """Upload the changes to Github."""
        params = {
            '_method': 'put',
            'login': self._username or USERNAME,
            'token': self._token or TOKEN,
        }
        names_map = self._renames
        original_names = names_map.values()
        index = len(original_names)
        for fn, content in self.files.items():
            ext = os.path.splitext(fn)[1] or '.txt'
            try:
                orig = names_map.pop(fn)
            except KeyError:
                # Find a unique filename
                while True:
                    orig = 'gistfile%s' % index
                    index += 1
                    if orig not in original_names:
                        break
            params.update({
                'file_name[%s]' % orig: fn,
                'file_ext[%s]' % orig: ext,
                'file_contents[%s]' % orig: content,
            })
        code, msg = self._post(params=params)
        if code == 200:  # OK
            # If successful, clear the cache
            self.reset()
        return code, msg

    @property
    def files(self):
        """Fetch Gist's files and store them in the 'files' property."""
        try:
            return self._files
        except AttributeError:
            self._files = _files = {}

        for fn in self._meta['files']:
            # Grab file contents
            _file_url = GIST_BASE % 'raw/%s/%s' % (self.id, urllib2.quote(fn))
            _files[fn] = cStringIO.StringIO()
            _files[fn].write(requests.get(_file_url).content)

        return _files

class Gists(object):
    """Gist API wrapper"""

    def __init__(self, username=None, token=None):
        # Token-based Authentication is unnecessary, gist api still in alpha
        self._username = username
        self._token = token

    @staticmethod
    def fetch_by_user(name):
        """Return a list of public Gist objects owned by
        the given GitHub username."""

        _url = GIST_JSON % 'gists/%s' % name

        # Return a list of Gist objects
        return [Gist(json=g)
                for g in json.loads(requests.get(_url).content)['gists']]

class GistComment(object):
    """Gist comments."""

    def __init__(self):
        self.body = None
        self.created_at = None
        self.gravatar_id = None
        self.id = None
        self.updated_at = None
        self.user = None

    def __repr__(self):
        return '<gist-comment %s>' % self.id

    @staticmethod
    def from_api(jsondict):
        """Returns new instance of GistComment containing given api dict."""
        comment = GistComment()

        comment.body = jsondict.get('body', None)
        comment.created_at = datetime.strptime(jsondict.get('created_at')[:-6], '%Y/%m/%d %H:%M:%S')
        comment.gravatar_id = jsondict.get('gravatar_id', None)
        comment.id = jsondict.get('id', None)
        comment.updated_at = datetime.strptime(jsondict.get('updated_at')[:-6], '%Y/%m/%d %H:%M:%S')
        comment.user = jsondict.get('user', None)

        return comment

<commit_msg>still need urllib2 for quoting<commit_after># -*- coding: utf-8 -*-

"""
GistAPI.py -- A Python wrapper for GitHub's Gist API
(c) 2011 Kenneth Reitz. MIT License.

Example usage:

>>> Gist('d4507e882a07ac6f9f92').repo
u'd4507e882a07ac6f9f92'

>>> Gist('d4507e882a07ac6f9f92').owner
u'kennethreitz'

>>> Gist('d4507e882a07ac6f9f92').description
u'Example Gist for gist.py'

>>> Gist('d4507e882a07ac6f9f92').created_at.isoformat()
'2010-05-16T10:51:15-07:00'

>>> Gist('d4507e882a07ac6f9f92').public
False

>>> Gist('d4507e882a07ac6f9f92').filenames
['exampleEmptyFile', 'exampleFile']

>>> Gists.fetch_by_user('kennethreitz')[-1].description
u'My .bashrc configuration'

>>> Gist(885658).comments[0].body
u'Great Stuff.'
"""

import cStringIO
import os.path

import requests
import urllib2
from datetime import datetime

try:
    import simplejson as json
except ImportError:
    import json

__all__ = ['Gist', 'Gists']

# Set your own credentials
USERNAME = 'YOUR GITHUB USERNAME'
TOKEN = 'YOUR GITHUB TOKEN'

GIST_BASE = 'http://gist.github.com/%s'
GIST_JSON = GIST_BASE % 'api/v1/json/%s'

class Gist(object):
    """Gist Object"""

    def __init__(self, id=None, json=None, username=None, token=None):
        self.id = id
        self._json = json
        self._username = username
        self._token = token

        # Map given repo id to gist id if none exists
        if self._json:
            self.id = json['repo']

        self.url = url = GIST_BASE % self.id
        self.embed_url = url + '.js'
        self.epic_embed_url = url + '.pibb'
        self.json_url = url + '.json'
        self.post_url = GIST_BASE % 'gists/%s' % self.id
        self.comments  = []

    def __repr__(self):
        return '<gist %s>' % self.id

    def __getattribute__(self, name):
        """Get attributes, but only if needed."""

        # Only make external API calls if needed
        if name in ('owner', 'description', 'created_at', 'public',
                    'files', 'filenames', 'repo', 'comments'):
            if not hasattr(self, '_meta'):
                self._meta = self._get_meta()

        return object.__getattribute__(self, name)

    def _get_meta(self):
        """Fetch Gist metadata."""

        # Use json data provided if available
        if self._json:
            _meta = self._json
            setattr(self, 'id', _meta['repo'])
        else:
            # Fetch Gist metadata
            _meta_url = GIST_JSON % self.id
            _meta = json.loads(requests.get(_meta_url).content)['gists'][0]

        for key, value in _meta.iteritems():

            if key == 'files':
                # Remap file key from API.
                setattr(self, 'filenames', value)
                # Store the {current_name: original_name} mapping.
                setattr(self, '_renames', dict((fn, fn) for fn in value))
            elif key == 'public':
                # Attach booleans
                setattr(self, key, value)
            elif key == 'created_at':
                # Attach datetime
                setattr(self, 'created_at', datetime.strptime(value[:-6], '%Y/%m/%d %H:%M:%S'))

            elif key == 'comments':
                _comments = []
                for comment in value:
                    c = GistComment().from_api(comment)
                    _comments.append(c)
                    setattr(self, 'comments', _comments)
            else:
                # Attach properties to object
                setattr(self, key, unicode(value))

        return _meta

    def _post(self, params, headers={}):
        """POST to the web form (internal method)."""
        r = requests.post(self.post_url, params, headers=headers)

        return r.status_code, r.content

    def reset(self):
        """Clear the local cache."""
        if hasattr(self, '_files'):
            del self._files
        if hasattr(self, '_meta'):
            del self._meta

    def auth(self, username, token):
        """Set credentials."""
        self._username = username
        self._token = token

    def rename(self, from_name, to_name):
        """Rename a file."""
        if from_name not in self.files:
            raise KeyError('File %r does not exist' % from_name)
        if to_name in self.files:
            raise KeyError('File %r already exist' % to_name)
        self.files[to_name] = self.files.pop(from_name)
        try:
            self._renames[to_name] = self._renames.pop(from_name)
        except KeyError:
            # New file
            pass

    def save(self):
        """Upload the changes to Github."""
        params = {
            '_method': 'put',
            'login': self._username or USERNAME,
            'token': self._token or TOKEN,
        }
        names_map = self._renames
        original_names = names_map.values()
        index = len(original_names)
        for fn, content in self.files.items():
            ext = os.path.splitext(fn)[1] or '.txt'
            try:
                orig = names_map.pop(fn)
            except KeyError:
                # Find a unique filename
                while True:
                    orig = 'gistfile%s' % index
                    index += 1
                    if orig not in original_names:
                        break
            params.update({
                'file_name[%s]' % orig: fn,
                'file_ext[%s]' % orig: ext,
                'file_contents[%s]' % orig: content,
            })
        code, msg = self._post(params=params)
        if code == 200:  # OK
            # If successful, clear the cache
            self.reset()
        return code, msg

    @property
    def files(self):
        """Fetch Gist's files and store them in the 'files' property."""
        try:
            return self._files
        except AttributeError:
            self._files = _files = {}

        for fn in self._meta['files']:
            # Grab file contents
            _file_url = GIST_BASE % 'raw/%s/%s' % (self.id, urllib2.quote(fn))
            _files[fn] = cStringIO.StringIO()
            _files[fn].write(requests.get(_file_url).content)

        return _files

class Gists(object):
    """Gist API wrapper"""

    def __init__(self, username=None, token=None):
        # Token-based Authentication is unnecessary, gist api still in alpha
        self._username = username
        self._token = token

    @staticmethod
    def fetch_by_user(name):
        """Return a list of public Gist objects owned by
        the given GitHub username."""

        _url = GIST_JSON % 'gists/%s' % name

        # Return a list of Gist objects
        return [Gist(json=g)
                for g in json.loads(requests.get(_url).content)['gists']]

class GistComment(object):
    """Gist comments."""

    def __init__(self):
        self.body = None
        self.created_at = None
        self.gravatar_id = None
        self.id = None
        self.updated_at = None
        self.user = None

    def __repr__(self):
        return '<gist-comment %s>' % self.id

    @staticmethod
    def from_api(jsondict):
        """Returns new instance of GistComment containing given api dict."""
        comment = GistComment()

        comment.body = jsondict.get('body', None)
        comment.created_at = datetime.strptime(jsondict.get('created_at')[:-6], '%Y/%m/%d %H:%M:%S')
        comment.gravatar_id = jsondict.get('gravatar_id', None)
        comment.id = jsondict.get('id', None)
        comment.updated_at = datetime.strptime(jsondict.get('updated_at')[:-6], '%Y/%m/%d %H:%M:%S')
        comment.user = jsondict.get('user', None)

        return comment

<|endoftext|>


<commit_before>import os
import pytest
import cPickle as pickle
from modules.beautify.beautify import beautify
from tests.constants import ARGS_DICT, BEAUTIFY_VF_SEROTYPE

def test_beautify():
    ## test vf & serotype json return
    vf_serotype_gene_dict = os.path.join('tests/refs', 'GCA_000005845.2_ASM584v2_genomic.fna_ectyper-vf_serotype.p')
    single_dict = dict(ARGS_DICT)
    single_dict.update({'i': vf_serotype_gene_dict})
    assert len(beautify(single_dict, vf_serotype_gene_dict)) == len(BEAUTIFY_VF_SEROTYPE)

    ## test serotype only json return
    # note: this is actually the same gene results file as above
    # we only differentiate what is returned to the user, because we want all analysis ran & added to the db
    # this mimicks user selection of serotype only
    single_dict.update({'options':{'vf': False, 'amr': False, 'serotype': True}})
    r = beautify(single_dict, serotype_gene_dict)
    print r
    assert len(r) == 1
<commit_msg>ADD: options showing up?<commit_after>import os
import pytest
import cPickle as pickle
from modules.beautify.beautify import beautify
from tests.constants import ARGS_DICT, BEAUTIFY_VF_SEROTYPE

def test_beautify():
    ## test vf & serotype json return
    vf_serotype_gene_dict = os.path.join('tests/refs', 'GCA_000005845.2_ASM584v2_genomic.fna_ectyper-vf_serotype.p')
    single_dict = dict(ARGS_DICT)
    single_dict.update({'i': vf_serotype_gene_dict})
    assert len(beautify(single_dict, vf_serotype_gene_dict)) == len(BEAUTIFY_VF_SEROTYPE)

    ## test serotype only json return
    # note: this is actually the same gene results file as above
    # we only differentiate what is returned to the user, because we want all analysis ran & added to the db
    # this mimicks user selection of serotype only
    single_dict.update({'options':{'vf': False, 'amr': False, 'serotype': True}})
    r = beautify(single_dict, vf_serotype_gene_dict)
    print r
    assert len(r) == 1
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

# Copyright © 2014 Puneeth Chaganti and others.
# See the LICENSE file for license rights and limitations (MIT).

""" A collection of helper functions for consuming the GitHub API. """

# Standard library
import base64
import json
import re

# 3rd party library
import requests

class GitHubUtils(object):
    """ A collection of helper functions for consuming the GitHUb API. """

    @staticmethod
    def commit(path, content, repo, token, extra_payload=None):
        """ Commit the given content to the given path in a repository. """

        if not GitHubUtils.exists(repo, path, token):

            url = 'repos/%s/contents/%s' % (repo, path)
            url = 'https://api.github.com/' + url

            headers = GitHubUtils.get_header(token)

            payload = {
                'path': path,
                'message': 'Adding %s (from statiki).' % path,
                'content': base64.standard_b64encode(content),
            }

            if extra_payload is not None:
                payload.update(extra_payload)

            response = requests.put(
                url, data=json.dumps(payload), headers=headers
            )

            committed = response.status_code == 201

        else:
            committed = False

        return committed

    @staticmethod
    def create_new_repository(full_name, token):
        """ Create a new repository given the name and a token.

        NOTE the token must have 'repo' scope, not just 'public_repo'.

        """

        user, user_type, name = GitHubUtils.get_user_and_repo(
            full_name, token
        )

        url = 'https://api.github.com/user/repos'
        headers = GitHubUtils.get_header(token)
        homepage = (
            name if GitHubUtils.is_user_pages(full_name)
            else 'http://%s.github.io/%s' % (user, name)
        )
        payload = {
            'name': name,
            'description': 'Website using Nikola, created from statiki',
            'homepage': homepage,
            'private': False,
            'has_issues': False,
            'has_wiki': False,
            'has_downloads': False,
        }

        response = requests.post(
            url, data=json.dumps(payload), headers=headers
        )

        return response.status_code == 201

    @staticmethod
    def exists(full_name, path, token):
        """ Return True if the given repository has the given path. """

        headers = GitHubUtils.get_header(token)

        url = 'repos/%s/contents/%s' % (full_name, path)
        url = 'https://api.github.com/' + url

        return requests.get(url, headers=headers).status_code == 200

    @staticmethod
    def get_header(token):
        """ Return a header with authorization info, given the token. """

        return {
            'Authorization': 'token %s' % token,
            'Content-Type': 'application/json'
        }

    @staticmethod
    def get_status():
        """ Return the server status of GitHub. """

        response = requests.get('https://status.github.com')
        pattern  = '(
(.|\s)*?
)'

        return re.findall(pattern, response.text)[0][0].strip()

    @staticmethod
    def get_user_and_repo(full_name, token):
        """ Return the username, user_type, repo_name, given a repository.

        user_type is one of User or Organization.

        """

        user, name = full_name.split('/')
        url = 'https://api.github.com/users/%s' % user
        headers = GitHubUtils.get_header(token)

        response = requests.get(url, headers=headers)

        user_type = response.json()['type']

        return user, user_type, name

    @staticmethod
    def is_valid_repository(full_name):
        """ Return True if such a repo exists on GitHub. """

        response = requests.get('https://github.com/%s' % full_name)
        return response.status_code == 200

    @staticmethod
    def is_user_pages(full_name):
        """ Return True if the repository is a user pages repository. """

        username, repo_name = full_name.split('/', 1)

        return (
            repo_name.startswith(username)
            and repo_name.endswith(('github.io', 'github.com'))
        )
<commit_msg>Fix the branch on which commits are made for user_pages repo.<commit_after># -*- coding: utf-8 -*-

# Copyright © 2014 Puneeth Chaganti and others.
# See the LICENSE file for license rights and limitations (MIT).

""" A collection of helper functions for consuming the GitHub API. """

# Standard library
import base64
import json
import re

# 3rd party library
import requests

class GitHubUtils(object):
    """ A collection of helper functions for consuming the GitHUb API. """

    @staticmethod
    def commit(path, content, repo, token, extra_payload=None):
        """ Commit the given content to the given path in a repository. """

        if not GitHubUtils.exists(repo, path, token):

            url = 'repos/%s/contents/%s' % (repo, path)
            url = 'https://api.github.com/' + url

            headers = GitHubUtils.get_header(token)
            branch = 'deploy' if GitHubUtils.is_user_pages(repo) else 'master'

            payload = {
                'path': path,
                'message': 'Adding %s (from statiki).' % path,
                'content': base64.standard_b64encode(content),
                'branch': branch
            }

            if extra_payload is not None:
                payload.update(extra_payload)

            response = requests.put(
                url, data=json.dumps(payload), headers=headers
            )

            committed = response.status_code == 201

        else:
            committed = False

        return committed

    @staticmethod
    def create_new_repository(full_name, token):
        """ Create a new repository given the name and a token.

        NOTE the token must have 'repo' scope, not just 'public_repo'.

        """

        user, user_type, name = GitHubUtils.get_user_and_repo(
            full_name, token
        )

        url = 'https://api.github.com/user/repos'
        headers = GitHubUtils.get_header(token)
        homepage = (
            name if GitHubUtils.is_user_pages(full_name)
            else 'http://%s.github.io/%s' % (user, name)
        )
        payload = {
            'name': name,
            'description': 'Website using Nikola, created from statiki',
            'homepage': homepage,
            'private': False,
            'has_issues': False,
            'has_wiki': False,
            'has_downloads': False,
        }

        response = requests.post(
            url, data=json.dumps(payload), headers=headers
        )

        return response.status_code == 201

    @staticmethod
    def exists(full_name, path, token):
        """ Return True if the given repository has the given path. """

        headers = GitHubUtils.get_header(token)

        url = 'repos/%s/contents/%s' % (full_name, path)
        url = 'https://api.github.com/' + url

        return requests.get(url, headers=headers).status_code == 200

    @staticmethod
    def get_header(token):
        """ Return a header with authorization info, given the token. """

        return {
            'Authorization': 'token %s' % token,
            'Content-Type': 'application/json'
        }

    @staticmethod
    def get_status():
        """ Return the server status of GitHub. """

        response = requests.get('https://status.github.com')
        pattern  = '(
(.|\s)*?
)'

        return re.findall(pattern, response.text)[0][0].strip()

    @staticmethod
    def get_user_and_repo(full_name, token):
        """ Return the username, user_type, repo_name, given a repository.

        user_type is one of User or Organization.

        """

        user, name = full_name.split('/')
        url = 'https://api.github.com/users/%s' % user
        headers = GitHubUtils.get_header(token)

        response = requests.get(url, headers=headers)

        user_type = response.json()['type']

        return user, user_type, name

    @staticmethod
    def is_valid_repository(full_name):
        """ Return True if such a repo exists on GitHub. """

        response = requests.get('https://github.com/%s' % full_name)
        return response.status_code == 200

    @staticmethod
    def is_user_pages(full_name):
        """ Return True if the repository is a user pages repository. """

        username, repo_name = full_name.split('/', 1)

        return (
            repo_name.startswith(username)
            and repo_name.endswith(('github.io', 'github.com'))
        )
<|endoftext|>


<commit_before>from datetime import datetime

MAGNITUDE_W = 1
"""
Magnitude value (1) for measures represented in W.
"""

MAGNITUDE_KW = 1000
"""
Magnitude value (1000) for measures represented in kW.
"""

class ValueWithTime(object):
    """
    Base class for values with time.
    """

    def _get_timestamp(self, name, element=None):
        """
        Formats a timestamp from the name of the value.

        :param name: a string with the name
        :param element: an lxml.objectify.StringElement, by default \
            self.objectified
        :return: a formatted string representing a timestamp \
            ('%Y-%m-%d %H:%M:%S')
        """
        if element is None:
            e = self.objectified
        else:
            e = element
        value = e.get(name)
        if len(value) > 15:
            date_value = value[0:14] + value[-1]
        else:
            date_value = value

        # Fix for SAGECOM which puts this timestamp when the period doesn't
        # affect the contracted tariff
        if date_value.upper() == 'FFFFFFFFFFFFFFW':
            date_value = '19000101000000W'

        time = datetime.strptime(date_value[:-1], '%Y%m%d%H%M%S')
        return time.strftime('%Y-%m-%d %H:%M:%S')

class Measure(ValueWithTime):
    """
    Base class for a set of measures.
    """

    def __init__(self, objectified_measure):
        """
        Create a Measure object.

        :param objectified_measure: an lxml.objectify.StringElement \
            representing a set of measures
        :return: a Measure object
        """
        self.objectified = objectified_measure

    @property
    def objectified(self):
        """
        The set of measures as an lxml.objectify.StringElement.

        :return: an lxml.objectify.StringElement representing a set of measures
        """
        return self._objectified

    @objectified.setter
    def objectified(self, value):
        """
        Stores an lxml.objectify.StringElement representing a set of measures.

        :param value: an lxml.objectify.StringElement representing a set of \
            measures
        :return:
        """
        self._objectified = value

class MeasureActiveReactive(Measure):
    """
    Base class for a set of measures with active and reactive measures.
    """

    def active_reactive(self, measure, measure_type):
        """
        Get the active and reactive measures.

        :param measure: an lxml.objectify.StringElement representing a set of \
            measures
        :param measure_type: the type of measure, added at the end of the \
            name of each measure ('a' or 'i')
        :return: a dict with the active and reactive measures
        """
        return {
            'ai': int(measure.get('AI{}'.format(measure_type))),
            'ae': int(measure.get('AE{}'.format(measure_type))),
            'r1': int(measure.get('R1{}'.format(measure_type))),
            'r2': int(measure.get('R2{}'.format(measure_type))),
            'r3': int(measure.get('R3{}'.format(measure_type))),
            'r4': int(measure.get('R4{}'.format(measure_type))),
        }

class Parameter(ValueWithTime):
    """
    Base class for a set of parameters.
    """

    def __init__(self, objectified_parameter, report_version):
        """
        Create a Parameter object.

        :param objectified_parameter: an lxml.objectify.StringElement \
            representing a set of parameters
        :return: a Measure object
        """
        self.objectified = objectified_parameter
        self.report_version = report_version

    @property
    def objectified(self):
        """
        The set of parameters as an lxml.objectify.StringElement.

        :return: an lxml.objectify.StringElement representing a set of \
            parameters
        """
        return self._objectified

    @objectified.setter
    def objectified(self, value):
        """
        Stores an lxml.objectify.StringElement representing a set of \
            parameters.

        :param value: an lxml.objectify.StringElement representing a set of \
            parameters
        :return:
        """
        self._objectified = value

    @property
    def report_version(self):
        """
        The version of the report.

        :return: a string with the version of the report
        """
        return self._report_version

    @report_version.setter
    def report_version(self, value):
        """
        Stores the report version.
        :param value: a string with the version of the report
        """
        self._report_version = value

    def get_boolean(self, name, element=None):
        """
        Gets a boolean value from the name of value.

        :param name: a string with the name
        :param element: an lxml.objectify.StringElement, by default \
            self.objectified
        :return:
        """
        if element is None:
            e = self.objectified
        else:
            e = element
        if e.get(name) == 'Y':
            returns = True
        else:
            returns = False
        return returns

    def to_integer(self, value):
        """
        Convert a value to an integer. If value is None then returns 0.

        :param value: a string with the value
        :return: an integer
        """
        if value is None:
            returns = 0
        else:
            returns = int(value)
        return returns

    def filter_integer(self, value):
        """
        Filter the provided value. Returns the value if is an integer or \
            returns 0 if not.

        :param value: an integer or not
        :return: an integer
        """
        if isinstance(value, int):
            returns = value
        else:
            returns = 0
        return returns

    @property
    def values(self):
        """
        Set of parameters.
        """
        raise NotImplementedError('This method is not implemented!')

class Meter(object):
    """
    Base class for a meter.
    """

    def __init__(self, objectified_meter):
        """
        Create a Meter object.

        :param objectified_meter: an lxml.objectify.StringElement \
            representing a meter
        :return: a Meter object
        """
        self.objectified = objectified_meter

    @property
    def objectified(self):
        """
        A meter as an lxml.objectify.StringElement.

        :return: an lxml.objectify.StringElement representing a meter
        """
        return self._objectified

    @objectified.setter
    def objectified(self, value):
        """
        Stores an lxml.objectify.StringElement representing a meter

        :param value: an lxml.objectify.StringElement representing a meter
        """
        self._objectified = value

    @property
    def errors(self):
        """
        The meter errors.

        :return: a dict with the meter errors
        """
        self._errors = {}
        if self.objectified.get('ErrCat'):
            self._errors = {
                'errcat': self.objectified.get('ErrCat'),
                'errcode': self.objectified.get('ErrCode')
            }
        return self._errors

    @property
    def name(self):
        """
        The name of the meter.

        :return: a string with the name of the meter
        """
        return self.objectified.get('Id')

    @property
    def magnitude(self):
        """
        The magnitude of the meter measures.

        :return: a int with the magnitude of the meter measures
        """
        return int(self.objectified.get('Magn'))

    @property
    def report_type(self):
        """
        The type of report. To implement in child classes.
        """
        raise NotImplementedError('This method is not implemented!')

    @property
    def measure_class(self):
        """
        The class to instance measures sets.

        :return: a class to instance measure sets
        """
        return Measure

    @property
    def measures(self):
        """
        Measure set objects of this meter.

        :return: a list of measure set objects
        """
        measures = []
        if hasattr(self.objectified, self.report_type):
            objectified = getattr(self.objectified, self.report_type)
            measures = map(self.measure_class, objectified)
        return measures

    @property
    def values(self):
        """
        Values of measure sets of this meter.

        :return: a list with the values of the measure sets
        """
        values = []
        for measure in self.measures:
            values.append(measure.value())
        return values

class MeterWithConcentratorName(Meter):
    """
    Base class for a meters of report that need the name of the concentrator \
        in the values, like S04 and S05.
    """

    def report_type(self):
        """
        The type of report. To implement in child classes.
        """
        raise NotImplementedError('This method is not implemented!')

    def __init__(self, objectified_meter, concentrator_name):
        """
        Create a Meter object using Meter constructor and adding the \
            concentrator name.

        :param objectified_meter: an lxml.objectify.StringElement \
            representing a meter
        :param concentrator_name: a string with the name of the concentrator
        :return: a Meter object
        """
        super(MeterWithConcentratorName, self).__init__(objectified_meter)
        self.concentrator_name = concentrator_name

    @property
    def concentrator_name(self):
        """
        A string with the concentrator name.

        :return: a string with the concentrator name
        """
        return self._concentrator_name

    @concentrator_name.setter
    def concentrator_name(self, value):
        """
        Stores a string with the concentrator name.

        :param value: a string with the concentrator name
        """
        self._concentrator_name = value

    @property
    def values(self):
        """
        Values of measure sets of this meter of report that need the name of \
            the concentrator and the meter,

        :return: a list with the values of the measure sets
        """
        values = []
        for measure in self.measures:
            for subvalue in measure.values:
                v = subvalue.copy()
                v['name'] = self.name
                v['cnc_name'] = self.concentrator_name
                values.append(v)
        return values

class Concentrator(object):
    """
    Base class for a concentrator.
    """

    def __init__(self, objectified_concentrator):
        """
        Create a Concentrator object.

        :param objectified_concentrator: an lxml.objectify.StringElement \
            representing a concentrator
        :return: a Concentrator object
        """
        self.objectified = objectified_concentrator

    @property
    def objectified(self):
        """
        A concentrator as an lxml.objectify.StringElement.

        :return: an lxml.objectify.StringElement representing a concentrator
        """
        return self._objectified

    @objectified.setter
    def objectified(self, value):
        """
        Stores a concentrator as an lxml.objectify.StringElement.

        :param value: an lxml.objectify.StringElement representing a \
            concentrator
        """
        self._objectified = value

    @property
    def name(self):
        """
        The name of the concentrator.

        :return: a string with the name of the concentrator
        """
        return self.objectified.get('Id')

class ConcentratorWithMeters(Concentrator):

    @property
    def meter_class(self):
        """
        The class to instance meters.

        :return: a class to instance meters
        """
        return Meter

    @property
    def meters(self):
        """
        Meter objects of this concentrator.

        :return: a list of meter objects
        """
        return map(self.meter_class, self.objectified.Cnt)

    @property
    def values(self):
        """
        Values of the meters of this concentrator.

        :return: a list with the values of the meters
        """
        values = []
        for meter in self.meters:
            values.extend(meter.values)
        return values

class ConcentratorWithMetersWithConcentratorName(ConcentratorWithMeters):
    """
    Base class for a concentrator of report that need the name of the \
        concentrator in the values, like S04 and S05.
    """

    @property
    def meters(self):
        """
        Meter objects of this concentrator. The name of concentrator is \
            passed to the meter.

        :return: a list of meter objects
        """
        meters = []
        for meter in self.objectified.Cnt:
            meters.append(self.meter_class(meter, self.name))
        return meters
<commit_msg>Create class that adds magnitude property to meter<commit_after>from datetime import datetime

MAGNITUDE_W = 1
"""
Magnitude value (1) for measures represented in W.
"""

MAGNITUDE_KW = 1000
"""
Magnitude value (1000) for measures represented in kW.
"""

class ValueWithTime(object):
    """
    Base class for values with time.
    """

    def _get_timestamp(self, name, element=None):
        """
        Formats a timestamp from the name of the value.

        :param name: a string with the name
        :param element: an lxml.objectify.StringElement, by default \
            self.objectified
        :return: a formatted string representing a timestamp \
            ('%Y-%m-%d %H:%M:%S')
        """
        if element is None:
            e = self.objectified
        else:
            e = element
        value = e.get(name)
        if len(value) > 15:
            date_value = value[0:14] + value[-1]
        else:
            date_value = value

        # Fix for SAGECOM which puts this timestamp when the period doesn't
        # affect the contracted tariff
        if date_value.upper() == 'FFFFFFFFFFFFFFW':
            date_value = '19000101000000W'

        time = datetime.strptime(date_value[:-1], '%Y%m%d%H%M%S')
        return time.strftime('%Y-%m-%d %H:%M:%S')

class Measure(ValueWithTime):
    """
    Base class for a set of measures.
    """

    def __init__(self, objectified_measure):
        """
        Create a Measure object.

        :param objectified_measure: an lxml.objectify.StringElement \
            representing a set of measures
        :return: a Measure object
        """
        self.objectified = objectified_measure

    @property
    def objectified(self):
        """
        The set of measures as an lxml.objectify.StringElement.

        :return: an lxml.objectify.StringElement representing a set of measures
        """
        return self._objectified

    @objectified.setter
    def objectified(self, value):
        """
        Stores an lxml.objectify.StringElement representing a set of measures.

        :param value: an lxml.objectify.StringElement representing a set of \
            measures
        :return:
        """
        self._objectified = value

class MeasureActiveReactive(Measure):
    """
    Base class for a set of measures with active and reactive measures.
    """

    def active_reactive(self, measure, measure_type):
        """
        Get the active and reactive measures.

        :param measure: an lxml.objectify.StringElement representing a set of \
            measures
        :param measure_type: the type of measure, added at the end of the \
            name of each measure ('a' or 'i')
        :return: a dict with the active and reactive measures
        """
        return {
            'ai': int(measure.get('AI{}'.format(measure_type))),
            'ae': int(measure.get('AE{}'.format(measure_type))),
            'r1': int(measure.get('R1{}'.format(measure_type))),
            'r2': int(measure.get('R2{}'.format(measure_type))),
            'r3': int(measure.get('R3{}'.format(measure_type))),
            'r4': int(measure.get('R4{}'.format(measure_type))),
        }

class Parameter(ValueWithTime):
    """
    Base class for a set of parameters.
    """

    def __init__(self, objectified_parameter, report_version):
        """
        Create a Parameter object.

        :param objectified_parameter: an lxml.objectify.StringElement \
            representing a set of parameters
        :return: a Measure object
        """
        self.objectified = objectified_parameter
        self.report_version = report_version

    @property
    def objectified(self):
        """
        The set of parameters as an lxml.objectify.StringElement.

        :return: an lxml.objectify.StringElement representing a set of \
            parameters
        """
        return self._objectified

    @objectified.setter
    def objectified(self, value):
        """
        Stores an lxml.objectify.StringElement representing a set of \
            parameters.

        :param value: an lxml.objectify.StringElement representing a set of \
            parameters
        :return:
        """
        self._objectified = value

    @property
    def report_version(self):
        """
        The version of the report.

        :return: a string with the version of the report
        """
        return self._report_version

    @report_version.setter
    def report_version(self, value):
        """
        Stores the report version.
        :param value: a string with the version of the report
        """
        self._report_version = value

    def get_boolean(self, name, element=None):
        """
        Gets a boolean value from the name of value.

        :param name: a string with the name
        :param element: an lxml.objectify.StringElement, by default \
            self.objectified
        :return:
        """
        if element is None:
            e = self.objectified
        else:
            e = element
        if e.get(name) == 'Y':
            returns = True
        else:
            returns = False
        return returns

    def to_integer(self, value):
        """
        Convert a value to an integer. If value is None then returns 0.

        :param value: a string with the value
        :return: an integer
        """
        if value is None:
            returns = 0
        else:
            returns = int(value)
        return returns

    def filter_integer(self, value):
        """
        Filter the provided value. Returns the value if is an integer or \
            returns 0 if not.

        :param value: an integer or not
        :return: an integer
        """
        if isinstance(value, int):
            returns = value
        else:
            returns = 0
        return returns

    @property
    def values(self):
        """
        Set of parameters.
        """
        raise NotImplementedError('This method is not implemented!')

class Meter(object):
    """
    Base class for a meter.
    """

    def __init__(self, objectified_meter):
        """
        Create a Meter object.

        :param objectified_meter: an lxml.objectify.StringElement \
            representing a meter
        :return: a Meter object
        """
        self.objectified = objectified_meter

    @property
    def objectified(self):
        """
        A meter as an lxml.objectify.StringElement.

        :return: an lxml.objectify.StringElement representing a meter
        """
        return self._objectified

    @objectified.setter
    def objectified(self, value):
        """
        Stores an lxml.objectify.StringElement representing a meter

        :param value: an lxml.objectify.StringElement representing a meter
        """
        self._objectified = value

    @property
    def errors(self):
        """
        The meter errors.

        :return: a dict with the meter errors
        """
        self._errors = {}
        if self.objectified.get('ErrCat'):
            self._errors = {
                'errcat': self.objectified.get('ErrCat'),
                'errcode': self.objectified.get('ErrCode')
            }
        return self._errors

    @property
    def name(self):
        """
        The name of the meter.

        :return: a string with the name of the meter
        """
        return self.objectified.get('Id')

    @property
    def report_type(self):
        """
        The type of report. To implement in child classes.
        """
        raise NotImplementedError('This method is not implemented!')

    @property
    def measure_class(self):
        """
        The class to instance measures sets.

        :return: a class to instance measure sets
        """
        return Measure

    @property
    def measures(self):
        """
        Measure set objects of this meter.

        :return: a list of measure set objects
        """
        measures = []
        if hasattr(self.objectified, self.report_type):
            objectified = getattr(self.objectified, self.report_type)
            measures = map(self.measure_class, objectified)
        return measures

    @property
    def values(self):
        """
        Values of measure sets of this meter.

        :return: a list with the values of the measure sets
        """
        values = []
        for measure in self.measures:
            values.append(measure.value())
        return values

class MeterWithConcentratorName(Meter):
    """
    Base class for a meters of report that need the name of the concentrator \
        in the values, like S04 and S05.
    """

    def report_type(self):
        """
        The type of report. To implement in child classes.
        """
        raise NotImplementedError('This method is not implemented!')

    def __init__(self, objectified_meter, concentrator_name):
        """
        Create a Meter object using Meter constructor and adding the \
            concentrator name.

        :param objectified_meter: an lxml.objectify.StringElement \
            representing a meter
        :param concentrator_name: a string with the name of the concentrator
        :return: a Meter object
        """
        super(MeterWithConcentratorName, self).__init__(objectified_meter)
        self.concentrator_name = concentrator_name

    @property
    def concentrator_name(self):
        """
        A string with the concentrator name.

        :return: a string with the concentrator name
        """
        return self._concentrator_name

    @concentrator_name.setter
    def concentrator_name(self, value):
        """
        Stores a string with the concentrator name.

        :param value: a string with the concentrator name
        """
        self._concentrator_name = value

    @property
    def values(self):
        """
        Values of measure sets of this meter of report that need the name of \
            the concentrator and the meter,

        :return: a list with the values of the measure sets
        """
        values = []
        for measure in self.measures:
            for subvalue in measure.values:
                v = subvalue.copy()
                v['name'] = self.name
                v['cnc_name'] = self.concentrator_name
                values.append(v)
        return values

class MeterWithMagnitude(MeterWithConcentratorName):

    @property
    def magnitude(self):
        """
        The magnitude of the meter measures.

        :return: a int with the magnitude of the meter measures
        """
        return int(self.objectified.get('Magn'))

class Concentrator(object):
    """
    Base class for a concentrator.
    """

    def __init__(self, objectified_concentrator):
        """
        Create a Concentrator object.

        :param objectified_concentrator: an lxml.objectify.StringElement \
            representing a concentrator
        :return: a Concentrator object
        """
        self.objectified = objectified_concentrator

    @property
    def objectified(self):
        """
        A concentrator as an lxml.objectify.StringElement.

        :return: an lxml.objectify.StringElement representing a concentrator
        """
        return self._objectified

    @objectified.setter
    def objectified(self, value):
        """
        Stores a concentrator as an lxml.objectify.StringElement.

        :param value: an lxml.objectify.StringElement representing a \
            concentrator
        """
        self._objectified = value

    @property
    def name(self):
        """
        The name of the concentrator.

        :return: a string with the name of the concentrator
        """
        return self.objectified.get('Id')

class ConcentratorWithMeters(Concentrator):

    @property
    def meter_class(self):
        """
        The class to instance meters.

        :return: a class to instance meters
        """
        return Meter

    @property
    def meters(self):
        """
        Meter objects of this concentrator.

        :return: a list of meter objects
        """
        return map(self.meter_class, self.objectified.Cnt)

    @property
    def values(self):
        """
        Values of the meters of this concentrator.

        :return: a list with the values of the meters
        """
        values = []
        for meter in self.meters:
            values.extend(meter.values)
        return values

class ConcentratorWithMetersWithConcentratorName(ConcentratorWithMeters):
    """
    Base class for a concentrator of report that need the name of the \
        concentrator in the values, like S04 and S05.
    """

    @property
    def meters(self):
        """
        Meter objects of this concentrator. The name of concentrator is \
            passed to the meter.

        :return: a list of meter objects
        """
        meters = []
        for meter in self.objectified.Cnt:
            meters.append(self.meter_class(meter, self.name))
        return meters
<|endoftext|>


<commit_before>from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Third-party
from astroML.utils import log_multivariate_gaussian
import numpy as np
from scipy.interpolate import splprep, splev
from scipy.misc import logsumexp

def ps1_data_to_X_cov(data, W):
    """
    Given Pan-STARRS1 photometry -- which must contain keys like
    'dered_g' and 'gErr', etc. -- and a mixing matrix, `W`, create
    a feature / data matrix `X` and a covariance matrix `V`.
    """
    X = np.vstack([data['dered_{}'.format(band)] for band in 'grizy']).T
    Xerr = np.vstack([data['{}Err'.format(band)] for band in 'grizy']).T

    X = np.einsum('nj,mj->nm', X, W)

    # compute error covariance with mixing matrix
    V = np.zeros(Xerr.shape + Xerr.shape[-1:])
    for i in range(Xerr.shape[1]):
        V[:,i,i] = Xerr[:,i]**2

    # each covariance C = WCW^T
    V = np.einsum('mj,njk->nmk', W, V)
    V = np.einsum('lk,nmk->nml', W, V)

    return X, V

def ps1_isoc_to_XCov(data, W, interpolate=False, n_interpolate=1024):
    """
    Given an isochrone in Pan-STARRS1 filters -- which must contain keys like
    'gP1', etc. -- and a mixing matrix, `W`, create a feature / data matrix `X`.
    """
    X = np.vstack([data['{}P1'.format(band)] for band in 'grizy']).T
    X = np.einsum('nj,mj->nm', X, W)

    if interpolate:
        tck, u = splprep(X[:,1:].T, u=X[:,0], k=3, s=1E-4)
        u_fine = np.linspace(u.min(), u.max(), n_interpolate)
        X = np.vstack([u_fine] + splev(u_fine, tck)).T

    return X

def likelihood_worker(allX, allCov, otherX, otherCov=None, smooth=None):
    if otherCov is not None:
        V = allCov[np.newaxis] + otherCov[:,np.newaxis]

        if smooth is not None:
            H = np.eye(allCov.shape[-1]) * smooth**2
            V += H[np.newaxis,np.newaxis]

        ll = log_multivariate_gaussian(allX[np.newaxis], otherX[:,np.newaxis], V)

    else:
        V = allCov
        if smooth is not None:
            H = np.eye(allCov.shape[-1]) * smooth**2
            V += H[np.newaxis]
        ll = log_multivariate_gaussian(allX[np.newaxis], otherX[:,np.newaxis], V[np.newaxis])

    ll = logsumexp(ll, axis=-1) # NOTE: could also max here
    return ll
<commit_msg>change index name<commit_after>from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Third-party
from astroML.utils import log_multivariate_gaussian
import numpy as np
from scipy.interpolate import splprep, splev
from scipy.misc import logsumexp

def ps1_data_to_X_cov(data, W):
    """
    Given Pan-STARRS1 photometry -- which must contain keys like
    'dered_g' and 'gErr', etc. -- and a mixing matrix, `W`, create
    a feature / data matrix `X` and a covariance matrix `V`.
    """
    X = np.vstack([data['dered_{}'.format(band)] for band in 'grizy']).T
    Xerr = np.vstack([data['{}Err'.format(band)] for band in 'grizy']).T

    X = np.einsum('ij,kj->ik', X, W)

    # compute error covariance with mixing matrix
    V = np.zeros(Xerr.shape + Xerr.shape[-1:])
    for i in range(Xerr.shape[1]):
        V[:,i,i] = Xerr[:,i]**2

    # each covariance C = WCW^T
    V = np.einsum('mj,njk->nmk', W, V)
    V = np.einsum('lk,nmk->nml', W, V)

    return X, V

def ps1_isoc_to_XCov(data, W, interpolate=False, n_interpolate=1024):
    """
    Given an isochrone in Pan-STARRS1 filters -- which must contain keys like
    'gP1', etc. -- and a mixing matrix, `W`, create a feature / data matrix `X`.
    """
    X = np.vstack([data['{}P1'.format(band)] for band in 'grizy']).T
    X = np.einsum('nj,mj->nm', X, W)

    if interpolate:
        tck, u = splprep(X[:,1:].T, u=X[:,0], k=3, s=1E-4)
        u_fine = np.linspace(u.min(), u.max(), n_interpolate)
        X = np.vstack([u_fine] + splev(u_fine, tck)).T

    return X

def likelihood_worker(allX, allCov, otherX, otherCov=None, smooth=None):
    if otherCov is not None:
        V = allCov[np.newaxis] + otherCov[:,np.newaxis]

        if smooth is not None:
            H = np.eye(allCov.shape[-1]) * smooth**2
            V += H[np.newaxis,np.newaxis]

        ll = log_multivariate_gaussian(allX[np.newaxis], otherX[:,np.newaxis], V)

    else:
        V = allCov
        if smooth is not None:
            H = np.eye(allCov.shape[-1]) * smooth**2
            V += H[np.newaxis]
        ll = log_multivariate_gaussian(allX[np.newaxis], otherX[:,np.newaxis], V[np.newaxis])

    ll = logsumexp(ll, axis=-1) # NOTE: could also max here
    return ll
<|endoftext|>


<commit_before>import tarfile
from gzip import GzipFile
from boto import connect_s3
from ConfigParser import SafeConfigParser, RawConfigParser
from datetime import datetime
from optparse import OptionParser
from os.path import expanduser
from os import unlink, fdopen
from psycopg2 import connect
from sys import stderr
from tempfile import mkstemp
from zc.lockfile import LockFile

def add_s3_options_to_parser(parser):
    parser.add_option("-b", "--bucket", dest="bucket", help="S3 Bucket")
    parser.add_option("-p", "--prefix", dest="prefix", help="Backup prefix (e.g. backup/dbname/)")
    parser.add_option("-a", "--access-key", dest='access_key', help="AWS access key")
    parser.add_option("-s", "--secret-key", dest='secret_key', help="NOT RECOMMENDED, use config files to set credentials. AWS secret key")

def update_config_for_s3(config,options,tool):
    keys = ['bucket', 'prefix', 'access_key', 'secret_key']
    for k in keys:
        if k not in config:
            config[k] = ''

    conf_parse = SafeConfigParser(config)
    conf_parse.read(['/etc/backuppsqls3.conf',expanduser('~/.backuppsqls3.conf'),options.config])

    for section in ["general", tool]:
        if conf_parse.has_section(section):
            for opt in config.keys():
                if conf_parse.has_option(section,opt):
                    config[opt] = conf_parse.get(section,opt)

                if getattr(options, opt) is not None:
                    config[opt] = getattr(options, opt)

def get_wal_config(tool,usage):

    parser = OptionParser(usage=usage)
    parser.add_option("-c", "--config", dest="config", default="", help="Configuration file")
    add_s3_options_to_parser(parser)
    (options, args) = parser.parse_args()

    if len(args) != 2:
        parser.print_usage()
        exit(1)

    config = {}
    update_config_for_s3(config,options,tool)
    return config, args

def backup_wal():
    config, args = get_wal_config("backup_wal", "usage: %prog [options] source_file file_name")
    source_file,file_name = args

    fi,t_file_name = mkstemp()
    fp = fdopen(fi,'w')
    gzf = GzipFile(fileobj=fp)
    gzf.writelines(open(source_file))
    gzf.close()
    fp.close()

    conn = connect_s3(config['access_key'], config['secret_key'])
    bucket = conn.get_bucket(config['bucket'])
    key = bucket.new_key(config['prefix']+file_name+".gz")
    key.set_contents_from_filename(t_file_name)

    unlink(t_file_name)

def restore_wal():
    config, args = get_wal_config("restore_wal", "usage: %prog [options] source_file target_path")
    source_file,target_path = args

    conn = connect_s3(config['access_key'], config['secret_key'])
    bucket = conn.get_bucket(config['bucket'])

    key = bucket.get_key(config['prefix']+source_file)
    if key is not None:
        key.get_contents_to_filename(target_path)
    else:
        print >> stderr, "File not found"
        exit(1)

def backup():
    usage = "usage: %prog [options]"
    parser = OptionParser(usage=usage)
    parser.add_option("-c", "--config", dest="config", default="", help="Configuration file")
    add_s3_options_to_parser(parser)
    parser.add_option("-l", "--lock-file", dest='lock_file', default='/tmp/backuppsqls3.lock')
    parser.add_option("--dump", dest="do_dump", default=False, action="store_true", help="Perform dump")
    parser.add_option("--dump-databases", dest="databases", default="ALL", help="Which databases to backup when a dump is done")
    parser.add_option("--base-backup", dest="do_base_backup", default=False, action="store_true", help="Perform base backup for PITR and check configuration")
    parser.add_option("--postgresql-conf", dest="postgresql_conf", help="PostgreSQL configuration file (used for base backup)")
    parser.add_option("--postgresql-user", dest="postgresql_user", help="PostgreSQL super user (default: current system user)")
    parser.add_option("--postgresql-password", dest="postgresql_password", help="Password for PostgreSQL user (default, use system ident)")

    (options, args) = parser.parse_args()
    config = {}

    update_config_for_s3(config,options,"backup")

    sec, postgres_conf = _read_postgres_config(options.postgresql_conf)
    test_archive_config(sec,postgres_conf,postgres_conf)

    config['lock_file'] = options.lock_file
    config['data_dir'] = postgres_conf.get(sec,'data_directory').replace('\'','').split('#')[0].strip()

    _do_base_backup(config)

def _read_postgres_config(config_file):
    fi, file_name = mkstemp()
    fp = fdopen(fi, 'w')
    sec = "sectionhead"
    print >> fp, "["+sec+"]"
    for line in open(config_file):
        print >> fp, line.strip()
    fp.close()
    s = RawConfigParser()
    s.read(file_name)
    unlink(file_name)
    return sec, s

def test_archive_config(sec,s,config_file):
    if not (s.has_option(sec,"archive_mode") and s.get(sec,"archive_mode").startswith('on') and s.has_option(sec, 'archive_command') and 'pg_backup_wal_s3' in s.get(sec, 'archive_command')):
        print "Please enable archiving in %s" % config_file
        print "Suggested entries:"
        print "  wal_level = archive"
        print "  archive_mode = on"
        print "  archive_command = 'pg_backup_wal_s3 %p %f'"
        exit(1)
    return True

def _do_base_backup(config):
    lock = LockFile(config['lock_file'])
    conn = connect("")
    cur = conn.cursor()
    label = datetime.now().strftime("%Y%m%d%H%M%S")
    cur.execute("SELECT pg_start_backup('%s');" % label)

    fi,file_name = mkstemp()
    fp = fdopen(fi,'w')
    tar_f = tarfile.open(fileobj=fp,mode='w:gz')
    tar_f.add(config['data_dir'],arcname="data",exclude=lambda x: '/pg_xlog' in x)
    tar_f.close()
    fp.close()

    conn = connect_s3(config['access_key'], config['secret_key'])
    bucket = conn.get_bucket(config['bucket'])
    key = bucket.new_key(config['prefix']+label+".tar.gz")
    key.set_contents_from_filename(file_name)

    unlink(file_name)
    cur.execute("SELECT pg_stop_backup();")
    lock.close()

def run():
    pass
<commit_msg>Also the dump part is working<commit_after>import tarfile
from gzip import GzipFile
from subprocess import  call
from boto import connect_s3
from ConfigParser import SafeConfigParser, RawConfigParser
from datetime import datetime
from optparse import OptionParser
import os
from os.path import expanduser
from os import unlink, fdopen
from psycopg2 import connect
from sys import stderr
from tempfile import mkstemp
from zc.lockfile import LockFile

def add_s3_options_to_parser(parser):
    parser.add_option("-b", "--bucket", dest="bucket", help="S3 Bucket")
    parser.add_option("-p", "--prefix", dest="prefix", help="Backup prefix (e.g. backup/dbname/)")
    parser.add_option("-a", "--access-key", dest='access_key', help="AWS access key")
    parser.add_option("-s", "--secret-key", dest='secret_key', help="NOT RECOMMENDED, use config files to set credentials. AWS secret key")

def update_config_for_s3(config,options,tool):
    keys = ['bucket', 'prefix', 'access_key', 'secret_key']
    for k in keys:
        if k not in config:
            config[k] = ''

    conf_parse = SafeConfigParser(config)
    conf_parse.read(['/etc/backuppsqls3.conf',expanduser('~/.backuppsqls3.conf'),options.config])

    for section in ["general", tool]:
        if conf_parse.has_section(section):
            for opt in config.keys():
                if conf_parse.has_option(section,opt):
                    config[opt] = conf_parse.get(section,opt)

                if getattr(options, opt) is not None:
                    config[opt] = getattr(options, opt)

def get_wal_config(tool,usage):

    parser = OptionParser(usage=usage)
    parser.add_option("-c", "--config", dest="config", default="", help="Configuration file")
    add_s3_options_to_parser(parser)
    (options, args) = parser.parse_args()

    if len(args) != 2:
        parser.print_usage()
        exit(1)

    config = {}
    update_config_for_s3(config,options,tool)
    return config, args

def backup_wal():
    config, args = get_wal_config("backup_wal", "usage: %prog [options] source_file file_name")
    source_file,file_name = args

    fi,t_file_name = mkstemp()
    fp = fdopen(fi,'w')
    gzf = GzipFile(fileobj=fp)
    gzf.writelines(open(source_file))
    gzf.close()
    fp.close()

    conn = connect_s3(config['access_key'], config['secret_key'])
    bucket = conn.get_bucket(config['bucket'])
    key = bucket.new_key(config['prefix']+file_name+".gz")
    key.set_contents_from_filename(t_file_name)

    unlink(t_file_name)

def restore_wal():
    config, args = get_wal_config("restore_wal", "usage: %prog [options] source_file target_path")
    source_file,target_path = args

    conn = connect_s3(config['access_key'], config['secret_key'])
    bucket = conn.get_bucket(config['bucket'])

    key = bucket.get_key(config['prefix']+source_file+".gz")
    if key is not None:
        fi,t_file_name = mkstemp()
        os.close(fi)
        key.get_contents_to_filename(t_file_name)
        gzf = GzipFile(t_file_name,'r')
        a = open(target_path,'w')
        a.write(gzf.read())
        a.close()
        gzf.close()
        unlink(t_file_name)
    else:
        print >> stderr, "File not found"
        exit(1)

def backup():
    usage = "usage: %prog [options]"
    parser = OptionParser(usage=usage)
    parser.add_option("-c", "--config", dest="config", default="", help="Configuration file")
    add_s3_options_to_parser(parser)
    parser.add_option("-l", "--lock-file", dest='lock_file', default='/tmp/backuppsqls3.lock')
    parser.add_option("--dump", dest="do_dump", default=False, action="store_true", help="Perform dump")
#    parser.add_option("--dump-databases", dest="databases", default="ALL", help="Which databases to backup when a dump is done")
    parser.add_option("--base-backup", dest="do_base_backup", default=False, action="store_true", help="Perform base backup for PITR and check configuration")
    parser.add_option("--postgresql-conf", dest="postgresql_conf", help="PostgreSQL configuration file (used for base backup)")
    parser.add_option("--postgresql-user", dest="postgresql_user", help="PostgreSQL super user (default: current system user)")
    parser.add_option("--postgresql-password", dest="postgresql_password", help="Password for PostgreSQL user (default, use system ident)")

    (options, args) = parser.parse_args()
    config = {}

    update_config_for_s3(config,options,"backup")

    sec, postgres_conf = _read_postgres_config(options.postgresql_conf)
    test_archive_config(sec,postgres_conf,postgres_conf)

    config['lock_file'] = options.lock_file
    config['data_dir'] = postgres_conf.get(sec,'data_directory').replace('\'','').split('#')[0].strip()

    if options.do_base_backup:
        _do_base_backup(config)
    if options.do_dump:
        _do_dump_backup(config)

def _read_postgres_config(config_file):
    fi, file_name = mkstemp()
    fp = fdopen(fi, 'w')
    sec = "sectionhead"
    print >> fp, "["+sec+"]"
    for line in open(config_file):
        print >> fp, line.strip()
    fp.close()
    s = RawConfigParser()
    s.read(file_name)
    unlink(file_name)
    return sec, s

def test_archive_config(sec,s,config_file):
    if not (s.has_option(sec,"archive_mode") and s.get(sec,"archive_mode").startswith('on') and s.has_option(sec, 'archive_command') and 'pg_backup_wal_s3' in s.get(sec, 'archive_command')):
        print "Please enable archiving in %s" % config_file
        print "Suggested entries:"
        print "  wal_level = archive"
        print "  archive_mode = on"
        print "  archive_command = 'pg_backup_wal_s3 %p %f'"
        exit(1)
    return True

def _do_base_backup(config):
    lock = LockFile(config['lock_file'])
    conn = connect("")
    cur = conn.cursor()
    label = datetime.now().strftime("%Y%m%d%H%M%S")
    cur.execute("SELECT pg_start_backup('%s');" % label)

    fi,file_name = mkstemp()
    fp = fdopen(fi,'w')
    tar_f = tarfile.open(fileobj=fp,mode='w:gz')
    tar_f.add(config['data_dir'],arcname="",exclude=lambda x: '/pg_xlog' in x)
    tar_f.close()
    fp.close()

    conn = connect_s3(config['access_key'], config['secret_key'])
    bucket = conn.get_bucket(config['bucket'])
    key = bucket.new_key(config['prefix']+"base_"+label+".tar.gz")
    key.set_contents_from_filename(file_name)

    unlink(file_name)
    cur.execute("SELECT pg_stop_backup();")
    lock.close()

def _do_dump_backup(config):
    label = datetime.now().strftime("%Y%m%d%H%M%S")

    fi,file_name = mkstemp()
    fp = fdopen(fi,'w')

    if call(["pg_dumpall"], stdout=fp) is not 0:
        print >> stderr, "Dump failed"
        exit(1)
    fp.close()

    fi2,file_name2 = mkstemp()

    f_in = open(file_name, 'rb')
    f_out = GzipFile(file_name2, 'wb')
    f_out.writelines(f_in)
    f_out.close()
    f_in.close()

    conn = connect_s3(config['access_key'], config['secret_key'])
    bucket = conn.get_bucket(config['bucket'])
    key = bucket.new_key(config['prefix']+"dump_"+label+".gz")
    key.set_contents_from_filename(file_name2)

    unlink(file_name)
    unlink(file_name2)

def run():
    pass
<|endoftext|>