<commit_before>from bson import ObjectId
from bson.errors import InvalidId
from collections import namedtuple
from mongoengine import Document, DateTimeField, BooleanField, StringField, DecimalField, ObjectIdField
from mongoengine.errors import InvalidQueryError
from mongoengine.queryset import Q
from pymongo.errors import DuplicateKeyError

FindParams = namedtuple('FindParams', 'projection sort limit request batch_size')
FindParams.__new__.__defaults__ = (None, None, 0, None, 0)

UpdateParams = namedtuple('UpdateByIdParams', 'unset upsert request')
UpdateParams.__new__.__defaults__ = ((), False, None)

class ModelPermissionException(Exception):
    pass

class BaseModel(Document):
    meta = {'abstract': True}

    last_updated = DateTimeField()
    deleted = BooleanField()

    @classmethod
    def allowed_find_query(cls, request):
        '''Returns a query that would find only records this user has permission to see'''
        raise NotImplementedError

    @classmethod
    def allowed_update_query(cls, request):
        '''Returns a query that would find only records this user has permission to edit'''
        raise NotImplementedError

    @classmethod
    def _get_lookup_query(cls, query, allow_deleted=False):
        if 'id' in query:
            query['_id'] = query.pop('id')
        if '_id' in query and isinstance(query['_id'], (str, unicode)):
            try:
                query['_id'] = ObjectId(query['_id'])
            except InvalidId:
                pass

        if not allow_deleted:
            query['deleted'] = None

        return query

    @classmethod
    def _get_lookup_query_find(cls, query, allow_deleted=False, request=None):
        query = cls._get_lookup_query(query, allow_deleted=allow_deleted)
        if request:
            query = {
                '$and': [query, cls.allowed_find_query(request)]
            }
        return query

    @classmethod
    def _get_lookup_query_update(cls, query, allow_deleted=False, request=None):
        query = cls._get_lookup_query(query, allow_deleted=allow_deleted)
        if request:
            query = {
                '$and': [query, cls.allowed_update_query(request)]
            }
        return query

    @classmethod
    def find(cls, params=FindParams(), **kwargs):
        query = cls._get_lookup_query_find(kwargs, request=params.request)
        return cls._get_collection().find(query, projection=params.projection, sort=params.sort,
                                          limit=params.limit, batch_size=params.batch_size)

    @classmethod
    def find_one(cls, params=FindParams(), **kwargs):
        query = cls._get_lookup_query_find(kwargs, request=params.request)
        return cls._get_collection().find_one(query, projection=params.projection)

    @classmethod
    def get_orm(cls, params=FindParams(), **kwargs):
        query = cls._get_lookup_query_find(kwargs, request=params.request)
        try:
            return cls.objects.no_dereference().get(**query)
        except InvalidQueryError:
            # Query contains $or or something
            return cls.objects.no_dereference().get(Q(__raw__=query))

    @classmethod
    def get_orm_by_id(cls, i, params=FindParams()):
        return cls.get_orm(params=params, _id=i)

    @classmethod
    def find_by_id(cls, i, params=FindParams()):
        return cls.find_one(_id=i, params=params)

    @classmethod
    def find_by_ids(cls, ids, params=FindParams()):
        ids = [ObjectId(i) for i in ids]
        return cls.find(_id={'$in': ids}, params=params)

    @classmethod
    def find_by_ids_ordered(cls, ids, params=FindParams(), strict=True):
        # pylint: disable=no-member
        if isinstance(cls.id, ObjectIdField):
            ids = [ObjectId(i) for i in ids]
        docs = {doc['_id']: doc for doc in cls.find(_id={'$in': ids}, params=params)}
        docs = [docs[i] for i in ids if i in docs]
        if strict and len(docs) != len(ids):
            raise Exception('Some docs not found')
        return docs

    @classmethod
    def count(cls, request=None, **kwargs):
        query = cls._get_lookup_query_find(kwargs, request=request)
        return cls._get_collection().count(query)

    @classmethod
    def exists(cls, request=None, **kwargs):
        params = FindParams(request=request, projection={'_id': 1})
        return bool(cls.find_one(params=params, **kwargs))

    @classmethod
    def _get_update_query(cls, unset=(), **kwargs):
        upd = {}
        if '$push' in kwargs:
            upd['$push'] = kwargs.pop('$push')
        if '$addToSet' in kwargs:
            upd['$addToSet'] = kwargs.pop('$addToSet')
        if '$pull' in kwargs:
            upd['$pull'] = kwargs.pop('$pull')
        if '$inc' in kwargs:
            upd['$inc'] = kwargs.pop('$inc')
        if kwargs:
            upd['$set'] = kwargs
        if unset:
            upd['$unset'] = {k: 1 for k in unset}
        return upd

    @classmethod
    def update_one(cls, lookup_dict, update_params=UpdateParams(), **kwargs):
        query = cls._get_lookup_query_update(lookup_dict, request=update_params.request)
        upd = cls._get_update_query(unset=update_params.unset, **kwargs)
        return cls._get_collection().update_one(query, upd, upsert=update_params.upsert)

    @classmethod
    def replace_one(cls, lookup_dict, update_params=UpdateParams(), **kwargs):
        query = cls._get_lookup_query_update(lookup_dict, request=update_params.request)
        return cls._get_collection().replace_one(query, kwargs, upsert=update_params.upsert)

    @classmethod
    def update_many(cls, lookup_dict, update_params=UpdateParams(), **kwargs):
        query = cls._get_lookup_query_update(lookup_dict, request=update_params.request)
        upd = cls._get_update_query(unset=update_params.unset, **kwargs)
        return cls._get_collection().update_many(query, upd)

    @classmethod
    def update_by_id(cls, _id, update_params=UpdateParams(), **kwargs):
        return cls.update_one({'_id': _id}, update_params=update_params, **kwargs)

    @classmethod
    def find_one_and_update(cls, lookup_dict, update, update_params=UpdateParams(), return_document=True, projection=None):
        query = cls._get_lookup_query_update(lookup_dict, request=update_params.request)
        upd = cls._get_update_query(unset=update_params.unset, **update)
        return cls._get_collection().find_one_and_update(query, upd, return_document=return_document,
                                                         projection=projection, upsert=update_params.upsert)

    @classmethod
    def delete_one(cls, request=None, **kwargs):
        query = cls._get_lookup_query_update(kwargs, request=request)
        return cls._get_collection().delete_one(query)

    @classmethod
    def delete_by_id(cls, _id, request=None):
        return cls.delete_one(request=request, _id=_id)

    @classmethod
    def delete_many(cls, request=None, **kwargs):
        query = cls._get_lookup_query_update(kwargs, request=request)
        return cls._get_collection().delete_many(query)

    @classmethod
    def insert_one(cls, doc):  # doc is not ** so insert_one can modify it
        has_id = '_id' in doc
        try:
            return cls._get_collection().insert_one(doc)
        except DuplicateKeyError:
            if not has_id:
                doc.pop('_id')  # Don't add an id if insert failed
            raise

    @classmethod
    def insert_many(cls, objs):
        return cls._get_collection().insert_many(objs)

    @classmethod
    def aggregate(cls, *args, **kwargs):
        return cls._get_collection().aggregate(*args, **kwargs)

    @classmethod
    def get_collection_name(cls):
        return cls._get_collection_name()

    @classmethod
    def msg404(cls, obj_id=None):
        return ("%s%s does not exist or you don't have permissions on it" %
                (cls._get_collection_name(), (' ' + str(obj_id)) if obj_id else ''))

    @classmethod
    def get_unique_indices(cls):
        indices = cls._meta['index_specs']
        return [idx for idx in indices if idx.get('unique')]

class LocationMixin(object):
    SERIALIZE_FIELDS = (
        'street',
        'city',
        'state_code',
        'region',
        'country_code',
        'zip_code',
        'lat',
        'lng',
    )

    street = StringField(max_length=64, required=True)
    city = StringField(max_length=32, required=True)
    state_code = StringField(max_length=2, min_length=2, required=True)
    region = StringField(max_length=32, required=True)
    country_code = StringField(max_length=3, min_length=2, required=True)
    zip_code = StringField(max_length=15, required=True)
    lat = DecimalField(precision=20, required=True, min_value=-90, max_value=90)
    lng = DecimalField(precision=20, required=True, min_value=-180, max_value=180)

    @staticmethod
    def get_location(doc):
        return {k: doc[k] for k in LocationMixin.SERIALIZE_FIELDS if k in doc}

    @staticmethod
    def display_location(doc):
        loc = doc.get('city', '')
        if 'state_code' in doc:
            loc += ', ' + doc['state_code']
        if 'country_code' in doc:
            loc += ', ' + doc['country_code']

        return loc.lstrip(', ')
<commit_msg>Support non-object id fields in find_by_ids<commit_after>from bson import ObjectId
from bson.errors import InvalidId
from collections import namedtuple
from mongoengine import Document, DateTimeField, BooleanField, StringField, DecimalField, ObjectIdField
from mongoengine.errors import InvalidQueryError
from mongoengine.queryset import Q
from pymongo.errors import DuplicateKeyError

FindParams = namedtuple('FindParams', 'projection sort limit request batch_size')
FindParams.__new__.__defaults__ = (None, None, 0, None, 0)

UpdateParams = namedtuple('UpdateByIdParams', 'unset upsert request')
UpdateParams.__new__.__defaults__ = ((), False, None)

class ModelPermissionException(Exception):
    pass

