<commit_before>"""
The MIT License (MIT)

Copyright (c) 2015-2016 Olof Kindgren <olof.kindgren@gmail.com>

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""

import yaml
import xml.etree.ElementTree as ET

from . import ipxact_yaml
import sys
if sys.version_info[0] > 2:
    UNICODE = 'unicode'
else:
    UNICODE = 'utf-8'

class IpxactInt(int):
    def __new__(cls, *args, **kwargs):
        if not args:
            return super(IpxactInt, cls).__new__(cls)

        expr = args[0]
        base = 10

        if len(expr) > 2 and expr[0:2] == '0x':
            # handle non-standard SystemVerilog (but commonly-used) syntax
            expr = expr[2:]
            base = 16
        elif "'" in expr:
            sep = expr.find("'")
            # ignore any bit size specified before ' because size is handled by other IP-XACT properties
            if expr[sep+1] in ["h", "H"]:
                base = 16
            elif expr[sep+1] in ["d", "D"]:
                base = 10
            elif expr[sep+1] in ["o", "O"]:
                base = 8
            elif expr[sep+1] in ["b", "B"]:
                base = 2
            else:
                raise ValueError("Could not convert expression to an integer: {}".format(args[0]))
            expr = expr[sep+2:]

        return super(IpxactInt, cls).__new__(cls, expr.replace('_', ''), base)

class IpxactBool(str):
    def __new__(cls, *args, **kwargs):
        if not args:
            return None
        elif args[0] in ['true', 'false']:
            return super(IpxactBool, cls).__new__(cls, args[0])
        else:
            raise Exception

class IpxactItem(object):
    nsmap = {'1.4'  : ('spirit' , 'http://www.spiritconsortium.org/XMLSchema/SPIRIT/1.4'),
             '1.5'  : ('spirit' , 'http://www.spiritconsortium.org/XMLSchema/SPIRIT/1.5'),
             '2009' : ('spirit' , 'http://www.spiritconsortium.org/XMLSchema/SPIRIT/1685-2009'),
             '2014' : ('ipxact' , 'http://www.accellera.org/XMLSchema/IPXACT/1685-2014'),
    }
    nsversion = '1.5'

    ATTRIBS = {}
    MEMBERS = {}
    CHILDREN = []
    CHILD = []
    def __init__(self, **kwargs):
        for k, v in kwargs.items():
            if k in self.MEMBERS:
                setattr(self, k, v)
            elif k in self.ATTRIBS:
                setattr(self, k, v)
            else:
                s = "{} has no member or attribute '{}'"
                raise AttributeError(s.format(self.__class__.__name__, k))

    def load(self, f):
        tree = ET.parse(f)
        root = tree.getroot()

        #Warning: Semi-horrible hack to find out which IP-Xact version that is used
        for key, value in self.nsmap.items():
            if root.tag[1:].startswith(value[1]):
                self.nsversion = key

        S = '{%s}' % self.nsmap[self.nsversion][1]
        if not (root.tag == S+self._tag):
            raise Exception

        self.parse_tree(root, self.nsmap[self.nsversion])

    def parse_tree(self, root, ns):
        if self.ATTRIBS:
            for _name, _type in self.ATTRIBS.items():
                _tagname = '{' + ns[1] + '}' + _name
                if _tagname in root.attrib:
                    setattr(self, _name, eval(_type)(root.attrib[_tagname]))
        for _name, _type in self.MEMBERS.items():
            tmp = root.find('./{}:{}'.format(ns[0], _name), {ns[0] : ns[1]})
            if tmp is not None and tmp.text is not None:
                setattr(self, _name, eval(_type)(tmp.text))
            else:
                setattr(self, _name, eval(_type)())

        for c in self.CHILDREN:
            for f in root.findall("./{}:{}".format(ns[0], c), {ns[0] : ns[1]}):
                child = getattr(self, c)[:]
                class_name = c[0].upper() + c[1:]
                t = __import__(self.__module__)
                t = getattr(t, 'ipyxact')
                t = getattr(t, class_name)()
                #t = eval(class_name)()
                t.parse_tree(f, ns)
                child.append(t)
                setattr(self, c, child)
        for c in self.CHILD:
            f = root.find("./{}:{}".format(ns[0], c), {ns[0] : ns[1]})
            if f is not None:
                child = getattr(self, c)
                class_name = c[0].upper() + c[1:]
                t = eval(class_name)()
                t.parse_tree(f, ns)
                setattr(self, c, t)

    def _write(self, root, S):
        for a in self.ATTRIBS:
            root.attrib[S+a] = getattr(self, a)

        for m in self.MEMBERS:
            tmp = getattr(self, m)
            if tmp is not None:
                ET.SubElement(root, S+m).text = str(tmp)

        for c in self.CHILDREN:
            for child_obj in getattr(self, c):
                subel = ET.SubElement(root, S+c)
                child_obj._write(subel, S)
        for c in self.CHILD:
            tmp = getattr(self, c)
            if tmp is not None:
                subel = ET.SubElement(root, S+c)
                tmp._write(subel, S)

    def write(self, f):
        ET.register_namespace(self.nsmap[self.nsversion][0], self.nsmap[self.nsversion][1])
        S = '{%s}' % self.nsmap[self.nsversion][1]
        root = ET.Element(S+self._tag)
        self._write(root, S)

        et = ET.ElementTree(root)
        et.write(f, xml_declaration=True, encoding=UNICODE)

def _generate_classes(j):
    for tag, _items in j.items():
        if 'ATTRIBS' in _items:
            for key, value in _items['ATTRIBS'].items():
                _items.update({key : eval(value)})
        if 'MEMBERS' in _items:
            for key, value in _items['MEMBERS'].items():
                _items.update({key : None})
        if 'CHILDREN' in _items:
            for key in _items['CHILDREN']:
                _items.update({key : []})
        if 'CHILD' in _items:
            for key in _items['CHILD']:
                _items.update({key : None})
        _items.update({'_tag' : tag})

        generatedClass = type(tag[0].upper()+tag[1:], (IpxactItem,), _items)
        globals()[generatedClass.__name__] = generatedClass

_generate_classes(yaml.load(ipxact_yaml.description))
<commit_msg>Return an int object to avoid errors with large numbers<commit_after>"""
The MIT License (MIT)

Copyright (c) 2015-2016 Olof Kindgren <olof.kindgren@gmail.com>

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.
"""

import yaml
import xml.etree.ElementTree as ET

from . import ipxact_yaml
import sys
if sys.version_info[0] > 2:
    UNICODE = 'unicode'
else:
    UNICODE = 'utf-8'

class IpxactInt(int):
    def __new__(cls, *args, **kwargs):
        if not args:
            return int()

        expr = args[0]
        base = 10

        if len(expr) > 2 and expr[0:2] == '0x':
            # handle non-standard SystemVerilog (but commonly-used) syntax
            expr = expr[2:]
            base = 16
        elif "'" in expr:
            sep = expr.find("'")
            # ignore any bit size specified before ' because size is handled by other IP-XACT properties
            if expr[sep+1] in ["h", "H"]:
                base = 16
            elif expr[sep+1] in ["d", "D"]:
                base = 10
            elif expr[sep+1] in ["o", "O"]:
                base = 8
            elif expr[sep+1] in ["b", "B"]:
                base = 2
            else:
                raise ValueError("Could not convert expression to an integer: {}".format(args[0]))
            expr = expr[sep+2:]

        return int(expr.replace('_', ''), base)

class IpxactBool(str):
    def __new__(cls, *args, **kwargs):
        if not args:
            return None
        elif args[0] in ['true', 'false']:
            return super(IpxactBool, cls).__new__(cls, args[0])
        else:
            raise Exception

class IpxactItem(object):
    nsmap = {'1.4'  : ('spirit' , 'http://www.spiritconsortium.org/XMLSchema/SPIRIT/1.4'),
             '1.5'  : ('spirit' , 'http://www.spiritconsortium.org/XMLSchema/SPIRIT/1.5'),
             '2009' : ('spirit' , 'http://www.spiritconsortium.org/XMLSchema/SPIRIT/1685-2009'),
             '2014' : ('ipxact' , 'http://www.accellera.org/XMLSchema/IPXACT/1685-2014'),
    }
    nsversion = '1.5'

    ATTRIBS = {}
    MEMBERS = {}
    CHILDREN = []
    CHILD = []
    def __init__(self, **kwargs):
        for k, v in kwargs.items():
            if k in self.MEMBERS:
                setattr(self, k, v)
            elif k in self.ATTRIBS:
                setattr(self, k, v)
            else:
                s = "{} has no member or attribute '{}'"
                raise AttributeError(s.format(self.__class__.__name__, k))

    def load(self, f):
        tree = ET.parse(f)
        root = tree.getroot()

        #Warning: Semi-horrible hack to find out which IP-Xact version that is used
        for key, value in self.nsmap.items():
            if root.tag[1:].startswith(value[1]):
                self.nsversion = key

        S = '{%s}' % self.nsmap[self.nsversion][1]
        if not (root.tag == S+self._tag):
            raise Exception

        self.parse_tree(root, self.nsmap[self.nsversion])

    def parse_tree(self, root, ns):
        if self.ATTRIBS:
            for _name, _type in self.ATTRIBS.items():
                _tagname = '{' + ns[1] + '}' + _name
                if _tagname in root.attrib:
                    setattr(self, _name, eval(_type)(root.attrib[_tagname]))
        for _name, _type in self.MEMBERS.items():
            tmp = root.find('./{}:{}'.format(ns[0], _name), {ns[0] : ns[1]})
            if tmp is not None and tmp.text is not None:
                setattr(self, _name, eval(_type)(tmp.text))
            else:
                setattr(self, _name, eval(_type)())

        for c in self.CHILDREN:
            for f in root.findall("./{}:{}".format(ns[0], c), {ns[0] : ns[1]}):
                child = getattr(self, c)[:]
                class_name = c[0].upper() + c[1:]
                t = __import__(self.__module__)
                t = getattr(t, 'ipyxact')
                t = getattr(t, class_name)()
                #t = eval(class_name)()
                t.parse_tree(f, ns)
                child.append(t)
                setattr(self, c, child)
        for c in self.CHILD:
            f = root.find("./{}:{}".format(ns[0], c), {ns[0] : ns[1]})
            if f is not None:
                child = getattr(self, c)
                class_name = c[0].upper() + c[1:]
                t = eval(class_name)()
                t.parse_tree(f, ns)
                setattr(self, c, t)

    def _write(self, root, S):
        for a in self.ATTRIBS:
            root.attrib[S+a] = getattr(self, a)

        for m in self.MEMBERS:
            tmp = getattr(self, m)
            if tmp is not None:
                ET.SubElement(root, S+m).text = str(tmp)

        for c in self.CHILDREN:
            for child_obj in getattr(self, c):
                subel = ET.SubElement(root, S+c)
                child_obj._write(subel, S)
        for c in self.CHILD:
            tmp = getattr(self, c)
            if tmp is not None:
                subel = ET.SubElement(root, S+c)
                tmp._write(subel, S)

    def write(self, f):
        ET.register_namespace(self.nsmap[self.nsversion][0], self.nsmap[self.nsversion][1])
        S = '{%s}' % self.nsmap[self.nsversion][1]
        root = ET.Element(S+self._tag)
        self._write(root, S)

        et = ET.ElementTree(root)
        et.write(f, xml_declaration=True, encoding=UNICODE)

def _generate_classes(j):
    for tag, _items in j.items():
        if 'ATTRIBS' in _items:
            for key, value in _items['ATTRIBS'].items():
                _items.update({key : eval(value)})
        if 'MEMBERS' in _items:
            for key, value in _items['MEMBERS'].items():
                _items.update({key : None})
        if 'CHILDREN' in _items:
            for key in _items['CHILDREN']:
                _items.update({key : []})
        if 'CHILD' in _items:
            for key in _items['CHILD']:
                _items.update({key : None})
        _items.update({'_tag' : tag})

        generatedClass = type(tag[0].upper()+tag[1:], (IpxactItem,), _items)
        globals()[generatedClass.__name__] = generatedClass

_generate_classes(yaml.load(ipxact_yaml.description))
<|endoftext|>


<commit_before>import os
import dj_database_url

print(type(os.environ['HEROKU']))

if os.environ['HEROKU'] is 'True':
    SECRET_KEY = os.environ['SECRET_KEY']
    TWITCH_API_KEY = os.environ['TWITCH_API_KEY']
    YOUTUBE_API_KEY = os.environ['YOUTUBE_API_KEY']
    DEBUG = False
else:
    from config.config import AdminConfig
    SECRET_KEY = AdminConfig.SECRET_KEY
    TWITCH_API_KEY = AdminConfig.TWITCH_API_KEY
    YOUTUBE_API_KEY = AdminConfig.YOUTUBE_API_KEY
    DEBUG = True

BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))

ALLOWED_HOSTS = [
    'jestemgraczem.pl',
    'localhost',
]

ADMINS = [
    ('otlet', 'otlet@otlet.pl')
]

INSTALLED_APPS = [
    'service.apps.ServiceConfig',
    'stream.apps.StreamConfig',
    'django.contrib.admin',
    'django.contrib.auth',
    'django.contrib.contenttypes',
    'django.contrib.sessions',
    'django.contrib.messages',
    'django.contrib.staticfiles',
    'rest_framework',
]

MIDDLEWARE = [
    'django.middleware.security.SecurityMiddleware',
    'whitenoise.middleware.WhiteNoiseMiddleware',
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.common.CommonMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
    'django.contrib.messages.middleware.MessageMiddleware',
    'django.middleware.clickjacking.XFrameOptionsMiddleware',
]

ROOT_URLCONF = 'JestemGraczem.urls'

TEMPLATES = [
    {
        'BACKEND': 'django.template.backends.django.DjangoTemplates',
        'DIRS': [
            BASE_DIR + '/templates/',
        ],
        'APP_DIRS': True,
        'OPTIONS': {
            'context_processors': [
                'django.template.context_processors.debug',
                'django.template.context_processors.request',
                'django.contrib.auth.context_processors.auth',
                'django.contrib.messages.context_processors.messages',
                'service.custom_template_variable.settings_processor'
            ],
            'debug': DEBUG,
        },
    },
]

WSGI_APPLICATION = 'JestemGraczem.wsgi.application'

DATABASES = {
    'default': {
        'ENGINE': 'django.db.backends.sqlite3',
        'NAME': os.path.join(BASE_DIR, 'db.sqlite3'),
    }
}

REST_FRAMEWORK = {
    'DEFAULT_PERMISSION_CLASSES': [
        'rest_framework.permissions.DjangoModelPermissionsOrAnonReadOnly',
    ],
    'DEFAULT_PAGINATION_CLASS': 'rest_framework.pagination.LimitOffsetPagination',
    'PAGE_SIZE': 10
}

AUTH_PASSWORD_VALIDATORS = [
    {
        'NAME': 'django.contrib.auth.password_validation.UserAttributeSimilarityValidator',
    },
    {
        'NAME': 'django.contrib.auth.password_validation.MinimumLengthValidator',
    },
    {
        'NAME': 'django.contrib.auth.password_validation.CommonPasswordValidator',
    },
    {
        'NAME': 'django.contrib.auth.password_validation.NumericPasswordValidator',
    },
]

LANGUAGE_CODE = 'pl-pl'
TIME_ZONE = 'Europe/Warsaw'
USE_I18N = True
USE_L10N = True
USE_TZ = True

db_from_env = dj_database_url.config(conn_max_age=500)

STATIC_ROOT = os.path.join(BASE_DIR, 'staticfiles')
STATIC_URL = '/static/'

SECURE_PROXY_SSL_HEADER = ('HTTP_X_FORWARDED_PROTO', 'https')

STATICFILES_DIRS = (
    os.path.join(BASE_DIR, 'static'),
)

STATICFILES_STORAGE = [
    'whitenoise.storage.CompressedManifestStaticFilesStorage',
]

CACHES = {
    'default': {
        'BACKEND': 'django.core.cache.backends.filebased.FileBasedCache',
        'LOCATION': os.path.join(BASE_DIR, 'cache').replace('\\', '/'),
        'TIMEOUT': 600,
        'OPTIONS': {
            'MAX_ENTRIES': 1000
        }
    }
}
<commit_msg>HEROKU DEPLOY<commit_after>import os
import dj_database_url

if os.environ['HEROKU'] == 'True':
    SECRET_KEY = os.environ['SECRET_KEY']
    TWITCH_API_KEY = os.environ['TWITCH_API_KEY']
    YOUTUBE_API_KEY = os.environ['YOUTUBE_API_KEY']
    DEBUG = False
else:
    from config.config import AdminConfig
    SECRET_KEY = AdminConfig.SECRET_KEY
    TWITCH_API_KEY = AdminConfig.TWITCH_API_KEY
    YOUTUBE_API_KEY = AdminConfig.YOUTUBE_API_KEY
    DEBUG = True

BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
PROJECT_ROOT = os.path.dirname(os.path.abspath(__file__))

ALLOWED_HOSTS = [
    'jestemgraczem.pl',
    'localhost',
]

ADMINS = [
    ('otlet', 'otlet@otlet.pl')
]

INSTALLED_APPS = [
    'service.apps.ServiceConfig',
    'stream.apps.StreamConfig',
    'django.contrib.admin',
    'django.contrib.auth',
    'django.contrib.contenttypes',
    'django.contrib.sessions',
    'django.contrib.messages',
    'django.contrib.staticfiles',
    'rest_framework',
]

MIDDLEWARE = [
    'django.middleware.security.SecurityMiddleware',
    'whitenoise.middleware.WhiteNoiseMiddleware',
    'django.contrib.sessions.middleware.SessionMiddleware',
    'django.middleware.common.CommonMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.contrib.auth.middleware.AuthenticationMiddleware',
    'django.contrib.messages.middleware.MessageMiddleware',
    'django.middleware.clickjacking.XFrameOptionsMiddleware',
]

ROOT_URLCONF = 'JestemGraczem.urls'

TEMPLATES = [
    {
        'BACKEND': 'django.template.backends.django.DjangoTemplates',
        'DIRS': [
            BASE_DIR + '/templates/',
        ],
        'APP_DIRS': True,
        'OPTIONS': {
            'context_processors': [
                'django.template.context_processors.debug',
                'django.template.context_processors.request',
                'django.contrib.auth.context_processors.auth',
                'django.contrib.messages.context_processors.messages',
                'service.custom_template_variable.settings_processor'
            ],
            'debug': DEBUG,
        },
    },
]

WSGI_APPLICATION = 'JestemGraczem.wsgi.application'

DATABASES = {
    'default': {
        'ENGINE': 'django.db.backends.sqlite3',
        'NAME': os.path.join(BASE_DIR, 'db.sqlite3'),
    }
}

REST_FRAMEWORK = {
    'DEFAULT_PERMISSION_CLASSES': [
        'rest_framework.permissions.DjangoModelPermissionsOrAnonReadOnly',
    ],
    'DEFAULT_PAGINATION_CLASS': 'rest_framework.pagination.LimitOffsetPagination',
    'PAGE_SIZE': 10
}

AUTH_PASSWORD_VALIDATORS = [
    {
        'NAME': 'django.contrib.auth.password_validation.UserAttributeSimilarityValidator',
    },
    {
        'NAME': 'django.contrib.auth.password_validation.MinimumLengthValidator',
    },
    {
        'NAME': 'django.contrib.auth.password_validation.CommonPasswordValidator',
    },
    {
        'NAME': 'django.contrib.auth.password_validation.NumericPasswordValidator',
    },
]

LANGUAGE_CODE = 'pl-pl'
TIME_ZONE = 'Europe/Warsaw'
USE_I18N = True
USE_L10N = True
USE_TZ = True

db_from_env = dj_database_url.config(conn_max_age=500)

STATIC_ROOT = os.path.join(BASE_DIR, 'staticfiles')
STATIC_URL = '/static/'

SECURE_PROXY_SSL_HEADER = ('HTTP_X_FORWARDED_PROTO', 'https')

STATICFILES_DIRS = (
    os.path.join(BASE_DIR, 'static'),
)

STATICFILES_STORAGE = [
    'whitenoise.storage.CompressedManifestStaticFilesStorage',
]

CACHES = {
    'default': {
        'BACKEND': 'django.core.cache.backends.filebased.FileBasedCache',
        'LOCATION': os.path.join(BASE_DIR, 'cache').replace('\\', '/'),
        'TIMEOUT': 600,
        'OPTIONS': {
            'MAX_ENTRIES': 1000
        }
    }
}
<|endoftext|>


<commit_before>import sys, os
import json

#This first if __name__ is needed if this file is test-run from the command line.
# path.append needs to happen before the first from import statement
if __name__ == '__main__':
    sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),os.pardir,os.pardir)))

import logging
from housinginsights.sources.base import BaseApiConn

class CensusApiConn(BaseApiConn):
    """

    """
    def __init__(self, proxies=None):
        #baseurl not actually used since we need the _urls property to hold many urls.
        #Needed to get call to super() to work correctly. TODO refactor so this is optional.
        baseurl = 'http://api.census.gov/data/'
        super(CensusApiConn, self).__init__(baseurl, proxies=None)
        secretsFileName =  os.path.abspath(os.path.join(os.path.dirname(__file__),
                                               os.pardir,"secrets.json"))
        self.api_keys = json.loads(open(secretsFileName).read())
        self.census_api_key = self.api_keys['census']['api_key']

        self._available_unique_data_ids = [
                        "acs5_2015",
                        "acs5_2015_moe"
                        ]

        self._urls =	{
                        "acs5_2015"     :"2015/acs5",
                        "acs5_2015_moe" :"2015/acs5"
                    	}
        self._fields = {
                        "acs5_2015"     : "NAME,B01003_001E,B25057_001E,B25058_001E,B25059_001E",
                        "acs5_2015_moe" : "NAME,B01003_001M,B25057_001M,B25058_001M,B25059_001M"
                       }

    def get_data(self, unique_data_ids=None, sample=False, output_type = 'csv', **kwargs):
        '''
        Gets all data sources associated with this ApiConn class.

        sample does not apply to this apiconn object because api doesn't let us request only some data.
        '''
        if unique_data_ids == None:
            unique_data_ids = self._available_unique_data_ids

        for u in unique_data_ids:
            if (u not in self._available_unique_data_ids):
                #TODO Change this error type to a more specific one that should be handled by calling function inproduction
                #We will want the calling function to continue with other data sources instead of erroring out.
                logging.info("  The unique_data_id '{}' is not supported by the CensusApiConn".format(u))
            else:
                result = self.get(self._urls[u], params={'key':self.census_api_key, 'get':self._fields[u], 'for': 'tract:*', 'in': 'state:11'})

                if result.status_code != 200:
                    err = "An error occurred during request: status {0}"
                    #TODO change this error type to be handleable by caller
                    raise Exception(err.format(result.status_code))

                content = result.text

                if output_type == 'stdout':
                    print(content)

                elif output_type == 'csv':
                    jsondata=json.loads(content)
                    self.result_to_csv(jsondata[0], jsondata[1:],  self.output_paths[u])

                #Can't yield content if we get multiple sources at once
                if len(unique_data_ids) == 1:
                    return content

if __name__ == '__main__':

    # Pushes everything from the logger to the command line output as well.
    log_path = os.path.abspath(os.path.join(os.path.dirname(__file__),os.pardir,os.pardir,"logs","sources.log"))
    logging.basicConfig(filename=log_path, level=logging.DEBUG)
    logging.getLogger().addHandler(logging.StreamHandler())

    api_conn = CensusApiConn()
    unique_data_ids = None #['crime_2013']
    api_conn.get_data(unique_data_ids)
<commit_msg>152-census data - added all years<commit_after>import sys, os
import json

#This first if __name__ is needed if this file is test-run from the command line.
# path.append needs to happen before the first from import statement
if __name__ == '__main__':
    sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__),os.pardir,os.pardir)))

import logging
from housinginsights.sources.base import BaseApiConn

class CensusApiConn(BaseApiConn):
    """

    """
    def __init__(self, proxies=None):
        #baseurl not actually used since we need the _urls property to hold many urls.
        #Needed to get call to super() to work correctly. TODO refactor so this is optional.
        baseurl = 'http://api.census.gov/data/'
        super(CensusApiConn, self).__init__(baseurl, proxies=None)
        secretsFileName =  os.path.abspath(os.path.join(os.path.dirname(__file__),
                                               os.pardir,"secrets.json"))
        self.api_keys = json.loads(open(secretsFileName).read())
        self.census_api_key = self.api_keys['census']['api_key']

        self._available_unique_data_ids = [
                        "acs5_2009",
                        "acs5_2010",
                        "acs5_2011",
                        "acs5_2012",
                        "acs5_2013",
                        "acs5_2014",
                        "acs5_2015",
                        "acs5_2009_moe",
                        "acs5_2010_moe",
                        "acs5_2011_moe",
                        "acs5_2012_moe",
                        "acs5_2013_moe",
                        "acs5_2014_moe",
                        "acs5_2015_moe"
                        ]

        self._urls =	{
                        "acs5_2009"     :"2009/acs5",
                        "acs5_2010"     :"2010/acs5",
                        "acs5_2011"     :"2011/acs5",
                        "acs5_2012"     :"2012/acs5",
                        "acs5_2013"     :"2013/acs5",
                        "acs5_2014"     :"2014/acs5",
                        "acs5_2015"     :"2015/acs5",
                        "acs5_2009_moe" :"2009/acs5",
                        "acs5_2010_moe" :"2010/acs5",
                        "acs5_2011_moe" :"2011/acs5",
                        "acs5_2012_moe" :"2012/acs5",
                        "acs5_2013_moe" :"2013/acs5",
                        "acs5_2014_moe" :"2014/acs5",
                        "acs5_2015_moe" :"2015/acs5"
                    	}
        self._fields = {
                        "acs5_2009"     : "NAME,B01003_001E,B25057_001E,B25058_001E,B25059_001E",
                        "acs5_2010"     : "NAME,B01003_001E,B25057_001E,B25058_001E,B25059_001E",
                        "acs5_2011"     : "NAME,B01003_001E,B25057_001E,B25058_001E,B25059_001E",
                        "acs5_2012"     : "NAME,B01003_001E,B25057_001E,B25058_001E,B25059_001E",
                        "acs5_2013"     : "NAME,B01003_001E,B25057_001E,B25058_001E,B25059_001E",
                        "acs5_2014"     : "NAME,B01003_001E,B25057_001E,B25058_001E,B25059_001E",
                        "acs5_2015"     : "NAME,B01003_001E,B25057_001E,B25058_001E,B25059_001E",
                        "acs5_2009_moe" : "NAME,B01003_001M,B25057_001M,B25058_001M,B25059_001M",
                        "acs5_2010_moe" : "NAME,B01003_001M,B25057_001M,B25058_001M,B25059_001M",
                        "acs5_2011_moe" : "NAME,B01003_001M,B25057_001M,B25058_001M,B25059_001M",
                        "acs5_2012_moe" : "NAME,B01003_001M,B25057_001M,B25058_001M,B25059_001M",
                        "acs5_2013_moe" : "NAME,B01003_001M,B25057_001M,B25058_001M,B25059_001M",
                        "acs5_2014_moe" : "NAME,B01003_001M,B25057_001M,B25058_001M,B25059_001M",
                        "acs5_2015_moe" : "NAME,B01003_001M,B25057_001M,B25058_001M,B25059_001M"
                       }

    def get_data(self, unique_data_ids=None, sample=False, output_type = 'csv', **kwargs):
        '''
        Gets all data sources associated with this ApiConn class.

        sample does not apply to this apiconn object because api doesn't let us request only some data.
        '''
        if unique_data_ids == None:
            unique_data_ids = self._available_unique_data_ids

        for u in unique_data_ids:
            if (u not in self._available_unique_data_ids):
                #TODO Change this error type to a more specific one that should be handled by calling function inproduction
                #We will want the calling function to continue with other data sources instead of erroring out.
                logging.info("  The unique_data_id '{}' is not supported by the CensusApiConn".format(u))
            else:
                result = self.get(self._urls[u], params={'key':self.census_api_key, 'get':self._fields[u], 'for': 'tract:*', 'in': 'state:11'})

                if result.status_code != 200:
                    err = "An error occurred during request: status {0}"
                    #TODO change this error type to be handleable by caller
                    raise Exception(err.format(result.status_code))

                content = result.text

                if output_type == 'stdout':
                    print(content)

                elif output_type == 'csv':
                    jsondata=json.loads(content)
                    self.result_to_csv(jsondata[0], jsondata[1:],  self.output_paths[u])

                #Can't yield content if we get multiple sources at once
                if len(unique_data_ids) == 1:
                    return content

if __name__ == '__main__':

    # Pushes everything from the logger to the command line output as well.
    log_path = os.path.abspath(os.path.join(os.path.dirname(__file__),os.pardir,os.pardir,"logs","sources.log"))
    logging.basicConfig(filename=log_path, level=logging.DEBUG)
    logging.getLogger().addHandler(logging.StreamHandler())

    api_conn = CensusApiConn()
    unique_data_ids = None #['crime_2013']
    api_conn.get_data(unique_data_ids)
<|endoftext|>


<commit_before>from sfa.util.sfalogging import logger
from sfa.util.xml import XpathFilter
from sfa.util.xrn import Xrn

from sfa.rspecs.elements.element import Element
from sfa.rspecs.elements.node import Node
from sfa.rspecs.elements.sliver import Sliver
from sfa.rspecs.elements.location import Location
from sfa.rspecs.elements.position_3d import Position3D
from sfa.rspecs.elements.hardware_type import HardwareType
from sfa.rspecs.elements.disk_image import DiskImage
from sfa.rspecs.elements.interface import Interface
from sfa.rspecs.elements.bwlimit import BWlimit
from sfa.rspecs.elements.pltag import PLTag
from sfa.rspecs.elements.versions.nitosv1Sliver import NITOSv1Sliver
from sfa.rspecs.elements.versions.nitosv1PLTag import NITOSv1PLTag
from sfa.rspecs.elements.versions.pgv2Services import PGv2Services

class NITOSv1Node:

    @staticmethod
    def add_nodes(xml, nodes):
        network_elems = xml.xpath('//network')
        if len(network_elems) > 0:
            network_elem = network_elems[0]
        elif len(nodes) > 0 and nodes[0].get('component_manager_id'):
            network_urn = nodes[0]['component_manager_id']
            network_elem = xml.add_element('network', name = Xrn(network_urn).get_hrn())
        else:
            network_elem = xml

        node_elems = []
        for node in nodes:
            node_fields = ['component_manager_id', 'component_id', 'boot_state']
            node_elem = network_elem.add_instance('node', node, node_fields)
            node_elems.append(node_elem)

            # determine network hrn
            network_hrn = None
            if 'component_manager_id' in node and node['component_manager_id']:
                network_hrn = Xrn(node['component_manager_id']).get_hrn()

            # set component_name attribute and  hostname element
            if 'component_id' in node and node['component_id']:
                component_name = Xrn(xrn=node['component_id']).get_leaf()
                node_elem.set('component_name', component_name)
                hostname_elem = node_elem.add_element('hostname')
                hostname_elem.set_text(component_name)

            # set site id
            if 'authority_id' in node and node['authority_id']:
                node_elem.set('site_id', node['authority_id'])

            # add locaiton
            location = node.get('location')
            if location:
                node_elem.add_instance('location', location, Location.fields)

            # add 3D Position of the node
            position_3d = node.get('position_3d')
            if position_3d:
                node_elem.add_instance('position_3d', position_3d, Position3D.fields)

            # all nitos nodes are exculsive
            exclusive_elem = node_elem.add_element('exclusive')
            exclusive_elem.set_text('TRUE')

            # add granularity of the reservation system
            granularity = node.get('granularity')['grain']
            if granularity:
                #node_elem.add_instance('granularity', granularity, granularity.fields)
                granularity_elem = node_elem.add_element('granularity')
                granularity_elem.set_text(str(granularity))
            # add hardware type
            #hardware_type = node.get('hardware_type')
            #if hardware_type:
            #    node_elem.add_instance('hardware_type', hardware_type)
            hardware_type_elem = node_elem.add_element('hardware_type')
            hardware_type_elem.set_text(node.get('hardware_type'))

            if isinstance(node.get('interfaces'), list):
                for interface in node.get('interfaces', []):
                    node_elem.add_instance('interface', interface, ['component_id', 'client_id', 'ipv4'])

            #if 'bw_unallocated' in node and node['bw_unallocated']:
            #    bw_unallocated = etree.SubElement(node_elem, 'bw_unallocated', units='kbps').text = str(int(node['bw_unallocated'])/1000)

            PGv2Services.add_services(node_elem, node.get('services', []))
            tags = node.get('tags', [])
            if tags:
                for tag in tags:
                    tag_elem = node_elem.add_element(tag['tagname'])
                    tag_elem.set_text(tag['value'])
            NITOSv1Sliver.add_slivers(node_elem, node.get('slivers', []))

    @staticmethod
    def add_slivers(xml, slivers):
        component_ids = []
        for sliver in slivers:
            filter = {}
            if isinstance(sliver, str):
                filter['component_id'] = '*%s*' % sliver
                sliver = {}
            elif 'component_id' in sliver and sliver['component_id']:
                filter['component_id'] = '*%s*' % sliver['component_id']
            if not filter:
                continue
            nodes = NITOSv1Node.get_nodes(xml, filter)
            if not nodes:
                continue
            node = nodes[0]
            NITOSv1Sliver.add_slivers(node, sliver)

    @staticmethod
    def remove_slivers(xml, hostnames):
        for hostname in hostnames:
            nodes = NITOSv1Node.get_nodes(xml, {'component_id': '*%s*' % hostname})
            for node in nodes:
                slivers = NITOSv1Sliver.get_slivers(node.element)
                for sliver in slivers:
                    node.element.remove(sliver.element)

    @staticmethod
    def get_nodes(xml, filter={}):
        xpath = '//node%s | //default:node%s' % (XpathFilter.xpath(filter), XpathFilter.xpath(filter))
        node_elems = xml.xpath(xpath)
        return NITOSv1Node.get_node_objs(node_elems)

    @staticmethod
    def get_nodes_with_slivers(xml):
        xpath = '//node[count(sliver)>0] | //default:node[count(default:sliver)>0]'
        node_elems = xml.xpath(xpath)
        return NITOSv1Node.get_node_objs(node_elems)

    @staticmethod
    def get_node_objs(node_elems):
        nodes = []
        for node_elem in node_elems:
            node = Node(node_elem.attrib, node_elem)
            if 'site_id' in node_elem.attrib:
                node['authority_id'] = node_elem.attrib['site_id']
            # get location
            location_elems = node_elem.xpath('./default:location | ./location')
            locations = [loc_elem.get_instance(Location) for loc_elem in location_elems]
            if len(locations) > 0:
                node['location'] = locations[0]
            # get bwlimit
            bwlimit_elems = node_elem.xpath('./default:bw_limit | ./bw_limit')
            bwlimits = [bwlimit_elem.get_instance(BWlimit) for bwlimit_elem in bwlimit_elems]
            if len(bwlimits) > 0:
                node['bwlimit'] = bwlimits[0]
            # get interfaces
            iface_elems = node_elem.xpath('./default:interface | ./interface')
            ifaces = [iface_elem.get_instance(Interface) for iface_elem in iface_elems]
            node['interfaces'] = ifaces
            # get services
            node['services'] = PGv2Services.get_services(node_elem)
            # get slivers
            node['slivers'] = NITOSv1Sliver.get_slivers(node_elem)
            # get tags
            node['tags'] =  NITOSv1PLTag.get_pl_tags(node_elem, ignore=Node.fields+["hardware_type"])
            # get hardware types
            hardware_type_elems = node_elem.xpath('./default:hardware_type | ./hardware_type')
            node['hardware_types'] = [hw_type.get_instance(HardwareType) for hw_type in hardware_type_elems]

            # temporary... play nice with old slice manager rspec
            if not node['component_name']:
                hostname_elem = node_elem.find("hostname")
                if hostname_elem != None:
                    node['component_name'] = hostname_elem.text

            nodes.append(node)
        return nodes

<commit_msg>Advertise the Nitos access gateway<commit_after>from sfa.util.sfalogging import logger
from sfa.util.xml import XpathFilter
from sfa.util.xrn import Xrn

from sfa.rspecs.elements.element import Element
from sfa.rspecs.elements.node import Node
from sfa.rspecs.elements.sliver import Sliver
from sfa.rspecs.elements.location import Location
from sfa.rspecs.elements.position_3d import Position3D
from sfa.rspecs.elements.hardware_type import HardwareType
from sfa.rspecs.elements.disk_image import DiskImage
from sfa.rspecs.elements.interface import Interface
from sfa.rspecs.elements.bwlimit import BWlimit
from sfa.rspecs.elements.pltag import PLTag
from sfa.rspecs.elements.versions.nitosv1Sliver import NITOSv1Sliver
from sfa.rspecs.elements.versions.nitosv1PLTag import NITOSv1PLTag
from sfa.rspecs.elements.versions.pgv2Services import PGv2Services

class NITOSv1Node:

    @staticmethod
    def add_nodes(xml, nodes):
        network_elems = xml.xpath('//network')
        if len(network_elems) > 0:
            network_elem = network_elems[0]
        elif len(nodes) > 0 and nodes[0].get('component_manager_id'):
            network_urn = nodes[0]['component_manager_id']
            network_elem = xml.add_element('network', name = Xrn(network_urn).get_hrn())
        else:
            network_elem = xml

        # needs to be improuved to retreive the gateway addr dynamically.
        gateway_addr = 'nitlab.inf.uth.gr'

        node_elems = []
        for node in nodes:
            node_fields = ['component_manager_id', 'component_id', 'boot_state']
            node_elem = network_elem.add_instance('node', node, node_fields)
            node_elems.append(node_elem)

            # determine network hrn
            network_hrn = None
            if 'component_manager_id' in node and node['component_manager_id']:
                network_hrn = Xrn(node['component_manager_id']).get_hrn()

            # set component_name attribute and  hostname element
            if 'component_id' in node and node['component_id']:
                component_name = Xrn(xrn=node['component_id']).get_leaf()
                node_elem.set('component_name', component_name)
                hostname_elem = node_elem.add_element('hostname')
                hostname_elem.set_text(component_name)

            # set site id
            if 'authority_id' in node and node['authority_id']:
                node_elem.set('site_id', node['authority_id'])

            # add locaiton
            location = node.get('location')
            if location:
                node_elem.add_instance('location', location, Location.fields)

            # add 3D Position of the node
            position_3d = node.get('position_3d')
            if position_3d:
                node_elem.add_instance('position_3d', position_3d, Position3D.fields)

            # all nitos nodes are exculsive
            exclusive_elem = node_elem.add_element('exclusive')
            exclusive_elem.set_text('TRUE')

            # In order to access nitos nodes, one need to pass through the nitos gateway
            # here we advertise Nitos access gateway address
            gateway_elem = node_elem.add_element('gateway')
            gateway_elem.set_text(gateway_addr)

            # add granularity of the reservation system
            granularity = node.get('granularity')['grain']
            if granularity:
                #node_elem.add_instance('granularity', granularity, granularity.fields)
                granularity_elem = node_elem.add_element('granularity')
                granularity_elem.set_text(str(granularity))
            # add hardware type
            #hardware_type = node.get('hardware_type')
            #if hardware_type:
            #    node_elem.add_instance('hardware_type', hardware_type)
            hardware_type_elem = node_elem.add_element('hardware_type')
            hardware_type_elem.set_text(node.get('hardware_type'))

            if isinstance(node.get('interfaces'), list):
                for interface in node.get('interfaces', []):
                    node_elem.add_instance('interface', interface, ['component_id', 'client_id', 'ipv4'])

            #if 'bw_unallocated' in node and node['bw_unallocated']:
            #    bw_unallocated = etree.SubElement(node_elem, 'bw_unallocated', units='kbps').text = str(int(node['bw_unallocated'])/1000)

            PGv2Services.add_services(node_elem, node.get('services', []))
            tags = node.get('tags', [])
            if tags:
                for tag in tags:
                    tag_elem = node_elem.add_element(tag['tagname'])
                    tag_elem.set_text(tag['value'])
            NITOSv1Sliver.add_slivers(node_elem, node.get('slivers', []))

    @staticmethod
    def add_slivers(xml, slivers):
        component_ids = []
        for sliver in slivers:
            filter = {}
            if isinstance(sliver, str):
                filter['component_id'] = '*%s*' % sliver
                sliver = {}
            elif 'component_id' in sliver and sliver['component_id']:
                filter['component_id'] = '*%s*' % sliver['component_id']
            if not filter:
                continue
            nodes = NITOSv1Node.get_nodes(xml, filter)
            if not nodes:
                continue
            node = nodes[0]
            NITOSv1Sliver.add_slivers(node, sliver)

    @staticmethod
    def remove_slivers(xml, hostnames):
        for hostname in hostnames:
            nodes = NITOSv1Node.get_nodes(xml, {'component_id': '*%s*' % hostname})
            for node in nodes:
                slivers = NITOSv1Sliver.get_slivers(node.element)
                for sliver in slivers:
                    node.element.remove(sliver.element)

    @staticmethod
    def get_nodes(xml, filter={}):
        xpath = '//node%s | //default:node%s' % (XpathFilter.xpath(filter), XpathFilter.xpath(filter))
        node_elems = xml.xpath(xpath)
        return NITOSv1Node.get_node_objs(node_elems)

    @staticmethod
    def get_nodes_with_slivers(xml):
        xpath = '//node[count(sliver)>0] | //default:node[count(default:sliver)>0]'
        node_elems = xml.xpath(xpath)
        return NITOSv1Node.get_node_objs(node_elems)

    @staticmethod
    def get_node_objs(node_elems):
        nodes = []
        for node_elem in node_elems:
            node = Node(node_elem.attrib, node_elem)
            if 'site_id' in node_elem.attrib:
                node['authority_id'] = node_elem.attrib['site_id']
            # get location
            location_elems = node_elem.xpath('./default:location | ./location')
            locations = [loc_elem.get_instance(Location) for loc_elem in location_elems]
            if len(locations) > 0:
                node['location'] = locations[0]
            # get bwlimit
            bwlimit_elems = node_elem.xpath('./default:bw_limit | ./bw_limit')
            bwlimits = [bwlimit_elem.get_instance(BWlimit) for bwlimit_elem in bwlimit_elems]
            if len(bwlimits) > 0:
                node['bwlimit'] = bwlimits[0]
            # get interfaces
            iface_elems = node_elem.xpath('./default:interface | ./interface')
            ifaces = [iface_elem.get_instance(Interface) for iface_elem in iface_elems]
            node['interfaces'] = ifaces
            # get services
            node['services'] = PGv2Services.get_services(node_elem)
            # get slivers
            node['slivers'] = NITOSv1Sliver.get_slivers(node_elem)
            # get tags
            node['tags'] =  NITOSv1PLTag.get_pl_tags(node_elem, ignore=Node.fields+["hardware_type"])
            # get hardware types
            hardware_type_elems = node_elem.xpath('./default:hardware_type | ./hardware_type')
            node['hardware_types'] = [hw_type.get_instance(HardwareType) for hw_type in hardware_type_elems]

            # temporary... play nice with old slice manager rspec
            if not node['component_name']:
                hostname_elem = node_elem.find("hostname")
                if hostname_elem != None:
                    node['component_name'] = hostname_elem.text

            nodes.append(node)
        return nodes

<|endoftext|>


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

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

import os
import yara
import logging
import unittest
from unittest.mock import create_autospec, Mock

from stoq import PayloadMeta, RequestMeta, Stoq, Payload
from stoq.exceptions import StoqPluginException
from stoq.data_classes import WorkerResponse, DispatcherResponse

class TestCore(unittest.TestCase):
    def setUp(self) -> None:
        self.plugin_name = 'yara'
        self.plugin_dir = os.path.join(os.getcwd(), 'yarascan')
        self.data_dir = os.path.join(os.getcwd(), 'tests', 'data')
        self.generic_data = b'testtesttest'
        logging.disable(logging.CRITICAL)

    def tearDown(self) -> None:
        logging.disable(logging.NOTSET)

    def test_scan(self) -> None:
        s = Stoq(
            plugin_dir_list=[self.plugin_dir],
            plugin_opts={
                self.plugin_name: {'worker_rules': f'{self.data_dir}/scan_rules.yar'}
            },
        )
        plugin = s.load_plugin(self.plugin_name)
        payload = Payload(self.generic_data)
        request_meta = RequestMeta(archive_payloads=False)
        response = plugin.scan(payload, request_meta)
        self.assertIsInstance(response, WorkerResponse)
        self.assertEqual('test_scan_rule', response.results['matches'][0]['rule'])

    def test_scan_meta_bytes(self) -> None:
        s = Stoq(
            plugin_dir_list=[self.plugin_dir],
            plugin_opts={
                self.plugin_name: {'worker_rules': f'{self.data_dir}/scan_rules.yar'}
            },
        )
        plugin = s.load_plugin(self.plugin_name)
        payload = Payload(b'meta_bytes')
        request_meta = RequestMeta(archive_payloads=False)
        response = plugin.scan(payload, request_meta)
        self.assertIsInstance(response, WorkerResponse)
        self.assertEqual(
            'test_scan_metadata_bytes', response.results['matches'][0]['rule']
        )
        self.assertEqual('ANeato', response.results['matches'][0]['meta']['bytes'])
        self.assertEqual(
            'Peter Rabbit', response.results['matches'][0]['meta']['author']
        )
        self.assertEqual('save_false', response.results['matches'][0]['meta']['plugin'])

    def test_scan_invalid_rule_file(self) -> None:
        s = Stoq(
            plugin_dir_list=[self.plugin_dir],
            plugin_opts={
                self.plugin_name: {'worker_rules': f'{self.data_dir}/nonexistent.yar'}
            },
        )
        with self.assertRaises(StoqPluginException):
            s.load_plugin(self.plugin_name)

    def test_scan_invalid_rules(self) -> None:
        s = Stoq(
            plugin_dir_list=[self.plugin_dir],
            plugin_opts={
                self.plugin_name: {'worker_rules': f'{self.data_dir}/invalid_rules.yar'}
            },
        )
        with self.assertRaises(yara.SyntaxError):
            s.load_plugin(self.plugin_name)

    def test_dispatcher(self) -> None:
        s = Stoq(
            plugin_dir_list=[self.plugin_dir],
            plugin_opts={
                self.plugin_name: {
                    'dispatch_rules': f'{self.data_dir}/dispatch_rules.yar'
                }
            },
        )
        plugin = s.load_plugin(self.plugin_name)
        payload = Payload(self.generic_data)
        request_meta = RequestMeta(archive_payloads=False)
        response = plugin.get_dispatches(payload, request_meta)
        self.assertIsInstance(response, DispatcherResponse)
        self.assertIn('test_dispatch_plugin', response.plugin_names)
        self.assertEqual(
            'test_dispatch_rule', response.meta['test_dispatch_plugin']['rule']
        )
        self.assertIn(
            'test_dispatch_plugin',
            response.meta['test_dispatch_plugin']['meta']['plugin'],
        )
        self.assertIn('True', response.meta['test_dispatch_plugin']['meta']['save'])
        self.assertEqual(
            ['tag1', 'tag2'], response.meta['test_dispatch_plugin']['tags']
        )

    def test_dispatcher_save_false(self) -> None:
        s = Stoq(
            plugin_dir_list=[self.plugin_dir],
            plugin_opts={
                self.plugin_name: {
                    'dispatch_rules': f'{self.data_dir}/dispatch_rules.yar'
                }
            },
        )
        plugin = s.load_plugin(self.plugin_name)
        payload = Payload(b'save_false')
        request_meta = RequestMeta(archive_payloads=False)
        response = plugin.get_dispatches(payload, request_meta)
        self.assertIsInstance(response, DispatcherResponse)
        self.assertIn('False', response.meta['save_false']['meta']['save'])
<commit_msg>Improve handling of paths to tests can be run anywhere<commit_after>#!/usr/bin/env python3

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

import os
import yara
import unittest

from pathlib import Path

from stoq import RequestMeta, Stoq, Payload
from stoq.exceptions import StoqPluginException
from stoq.data_classes import WorkerResponse, DispatcherResponse

class TestCore(unittest.TestCase):
    def setUp(self) -> None:
        self.plugin_name = 'yara'
        self.base_dir = Path(os.path.realpath(__file__)).parent
        self.data_dir = os.path.join(self.base_dir, 'data')
        self.plugin_dir = os.path.join(self.base_dir.parent, 'yarascan')
        self.generic_data = b'testtesttest'

    def tearDown(self) -> None:
        pass

    def test_scan(self) -> None:
        s = Stoq(
            plugin_dir_list=[self.plugin_dir],
            plugin_opts={
                self.plugin_name: {'worker_rules': f'{self.data_dir}/scan_rules.yar'}
            },
        )
        plugin = s.load_plugin(self.plugin_name)
        payload = Payload(self.generic_data)
        request_meta = RequestMeta(archive_payloads=False)
        response = plugin.scan(payload, request_meta)
        self.assertIsInstance(response, WorkerResponse)
        self.assertEqual('test_scan_rule', response.results['matches'][0]['rule'])

    def test_scan_meta_bytes(self) -> None:
        s = Stoq(
            plugin_dir_list=[self.plugin_dir],
            plugin_opts={
                self.plugin_name: {'worker_rules': f'{self.data_dir}/scan_rules.yar'}
            },
        )
        plugin = s.load_plugin(self.plugin_name)
        payload = Payload(b'meta_bytes')
        request_meta = RequestMeta(archive_payloads=False)
        response = plugin.scan(payload, request_meta)
        self.assertIsInstance(response, WorkerResponse)
        self.assertEqual(
            'test_scan_metadata_bytes', response.results['matches'][0]['rule']
        )
        self.assertEqual('ANeato', response.results['matches'][0]['meta']['bytes'])
        self.assertEqual(
            'Peter Rabbit', response.results['matches'][0]['meta']['author']
        )
        self.assertEqual('save_false', response.results['matches'][0]['meta']['plugin'])

    def test_scan_invalid_rule_file(self) -> None:
        s = Stoq(
            plugin_dir_list=[self.plugin_dir],
            plugin_opts={
                self.plugin_name: {'worker_rules': f'{self.data_dir}/nonexistent.yar'}
            },
        )
        with self.assertRaises(StoqPluginException):
            s.load_plugin(self.plugin_name)

    def test_scan_invalid_rules(self) -> None:
        s = Stoq(
            plugin_dir_list=[self.plugin_dir],
            plugin_opts={
                self.plugin_name: {'worker_rules': f'{self.data_dir}/invalid_rules.yar'}
            },
        )
        with self.assertRaises(yara.SyntaxError):
            s.load_plugin(self.plugin_name)

    def test_dispatcher(self) -> None:
        s = Stoq(
            plugin_dir_list=[self.plugin_dir],
            plugin_opts={
                self.plugin_name: {
                    'dispatch_rules': f'{self.data_dir}/dispatch_rules.yar'
                }
            },
        )
        plugin = s.load_plugin(self.plugin_name)
        payload = Payload(self.generic_data)
        request_meta = RequestMeta(archive_payloads=False)
        response = plugin.get_dispatches(payload, request_meta)
        self.assertIsInstance(response, DispatcherResponse)
        self.assertIn('test_dispatch_plugin', response.plugin_names)
        self.assertEqual(
            'test_dispatch_rule', response.meta['test_dispatch_plugin']['rule']
        )
        self.assertIn(
            'test_dispatch_plugin',
            response.meta['test_dispatch_plugin']['meta']['plugin'],
        )
        self.assertIn('True', response.meta['test_dispatch_plugin']['meta']['save'])
        self.assertEqual(
            ['tag1', 'tag2'], response.meta['test_dispatch_plugin']['tags']
        )

    def test_dispatcher_save_false(self) -> None:
        s = Stoq(
            plugin_dir_list=[self.plugin_dir],
            plugin_opts={
                self.plugin_name: {
                    'dispatch_rules': f'{self.data_dir}/dispatch_rules.yar'
                }
            },
        )
        plugin = s.load_plugin(self.plugin_name)
        payload = Payload(b'save_false')
        request_meta = RequestMeta(archive_payloads=False)
        response = plugin.get_dispatches(payload, request_meta)
        self.assertIsInstance(response, DispatcherResponse)
        self.assertIn('False', response.meta['save_false']['meta']['save'])
<|endoftext|>


<commit_before><commit_msg>Split out alarm sounding status into two events.<commit_after><|endoftext|>


<commit_before>"""Tools to create quicklook visualisations of VST/OmegaCAM data."""
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import os
import numpy as np
import imageio

import matplotlib
import matplotlib.pyplot as pl
from matplotlib.colors import LogNorm
import matplotlib.patheffects as path_effects
import matplotlib.image as mimg

from astropy import log
from astropy.io import fits
from astropy.wcs import WCS
from astropy.visualization import ManualInterval, AsymmetricPercentileInterval, SqrtStretch, LogStretch, AsinhStretch
from astropy.utils.console import ProgressBar

from . import VPHAS_DATA_PATH, OMEGACAM_CCD_ARRANGEMENT
from .footprint import VphasOffset, NotObservedException

##########
# CLASSES
##########

class VphasDataException(Exception):
    pass

class LuptonColorStretch():
    """
    Stretch a MxNx3 (RGB) array using the colour-preserving method by
    Lupton et al. (arXiv:astro-ph/0312483).

    Parameters
    ----------
    intensity_stretch : object derived from `BaseStretch`
        The stretch to apply on the integrated intensity, e.g. `LogStretch`.
    """
    def __init__(self, intensity_stretch=None):
        if intensity_stretch is None:
            self.intensity_stretch = AsinhStretch(a=0.03)
        else:
            self.intensity_stretch = intensity_stretch

    def __call__(self, values):
        """Transform the image using this stretch.

        Parameters
        ----------
        values : `~numpy.ndarray` of shape MxNx3 (RGB)
            The input image, which should already be normalized to the [0:1]
            range.

        Returns
        -------
        new_values : `~numpy.ndarray`
            The transformed values.
        """
        intensity = values.sum(axis=2) / 3.
        old_settings = np.seterr(divide='ignore', invalid='ignore')
        for chn in range(3):
            values[:, :, chn] *= self.intensity_stretch(intensity) / intensity
            values[:, :, chn][intensity == 0] = 0
        np.seterr(**old_settings)
        maxrgb = values.max(axis=2)
        maxrgb_gt_1 = maxrgb > 1
        for chn in range(3):
            values[:, :, chn][maxrgb_gt_1] /= maxrgb[maxrgb_gt_1]
        return values

class VphasColourFrame():
    """Create a pretty colour quicklook from a VPHAS frame.

    Parameters
    ----------
    offset : str
        Name of the VPHAS offset, e.g. "0001a".

    ccd : int
        CCD number, referring to the FITS extension number.
    """
    def __init__(self, offset, ccd):
        self.offset = offset
        self.ccd = ccd
        self.filenames = VphasOffset(offset).get_filenames()
        if len(self.filenames) < 6:
            raise NotObservedException('Not all six bands have been observed'
                                       ' for offset {0}.'.format(offset))

    def as_array(self,
                 interval_r=AsymmetricPercentileInterval(2.5, 99.0),
                 interval_g=AsymmetricPercentileInterval(5., 99.2),
                 interval_b=AsymmetricPercentileInterval(10., 99.2)):
        """Returns the colour image as a MxNx3 (RGB) array."""
        # First we determine the shifts
        cx, cy = 2050, 1024  # central coordinates of the image
        maxshiftx, maxshifty = 0, 0
        aligned_imgs = {}
        for idx, band in enumerate(self.filenames.keys()):
            fn = self.filenames[band]
            hdu = fits.open(os.path.join(VPHAS_DATA_PATH, fn))[self.ccd]
            wcs = WCS(hdu.header)
            img = hdu.data
            if idx == 0:  # The first image acts as reference
                cra, cdec = wcs.wcs_pix2world(cx, cy, 1)
            else:  # For all subsequent images, compute the shift using the WCS
                refx, refy = wcs.wcs_world2pix(cra, cdec, 1)
                shiftx = int(refx - cx)
                shifty = int(refy - cy)
                # Now apply the required shift to the image
                if shiftx > 0:
                    img = np.pad(img, ((0, 0), (0, shiftx)),
                                 mode=str('constant'))[:, shiftx:]
                elif shiftx < 0:
                    img = np.pad(img, ((0, 0), (-shiftx, 0)),
                                 mode=str('constant'))[:, :shiftx]
                if shifty > 0:
                    img = np.pad(img, ((0, shifty), (0, 0)),
                                 mode=str('constant'))[shifty:, :]
                elif shifty < 0:
                    img = np.pad(img, ((-shifty, 0), (0, 0)),
                                 mode=str('constant'))[:shifty, :]
                # The maximum shift applied will determine the final img shape
                maxshiftx = max(abs(shiftx), maxshiftx)
                maxshifty = max(abs(shifty), maxshifty)
            aligned_imgs[band] = img

        if maxshiftx > cx or maxshifty > cy:
            raise VphasDataException('{0}-{1}: bands do not overlap'.format(self.offset, self.ccd))
        # New stretch, scale, and stack the data into an MxNx3 array
        r = aligned_imgs['i'] + aligned_imgs['ha']
        g = aligned_imgs['g'] + aligned_imgs['r'] + aligned_imgs['r2']
        b = aligned_imgs['u'] + 2 * aligned_imgs['g']
        r, g, b = 1.5*r, 0.8*g, 2.2*b
        vmin_r, vmax_r = np.percentile(r, [1., 99.5])
        vmin_g, vmax_g = np.percentile(g, [5., 99.5])
        vmin_b, vmax_b = np.percentile(b, [10., 99.5])
        #log.info((vmin_r, vmin_g, vmin_b))
        #log.info((vmax_r, vmax_g, vmax_b))
        #if vmin_b < 100:
        #    vmin_b = 100
        minrange = np.max((1250., vmax_g-vmin_g, vmax_g-vmin_g, vmax_b-vmin_b))
        if (vmax_r - vmin_r) < minrange:
            vmax_r = vmin_r + minrange
        if (vmax_g - vmin_g) < minrange:
            vmax_g = vmin_g + minrange
        if (vmax_b - vmin_b) < minrange:
            vmax_b = vmin_b + minrange
        interval_r = ManualInterval(vmin_r, vmax_r)
        interval_g = ManualInterval(vmin_g, vmax_g)
        interval_b = ManualInterval(vmin_b, vmax_b)

        r = interval_r(r)
        g = interval_g(g)
        b = interval_b(b)
        stacked = np.dstack((r[maxshifty:-maxshifty, maxshiftx:-maxshiftx],
                             g[maxshifty:-maxshifty, maxshiftx:-maxshiftx],
                             b[maxshifty:-maxshifty, maxshiftx:-maxshiftx]))
        return LuptonColorStretch()(stacked)

    def imsave(self, out_fn=None):
        if out_fn is None:
            out_fn = 'vphas-{0}-{1}.jpg'.format(self.offset, self.ccd)
        log.info('Writing {0}'.format(out_fn))
        #mimg.imsave(out_fn, np.rot90(self.as_array()), origin='lower')
        img = np.rot90(self.as_array())
        imageio.imsave(out_fn, img, quality=90, optimize=True)

############
# FUNCTIONS
############

def vphas_quicklook(offset, ccd=None, out_fn=None):
    VphasColourFrame(offset, ccd).imsave(out_fn)

def vphas_quicklook_main(args=None):
    import argparse

    parser = argparse.ArgumentParser(
        description='Create a beautiful color image from single VPHAS frames.')
    parser.add_argument('-o', '--output', metavar='filename',
                        type=str, default=None,
                        help='Filename for the output image (Default is a '
                             'JPG file named vphas-offset-ccd.jpg)')
    parser.add_argument('-c', '--ccd', nargs='+', type=int, default=None,
                        help='CCD number between 1 and 32.'
                             '(Default is to save all CCDs.)')
    parser.add_argument('--min_percent_r', type=float, default=1.0,
                        help=('The percentile value used to determine the '
                              'minimum cut level for the red channel'))
    parser.add_argument('--max_percent_r', type=float, default=99.5,
                        help=('The percentile value used to determine the '
                              'maximum cut level for the red channel'))
    parser.add_argument('offset', nargs='+',
                        help='Name of the VPHAS offset pointing.')
    args = parser.parse_args(args)

    if args.ccd is None:
        args.ccd = range(1, 33)
    for offset in args.offset:
        try:
            for ccd in args.ccd:
                vphas_quicklook(offset,
                                ccd=ccd,
                                out_fn=args.output)
        except NotObservedException as e:
            log.error(e)

def vst_pawplot(filename, out_fn=None, dpi=100,
                min_percent=1.0, max_percent=99.5,
                cmap='gist_heat', show_hdu=False):
    """Plots the 32-CCD OmegaCam mosaic as a pretty bitmap.

    Parameters
    ----------
    filename : str
        The filename of the FITS file.
    out_fn : str
        The filename of the output bitmap image.  The type of bitmap
        is determined by the filename extension (e.g. '.jpg', '.png').
        The default is a PNG file with the same name as the FITS file.
    dpi : float, optional [dots per inch]
        Resolution of the output 10-by-9 inch output bitmap.
        The default is 100.
    min_percent : float, optional
        The percentile value used to determine the pixel value of
        minimum cut level.  The default is 1.0.
    max_percent : float, optional
        The percentile value used to determine the pixel value of
        maximum cut level.  The default is 99.5.
    cmap : str, optional
        The matplotlib color map name.  The default is 'gist_heat',
        can also be e.g. 'Greys_r'.
    show_hdu : boolean, optional
        Plot the HDU extension number if True (default: False).
    """
    if out_fn is None:
        out_fn = filename + '.png'
    if cmap not in pl.cm.datad.keys():
        raise ValueError('{0} is not a valid matplotlib colormap '
                         'name'.format(cmap))
    log.info('Writing {0}'.format(out_fn))
    # Prepare the plot
    f = fits.open(filename)
    pl.interactive(False)
    fig = pl.figure(figsize=(10, 9))
    # Determine vmin/vmax based on a sample of pixels across the mosaic
    sample = np.concatenate((f[1].data[::20, ::10],
                             f[11].data[::20, ::10],
                             f[22].data[::20, ::10],
                             f[32].data[::20, ::10]))
    vmin, vmax = np.percentile(sample, [min_percent, max_percent])
    del sample  # save memory
    log.debug('vst_pawplot: vmin={0}, vmax={1}'.format(vmin, vmax))
    # Plot the extensions
    idx = 0
    for hduno in ProgressBar(OMEGACAM_CCD_ARRANGEMENT):
        idx += 1
        log.debug('vst_pawplot: adding HDU #{0}'.format(hduno))
        ax = fig.add_subplot(4, 8, idx)
        sampling = int(500 / dpi)
        im = ax.matshow(f[hduno].data[::sampling, ::-sampling],
                        norm=LogNorm(vmin=vmin, vmax=vmax),
                        cmap=cmap, origin='lower')
        ax.set_xticks([])
        ax.set_yticks([])
        ax.axis('off')
        if show_hdu:
            # Show the HDU extension number on top of the image
            txt = ax.text(0.05, 0.97, hduno, fontsize=14, color='white',
                          ha='left', va='top', transform=ax.transAxes)
            # Add a black border around the text
            txt.set_path_effects([path_effects.Stroke(linewidth=2, foreground='black'),
                                  path_effects.Normal()])

    # Aesthetics and colorbar
    fig.subplots_adjust(left=0.04, right=0.85,
                        top=0.93, bottom=0.05,
                        wspace=0.02, hspace=0.02)
    cbar_ax = fig.add_axes([0.9, 0.06, 0.02, 0.86])
    t = np.logspace(np.log10(vmin), np.log10(vmax), num=10)
    cb = fig.colorbar(im, cax=cbar_ax, ticks=t, format='%.0f')
    cb.set_label('Pixel count [ADU]')
    # Title and footer text
    fig.text(0.05, 0.95, filename, fontsize=16, ha='left')
    try:
        filterfooter = '{0} ({1}/{2})'.format(
                       f[0].header['ESO INS FILT1 NAME'],
                       f[0].header['ESO TPL EXPNO'],
                       f[0].header['ESO TPL NEXP'])
        fig.text(0.04, 0.02, f[0].header['OBJECT'], fontsize=12, ha='left')
        fig.text(0.50, 0.02, filterfooter, fontsize=12, ha='right')
        fig.text(0.85, 0.02, f[0].header['DATE-OBS'][0:19],
                 fontsize=12, ha='right')
    except KeyError as e:
        log.warning('Could not write footer text: {0}'.format(e))
        pass
    # We're done
    fig.savefig(out_fn, dpi=dpi)
    pl.close()
    del f

def vst_pawplot_main(args=None):
    import argparse

    parser = argparse.ArgumentParser(
        description='Create a beautiful plot showing the data for all 32 CCDs '
                    'in a single image.')
    parser.add_argument('-o', metavar='filename',
                        type=str, default=None,
                        help='Filename for the output image (Default is a '
                        'PNG file with the same name as the FITS file)')
    parser.add_argument('--dpi', type=float, default=100.,
                        help=('The resolution of the output image.'))
    parser.add_argument('--min_percent', type=float, default=1.0,
                        help=('The percentile value used to determine the '
                              'minimum cut level'))
    parser.add_argument('--max_percent', type=float, default=99.5,
                        help=('The percentile value used to determine the '
                              'maximum cut level'))
    parser.add_argument('--cmap', metavar='colormap_name', type=str,
                        default='gist_heat', help='matplotlib color map name')
    parser.add_argument('--show-hdu', action='store_true',
                        help='Display the HDU extension numbers.')
    parser.add_argument('filename', nargs='+',
                        help='Path to one or more FITS files to convert')
    args = parser.parse_args(args)

    for filename in args.filename:
        vst_pawplot(filename,
                    out_fn=args.o,
                    dpi=args.dpi,
                    min_percent=args.min_percent,
                    max_percent=args.max_percent,
                    cmap=args.cmap,
                    show_hdu=args.show_hdu)

if __name__ == '__main__':
    # Example use:
    vst_pawplot(os.path.join(VPHAS_DATA_PATH, 'o20120918_00025.fit'),
                '/tmp/test.jpg')
    vphas_quicklook("0778a", 24, "/tmp/test2.jpg")
<commit_msg>Print decimal digits on vst-pawplot colorbar<commit_after>"""Tools to create quicklook visualisations of VST/OmegaCAM data."""
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import os
import numpy as np
import imageio

import matplotlib
import matplotlib.pyplot as pl
from matplotlib.colors import LogNorm
import matplotlib.patheffects as path_effects
import matplotlib.image as mimg

from astropy import log
from astropy.io import fits
from astropy.wcs import WCS
from astropy.visualization import ManualInterval, AsymmetricPercentileInterval, SqrtStretch, LogStretch, AsinhStretch
from astropy.utils.console import ProgressBar

from . import VPHAS_DATA_PATH, OMEGACAM_CCD_ARRANGEMENT
from .footprint import VphasOffset, NotObservedException

##########
# CLASSES
##########

class VphasDataException(Exception):
    pass

class LuptonColorStretch():
    """
    Stretch a MxNx3 (RGB) array using the colour-preserving method by
    Lupton et al. (arXiv:astro-ph/0312483).

    Parameters
    ----------
    intensity_stretch : object derived from `BaseStretch`
        The stretch to apply on the integrated intensity, e.g. `LogStretch`.
    """
    def __init__(self, intensity_stretch=None):
        if intensity_stretch is None:
            self.intensity_stretch = AsinhStretch(a=0.03)
        else:
            self.intensity_stretch = intensity_stretch

    def __call__(self, values):
        """Transform the image using this stretch.

        Parameters
        ----------
        values : `~numpy.ndarray` of shape MxNx3 (RGB)
            The input image, which should already be normalized to the [0:1]
            range.

        Returns
        -------
        new_values : `~numpy.ndarray`
            The transformed values.
        """
        intensity = values.sum(axis=2) / 3.
        old_settings = np.seterr(divide='ignore', invalid='ignore')
        for chn in range(3):
            values[:, :, chn] *= self.intensity_stretch(intensity) / intensity
            values[:, :, chn][intensity == 0] = 0
        np.seterr(**old_settings)
        maxrgb = values.max(axis=2)
        maxrgb_gt_1 = maxrgb > 1
        for chn in range(3):
            values[:, :, chn][maxrgb_gt_1] /= maxrgb[maxrgb_gt_1]
        return values

class VphasColourFrame():
    """Create a pretty colour quicklook from a VPHAS frame.

    Parameters
    ----------
    offset : str
        Name of the VPHAS offset, e.g. "0001a".

    ccd : int
        CCD number, referring to the FITS extension number.
    """
    def __init__(self, offset, ccd):
        self.offset = offset
        self.ccd = ccd
        self.filenames = VphasOffset(offset).get_filenames()
        if len(self.filenames) < 6:
            raise NotObservedException('Not all six bands have been observed'
                                       ' for offset {0}.'.format(offset))

    def as_array(self,
                 interval_r=AsymmetricPercentileInterval(2.5, 99.0),
                 interval_g=AsymmetricPercentileInterval(5., 99.2),
                 interval_b=AsymmetricPercentileInterval(10., 99.2)):
        """Returns the colour image as a MxNx3 (RGB) array."""
        # First we determine the shifts
        cx, cy = 2050, 1024  # central coordinates of the image
        maxshiftx, maxshifty = 0, 0
        aligned_imgs = {}
        for idx, band in enumerate(self.filenames.keys()):
            fn = self.filenames[band]
            hdu = fits.open(os.path.join(VPHAS_DATA_PATH, fn))[self.ccd]
            wcs = WCS(hdu.header)
            img = hdu.data
            if idx == 0:  # The first image acts as reference
                cra, cdec = wcs.wcs_pix2world(cx, cy, 1)
            else:  # For all subsequent images, compute the shift using the WCS
                refx, refy = wcs.wcs_world2pix(cra, cdec, 1)
                shiftx = int(refx - cx)
                shifty = int(refy - cy)
                # Now apply the required shift to the image
                if shiftx > 0:
                    img = np.pad(img, ((0, 0), (0, shiftx)),
                                 mode=str('constant'))[:, shiftx:]
                elif shiftx < 0:
                    img = np.pad(img, ((0, 0), (-shiftx, 0)),
                                 mode=str('constant'))[:, :shiftx]
                if shifty > 0:
                    img = np.pad(img, ((0, shifty), (0, 0)),
                                 mode=str('constant'))[shifty:, :]
                elif shifty < 0:
                    img = np.pad(img, ((-shifty, 0), (0, 0)),
                                 mode=str('constant'))[:shifty, :]
                # The maximum shift applied will determine the final img shape
                maxshiftx = max(abs(shiftx), maxshiftx)
                maxshifty = max(abs(shifty), maxshifty)
            aligned_imgs[band] = img

        if maxshiftx > cx or maxshifty > cy:
            raise VphasDataException('{0}-{1}: bands do not overlap'.format(self.offset, self.ccd))
        # New stretch, scale, and stack the data into an MxNx3 array
        r = aligned_imgs['i'] + aligned_imgs['ha']
        g = aligned_imgs['g'] + aligned_imgs['r'] + aligned_imgs['r2']
        b = aligned_imgs['u'] + 2 * aligned_imgs['g']
        r, g, b = 1.5*r, 0.8*g, 2.2*b
        vmin_r, vmax_r = np.percentile(r, [1., 99.5])
        vmin_g, vmax_g = np.percentile(g, [5., 99.5])
        vmin_b, vmax_b = np.percentile(b, [10., 99.5])
        #log.info((vmin_r, vmin_g, vmin_b))
        #log.info((vmax_r, vmax_g, vmax_b))
        #if vmin_b < 100:
        #    vmin_b = 100
        minrange = np.max((1250., vmax_g-vmin_g, vmax_g-vmin_g, vmax_b-vmin_b))
        if (vmax_r - vmin_r) < minrange:
            vmax_r = vmin_r + minrange
        if (vmax_g - vmin_g) < minrange:
            vmax_g = vmin_g + minrange
        if (vmax_b - vmin_b) < minrange:
            vmax_b = vmin_b + minrange
        interval_r = ManualInterval(vmin_r, vmax_r)
        interval_g = ManualInterval(vmin_g, vmax_g)
        interval_b = ManualInterval(vmin_b, vmax_b)

        r = interval_r(r)
        g = interval_g(g)
        b = interval_b(b)
        stacked = np.dstack((r[maxshifty:-maxshifty, maxshiftx:-maxshiftx],
                             g[maxshifty:-maxshifty, maxshiftx:-maxshiftx],
                             b[maxshifty:-maxshifty, maxshiftx:-maxshiftx]))
        return LuptonColorStretch()(stacked)

    def imsave(self, out_fn=None):
        if out_fn is None:
            out_fn = 'vphas-{0}-{1}.jpg'.format(self.offset, self.ccd)
        log.info('Writing {0}'.format(out_fn))
        #mimg.imsave(out_fn, np.rot90(self.as_array()), origin='lower')
        img = np.rot90(self.as_array())
        imageio.imsave(out_fn, img, quality=90, optimize=True)

############
# FUNCTIONS
############

def vphas_quicklook(offset, ccd=None, out_fn=None):
    VphasColourFrame(offset, ccd).imsave(out_fn)

def vphas_quicklook_main(args=None):
    import argparse

    parser = argparse.ArgumentParser(
        description='Create a beautiful color image from single VPHAS frames.')
    parser.add_argument('-o', '--output', metavar='filename',
                        type=str, default=None,
                        help='Filename for the output image (Default is a '
                             'JPG file named vphas-offset-ccd.jpg)')
    parser.add_argument('-c', '--ccd', nargs='+', type=int, default=None,
                        help='CCD number between 1 and 32.'
                             '(Default is to save all CCDs.)')
    parser.add_argument('--min_percent_r', type=float, default=1.0,
                        help=('The percentile value used to determine the '
                              'minimum cut level for the red channel'))
    parser.add_argument('--max_percent_r', type=float, default=99.5,
                        help=('The percentile value used to determine the '
                              'maximum cut level for the red channel'))
    parser.add_argument('offset', nargs='+',
                        help='Name of the VPHAS offset pointing.')
    args = parser.parse_args(args)

    if args.ccd is None:
        args.ccd = range(1, 33)
    for offset in args.offset:
        try:
            for ccd in args.ccd:
                vphas_quicklook(offset,
                                ccd=ccd,
                                out_fn=args.output)
        except NotObservedException as e:
            log.error(e)

def vst_pawplot(filename, out_fn=None, dpi=100,
                min_percent=1.0, max_percent=99.5,
                cmap='gist_heat', show_hdu=False):
    """Plots the 32-CCD OmegaCam mosaic as a pretty bitmap.

    Parameters
    ----------
    filename : str
        The filename of the FITS file.
    out_fn : str
        The filename of the output bitmap image.  The type of bitmap
        is determined by the filename extension (e.g. '.jpg', '.png').
        The default is a JPG file with the same name as the FITS file.
    dpi : float, optional [dots per inch]
        Resolution of the output 10-by-9 inch output bitmap.
        The default is 100.
    min_percent : float, optional
        The percentile value used to determine the pixel value of
        minimum cut level.  The default is 1.0.
    max_percent : float, optional
        The percentile value used to determine the pixel value of
        maximum cut level.  The default is 99.5.
    cmap : str, optional
        The matplotlib color map name.  The default is 'gist_heat',
        can also be e.g. 'Greys_r'.
    show_hdu : boolean, optional
        Plot the HDU extension number if True (default: False).
    """
    if out_fn is None:
        out_fn = filename + '.jpg'
    if cmap not in pl.cm.datad.keys():
        raise ValueError('{0} is not a valid matplotlib colormap '
                         'name'.format(cmap))
    log.info('Writing {0}'.format(out_fn))
    # Prepare the plot
    f = fits.open(filename)
    pl.interactive(False)
    fig = pl.figure(figsize=(10, 9))
    # Determine vmin/vmax based on a sample of pixels across the mosaic
    sample = np.concatenate((f[1].data[::20, ::10],
                             f[11].data[::20, ::10],
                             f[22].data[::20, ::10],
                             f[32].data[::20, ::10]))
    vmin, vmax = np.percentile(sample, [min_percent, max_percent])
    del sample  # save memory
    log.debug('vst_pawplot: vmin={0}, vmax={1}'.format(vmin, vmax))
    # Plot the extensions
    idx = 0
    for hduno in ProgressBar(OMEGACAM_CCD_ARRANGEMENT):
        idx += 1
        log.debug('vst_pawplot: adding HDU #{0}'.format(hduno))
        ax = fig.add_subplot(4, 8, idx)
        sampling = int(500 / dpi)
        im = ax.matshow(f[hduno].data[::sampling, ::-sampling],
                        norm=LogNorm(vmin=vmin, vmax=vmax),
                        cmap=cmap, origin='lower')
        ax.set_xticks([])
        ax.set_yticks([])
        ax.axis('off')
        if show_hdu:
            # Show the HDU extension number on top of the image
            txt = ax.text(0.05, 0.97, hduno, fontsize=14, color='white',
                          ha='left', va='top', transform=ax.transAxes)
            # Add a black border around the text
            txt.set_path_effects([path_effects.Stroke(linewidth=2, foreground='black'),
                                  path_effects.Normal()])

    # Aesthetics and colorbar
    fig.subplots_adjust(left=0.04, right=0.85,
                        top=0.93, bottom=0.05,
                        wspace=0.02, hspace=0.02)
    cbar_ax = fig.add_axes([0.9, 0.06, 0.02, 0.86])
    t = np.logspace(np.log10(vmin), np.log10(vmax), num=10)
    if (vmax - vmin) > 10:
        fmt = '%.0f'
    elif (vmax - vmin) > 1:
        fmt = '%.1f'
    else:
        fmt = '%.2f'
    cb = fig.colorbar(im, cax=cbar_ax, ticks=t, format=fmt)
    cb.set_label('Pixel count [ADU]')
    # Title and footer text
    fig.text(0.05, 0.95, filename, fontsize=16, ha='left')
    try:
        filterfooter = '{0} ({1}/{2})'.format(
                       f[0].header['ESO INS FILT1 NAME'],
                       f[0].header['ESO TPL EXPNO'],
                       f[0].header['ESO TPL NEXP'])
        fig.text(0.04, 0.02, f[0].header['OBJECT'], fontsize=12, ha='left')
        fig.text(0.50, 0.02, filterfooter, fontsize=12, ha='right')
        fig.text(0.85, 0.02, f[0].header['DATE-OBS'][0:19],
                 fontsize=12, ha='right')
    except KeyError as e:
        log.warning('Could not write footer text: {0}'.format(e))
        pass
    # We're done
    fig.savefig(out_fn, dpi=dpi)
    pl.close()
    del f

def vst_pawplot_main(args=None):
    import argparse

    parser = argparse.ArgumentParser(
        description='Create a beautiful plot showing the data for all 32 CCDs '
                    'in a single image.')
    parser.add_argument('-o', metavar='filename',
                        type=str, default=None,
                        help='Filename for the output image (Default is a '
                        'PNG file with the same name as the FITS file)')
    parser.add_argument('--dpi', type=float, default=100.,
                        help=('The resolution of the output image.'))
    parser.add_argument('--min_percent', type=float, default=1.0,
                        help=('The percentile value used to determine the '
                              'minimum cut level'))
    parser.add_argument('--max_percent', type=float, default=99.5,
                        help=('The percentile value used to determine the '
                              'maximum cut level'))
    parser.add_argument('--cmap', metavar='colormap_name', type=str,
                        default='gist_heat', help='matplotlib color map name')
    parser.add_argument('--show-hdu', action='store_true',
                        help='Display the HDU extension numbers.')
    parser.add_argument('filename', nargs='+',
                        help='Path to one or more FITS files to convert')
    args = parser.parse_args(args)

    for filename in args.filename:
        vst_pawplot(filename,
                    out_fn=args.o,
                    dpi=args.dpi,
                    min_percent=args.min_percent,
                    max_percent=args.max_percent,
                    cmap=args.cmap,
                    show_hdu=args.show_hdu)

if __name__ == '__main__':
    # Example use:
    vst_pawplot(os.path.join(VPHAS_DATA_PATH, 'o20120918_00025.fit'),
                '/tmp/test.jpg')
    vphas_quicklook("0778a", 24, "/tmp/test2.jpg")
<|endoftext|>