class BaseModel(Document):
    meta = {'abstract': True}

    last_updated = DateTimeField()
    deleted = BooleanField()

    @classmethod
    def allowed_find_query(cls, request):
        '''Returns a query that would find only records this user has permission to see'''
        raise NotImplementedError

    @classmethod
    def allowed_update_query(cls, request):
        '''Returns a query that would find only records this user has permission to edit'''
        raise NotImplementedError

    @classmethod
    def _get_lookup_query(cls, query, allow_deleted=False):
        if 'id' in query:
            query['_id'] = query.pop('id')
        if '_id' in query and isinstance(query['_id'], (str, unicode)):
            try:
                query['_id'] = ObjectId(query['_id'])
            except InvalidId:
                pass

        if not allow_deleted:
            query['deleted'] = None

        return query

    @classmethod
    def _get_lookup_query_find(cls, query, allow_deleted=False, request=None):
        query = cls._get_lookup_query(query, allow_deleted=allow_deleted)
        if request:
            query = {
                '$and': [query, cls.allowed_find_query(request)]
            }
        return query

    @classmethod
    def _get_lookup_query_update(cls, query, allow_deleted=False, request=None):
        query = cls._get_lookup_query(query, allow_deleted=allow_deleted)
        if request:
            query = {
                '$and': [query, cls.allowed_update_query(request)]
            }
        return query

    @classmethod
    def find(cls, params=FindParams(), **kwargs):
        query = cls._get_lookup_query_find(kwargs, request=params.request)
        return cls._get_collection().find(query, projection=params.projection, sort=params.sort,
                                          limit=params.limit, batch_size=params.batch_size)

    @classmethod
    def find_one(cls, params=FindParams(), **kwargs):
        query = cls._get_lookup_query_find(kwargs, request=params.request)
        return cls._get_collection().find_one(query, projection=params.projection)

    @classmethod
    def get_orm(cls, params=FindParams(), **kwargs):
        query = cls._get_lookup_query_find(kwargs, request=params.request)
        try:
            return cls.objects.no_dereference().get(**query)
        except InvalidQueryError:
            # Query contains $or or something
            return cls.objects.no_dereference().get(Q(__raw__=query))

    @classmethod
    def get_orm_by_id(cls, i, params=FindParams()):
        return cls.get_orm(params=params, _id=i)

    @classmethod
    def find_by_id(cls, i, params=FindParams()):
        return cls.find_one(_id=i, params=params)

    @classmethod
    def find_by_ids(cls, ids, params=FindParams()):
        if isinstance(cls.id, ObjectIdField):
            ids = [ObjectId(i) for i in ids]
        return cls.find(_id={'$in': ids}, params=params)

    @classmethod
    def find_by_ids_ordered(cls, ids, params=FindParams(), strict=True):
        # pylint: disable=no-member
        if isinstance(cls.id, ObjectIdField):
            ids = [ObjectId(i) for i in ids]
        docs = {doc['_id']: doc for doc in cls.find(_id={'$in': ids}, params=params)}
        docs = [docs[i] for i in ids if i in docs]
        if strict and len(docs) != len(ids):
            raise Exception('Some docs not found')
        return docs

    @classmethod
    def count(cls, request=None, **kwargs):
        query = cls._get_lookup_query_find(kwargs, request=request)
        return cls._get_collection().count(query)

    @classmethod
    def exists(cls, request=None, **kwargs):
        params = FindParams(request=request, projection={'_id': 1})
        return bool(cls.find_one(params=params, **kwargs))

    @classmethod
    def _get_update_query(cls, unset=(), **kwargs):
        upd = {}
        if '$push' in kwargs:
            upd['$push'] = kwargs.pop('$push')
        if '$addToSet' in kwargs:
            upd['$addToSet'] = kwargs.pop('$addToSet')
        if '$pull' in kwargs:
            upd['$pull'] = kwargs.pop('$pull')
        if '$inc' in kwargs:
            upd['$inc'] = kwargs.pop('$inc')
        if kwargs:
            upd['$set'] = kwargs
        if unset:
            upd['$unset'] = {k: 1 for k in unset}
        return upd

    @classmethod
    def update_one(cls, lookup_dict, update_params=UpdateParams(), **kwargs):
        query = cls._get_lookup_query_update(lookup_dict, request=update_params.request)
        upd = cls._get_update_query(unset=update_params.unset, **kwargs)
        return cls._get_collection().update_one(query, upd, upsert=update_params.upsert)

    @classmethod
    def replace_one(cls, lookup_dict, update_params=UpdateParams(), **kwargs):
        query = cls._get_lookup_query_update(lookup_dict, request=update_params.request)
        return cls._get_collection().replace_one(query, kwargs, upsert=update_params.upsert)

    @classmethod
    def update_many(cls, lookup_dict, update_params=UpdateParams(), **kwargs):
        query = cls._get_lookup_query_update(lookup_dict, request=update_params.request)
        upd = cls._get_update_query(unset=update_params.unset, **kwargs)
        return cls._get_collection().update_many(query, upd)

    @classmethod
    def update_by_id(cls, _id, update_params=UpdateParams(), **kwargs):
        return cls.update_one({'_id': _id}, update_params=update_params, **kwargs)

    @classmethod
    def find_one_and_update(cls, lookup_dict, update, update_params=UpdateParams(), return_document=True, projection=None):
        query = cls._get_lookup_query_update(lookup_dict, request=update_params.request)
        upd = cls._get_update_query(unset=update_params.unset, **update)
        return cls._get_collection().find_one_and_update(query, upd, return_document=return_document,
                                                         projection=projection, upsert=update_params.upsert)

    @classmethod
    def delete_one(cls, request=None, **kwargs):
        query = cls._get_lookup_query_update(kwargs, request=request)
        return cls._get_collection().delete_one(query)

    @classmethod
    def delete_by_id(cls, _id, request=None):
        return cls.delete_one(request=request, _id=_id)

    @classmethod
    def delete_many(cls, request=None, **kwargs):
        query = cls._get_lookup_query_update(kwargs, request=request)
        return cls._get_collection().delete_many(query)

    @classmethod
    def insert_one(cls, doc):  # doc is not ** so insert_one can modify it
        has_id = '_id' in doc
        try:
            return cls._get_collection().insert_one(doc)
        except DuplicateKeyError:
            if not has_id:
                doc.pop('_id')  # Don't add an id if insert failed
            raise

    @classmethod
    def insert_many(cls, objs):
        return cls._get_collection().insert_many(objs)

    @classmethod
    def aggregate(cls, *args, **kwargs):
        return cls._get_collection().aggregate(*args, **kwargs)

    @classmethod
    def get_collection_name(cls):
        return cls._get_collection_name()

    @classmethod
    def msg404(cls, obj_id=None):
        return ("%s%s does not exist or you don't have permissions on it" %
                (cls._get_collection_name(), (' ' + str(obj_id)) if obj_id else ''))

    @classmethod
    def get_unique_indices(cls):
        indices = cls._meta['index_specs']
        return [idx for idx in indices if idx.get('unique')]

class LocationMixin(object):
    SERIALIZE_FIELDS = (
        'street',
        'city',
        'state_code',
        'region',
        'country_code',
        'zip_code',
        'lat',
        'lng',
    )

    street = StringField(max_length=64, required=True)
    city = StringField(max_length=32, required=True)
    state_code = StringField(max_length=2, min_length=2, required=True)
    region = StringField(max_length=32, required=True)
    country_code = StringField(max_length=3, min_length=2, required=True)
    zip_code = StringField(max_length=15, required=True)
    lat = DecimalField(precision=20, required=True, min_value=-90, max_value=90)
    lng = DecimalField(precision=20, required=True, min_value=-180, max_value=180)

    @staticmethod
    def get_location(doc):
        return {k: doc[k] for k in LocationMixin.SERIALIZE_FIELDS if k in doc}

    @staticmethod
    def display_location(doc):
        loc = doc.get('city', '')
        if 'state_code' in doc:
            loc += ', ' + doc['state_code']
        if 'country_code' in doc:
            loc += ', ' + doc['country_code']

        return loc.lstrip(', ')
<|endoftext|>


<commit_before>import re
import urllib

from django.conf import settings
from django.contrib.auth import (
    REDIRECT_FIELD_NAME, authenticate, login as auth_login)
from django.core.urlresolvers import reverse
from django.http import HttpResponse, HttpResponseRedirect
from django.shortcuts import render_to_response
from django.template import RequestContext
from django.utils.safestring import SafeString

from openid.consumer.consumer import (
    Consumer, SUCCESS, CANCEL, FAILURE)
from openid.consumer.discover import DiscoveryFailure
from openid.extensions import sreg

from django_openid_auth import teams
from django_openid_auth.forms import OpenIDLoginForm
from django_openid_auth.store import DjangoOpenIDStore

next_url_re = re.compile('^/[-\w/]+$')

def is_valid_next_url(next):
    # When we allow this:
    #   /openid/?next=/welcome/
    # For security reasons we want to restrict the next= bit to being a local
    # path, not a complete URL.
    return bool(next_url_re.match(next))

def sanitise_redirect_url(redirect_to):
    """Sanitise the redirection URL."""
    # Light security check -- make sure redirect_to isn't garbage.
    if not redirect_to or '//' in redirect_to or ' ' in redirect_to:
        redirect_to = settings.LOGIN_REDIRECT_URL
    return redirect_to

def make_consumer(request):
    """Create an OpenID Consumer object for the given Django request."""
    # Give the OpenID library its own space in the session object.
    session = request.session.setdefault('OPENID', {})
    store = DjangoOpenIDStore()
    return Consumer(session, store)

def render_openid_request(request, openid_request, return_to, trust_root=None):
    """Render an OpenID authentication request."""
    if trust_root is None:
        trust_root = getattr(settings, 'OPENID_TRUST_ROOT',
                             request.build_absolute_uri('/'))

    if openid_request.shouldSendRedirect():
        redirect_url = openid_request.redirectURL(
            trust_root, return_to)
        return HttpResponseRedirect(redirect_url)
    else:
        form_html = openid_request.htmlMarkup(
            trust_root, return_to, form_tag_attrs={'id': 'openid_message'})
        return HttpResponse(form_html, content_type='text/html;charset=UTF-8')

def parse_openid_response(request):
    """Parse an OpenID response from a Django request."""
    # Short cut if there is no request parameters.
    #if len(request.REQUEST) == 0:
    #    return None

    current_url = request.build_absolute_uri()

    consumer = make_consumer(request)
    return consumer.complete(dict(request.REQUEST.items()), current_url)

def login_begin(request, template_name='openid/login.html',
                redirect_field_name=REDIRECT_FIELD_NAME):
    """Begin an OpenID login request, possibly asking for an identity URL."""
    redirect_to = request.REQUEST.get(redirect_field_name, '')

    # Get the OpenID URL to try.  First see if we've been configured
    # to use a fixed server URL.
    openid_url = getattr(settings, 'OPENID_SSO_SERVER_URL', None)

    if openid_url is None:
        if request.POST:
            login_form = OpenIDLoginForm(data=request.POST)
            if login_form.is_valid():
                openid_url = login_form.cleaned_data['openid_identifier']
        else:
            login_form = OpenIDLoginForm()

        # Invalid or no form data:
        if openid_url is None:
            return render_to_response(template_name, {
                    'form': login_form,
                    redirect_field_name: redirect_to
                    }, context_instance=RequestContext(request))

    error = None
    consumer = make_consumer(request)
    try:
        openid_request = consumer.begin(openid_url)
    except DiscoveryFailure, exc:
        # XXX: make this a proper error page.
        return HttpResponse("OpenID discovery error: %s" % (str(exc),))

    # Request some user details.
    openid_request.addExtension(
        sreg.SRegRequest(optional=['email', 'fullname', 'nickname']))

    # Request team info
    launchpad_teams = getattr(settings, 'OPENID_LAUNCHPAD_TEAMS_MAPPING',
                              {})
    openid_request.addExtension(teams.TeamsRequest(launchpad_teams.keys()))

    # Construct the request completion URL, including the page we
    # should redirect to.
    return_to = request.build_absolute_uri(reverse(login_complete))
    if redirect_to:
        if '?' in return_to:
            return_to += '&'
        else:
            return_to += '?'
        return_to += urllib.urlencode({redirect_field_name: redirect_to})

    return render_openid_request(request, openid_request, return_to)

def login_complete(request, redirect_field_name=REDIRECT_FIELD_NAME):
    redirect_to = request.REQUEST.get(redirect_field_name, '')

    openid_response = parse_openid_response(request)
    if not openid_response:
        return HttpResponse('This is an OpenID relying party endpoint.')

    if openid_response.status == SUCCESS:
        user = authenticate(openid_response=openid_response)
        if user is not None:
            if user.is_active:
                auth_login(request, user)
                return HttpResponseRedirect(sanitise_redirect_url(redirect_to))
            else:
                return HttpResponse('Disabled account')
        else:
            return HttpResponse('invalid user')
    elif openid_response.status == FAILURE:
        return HttpResponse("Failure: %s" % openid_response.message)
    elif openid_response.status == CANCEL:
        return HttpResponse("Cancel")
    else:
        assert False, (
            "Unknown OpenID response type: %r" % openid_response.status)

def logo(request):
    return HttpResponse(
        OPENID_LOGO_BASE_64.decode('base64'), mimetype='image/gif'
    )

# Logo from http://openid.net/login-bg.gif
# Embedded here for convenience; you should serve this as a static file
OPENID_LOGO_BASE_64 = """
R0lGODlhEAAQAMQAAO3t7eHh4srKyvz8/P5pDP9rENLS0v/28P/17tXV1dHEvPDw8M3Nzfn5+d3d
3f5jA97Syvnv6MfLzcfHx/1mCPx4Kc/S1Pf189C+tP+xgv/k1N3OxfHy9NLV1/39/f///yH5BAAA
AAAALAAAAAAQABAAAAVq4CeOZGme6KhlSDoexdO6H0IUR+otwUYRkMDCUwIYJhLFTyGZJACAwQcg
EAQ4kVuEE2AIGAOPQQAQwXCfS8KQGAwMjIYIUSi03B7iJ+AcnmclHg4TAh0QDzIpCw4WGBUZeikD
Fzk0lpcjIQA7
"""
<commit_msg>Only add the teams extension to the request if we have teams to check.<commit_after>import re
import urllib

from django.conf import settings
from django.contrib.auth import (
    REDIRECT_FIELD_NAME, authenticate, login as auth_login)
from django.core.urlresolvers import reverse
from django.http import HttpResponse, HttpResponseRedirect
from django.shortcuts import render_to_response
from django.template import RequestContext
from django.utils.safestring import SafeString

from openid.consumer.consumer import (
    Consumer, SUCCESS, CANCEL, FAILURE)
from openid.consumer.discover import DiscoveryFailure
from openid.extensions import sreg

from django_openid_auth import teams
from django_openid_auth.forms import OpenIDLoginForm
from django_openid_auth.store import DjangoOpenIDStore

next_url_re = re.compile('^/[-\w/]+$')

def is_valid_next_url(next):
    # When we allow this:
    #   /openid/?next=/welcome/
    # For security reasons we want to restrict the next= bit to being a local
    # path, not a complete URL.
    return bool(next_url_re.match(next))

def sanitise_redirect_url(redirect_to):
    """Sanitise the redirection URL."""
    # Light security check -- make sure redirect_to isn't garbage.
    if not redirect_to or '//' in redirect_to or ' ' in redirect_to:
        redirect_to = settings.LOGIN_REDIRECT_URL
    return redirect_to

def make_consumer(request):
    """Create an OpenID Consumer object for the given Django request."""
    # Give the OpenID library its own space in the session object.
    session = request.session.setdefault('OPENID', {})
    store = DjangoOpenIDStore()
    return Consumer(session, store)

def render_openid_request(request, openid_request, return_to, trust_root=None):
    """Render an OpenID authentication request."""
    if trust_root is None:
        trust_root = getattr(settings, 'OPENID_TRUST_ROOT',
                             request.build_absolute_uri('/'))

    if openid_request.shouldSendRedirect():
        redirect_url = openid_request.redirectURL(
            trust_root, return_to)
        return HttpResponseRedirect(redirect_url)
    else:
        form_html = openid_request.htmlMarkup(
            trust_root, return_to, form_tag_attrs={'id': 'openid_message'})
        return HttpResponse(form_html, content_type='text/html;charset=UTF-8')

def parse_openid_response(request):
    """Parse an OpenID response from a Django request."""
    # Short cut if there is no request parameters.
    #if len(request.REQUEST) == 0:
    #    return None

    current_url = request.build_absolute_uri()

    consumer = make_consumer(request)
    return consumer.complete(dict(request.REQUEST.items()), current_url)

def login_begin(request, template_name='openid/login.html',
                redirect_field_name=REDIRECT_FIELD_NAME):
    """Begin an OpenID login request, possibly asking for an identity URL."""
    redirect_to = request.REQUEST.get(redirect_field_name, '')

    # Get the OpenID URL to try.  First see if we've been configured
    # to use a fixed server URL.
    openid_url = getattr(settings, 'OPENID_SSO_SERVER_URL', None)

    if openid_url is None:
        if request.POST:
            login_form = OpenIDLoginForm(data=request.POST)
            if login_form.is_valid():
                openid_url = login_form.cleaned_data['openid_identifier']
        else:
            login_form = OpenIDLoginForm()

        # Invalid or no form data:
        if openid_url is None:
            return render_to_response(template_name, {
                    'form': login_form,
                    redirect_field_name: redirect_to
                    }, context_instance=RequestContext(request))

    error = None
    consumer = make_consumer(request)
    try:
        openid_request = consumer.begin(openid_url)
    except DiscoveryFailure, exc:
        # XXX: make this a proper error page.
        return HttpResponse("OpenID discovery error: %s" % (str(exc),))

    # Request some user details.
    openid_request.addExtension(
        sreg.SRegRequest(optional=['email', 'fullname', 'nickname']))

    # Request team info
    launchpad_teams = getattr(settings, 'OPENID_LAUNCHPAD_TEAMS_MAPPING', {})
    if launchpad_teams:
        openid_request.addExtension(teams.TeamsRequest(launchpad_teams.keys()))

    # Construct the request completion URL, including the page we
    # should redirect to.
    return_to = request.build_absolute_uri(reverse(login_complete))
    if redirect_to:
        if '?' in return_to:
            return_to += '&'
        else:
            return_to += '?'
        return_to += urllib.urlencode({redirect_field_name: redirect_to})

    return render_openid_request(request, openid_request, return_to)

def login_complete(request, redirect_field_name=REDIRECT_FIELD_NAME):
    redirect_to = request.REQUEST.get(redirect_field_name, '')

    openid_response = parse_openid_response(request)
    if not openid_response:
        return HttpResponse('This is an OpenID relying party endpoint.')

    if openid_response.status == SUCCESS:
        user = authenticate(openid_response=openid_response)
        if user is not None:
            if user.is_active:
                auth_login(request, user)
                return HttpResponseRedirect(sanitise_redirect_url(redirect_to))
            else:
                return HttpResponse('Disabled account')
        else:
            return HttpResponse('invalid user')
    elif openid_response.status == FAILURE:
        return HttpResponse("Failure: %s" % openid_response.message)
    elif openid_response.status == CANCEL:
        return HttpResponse("Cancel")
    else:
        assert False, (
            "Unknown OpenID response type: %r" % openid_response.status)

def logo(request):
    return HttpResponse(
        OPENID_LOGO_BASE_64.decode('base64'), mimetype='image/gif'
    )

# Logo from http://openid.net/login-bg.gif
# Embedded here for convenience; you should serve this as a static file
OPENID_LOGO_BASE_64 = """
R0lGODlhEAAQAMQAAO3t7eHh4srKyvz8/P5pDP9rENLS0v/28P/17tXV1dHEvPDw8M3Nzfn5+d3d
3f5jA97Syvnv6MfLzcfHx/1mCPx4Kc/S1Pf189C+tP+xgv/k1N3OxfHy9NLV1/39/f///yH5BAAA
AAAALAAAAAAQABAAAAVq4CeOZGme6KhlSDoexdO6H0IUR+otwUYRkMDCUwIYJhLFTyGZJACAwQcg
EAQ4kVuEE2AIGAOPQQAQwXCfS8KQGAwMjIYIUSi03B7iJ+AcnmclHg4TAh0QDzIpCw4WGBUZeikD
Fzk0lpcjIQA7
"""
<|endoftext|>


<commit_before>__version__ = '0.2.0'
<commit_msg>Bump to 0.3.0<commit_after># -*- coding: utf-8 -*-

__version__ = '0.3.0'
<|endoftext|>


<commit_before>__version__ = '0.1.14'

default_app_config = 'djangocms_forms.apps.DjangoCMSFormsConfig'
<commit_msg>Bump Version<commit_after>__version__ = '0.1.15'

default_app_config = 'djangocms_forms.apps.DjangoCMSFormsConfig'
<|endoftext|>


<commit_before>#!/usr/bin/env python
# encoding: utf-8
# pylint: disable=no-member, no-init, too-many-public-methods
# pylint: disable=attribute-defined-outside-init
# pylint: disable=missing-docstring, unused-import, invalid-name, import-error, super-on-old-class
"""Tests for final submissions."""

import os
os.environ['FLASK_CONF'] = 'TEST'
import datetime

from test_base import APIBaseTestCase, unittest

from app import models, utils, constants

from google.appengine.ext import ndb
from google.appengine.ext import deferred
from google.appengine.ext import testbed

class FinalSubmissionTest(APIBaseTestCase):
    def get_accounts(self):
        """
        Returns the accounts you want to exist in your system.
        """
        keys = ["student0", "student1", "student2", "staff", "admin"]
        rval = {key: models.User(email=[key+"@b.edu"]) for key in keys}
        for k, val in rval.items():
            if 'admin' in k:
                val.is_admin = True
        return rval

    def setUp(self):
        super(APIBaseTestCase, self).setUp()
        self.accounts = self.get_accounts()
        for user in self.accounts.values():
            user.put()

        self.courses = {
            "first": models.Course(
                institution="UC Awesome",
                display_name="First Course",
                instructor=[self.accounts['admin'].key]),
            "second": models.Course(
                institution="UC Awesome",
                display_name="Second Course",
                instructor=[self.accounts['admin'].key]),
            }

        for course in self.courses.values():
            course.put()

        for student in ["student0", "student1", "student2"]:
            models.Participant.add_role(
                self.accounts[student], self.courses['first'], constants.STUDENT_ROLE)

        self.assignments = {
            "first": models.Assignment(
                name="first",
                points=3,
                creator=self.accounts["admin"].key,
                course=self.courses['first'].key,
                display_name="first display",
                templates="{}",
                max_group_size=3,
                due_date=datetime.datetime.now()
                ),
            }
        for assign in self.assignments.values():
            assign.put()
        self.assign = self.assignments["first"]

    def run_deferred(self):
        """Execute all deferred tasks."""
        for task in self.taskqueue_stub.get_filtered_tasks():
            deferred.run(task.payload)

    def test_one_final(self):
        """An invite/accept/exit/invite sequence keeps a final submission."""
        self.user = self.accounts['student0']

        # Submit
        messages = {'file_contents': {'submit': True, 'trends.py': 'hi!'}}
        self.post_json('/submission?{}={}'.format(
            'access_token', self.accounts['admin'].email[0]),
            data={'assignment': self.assign.name,
                  'submitter': self.user.key.id(),
                  'messages': messages,
                  'access_token': self.user.email[0]})
        self.run_deferred()

        finals = list(models.FinalSubmission.query().fetch())
        self.assertEqual(1, len(finals))
        final = finals[0]
        subm = final.submission.get()
        self.assertIsNotNone(subm)
        backup = subm.backup.get()
        self.assertIsNotNone(backup)
        self.assertEqual(backup.submitter, self.user.key)
        # TODO Not sure how to make/verify this final_submission get request...
        # self.assertEqual(final, self.user.get_final_submission(self.assign))
        # self.get('/user/{}/final_submission'.format(self.user.email[0]),
        #          data={'assignment': self.assign.key.id()})

        # Invite
        #invited = self.accounts['student1']
        # TODO This post is being made with admin as the user; not sure why...
        #self.post_json('/assignment/{}/invite'.format(self.assign.key.id()),
        #               data={'email': invited.email[0]})
        # TODO Check final submissions

        # Accept
        # TODO

        # Exit
        # TODO

        # Invite
        # TODO

if __name__ == "__main__":
    unittest.main()
<commit_msg>Added more verbose test failure.<commit_after>#!/usr/bin/env python
# encoding: utf-8
# pylint: disable=no-member, no-init, too-many-public-methods
# pylint: disable=attribute-defined-outside-init
# pylint: disable=missing-docstring, unused-import, invalid-name, import-error, super-on-old-class
"""Tests for final submissions."""

import os
os.environ['FLASK_CONF'] = 'TEST'
import datetime

from test_base import APIBaseTestCase, unittest

from app import models, utils, constants

from google.appengine.ext import ndb
from google.appengine.ext import deferred
from google.appengine.ext import testbed

class FinalSubmissionTest(APIBaseTestCase):
    def get_accounts(self):
        """
        Returns the accounts you want to exist in your system.
        """
        keys = ["student0", "student1", "student2", "staff", "admin"]
        rval = {key: models.User(email=[key+"@b.edu"]) for key in keys}
        for k, val in rval.items():
            if 'admin' in k:
                val.is_admin = True
        return rval

    def setUp(self):
        super(APIBaseTestCase, self).setUp()
        self.accounts = self.get_accounts()
        for user in self.accounts.values():
            user.put()

        self.courses = {
            "first": models.Course(
                institution="UC Awesome",
                display_name="First Course",
                instructor=[self.accounts['admin'].key]),
            "second": models.Course(
                institution="UC Awesome",
                display_name="Second Course",
                instructor=[self.accounts['admin'].key]),
            }

        for course in self.courses.values():
            course.put()

        for student in ["student0", "student1", "student2"]:
            models.Participant.add_role(
                self.accounts[student], self.courses['first'], constants.STUDENT_ROLE)

        self.assignments = {
            "first": models.Assignment(
                name="first",
                points=3,
                creator=self.accounts["admin"].key,
                course=self.courses['first'].key,
                display_name="first display",
                templates="{}",
                max_group_size=3,
                due_date=datetime.datetime.now()
                ),
            }
        for assign in self.assignments.values():
            assign.put()
        self.assign = self.assignments["first"]

    def run_deferred(self):
        """Execute all deferred tasks."""
        for task in self.taskqueue_stub.get_filtered_tasks():
            deferred.run(task.payload)

    def test_one_final(self):
        """An invite/accept/exit/invite sequence keeps a final submission."""
        self.user = self.accounts['student0']

        # Submit
        messages = {'file_contents': {'submit': True, 'trends.py': 'hi!'}}
        self.post_json('/submission?{}={}'.format(
            'access_token', self.accounts['admin'].email[0]),
            data={'assignment': self.assign.name,
                  'submitter': self.user.key.id(),
                  'messages': messages,
                  'access_token': self.user.email[0]})
        self.run_deferred()

        finals = list(models.FinalSubmission.query().fetch())
        self.assertEqual(1, len(finals))
        final = finals[0]
        subm = final.submission.get()
        self.assertIsNotNone(subm)
        backup = subm.backup.get()
        self.assertIsNotNone(backup)
        self.assertEqual(backup.submitter, self.user.key,
                "{} != {}".format(backup.submitter.get(), self.user))
        # TODO Not sure how to make/verify this final_submission get request...
        # self.assertEqual(final, self.user.get_final_submission(self.assign))
        # self.get('/user/{}/final_submission'.format(self.user.email[0]),
        #          data={'assignment': self.assign.key.id()})

        # Invite
        #invited = self.accounts['student1']
        # TODO This post is being made with admin as the user; not sure why...
        #self.post_json('/assignment/{}/invite'.format(self.assign.key.id()),
        #               data={'email': invited.email[0]})
        # TODO Check final submissions

        # Accept
        # TODO

        # Exit
        # TODO

        # Invite
        # TODO

if __name__ == "__main__":
    unittest.main()
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os, sys
# env
sys.path.append('/usr/lib/python2.7/dist-packages/')
sys.path.append('/usr/lib/python2.7/')
sys.path.append('/usr/local/lib/python2.7/dist-packages/')
sys.path.append('/data2/django_1.11/')
sys.path.append('/data2/django_projects/')
sys.path.append('/data2/django_third/')
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'djequis.settings')

import django
django.setup()

from django.conf import settings

# informix environment
os.environ['INFORMIXSERVER'] = settings.INFORMIXSERVER
os.environ['DBSERVERNAME'] = settings.DBSERVERNAME
os.environ['INFORMIXDIR'] = settings.INFORMIXDIR
os.environ['ODBCINI'] = settings.ODBCINI
os.environ['ONCONFIG'] = settings.ONCONFIG
os.environ['INFORMIXSQLHOSTS'] = settings.INFORMIXSQLHOSTS
os.environ['LD_LIBRARY_PATH'] = settings.LD_LIBRARY_PATH
os.environ['LD_RUN_PATH'] = settings.LD_RUN_PATH

from djequis.sql.provisioning import INSERT_EMAIL_RECORD
from djequis.sql.provisioning import INSERT_CVID_RECORD
from djequis.sql.provisioning import SELECT_NEW_PEOPLE

from djzbar.utils.informix import do_sql
from djzbar.settings import INFORMIX_EARL_TEST
from djzbar.settings import INFORMIX_EARL_PROD

from openpyxl import Workbook
from openpyxl import load_workbook

import csv
import time
import argparse
import logging

"""
provising framework for new students and employees
"""

# set up command-line options
desc = """
Accepts as input a database name
"""

parser = argparse.ArgumentParser(description=desc)

parser.add_argument(
    '-d', '--database',
    required=True,
    help="Database name (cars or train).",
    dest='database'
)
parser.add_argument(
    '-f', '--filetype',
    required=True,
    help="File type (csv or xlsx).",
    dest='filetype'
)
parser.add_argument(
    '--test',
    action='store_true',
    help="Dry run?",
    dest='test'
)

TIMESTAMP = time.strftime("%Y%m%d%H%M%S")

# create logger for collecting failed inserts
logger = logging.getLogger(__name__)
# create handler and set level to info
handler = logging.FileHandler('{}provisioning.log'.format(
    settings.LOG_FILEPATH)
)
handler.setLevel(logging.INFO)
formatter = logging.Formatter(
    '%(asctime)s: %(levelname)s: %(message)s', datefmt='%m/%d/%Y %I:%M:%S %p'
)
handler.setFormatter(formatter)
logger.addHandler(handler)

def _generate_files(results, filetype, group):
    """
    Active Directory required fields in order:

        loginID, lastName, firstName, nameID,
        [facultyStatus, staffStatus, studentStatus, retireStatus,]
        dob, zip, acct-types, proxID, phoneExt, depts

    at least one of the Status fields must be populated
    """

    if test:
        sendero = settings.PROVISIONING_DATA_DIRECTORY_TEST
    else:
        sendero = settings.PROVISIONING_DATA_DIRECTORY

    root = '{}{}_{}'.format(sendero, group, TIMESTAMP)

    if filetype == 'csv':
        # create .csv file
        csvphile = ('{}.csv'.format(root))
        phile = open(csvphile,"w")
        output = csv.writer(phile, quoting=csv.QUOTE_ALL)

        for result in results:
            output.writerow(result)

        # close the csv file
        phile.close()

    elif filetype == 'xlsx':
        # load our XLSX template
        phile = load_workbook(
            '{}/static/xml/{}.xlsx'.format(settings.ROOT_DIR, group)
        )
        # obtain the active worksheet
        ws = phile.active

        for result in results:
            row = []
            for r in result:
                row.append(r)
            ws.append(row)

        # Save the xml file
        phile.save('{}.xlsx'.format(root))

    else:
        print("filetype must be: 'csv' or 'xlsx'\n")
        phile = None
        parser.print_help()
        exit(-1)

    return phile

def main():
    """
    main function
    """

    key = 'debug'

    if filetype not in ['csv','xlsx']:
        print("filetype must be: 'csv' or 'xlsx'\n")
        parser.print_help()
        exit(-1)

    if database == 'train':
        EARL = INFORMIX_EARL_TEST
    elif database == 'cars':
        EARL = INFORMIX_EARL_PROD
    else:
        print("database must be: 'cars' or 'train'\n")
        parser.print_help()
        exit(-1)

    sql = SELECT_NEW_PEOPLE
    if test:
        print("new people sql")
        print("sql = {}".format(sql))

    people = []
    objects = do_sql(sql, key=key, earl=EARL)

    # the people list allows us to iterate over the result set more
    # than once, whereas just using objects result would throw an
    # error after the first iteration.
    for o in objects:
        if test:
            print(o)
        people.append(o)

    if people:
        response = _generate_files(people, filetype, 'new_people')

        if not response:
            print("no response")
        else:
            for p in people:
                try:
                    sql = INSERT_EMAIL_RECORD(cid=p.id, ldap=p.loginid)
                    if test:
                        print(sql)
                    else:
                        do_sql(sql, key=key, earl=EARL)
                except:
                    logger.info("failed insert = {}".format(p))

                try:
                    sql = INSERT_CVID_RECORD(cid=p.id, ldap=p.loginid)
                    if test:
                        print(sql)
                    else:
                        do_sql(sql, key=key, earl=EARL)
                except:
                    logger.info("failed insert = {}".format(p))
    else:
        print("No objects found for provisioning.")

######################
# shell command line
######################

if __name__ == '__main__':
    args = parser.parse_args()
    test = args.test
    database = args.database.lower()
    filetype = args.filetype.lower()
    sys.exit(main())
<commit_msg>added output for when we find folks who need provisioning and not using --test flag<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os, sys
# env
sys.path.append('/usr/lib/python2.7/dist-packages/')
sys.path.append('/usr/lib/python2.7/')
sys.path.append('/usr/local/lib/python2.7/dist-packages/')
sys.path.append('/data2/django_1.11/')
sys.path.append('/data2/django_projects/')
sys.path.append('/data2/django_third/')
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'djequis.settings')

import django
django.setup()

from django.conf import settings

# informix environment
os.environ['INFORMIXSERVER'] = settings.INFORMIXSERVER
os.environ['DBSERVERNAME'] = settings.DBSERVERNAME
os.environ['INFORMIXDIR'] = settings.INFORMIXDIR
os.environ['ODBCINI'] = settings.ODBCINI
os.environ['ONCONFIG'] = settings.ONCONFIG
os.environ['INFORMIXSQLHOSTS'] = settings.INFORMIXSQLHOSTS
os.environ['LD_LIBRARY_PATH'] = settings.LD_LIBRARY_PATH
os.environ['LD_RUN_PATH'] = settings.LD_RUN_PATH

from djequis.sql.provisioning import INSERT_EMAIL_RECORD
from djequis.sql.provisioning import INSERT_CVID_RECORD
from djequis.sql.provisioning import SELECT_NEW_PEOPLE

from djzbar.utils.informix import do_sql
from djzbar.settings import INFORMIX_EARL_TEST
from djzbar.settings import INFORMIX_EARL_PROD

from openpyxl import Workbook
from openpyxl import load_workbook

import csv
import time
import argparse
import logging

"""
provising framework for new students and employees
"""

# set up command-line options
desc = """
Accepts as input a database name
"""

parser = argparse.ArgumentParser(description=desc)

parser.add_argument(
    '-d', '--database',
    required=True,
    help="Database name (cars or train).",
    dest='database'
)
parser.add_argument(
    '-f', '--filetype',
    required=True,
    help="File type (csv or xlsx).",
    dest='filetype'
)
parser.add_argument(
    '--test',
    action='store_true',
    help="Dry run?",
    dest='test'
)

TIMESTAMP = time.strftime("%Y%m%d%H%M%S")

# create logger for collecting failed inserts
logger = logging.getLogger(__name__)
# create handler and set level to info
handler = logging.FileHandler('{}provisioning.log'.format(
    settings.LOG_FILEPATH)
)
handler.setLevel(logging.INFO)
formatter = logging.Formatter(
    '%(asctime)s: %(levelname)s: %(message)s', datefmt='%m/%d/%Y %I:%M:%S %p'
)
handler.setFormatter(formatter)
logger.addHandler(handler)

def _generate_files(results, filetype, group):
    """
    Active Directory required fields in order:

        loginID, lastName, firstName, nameID,
        [facultyStatus, staffStatus, studentStatus, retireStatus,]
        dob, zip, acct-types, proxID, phoneExt, depts

    at least one of the Status fields must be populated
    """

    if test:
        sendero = settings.PROVISIONING_DATA_DIRECTORY_TEST
    else:
        sendero = settings.PROVISIONING_DATA_DIRECTORY

    root = '{}{}_{}'.format(sendero, group, TIMESTAMP)

    if filetype == 'csv':
        # create .csv file
        csvphile = ('{}.csv'.format(root))
        phile = open(csvphile,"w")
        output = csv.writer(phile, quoting=csv.QUOTE_ALL)

        for result in results:
            output.writerow(result)

        # close the csv file
        phile.close()

    elif filetype == 'xlsx':
        # load our XLSX template
        phile = load_workbook(
            '{}/static/xml/{}.xlsx'.format(settings.ROOT_DIR, group)
        )
        # obtain the active worksheet
        ws = phile.active

        for result in results:
            row = []
            for r in result:
                row.append(r)
            ws.append(row)

        # Save the xml file
        phile.save('{}.xlsx'.format(root))

    else:
        print("filetype must be: 'csv' or 'xlsx'\n")
        phile = None
        parser.print_help()
        exit(-1)

    return phile

def main():
    """
    main function
    """

    key = 'debug'

    if filetype not in ['csv','xlsx']:
        print("filetype must be: 'csv' or 'xlsx'\n")
        parser.print_help()
        exit(-1)

    if database == 'train':
        EARL = INFORMIX_EARL_TEST
    elif database == 'cars':
        EARL = INFORMIX_EARL_PROD
    else:
        print("database must be: 'cars' or 'train'\n")
        parser.print_help()
        exit(-1)

    sql = SELECT_NEW_PEOPLE
    if test:
        print("new people sql")
        print("sql = {}".format(sql))

    people = []
    objects = do_sql(sql, key=key, earl=EARL)

    # the people list allows us to iterate over the result set more
    # than once, whereas just using objects result would throw an
    # error after the first iteration.
    for o in objects:
        if test:
            print(o)
        people.append(o)

    if people:
        response = _generate_files(people, filetype, 'new_people')

        if not response:
            print("no response")
        else:
            print("{} object(s) found for provisioning.".format(len(people))
            for p in people:
                try:
                    sql = INSERT_EMAIL_RECORD(cid=p.id, ldap=p.loginid)
                    if test:
                        print(sql)
                    else:
                        do_sql(sql, key=key, earl=EARL)
                except:
                    logger.info("failed insert = {}".format(p))

                try:
                    sql = INSERT_CVID_RECORD(cid=p.id, ldap=p.loginid)
                    if test:
                        print(sql)
                    else:
                        do_sql(sql, key=key, earl=EARL)
                except:
                    logger.info("failed insert = {}".format(p))
    else:
        print("No objects found for provisioning.")

######################
# shell command line
######################

if __name__ == '__main__':
    args = parser.parse_args()
    test = args.test
    database = args.database.lower()
    filetype = args.filetype.lower()
    sys.exit(main())
<|endoftext|>


<commit_before># Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
#
# 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.
"""Tests for dm_pix._src.augment."""

import functools

from absl.testing import absltest
from absl.testing import parameterized
from dm_pix._src import augment
import jax
import numpy as np
import tensorflow as tf
import inspect

_IMG_SHAPE = (131, 111, 3)
_RAND_FLOATS_IN_RANGE = list(
    np.random.uniform(0., 1., size=(10,) + _IMG_SHAPE).astype(np.float32))
_RAND_FLOATS_OUT_OF_RANGE = list(
    np.random.uniform(-0.5, 1.5, size=(10,) + _IMG_SHAPE).astype(np.float32))
_KERNEL_SIZE = _IMG_SHAPE[0] / 10.

class _ImageAugmentationTest(parameterized.TestCase):
  """Runs tests for the various augments with the correct arguments."""

  def _test_fn_with_random_arg(self, images_list, jax_fn, tf_fn, **kw_range):
    pass

  def _test_fn(self, images_list, jax_fn, tf_fn):
    pass

  def assertAllCloseTolerant(self, x, y):
    # Increase tolerance on TPU due to lower precision.
    tol = 1e-2 if jax.local_devices()[0].platform == "tpu" else 1e-4
    np.testing.assert_allclose(x, y, rtol=tol, atol=tol)
    self.assertEqual(x.dtype, y.dtype)

  @parameterized.named_parameters(("in_range", _RAND_FLOATS_IN_RANGE),
                                  ("out_of_range", _RAND_FLOATS_OUT_OF_RANGE))
  def test_adjust_brightness(self, images_list):
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=augment.adjust_brightness,
        tf_fn=tf.image.adjust_brightness,
        delta=(-0.5, 0.5))

    key = jax.random.PRNGKey(0)
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=functools.partial(augment.random_brightness, key),
        tf_fn=None,
        max_delta=(0, 0.5))

  @parameterized.named_parameters(("in_range", _RAND_FLOATS_IN_RANGE),
                                  ("out_of_range", _RAND_FLOATS_OUT_OF_RANGE))
  def test_adjust_contrast(self, images_list):
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=augment.adjust_contrast,
        tf_fn=tf.image.adjust_contrast,
        factor=(0.5, 1.5))
    key = jax.random.PRNGKey(0)
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=functools.partial(augment.random_contrast, key, upper=1),
        tf_fn=None,
        lower=(0, 0.9))

  # Doesn't make sense outside of [0, 1].
  @parameterized.named_parameters(("in_range", _RAND_FLOATS_IN_RANGE))
  def test_adjust_gamma(self, images_list):
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=augment.adjust_gamma,
        tf_fn=tf.image.adjust_gamma,
        gamma=(0.5, 1.5))

  @parameterized.named_parameters(("in_range", _RAND_FLOATS_IN_RANGE),
                                  ("out_of_range", _RAND_FLOATS_OUT_OF_RANGE))
  def test_adjust_saturation(self, images_list):
    # tf.image.adjust_saturation has a buggy implementation when the green and
    # blue channels have very close values that don't match the red channel.
    # This is due to a rounding error in http://shortn/_ETSJsEwUj5
    # if (g - b) < 0 but small enough that (hh + 1) == 1.
    # Eg: tf.image.adjust_saturation([[[0.75, 0.0369078, 0.0369079]]], 1.0)
    #     -> [[[0.03690779, 0.03690779, 0.03690779]]]
    # Perturb the inputs slightly so that this doesn't happen.
    def perturb(rgb):
      rgb_new = np.copy(rgb)
      rgb_new[..., 1] += 0.001 * (np.abs(rgb[..., 2] - rgb[..., 1]) < 1e-3)
      return rgb_new

    images_list = list(map(perturb, images_list))
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=augment.adjust_saturation,
        tf_fn=tf.image.adjust_saturation,
        factor=(0.5, 1.5))
    key = jax.random.PRNGKey(0)
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=functools.partial(augment.random_saturation, key, upper=1),
        tf_fn=None,
        lower=(0, 0.9))

  # CPU TF uses a different hue adjustment method outside of the [0, 1] range.
  # Disable out-of-range tests.
  @parameterized.named_parameters(
      ("in_range", _RAND_FLOATS_IN_RANGE),)
  def test_adjust_hue(self, images_list):
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=augment.adjust_hue,
        tf_fn=tf.image.adjust_hue,
        delta=(-0.5, 0.5))
    key = jax.random.PRNGKey(0)
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=functools.partial(augment.random_hue, key),
        tf_fn=None,
        max_delta=(0, 0.5))

  @parameterized.named_parameters(("in_range", _RAND_FLOATS_IN_RANGE),
                                  ("out_of_range", _RAND_FLOATS_OUT_OF_RANGE))
  def test_rot90(self, images_list):
    self._test_fn(
        images_list,
        jax_fn=lambda img: augment.rot90(img, k=1),
        tf_fn=lambda img: tf.image.rot90(img, k=1))
    self._test_fn(
        images_list,
        jax_fn=lambda img: augment.rot90(img, k=2),
        tf_fn=lambda img: tf.image.rot90(img, k=2))
    self._test_fn(
        images_list,
        jax_fn=lambda img: augment.rot90(img, k=3),
        tf_fn=lambda img: tf.image.rot90(img, k=3))

  # The functions below don't have a TF equivalent to compare to, we just check
  # that they run.
  @parameterized.named_parameters(("in_range", _RAND_FLOATS_IN_RANGE),
                                  ("out_of_range", _RAND_FLOATS_OUT_OF_RANGE))
  def test_flip(self, images_list):
    self._test_fn(
        images_list,
        jax_fn=augment.flip_left_right,
        tf_fn=tf.image.flip_left_right)
    self._test_fn(
        images_list, jax_fn=augment.flip_up_down, tf_fn=tf.image.flip_up_down)
    key = jax.random.PRNGKey(0)
    self._test_fn(
        images_list,
        jax_fn=functools.partial(augment.random_flip_left_right, key),
        tf_fn=None)
    self._test_fn(
        images_list,
        jax_fn=functools.partial(augment.random_flip_up_down, key),
        tf_fn=None)
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=functools.partial(augment.random_flip_left_right, key),
        tf_fn=None,
        probability=(0., 1.))
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=functools.partial(augment.random_flip_up_down, key),
        tf_fn=None,
        probability=(0., 1.))

  @parameterized.named_parameters(("in_range", _RAND_FLOATS_IN_RANGE),
                                  ("out_of_range", _RAND_FLOATS_OUT_OF_RANGE))
  def test_solarize(self, images_list):
    self._test_fn_with_random_arg(
        images_list, jax_fn=augment.solarize, tf_fn=None, threshold=(0., 1.))

  @parameterized.named_parameters(("in_range", _RAND_FLOATS_IN_RANGE),
                                  ("out_of_range", _RAND_FLOATS_OUT_OF_RANGE))
  def test_gaussian_blur(self, images_list):
    blur_fn = functools.partial(augment.gaussian_blur, kernel_size=_KERNEL_SIZE)
    self._test_fn_with_random_arg(
        images_list, jax_fn=blur_fn, tf_fn=None, sigma=(0.1, 2.0))

  @parameterized.named_parameters(("in_range", _RAND_FLOATS_IN_RANGE),
                                  ("out_of_range", _RAND_FLOATS_OUT_OF_RANGE))
  def test_random_crop(self, images_list):
    key = jax.random.PRNGKey(43)
    crop_fn = lambda img: augment.random_crop(key, img, (100, 100, 3))
    self._test_fn(images_list, jax_fn=crop_fn, tf_fn=None)

class TestMatchTensorflow(_ImageAugmentationTest):

  def _test_fn_with_random_arg(self, images_list, jax_fn, tf_fn, **kw_range):
    if tf_fn is None:
      return
    assert len(kw_range) == 1
    kw_name, (random_min, random_max) = list(kw_range.items())[0]
    for image_rgb in images_list:
      argument = np.random.uniform(random_min, random_max, size=())
      adjusted_jax = jax_fn(image_rgb, **{kw_name: argument})
      adjusted_tf = tf_fn(image_rgb, argument).numpy()
      self.assertAllCloseTolerant(adjusted_jax, adjusted_tf)

  def _test_fn(self, images_list, jax_fn, tf_fn):
    if tf_fn is None:
      return
    for image_rgb in images_list:
      adjusted_jax = jax_fn(image_rgb)
      adjusted_tf = tf_fn(image_rgb).numpy()
      self.assertAllCloseTolerant(adjusted_jax, adjusted_tf)

class TestVmap(_ImageAugmentationTest):

  def _test_fn_with_random_arg(self, images_list, jax_fn, tf_fn, **kw_range):
    del tf_fn  # unused.
    assert len(kw_range) == 1
    kw_name, (random_min, random_max) = list(kw_range.items())[0]
    arguments = [
        np.random.uniform(random_min, random_max, size=()) for _ in images_list
    ]
    fn_vmap = jax.vmap(jax_fn)
    outputs_vmaped = list(
        fn_vmap(np.stack(images_list, axis=0),
        **{kw_name: np.stack(arguments, axis=0)}))
    assert len(images_list) == len(outputs_vmaped)
    assert len(images_list) == len(arguments)
    for image_rgb, argument, adjusted_vmap in zip(images_list, arguments,
                                                  outputs_vmaped):
      adjusted_jax = jax_fn(image_rgb, **{kw_name: argument})
      self.assertAllCloseTolerant(adjusted_jax, adjusted_vmap)

  def _test_fn(self, images_list, jax_fn, tf_fn):
    del tf_fn  # unused.
    fn_vmap = jax.vmap(jax_fn)
    outputs_vmaped = list(fn_vmap(np.stack(images_list, axis=0)))
    assert len(images_list) == len(outputs_vmaped)
    for image_rgb, adjusted_vmap in zip(images_list, outputs_vmaped):
      adjusted_jax = jax_fn(image_rgb)
      self.assertAllCloseTolerant(adjusted_jax, adjusted_vmap)

class TestJit(_ImageAugmentationTest):

  def _test_fn_with_random_arg(self, images_list, jax_fn, tf_fn, **kw_range):
    del tf_fn  # unused.
    assert len(kw_range) == 1
    kw_name, (random_min, random_max) = list(kw_range.items())[0]
    jax_fn_jitted = jax.jit(jax_fn)
    for image_rgb in images_list:
      argument = np.random.uniform(random_min, random_max, size=())
      adjusted_jax = jax_fn(image_rgb, **{kw_name: argument})
      adjusted_jit = jax_fn_jitted(image_rgb, **{kw_name: argument})
      self.assertAllCloseTolerant(adjusted_jax, adjusted_jit)

  def _test_fn(self, images_list, jax_fn, tf_fn):
    del tf_fn  # unused.
    jax_fn_jitted = jax.jit(jax_fn)
    for image_rgb in images_list:
      adjusted_jax = jax_fn(image_rgb)
      adjusted_jit = jax_fn_jitted(image_rgb)
      self.assertAllCloseTolerant(adjusted_jax, adjusted_jit)

if __name__ == "__main__":
  absltest.main()
<commit_msg><commit_after># Copyright 2020 DeepMind Technologies Limited. All Rights Reserved.
#
# 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.
"""Tests for dm_pix._src.augment."""

import functools

from absl.testing import absltest
from absl.testing import parameterized
from dm_pix._src import augment
import jax
import numpy as np
import tensorflow as tf

_IMG_SHAPE = (131, 111, 3)
_RAND_FLOATS_IN_RANGE = list(
    np.random.uniform(0., 1., size=(10,) + _IMG_SHAPE).astype(np.float32))
_RAND_FLOATS_OUT_OF_RANGE = list(
    np.random.uniform(-0.5, 1.5, size=(10,) + _IMG_SHAPE).astype(np.float32))
_KERNEL_SIZE = _IMG_SHAPE[0] / 10.

class _ImageAugmentationTest(parameterized.TestCase):
  """Runs tests for the various augments with the correct arguments."""

  def _test_fn_with_random_arg(self, images_list, jax_fn, tf_fn, **kw_range):
    pass

  def _test_fn(self, images_list, jax_fn, tf_fn):
    pass

  def assertAllCloseTolerant(self, x, y):
    # Increase tolerance on TPU due to lower precision.
    tol = 1e-2 if jax.local_devices()[0].platform == "tpu" else 1e-4
    np.testing.assert_allclose(x, y, rtol=tol, atol=tol)
    self.assertEqual(x.dtype, y.dtype)

  @parameterized.named_parameters(("in_range", _RAND_FLOATS_IN_RANGE),
                                  ("out_of_range", _RAND_FLOATS_OUT_OF_RANGE))
  def test_adjust_brightness(self, images_list):
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=augment.adjust_brightness,
        tf_fn=tf.image.adjust_brightness,
        delta=(-0.5, 0.5))

    key = jax.random.PRNGKey(0)
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=functools.partial(augment.random_brightness, key),
        tf_fn=None,
        max_delta=(0, 0.5))

  @parameterized.named_parameters(("in_range", _RAND_FLOATS_IN_RANGE),
                                  ("out_of_range", _RAND_FLOATS_OUT_OF_RANGE))
  def test_adjust_contrast(self, images_list):
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=augment.adjust_contrast,
        tf_fn=tf.image.adjust_contrast,
        factor=(0.5, 1.5))
    key = jax.random.PRNGKey(0)
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=functools.partial(augment.random_contrast, key, upper=1),
        tf_fn=None,
        lower=(0, 0.9))

  # Doesn't make sense outside of [0, 1].
  @parameterized.named_parameters(("in_range", _RAND_FLOATS_IN_RANGE))
  def test_adjust_gamma(self, images_list):
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=augment.adjust_gamma,
        tf_fn=tf.image.adjust_gamma,
        gamma=(0.5, 1.5))

  @parameterized.named_parameters(("in_range", _RAND_FLOATS_IN_RANGE),
                                  ("out_of_range", _RAND_FLOATS_OUT_OF_RANGE))
  def test_adjust_saturation(self, images_list):
    # tf.image.adjust_saturation has a buggy implementation when the green and
    # blue channels have very close values that don't match the red channel.
    # This is due to a rounding error in http://shortn/_ETSJsEwUj5
    # if (g - b) < 0 but small enough that (hh + 1) == 1.
    # Eg: tf.image.adjust_saturation([[[0.75, 0.0369078, 0.0369079]]], 1.0)
    #     -> [[[0.03690779, 0.03690779, 0.03690779]]]
    # Perturb the inputs slightly so that this doesn't happen.
    def perturb(rgb):
      rgb_new = np.copy(rgb)
      rgb_new[..., 1] += 0.001 * (np.abs(rgb[..., 2] - rgb[..., 1]) < 1e-3)
      return rgb_new

    images_list = list(map(perturb, images_list))
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=augment.adjust_saturation,
        tf_fn=tf.image.adjust_saturation,
        factor=(0.5, 1.5))
    key = jax.random.PRNGKey(0)
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=functools.partial(augment.random_saturation, key, upper=1),
        tf_fn=None,
        lower=(0, 0.9))

  # CPU TF uses a different hue adjustment method outside of the [0, 1] range.
  # Disable out-of-range tests.
  @parameterized.named_parameters(
      ("in_range", _RAND_FLOATS_IN_RANGE),)
  def test_adjust_hue(self, images_list):
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=augment.adjust_hue,
        tf_fn=tf.image.adjust_hue,
        delta=(-0.5, 0.5))
    key = jax.random.PRNGKey(0)
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=functools.partial(augment.random_hue, key),
        tf_fn=None,
        max_delta=(0, 0.5))

  @parameterized.named_parameters(("in_range", _RAND_FLOATS_IN_RANGE),
                                  ("out_of_range", _RAND_FLOATS_OUT_OF_RANGE))
  def test_rot90(self, images_list):
    self._test_fn(
        images_list,
        jax_fn=lambda img: augment.rot90(img, k=1),
        tf_fn=lambda img: tf.image.rot90(img, k=1))
    self._test_fn(
        images_list,
        jax_fn=lambda img: augment.rot90(img, k=2),
        tf_fn=lambda img: tf.image.rot90(img, k=2))
    self._test_fn(
        images_list,
        jax_fn=lambda img: augment.rot90(img, k=3),
        tf_fn=lambda img: tf.image.rot90(img, k=3))

  # The functions below don't have a TF equivalent to compare to, we just check
  # that they run.
  @parameterized.named_parameters(("in_range", _RAND_FLOATS_IN_RANGE),
                                  ("out_of_range", _RAND_FLOATS_OUT_OF_RANGE))
  def test_flip(self, images_list):
    self._test_fn(
        images_list,
        jax_fn=augment.flip_left_right,
        tf_fn=tf.image.flip_left_right)
    self._test_fn(
        images_list, jax_fn=augment.flip_up_down, tf_fn=tf.image.flip_up_down)
    key = jax.random.PRNGKey(0)
    self._test_fn(
        images_list,
        jax_fn=functools.partial(augment.random_flip_left_right, key),
        tf_fn=None)
    self._test_fn(
        images_list,
        jax_fn=functools.partial(augment.random_flip_up_down, key),
        tf_fn=None)
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=functools.partial(augment.random_flip_left_right, key),
        tf_fn=None,
        probability=(0., 1.))
    self._test_fn_with_random_arg(
        images_list,
        jax_fn=functools.partial(augment.random_flip_up_down, key),
        tf_fn=None,
        probability=(0., 1.))

  @parameterized.named_parameters(("in_range", _RAND_FLOATS_IN_RANGE),
                                  ("out_of_range", _RAND_FLOATS_OUT_OF_RANGE))
  def test_solarize(self, images_list):
    self._test_fn_with_random_arg(
        images_list, jax_fn=augment.solarize, tf_fn=None, threshold=(0., 1.))

  @parameterized.named_parameters(("in_range", _RAND_FLOATS_IN_RANGE),
                                  ("out_of_range", _RAND_FLOATS_OUT_OF_RANGE))
  def test_gaussian_blur(self, images_list):
    blur_fn = functools.partial(augment.gaussian_blur, kernel_size=_KERNEL_SIZE)
    self._test_fn_with_random_arg(
        images_list, jax_fn=blur_fn, tf_fn=None, sigma=(0.1, 2.0))

  @parameterized.named_parameters(("in_range", _RAND_FLOATS_IN_RANGE),
                                  ("out_of_range", _RAND_FLOATS_OUT_OF_RANGE))
  def test_random_crop(self, images_list):
    key = jax.random.PRNGKey(43)
    crop_fn = lambda img: augment.random_crop(key, img, (100, 100, 3))
    self._test_fn(images_list, jax_fn=crop_fn, tf_fn=None)

class TestMatchTensorflow(_ImageAugmentationTest):

  def _test_fn_with_random_arg(self, images_list, jax_fn, tf_fn, **kw_range):
    if tf_fn is None:
      return
    assert len(kw_range) == 1
    kw_name, (random_min, random_max) = list(kw_range.items())[0]
    for image_rgb in images_list:
      argument = np.random.uniform(random_min, random_max, size=())
      adjusted_jax = jax_fn(image_rgb, **{kw_name: argument})
      adjusted_tf = tf_fn(image_rgb, argument).numpy()
      self.assertAllCloseTolerant(adjusted_jax, adjusted_tf)

  def _test_fn(self, images_list, jax_fn, tf_fn):
    if tf_fn is None:
      return
    for image_rgb in images_list:
      adjusted_jax = jax_fn(image_rgb)
      adjusted_tf = tf_fn(image_rgb).numpy()
      self.assertAllCloseTolerant(adjusted_jax, adjusted_tf)

class TestVmap(_ImageAugmentationTest):

  def _test_fn_with_random_arg(self, images_list, jax_fn, tf_fn, **kw_range):
    del tf_fn  # unused.
    assert len(kw_range) == 1
    kw_name, (random_min, random_max) = list(kw_range.items())[0]
    arguments = [
        np.random.uniform(random_min, random_max, size=()) for _ in images_list
    ]
    fn_vmap = jax.vmap(jax_fn)
    outputs_vmaped = list(
        fn_vmap(np.stack(images_list, axis=0),
        **{kw_name: np.stack(arguments, axis=0)}))
    assert len(images_list) == len(outputs_vmaped)
    assert len(images_list) == len(arguments)
    for image_rgb, argument, adjusted_vmap in zip(images_list, arguments,
                                                  outputs_vmaped):
      adjusted_jax = jax_fn(image_rgb, **{kw_name: argument})
      self.assertAllCloseTolerant(adjusted_jax, adjusted_vmap)

  def _test_fn(self, images_list, jax_fn, tf_fn):
    del tf_fn  # unused.
    fn_vmap = jax.vmap(jax_fn)
    outputs_vmaped = list(fn_vmap(np.stack(images_list, axis=0)))
    assert len(images_list) == len(outputs_vmaped)
    for image_rgb, adjusted_vmap in zip(images_list, outputs_vmaped):
      adjusted_jax = jax_fn(image_rgb)
      self.assertAllCloseTolerant(adjusted_jax, adjusted_vmap)

class TestJit(_ImageAugmentationTest):

  def _test_fn_with_random_arg(self, images_list, jax_fn, tf_fn, **kw_range):
    del tf_fn  # unused.
    assert len(kw_range) == 1
    kw_name, (random_min, random_max) = list(kw_range.items())[0]
    jax_fn_jitted = jax.jit(jax_fn)
    for image_rgb in images_list:
      argument = np.random.uniform(random_min, random_max, size=())
      adjusted_jax = jax_fn(image_rgb, **{kw_name: argument})
      adjusted_jit = jax_fn_jitted(image_rgb, **{kw_name: argument})
      self.assertAllCloseTolerant(adjusted_jax, adjusted_jit)

  def _test_fn(self, images_list, jax_fn, tf_fn):
    del tf_fn  # unused.
    jax_fn_jitted = jax.jit(jax_fn)
    for image_rgb in images_list:
      adjusted_jax = jax_fn(image_rgb)
      adjusted_jit = jax_fn_jitted(image_rgb)
      self.assertAllCloseTolerant(adjusted_jax, adjusted_jit)

if __name__ == "__main__":
  absltest.main()
<|endoftext|>


<commit_before>import os

import numpy as np
import pandas as pd
from cea.demand.metamodel.nn_generator.nn_trainer_resume import nn_model_collector

def get_nn_estimations(model, scalerT, scalerX, urban_input_matrix, locator):
    input_NN_x = urban_input_matrix
    inputs_x = scalerX.transform(input_NN_x)

    model_estimates = model.predict(inputs_x)
    filtered_predict = scalerT.inverse_transform(model_estimates)

    # filter_logic = np.isin(targets_t, 0)
    # target_anomalies = np.asarray(np.where(filter_logic), dtype=np.int)
    # t_anomalies_rows, t_anomalies_cols = target_anomalies.shape
    # anomalies_replacements = np.zeros(t_anomalies_cols)
    # filtered_predict[target_anomalies, 0] = anomalies_replacements

    model_estimates = locator.get_neural_network_estimates()
    filtered_predict = pd.DataFrame(filtered_predict)
    filtered_predict.to_csv(model_estimates, index=False, header=False, float_format='%.3f', decimal='.')

def test_sample_collector(locator):
    test_sample_path = locator.get_nn_inout_folder()
    file_path_inputs = os.path.join(test_sample_path, "input_predict.csv")
    urban_input_matrix = np.asarray(pd.read_csv(file_path_inputs))

    return urban_input_matrix

def test_nn_performance(locator):
    urban_input_matrix = test_sample_collector(locator)
    model, scalerT, scalerX = nn_model_collector(locator)
    get_nn_estimations(model, scalerT, scalerX, urban_input_matrix, locator)

def run_as_script():
    import cea.config
    config = cea.config.Configuration()
    locator = cea.inputlocator.InputLocator(scenario_path=config.scenario)

    test_nn_performance(locator)

if __name__ == '__main__':
    run_as_script()
<commit_msg>correcting the estimate<commit_after>import os

import numpy as np
import pandas as pd
from cea.demand.metamodel.nn_generator.nn_trainer_resume import nn_model_collector

def get_nn_estimations(model, scalerT, scalerX, urban_input_matrix, locator):
    input_NN_x = urban_input_matrix
    inputs_x = scalerX.transform(input_NN_x)

    model_estimates = model.predict(inputs_x)
    filtered_predict = scalerT.inverse_transform(model_estimates)

    # filter_logic = np.isin(targets_t, 0)
    # target_anomalies = np.asarray(np.where(filter_logic), dtype=np.int)
    # t_anomalies_rows, t_anomalies_cols = target_anomalies.shape
    # anomalies_replacements = np.zeros(t_anomalies_cols)
    # filtered_predict[target_anomalies, 0] = anomalies_replacements

    model_estimates = locator.get_neural_network_estimates()
    filtered_predict = pd.DataFrame(filtered_predict)
    filtered_predict.to_csv(model_estimates, index=False, header=False, float_format='%.3f', decimal='.')

def test_sample_collector(locator):
    test_sample_path = locator.get_nn_inout_folder()
    overrides_path=locator.get_building_overrides()
    os.remove(overrides_path)

    file_path_inputs = os.path.join(test_sample_path, "input_predict.csv")
    urban_input_matrix = np.asarray(pd.read_csv(file_path_inputs))

    return urban_input_matrix

def input_prepare_estimate(list_building_names, locator, target_parameters, gv):
    '''
    this function prepares the inputs and targets for the neural net by splitting the jobs between different processors
    :param list_building_names: a list of building names
    :param locator: points to the variables
    :param target_parameters: (imported from 'nn_settings.py') a list containing the name of desirable outputs
    :param gv: global variables
    :return: inputs and targets for the whole dataset (urban_input_matrix, urban_taget_matrix)
    '''

    #   open multiprocessing pool
    pool = mp.Pool()
    #   count number of CPUs
    gv.log("Using %i CPU's" % mp.cpu_count())
    #   creat an empty job list to be filled later
    joblist = []
    #   create one job for each data preparation task i.e. each building
    for building_name in list_building_names:
        job = pool.apply_async(input_matrix.input_prepare_multi_processing,
                               [building_name, gv, locator, target_parameters])
        joblist.append(job)
    #   run the input/target preperation for all buildings in the list (here called jobs)
    for i, job in enumerate(joblist):
        NN_input_ready , NN_target_ready=job.get(240)
        #   remove buildings that have "NaN" in their input (e.g. if heating/cooling is off, the indoor temperature
        #   will be returned as "NaN"). Afterwards, stack the inputs/targets of all buildings
        check_nan=1*(np.isnan(np.sum(NN_input_ready)))
        if check_nan == 0:
            if i == 0:
                urban_input_matrix = NN_input_ready
                urban_taget_matrix = NN_target_ready
            else:
                urban_input_matrix = np.concatenate((urban_input_matrix, NN_input_ready))
                urban_taget_matrix = np.concatenate((urban_taget_matrix, NN_target_ready))

    #   close the multiprocessing
    pool.close()

    return urban_input_matrix, urban_taget_matrix

def test_nn_performance(locator):
    urban_input_matrix = test_sample_collector(locator)
    model, scalerT, scalerX = nn_model_collector(locator)
    get_nn_estimations(model, scalerT, scalerX, urban_input_matrix, locator)

def run_as_script():
    import cea.config
    config = cea.config.Configuration()
    locator = cea.inputlocator.InputLocator(scenario_path=config.scenario)

    test_nn_performance(locator)

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