<commit_before># -*- coding: utf-8 -*-
"""
Common classes and methods for PyLTI module
"""

from __future__ import absolute_import
import logging
import oauth2
import oauth.oauth as oauth
from lxml import etree

log = logging.getLogger('pylti.common')  # pylint: disable=invalid-name

"""
Classes to handle oauth portion of LTI
"""

class LTIOAuthDataStore(oauth.OAuthDataStore):
    """
    Largely taken from reference implementation
    for app engine at https://code.google.com/p/ims-dev/
    """

    def __init__(self, consumers):
        """
        Create OAuth store
        """
        oauth.OAuthDataStore.__init__(self)
        self.consumers = consumers

    def lookup_consumer(self, key):
        """
        Search through keys
        """
        if not self.consumers:
            log.critical(("No consumers defined in settings."
                          "Have you created a configuration file?"))
            return None

        consumer = self.consumers.get(key)
        if not consumer:
            log.info("Did not find consumer, using key: %s ", key)
            return None

        secret = consumer.get('secret', None)
        if not secret:
            log.critical(('Consumer %s, is missing secret'
                          'in settings file, and needs correction.'), key)
            return None
        return oauth.OAuthConsumer(key, secret)

    def lookup_nonce(self, oauth_consumer, oauth_token, nonce):
        return None

class LTIException(Exception):
    """
    Custom LTI exception for proper handling
    of LTI specific errors
    """
    pass

class LTINotInSessionException(LTIException):
    """
    Custom LTI exception for proper handling
    of LTI specific errors
    """
    pass

class LTIRoleException(LTIException):
    """
    Exception class for when LTI user doesn't have the
    right role.
    """
    pass

"""
Decorators to handle LTI session management,
authentication, etc.
"""

LTI_PROPERTY_LIST = [
    'oauth_consumer_key',
    'launch_presentation_return_url',
    'user_id',
    'oauth_nonce',
    'context_label',
    'context_id',
    'resource_link_title',
    'resource_link_id',
    'lis_person_contact_email_primary',
    'lis_person_contact_emailprimary',
    'lis_person_name_full',
    'lis_person_name_family',
    'lis_person_name_given',
    'lis_result_sourcedid',
    'launch_type',
    'lti_message',
    'lti_version',
    'roles',
    'lis_outcome_service_url'
]

LTI_STAFF_ROLES = ['Instructor', 'Administrator', ]

LTI_SESSION_KEY = 'lti_authenticated'

def _post_patched_request(body, client, url):
    """
    Authorization header needs to be capitalized for some LTI clients
    this function ensures that header is capitalized
    :param body: body of the call
    :param client: OAuth Client
    :param url: outcome url
    :return: response
    """
    monkey_patch_function = None
    import httplib2

    http = httplib2.Http
    # pylint: disable=protected-access
    normalize = http._normalize_headers

    def my_normalize(self, headers):
        """ This function patches Authorization header """
        ret = normalize(self, headers)
        if ret.has_key('authorization'):
            ret['Authorization'] = ret.pop('authorization')
        log.debug("headers")
        log.debug(headers)
        return ret

    http._normalize_headers = my_normalize
    monkey_patch_function = normalize

    # pylint: disable=unused-variable
    response, content = client.request(
        url,
        'POST',
        body=body,
        headers={'Content-Type': 'application/xml'})

    http = httplib2.Http
    #pylint: disable=protected-access
    http._normalize_headers = monkey_patch_function

    return response

def post_message(consumers, lti_key, url, body):
    """
        Posts a signed message to LTI consumer
    :param consumers: consumers from config
    :param lti_key: key to find appropriate consumer
    :param url: post url
    :param body: xml body
    :return: success
    """
    oauth_store = LTIOAuthDataStore(consumers)
    oauth_server = oauth.OAuthServer(oauth_store)
    oauth_server.add_signature_method(oauth.OAuthSignatureMethod_HMAC_SHA1())
    lti_consumer = oauth_store.lookup_consumer(lti_key)
    secret = lti_consumer.secret
    consumer = oauth2.Consumer(key=lti_key, secret=secret)
    client = oauth2.Client(consumer)
    # monkey_patch_headers ensures that Authorization header is NOT lower cased
    response = _post_patched_request(body, client, url)
    #TODO: inspect content and return True if success
    log.debug("key {}".format(lti_key))
    log.debug("secret {}".format(secret))
    log.debug("url {}".format(url))
    log.debug(body)
    log.debug(response)
    # log.debug(content)
    return True

def verify_request_common(consumers, url, method, headers, params):
    """
    :param consumers: consumers from config file
    :param url: request url
    :param method: request method
    :param headers: request headers
    :param params: request params
    :return: is request valid
    """

    log.debug("consumers {}".format(consumers))
    log.debug("url {}".format(url))
    log.debug("method {}".format(method))
    log.debug("headers {}".format(headers))
    log.debug("params {}".format(params))
    oauth_store = LTIOAuthDataStore(consumers)
    oauth_server = oauth.OAuthServer(oauth_store)
    oauth_server.add_signature_method(
        oauth.OAuthSignatureMethod_PLAINTEXT())
    oauth_server.add_signature_method(
        oauth.OAuthSignatureMethod_HMAC_SHA1())

    # Check header for SSL before selecting the url
    if headers.get('X-Forwarded-Proto', 'http') == 'https':
        url = url.replace('http', 'https', 1)

    oauth_request = oauth.OAuthRequest.from_request(
        method,
        url,
        headers=dict(headers),
        parameters=params
    )

    if not oauth_request:
        log.info('Received non oauth request on oauth protected page')
        raise LTIException('This page requires a valid oauth session '
                           'or request')
    try:
        #pylint: disable=protected-access
        consumer = oauth_server._get_consumer(oauth_request)
        oauth_server._check_signature(oauth_request, consumer, None)
    except oauth.OAuthError as err:
        # Rethrow our own for nice error handling (don't print
        # error message as it will contain the key
        raise LTIException("OAuth error: Please check your key and secret")

    return True

def generate_request_xml(message_identifier_id, operation, \
                         lis_result_sourcedid, score):
    root = etree.Element('imsx_POXEnvelopeRequest', xmlns= \
        'http://www.imsglobal.org/services/ltiv1p1/xsd/imsoms_v1p0')

    header = etree.SubElement(root, 'imsx_POXHeader')
    header_info = etree.SubElement(header, 'imsx_POXRequestHeaderInfo')
    version = etree.SubElement(header_info, 'imsx_version')
    version.text = 'V1.0'
    message_identifier = etree.SubElement(header_info,
                                          'imsx_messageIdentifier')
    message_identifier.text = message_identifier_id
    body = etree.SubElement(root, 'imsx_POXBody')
    xml_request = etree.SubElement(body, '%s%s' % (operation, 'Request'))
    record = etree.SubElement(xml_request, 'resultRecord')

    guid = etree.SubElement(record, 'sourcedGUID')

    sourcedid = etree.SubElement(guid, 'sourcedId')
    sourcedid.text = lis_result_sourcedid
    if not score is None:
        result = etree.SubElement(record, 'result')
        result_score = etree.SubElement(result, 'resultScore')
        language = etree.SubElement(result_score, 'language')
        language.text = 'en'
        text_string = etree.SubElement(result_score, 'textString')
        text_string.text = score.__str__()
    log.debug("XML Response: \n{}".format(\
        etree.tostring(root, xml_declaration=True, encoding='utf-8')))
    return etree.tostring(root, xml_declaration=True, encoding='utf-8')

<commit_msg>Moving constants to the top.<commit_after># -*- coding: utf-8 -*-
"""
Common classes and methods for PyLTI module
"""

from __future__ import absolute_import
import logging
import oauth2
import oauth.oauth as oauth
from lxml import etree

log = logging.getLogger('pylti.common')  # pylint: disable=invalid-name

LTI_PROPERTY_LIST = [
    'oauth_consumer_key',
    'launch_presentation_return_url',
    'user_id',
    'oauth_nonce',
    'context_label',
    'context_id',
    'resource_link_title',
    'resource_link_id',
    'lis_person_contact_email_primary',
    'lis_person_contact_emailprimary',
    'lis_person_name_full',
    'lis_person_name_family',
    'lis_person_name_given',
    'lis_result_sourcedid',
    'launch_type',
    'lti_message',
    'lti_version',
    'roles',
    'lis_outcome_service_url'
]

LTI_STAFF_ROLES = ['Instructor', 'Administrator', ]

LTI_SESSION_KEY = 'lti_authenticated'

"""
Classes to handle oauth portion of LTI
"""

class LTIOAuthDataStore(oauth.OAuthDataStore):
    """
    Largely taken from reference implementation
    for app engine at https://code.google.com/p/ims-dev/
    """

    def __init__(self, consumers):
        """
        Create OAuth store
        """
        oauth.OAuthDataStore.__init__(self)
        self.consumers = consumers

    def lookup_consumer(self, key):
        """
        Search through keys
        """
        if not self.consumers:
            log.critical(("No consumers defined in settings."
                          "Have you created a configuration file?"))
            return None

        consumer = self.consumers.get(key)
        if not consumer:
            log.info("Did not find consumer, using key: %s ", key)
            return None

        secret = consumer.get('secret', None)
        if not secret:
            log.critical(('Consumer %s, is missing secret'
                          'in settings file, and needs correction.'), key)
            return None
        return oauth.OAuthConsumer(key, secret)

    def lookup_nonce(self, oauth_consumer, oauth_token, nonce):
        return None

class LTIException(Exception):
    """
    Custom LTI exception for proper handling
    of LTI specific errors
    """
    pass

class LTINotInSessionException(LTIException):
    """
    Custom LTI exception for proper handling
    of LTI specific errors
    """
    pass

class LTIRoleException(LTIException):
    """
    Exception class for when LTI user doesn't have the
    right role.
    """
    pass

def _post_patched_request(body, client, url):
    """
    Authorization header needs to be capitalized for some LTI clients
    this function ensures that header is capitalized
    :param body: body of the call
    :param client: OAuth Client
    :param url: outcome url
    :return: response
    """
    monkey_patch_function = None
    import httplib2

    http = httplib2.Http
    # pylint: disable=protected-access
    normalize = http._normalize_headers

    def my_normalize(self, headers):
        """ This function patches Authorization header """
        ret = normalize(self, headers)
        if ret.has_key('authorization'):
            ret['Authorization'] = ret.pop('authorization')
        log.debug("headers")
        log.debug(headers)
        return ret

    http._normalize_headers = my_normalize
    monkey_patch_function = normalize

    # pylint: disable=unused-variable
    response, content = client.request(
        url,
        'POST',
        body=body,
        headers={'Content-Type': 'application/xml'})

    http = httplib2.Http
    #pylint: disable=protected-access
    http._normalize_headers = monkey_patch_function

    return response

def post_message(consumers, lti_key, url, body):
    """
        Posts a signed message to LTI consumer
    :param consumers: consumers from config
    :param lti_key: key to find appropriate consumer
    :param url: post url
    :param body: xml body
    :return: success
    """
    oauth_store = LTIOAuthDataStore(consumers)
    oauth_server = oauth.OAuthServer(oauth_store)
    oauth_server.add_signature_method(oauth.OAuthSignatureMethod_HMAC_SHA1())
    lti_consumer = oauth_store.lookup_consumer(lti_key)
    secret = lti_consumer.secret
    consumer = oauth2.Consumer(key=lti_key, secret=secret)
    client = oauth2.Client(consumer)
    # monkey_patch_headers ensures that Authorization header is NOT lower cased
    response = _post_patched_request(body, client, url)
    #TODO: inspect content and return True if success
    log.debug("key {}".format(lti_key))
    log.debug("secret {}".format(secret))
    log.debug("url {}".format(url))
    log.debug(body)
    log.debug(response)
    # log.debug(content)
    return True

def verify_request_common(consumers, url, method, headers, params):
    """
    :param consumers: consumers from config file
    :param url: request url
    :param method: request method
    :param headers: request headers
    :param params: request params
    :return: is request valid
    """

    log.debug("consumers {}".format(consumers))
    log.debug("url {}".format(url))
    log.debug("method {}".format(method))
    log.debug("headers {}".format(headers))
    log.debug("params {}".format(params))
    oauth_store = LTIOAuthDataStore(consumers)
    oauth_server = oauth.OAuthServer(oauth_store)
    oauth_server.add_signature_method(
        oauth.OAuthSignatureMethod_PLAINTEXT())
    oauth_server.add_signature_method(
        oauth.OAuthSignatureMethod_HMAC_SHA1())

    # Check header for SSL before selecting the url
    if headers.get('X-Forwarded-Proto', 'http') == 'https':
        url = url.replace('http', 'https', 1)

    oauth_request = oauth.OAuthRequest.from_request(
        method,
        url,
        headers=dict(headers),
        parameters=params
    )

    if not oauth_request:
        log.info('Received non oauth request on oauth protected page')
        raise LTIException('This page requires a valid oauth session '
                           'or request')
    try:
        #pylint: disable=protected-access
        consumer = oauth_server._get_consumer(oauth_request)
        oauth_server._check_signature(oauth_request, consumer, None)
    except oauth.OAuthError as err:
        # Rethrow our own for nice error handling (don't print
        # error message as it will contain the key
        raise LTIException("OAuth error: Please check your key and secret")

    return True

def generate_request_xml(message_identifier_id, operation, \
                         lis_result_sourcedid, score):
    root = etree.Element('imsx_POXEnvelopeRequest', xmlns= \
        'http://www.imsglobal.org/services/ltiv1p1/xsd/imsoms_v1p0')

    header = etree.SubElement(root, 'imsx_POXHeader')
    header_info = etree.SubElement(header, 'imsx_POXRequestHeaderInfo')
    version = etree.SubElement(header_info, 'imsx_version')
    version.text = 'V1.0'
    message_identifier = etree.SubElement(header_info,
                                          'imsx_messageIdentifier')
    message_identifier.text = message_identifier_id
    body = etree.SubElement(root, 'imsx_POXBody')
    xml_request = etree.SubElement(body, '%s%s' % (operation, 'Request'))
    record = etree.SubElement(xml_request, 'resultRecord')

    guid = etree.SubElement(record, 'sourcedGUID')

    sourcedid = etree.SubElement(guid, 'sourcedId')
    sourcedid.text = lis_result_sourcedid
    if not score is None:
        result = etree.SubElement(record, 'result')
        result_score = etree.SubElement(result, 'resultScore')
        language = etree.SubElement(result_score, 'language')
        language.text = 'en'
        text_string = etree.SubElement(result_score, 'textString')
        text_string.text = score.__str__()
    log.debug("XML Response: \n{}".format(\
        etree.tostring(root, xml_declaration=True, encoding='utf-8')))
    return etree.tostring(root, xml_declaration=True, encoding='utf-8')

<|endoftext|>


<commit_before>import os
import re
import csv
import errno
# import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

from sklearn.preprocessing import StandardScaler, LabelEncoder
from collections import OrderedDict

def sizeof_file(name, suffix='B'):
    num = os.stat(name).st_size
    for unit in ['','Ki','Mi','Gi','Ti','Pi','Ei','Zi']:
        if abs(num) < 1024.0:
            return "%3.1f%s%s" % (num, unit, suffix)
        num /= 1024.0
    return "%.1f%s%s" % (num, 'Yi', suffix)

def path_exists(path):
    try:
        os.makedirs(path)
    except OSError as exception:
        if exception.errno != errno.EEXIST:
            raise

def localization():
    APP_PATH = os.path.dirname(os.path.abspath(__file__))
    UPLOAD_FOLDER = os.path.join(APP_PATH, 'uploads')
    devices_file = 'devices'
    devices_path = os.path.join(UPLOAD_FOLDER, devices_file)
    files = [f for f in os.listdir(UPLOAD_FOLDER) if re.match(r'p[0-9]+', f)]
    len_files = len(files)
    print(devices_path)
    print('len_file', len_files)
    print('files', files)

    if len_files != 0:
        for i in range(1, len_files+1):
            NUMBER = i
            name_file = 'p'+str(NUMBER)
            path_file = os.path.join(UPLOAD_FOLDER, name_file)
            with open(path_file, 'r') as in_file:
                stripped = (line.strip() for line in in_file)
                lines = (line for line in stripped if line)
                grouped = zip(*[lines] * 2)
                with open(path_file+'.csv', 'w') as out_file:
                    writer = csv.writer(out_file)
                    writer.writerow(('Beacon', 'RSSI'))
                    writer.writerows(grouped)

            df = pd.read_csv(path_file+'.csv')
            os.remove(path_file+'.csv')
            condition = ''

            with open(devices_path, 'r') as inFile:
                stripped = (line.strip() for line in inFile)

                for x in stripped:
                    condition += '(Beacon == "' + str(x) + '")' + ' | '
                condition = condition[:-3]
            df = df.query(condition)

            encoder = LabelEncoder()
            df['Beacon'] = encoder.fit_transform(df['Beacon'])

            mapped_df = {}
            for b, data in df.groupby('Beacon'):
                mapped_df[b] = data

            dict_to_df = OrderedDict()
            for index in range(len(mapped_df)):
                dict_to_df[index] = mapped_df[index]['RSSI']

            df_outliers = pd.DataFrame(dict_to_df)

            #Dealing with nan and duplicates
            df_outliers = df_outliers.fillna(method='ffill')
            df_outliers = df_outliers.dropna()
            df_outliers = df_outliers.drop_duplicates()

            #Scaling
            scaler = StandardScaler()
            df_outliers_t = scaler.fit_transform(df_outliers.ix[:, 0:len(df_outliers.columns)-1])
            df_outliers_t = pd.DataFrame(df_outliers_t)

            #Removing outliers aplying the mean
            df_no_outliers = df_outliers_t[df_outliers_t.apply(lambda x: np.abs(x - x.median()) / x.std() < 3).all(axis=1)]
            df_no_outliers = scaler.inverse_transform(df_no_outliers)
            df_no_outliers = pd.DataFrame(df_no_outliers)
            vector_no_out = NUMBER*np.ones((len(df_no_outliers),), dtype=np.int)
            vector_out = NUMBER*np.ones((len(df_outliers),), dtype=np.int)
            df_no_outliers = df_no_outliers.assign(position=pd.Series(vector_no_out).values)
            df_outliers = df_outliers.assign(position=pd.Series(vector_out).values)

            #Saving files with outliers and without it.
            df_outliers.to_csv(os.path.join(UPLOAD_FOLDER, name_file+'_outliers'+'.csv'), index=False)
            df_no_outliers.to_csv(os.path.join(UPLOAD_FOLDER, name_file+'_no_outliers'+'.csv'), index=False)

        for file in ["outliers.csv", "no_outliers.csv"]:
            file_out = os.path.join(UPLOAD_FOLDER, "Localization_"+file)
            with open(file_out,"a") as fout:
                name_first_file = os.path.join(UPLOAD_FOLDER, "p1_"+file)
                with open(name_first_file) as f_in:
                    for line in f_in:
                        fout.write(line)
                os.remove(name_first_file)
                for num in range(2,len_files+1):
                    name = os.path.join(UPLOAD_FOLDER,"p"+str(num)+"_"+file)
                    with open(name) as f:
                        f.readline() # skip the header
                        for line in f:
                             fout.write(line)
                    os.remove(name)
    else:
        print("There are no files for localization")
<commit_msg>Improve hyper<commit_after>import os
import re
import csv
import errno
# import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

from sklearn.preprocessing import StandardScaler, LabelEncoder
from collections import OrderedDict

def sizeof_file(name, suffix='B'):
    num = os.stat(name).st_size
    for unit in ['','Ki','Mi','Gi','Ti','Pi','Ei','Zi']:
        if abs(num) < 1024.0:
            return "%3.1f%s%s" % (num, unit, suffix)
        num /= 1024.0
    return "%.1f%s%s" % (num, 'Yi', suffix)

def path_exists(path):
    try:
        os.makedirs(path)
    except OSError as exception:
        if exception.errno != errno.EEXIST:
            raise

def mean_error_localization(y_real, y_test):
    pos = {}
    # pos['1'] = (0,0)
    # pos['2'] = (1.5,0)
    # pos['3'] = (3,0)
    # pos['4'] = (0,1.5)
    # pos['5'] = (1.5,1.5)
    # pos['6'] = (3,1.5)

    y = -1.5
    for e in range(1,16):
        if (e % 3) == 1:
            y += 1.5

        pos[str(e)] = (((e-1)%3)*1.5, y)

    error = []
    for a,b in zip(y_real, y_test):
        error.append(np.linalg.norm(np.array(pos[str(a)]) - np.array(pos[str(b)])))

    return np.mean(np.array(error))

def localization():
    APP_PATH = os.path.dirname(os.path.abspath(__file__))
    UPLOAD_FOLDER = os.path.join(APP_PATH, 'uploads')
    devices_file = 'devices'
    devices_path = os.path.join(UPLOAD_FOLDER, devices_file)
    files = [f for f in os.listdir(UPLOAD_FOLDER) if re.match(r'p[0-9]+', f)]
    len_files = len(files)
    print(devices_path)
    print('len_file', len_files)
    print('files', files)

    if len_files != 0:
        for i in range(1, len_files+1):
            NUMBER = i
            name_file = 'p'+str(NUMBER)
            path_file = os.path.join(UPLOAD_FOLDER, name_file)
            with open(path_file, 'r') as in_file:
                stripped = (line.strip() for line in in_file)
                lines = (line for line in stripped if line)
                grouped = zip(*[lines] * 2)
                with open(path_file+'.csv', 'w') as out_file:
                    writer = csv.writer(out_file)
                    writer.writerow(('Beacon', 'RSSI'))
                    writer.writerows(grouped)

            df = pd.read_csv(path_file+'.csv')
            os.remove(path_file+'.csv')
            condition = ''

            with open(devices_path, 'r') as inFile:
                stripped = (line.strip() for line in inFile)

                for x in stripped:
                    condition += '(Beacon == "' + str(x) + '")' + ' | '
                condition = condition[:-3]
            df = df.query(condition)

            encoder = LabelEncoder()
            df['Beacon'] = encoder.fit_transform(df['Beacon'])

            mapped_df = {}
            for b, data in df.groupby('Beacon'):
                mapped_df[b] = data

            dict_to_df = OrderedDict()
            for index in range(len(mapped_df)):
                dict_to_df[index] = mapped_df[index]['RSSI']

            df_outliers = pd.DataFrame(dict_to_df)

            #Dealing with nan and duplicates
            df_outliers = df_outliers.fillna(method='ffill')
            df_outliers = df_outliers.dropna()
            df_outliers = df_outliers.drop_duplicates()

            #Scaling
            scaler = StandardScaler()
            df_outliers_t = scaler.fit_transform(df_outliers.ix[:, 0:len(df_outliers.columns)-1])
            df_outliers_t = pd.DataFrame(df_outliers_t)

            #Removing outliers aplying the mean
            df_no_outliers = df_outliers_t[df_outliers_t.apply(lambda x: np.abs(x - x.median()) / x.std() < 3).all(axis=1)]
            df_no_outliers = scaler.inverse_transform(df_no_outliers)
            df_no_outliers = pd.DataFrame(df_no_outliers)
            vector_no_out = NUMBER*np.ones((len(df_no_outliers),), dtype=np.int)
            vector_out = NUMBER*np.ones((len(df_outliers),), dtype=np.int)
            df_no_outliers = df_no_outliers.assign(position=pd.Series(vector_no_out).values)
            df_outliers = df_outliers.assign(position=pd.Series(vector_out).values)

            #Saving files with outliers and without it.
            df_outliers.to_csv(os.path.join(UPLOAD_FOLDER, name_file+'_outliers'+'.csv'), index=False)
            df_no_outliers.to_csv(os.path.join(UPLOAD_FOLDER, name_file+'_no_outliers'+'.csv'), index=False)

        for file in ["outliers.csv", "no_outliers.csv"]:
            file_out = os.path.join(UPLOAD_FOLDER, "Localization_"+file)
            with open(file_out,"a") as fout:
                name_first_file = os.path.join(UPLOAD_FOLDER, "p1_"+file)
                with open(name_first_file) as f_in:
                    for line in f_in:
                        fout.write(line)
                os.remove(name_first_file)
                for num in range(2,len_files+1):
                    name = os.path.join(UPLOAD_FOLDER,"p"+str(num)+"_"+file)
                    with open(name) as f:
                        f.readline() # skip the header
                        for line in f:
                             fout.write(line)
                    os.remove(name)
    else:
        print("There are no files for localization")
<|endoftext|>


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

"""Tools for dealing with Mongo's BSON data representation.

Generally not needed to be used by application developers."""

import types
import struct
import random
import re
import datetime
import calendar
import logging

from binary import Binary
from code import Code
from objectid import ObjectId
from dbref import DBRef
from son import SON
from errors import InvalidBSON, InvalidDocument, UnsupportedTag

try:
    import _cbson
    _use_c = True
except ImportError:
    _use_c = False

_logger = logging.getLogger("pymongo.bson")
# _logger.setLevel(logging.DEBUG)
# _logger.addHandler(logging.StreamHandler())

def _get_int(data):
    try:
        value = struct.unpack("<i", data[:4])[0]
    except struct.error:
        raise InvalidBSON()

    return (value, data[4:])

def _get_c_string(data):
    try:
        end = data.index("\x00")
    except ValueError:
        raise InvalidBSON()

    return (unicode(data[:end], "utf-8"), data[end + 1:])

def _make_c_string(string, encode=True):
    if encode and isinstance(string, unicode):
        return string.encode("utf-8") + "\x00"
    return string + "\x00"

def _validate_number(data):
    assert len(data) >= 8
    return data[8:]

def _validate_string(data):
    (length, data) = _get_int(data)
    assert len(data) >= length
    assert data[length - 1] == "\x00"
    return data[length:]

def _validate_object(data):
    return _validate_document(data, None)

_valid_array_name = re.compile("^\d+$")
def _validate_array(data):
    return _validate_document(data, _valid_array_name)

def _validate_binary(data):
    (length, data) = _get_int(data)
    # + 1 for the subtype byte
    assert len(data) >= length + 1
    return data[length + 1:]

def _validate_undefined(data):
    return data

_OID_SIZE = 12
def _validate_oid(data):
    assert len(data) >= _OID_SIZE
    return data[_OID_SIZE:]

def _validate_boolean(data):
    assert len(data) >= 1
    return data[1:]

_DATE_SIZE = 8
def _validate_date(data):
    assert len(data) >= _DATE_SIZE
    return data[_DATE_SIZE:]

_validate_null = _validate_undefined

def _validate_regex(data):
    (regex, data) = _get_c_string(data)
    (options, data) = _get_c_string(data)
    return data

def _validate_ref(data):
    data = _validate_string(data)
    return _validate_oid(data)

_validate_code = _validate_string

# still not sure what is actually stored here, but i know how big it is...
def _validate_code_w_scope(data):
    (length, data) = _get_int(data)
    assert len(data) >= length + 1
    return data[length + 1:]

_validate_symbol = _validate_string

def _validate_number_int(data):
    assert len(data) >= 4
    return data[4:]

_element_validator = {
    "\x01": _validate_number,
    "\x02": _validate_string,
    "\x03": _validate_object,
    "\x04": _validate_array,
    "\x05": _validate_binary,
    "\x06": _validate_undefined,
    "\x07": _validate_oid,
    "\x08": _validate_boolean,
    "\x09": _validate_date,
    "\x0A": _validate_null,
    "\x0B": _validate_regex,
    "\x0C": _validate_ref,
    "\x0D": _validate_code,
    "\x0E": _validate_symbol,
    "\x0F": _validate_code_w_scope,
    "\x10": _validate_number_int,
}

def _validate_element_data(type, data):
    try:
        return _element_validator[type](data)
    except KeyError:
        raise InvalidBSON("unrecognized type: %s" % type)

def _validate_element(data, valid_name):
    element_type = data[0]
    (element_name, data) = _get_c_string(data[1:])
    if valid_name:
        assert valid_name.match(element_name), "%r doesn't match %s" % (element_name, valid_name.pattern)
    return _validate_element_data(element_type, data)

def _validate_elements(data, valid_name):
    while data:
        data = _validate_element(data, valid_name)

def _validate_document(data, valid_name=None):
    try:
        obj_size = struct.unpack("<i", data[:4])[0]
    except struct.error:
        raise InvalidBSON()

    assert obj_size <= len(data)
    obj = data[4:obj_size]
    assert len(obj)

    eoo = obj[-1]
    assert eoo == "\x00"

    elements = obj[:-1]
    _validate_elements(elements, valid_name)

    return data[obj_size:]

def _get_number(data):
    return (struct.unpack("<d", data[:8])[0], data[8:])

def _get_string(data):
    return _get_c_string(data[4:])

def _get_object(data):
    return _document_to_dict(data)

def _get_array(data):
    (obj, data) = _get_object(data)
    result = []
    i = 0
    while True:
        try:
            result.append(obj[str(i)])
            i += 1
        except KeyError:
            break
    return (result, data)

def _get_binary(data):
    (length, data) = _get_int(data)
    subtype = ord(data[0])
    data = data[1:]
    if subtype == 2:
        (length2, data) = _get_int(data)
        if length2 != length - 4:
            raise InvalidBSON("invalid binary (st 2) - lengths don't match!")
        length = length2
    return (Binary(data[:length], subtype), data[length:])

def _get_oid(data):
    oid = _shuffle_oid(data[:12])
    return (ObjectId(oid), data[12:])

def _get_boolean(data):
    return (data[0] == "\x01", data[1:])

def _get_date(data):
    seconds = float(struct.unpack("<q", data[:8])[0]) / 1000.0
    return (datetime.datetime.utcfromtimestamp(seconds), data[8:])

def _get_null(data):
    return (None, data)

def _get_regex(data):
    (pattern, data) = _get_c_string(data)
    (bson_flags, data) = _get_c_string(data)
    flags = 0
    if "i" in bson_flags:
        flags |= re.IGNORECASE
    if "l" in bson_flags:
        flags |= re.LOCALE
    if "m" in bson_flags:
        flags |= re.MULTILINE
    if "s" in bson_flags:
        flags |= re.DOTALL
    if "u" in bson_flags:
        flags |= re.UNICODE
    if "x" in bson_flags:
        flags |= re.VERBOSE
    return (re.compile(pattern, flags), data)

def _get_ref(data):
    (collection, data) = _get_c_string(data[4:])
    (oid, data) = _get_oid(data)
    return (DBRef(collection, oid), data)

_element_getter = {
    "\x01": _get_number,
    "\x02": _get_string,
    "\x03": _get_object,
    "\x04": _get_array,
    "\x05": _get_binary,
    "\x06": _get_null, # undefined
    "\x07": _get_oid,
    "\x08": _get_boolean,
    "\x09": _get_date,
    "\x0A": _get_null,
    "\x0B": _get_regex,
    "\x0C": _get_ref,
    "\x0D": _get_string, # code
    "\x0E": _get_string, # symbol
#    "\x0F": _get_code_w_scope
    "\x10": _get_int, # number_int
}

def _element_to_dict(data):
    element_type = data[0]
    (element_name, data) = _get_c_string(data[1:])
    (value, data) = _element_getter[element_type](data)
    return (element_name, value, data)

def _elements_to_dict(data):
    result = SON()
    while data:
        (key, value, data) = _element_to_dict(data)
        result[key] = value
    return result

def _document_to_dict(data):
    obj_size = struct.unpack("<i", data[:4])[0]
    elements = data[4:obj_size - 1]
    return (_elements_to_dict(elements), data[obj_size:])

def _shuffle_oid(data):
    return data[7::-1] + data[:7:-1]
if _use_c:
    _shuffle_oid = _cbson._shuffle_oid

_RE_TYPE = type(_valid_array_name)
def _element_to_bson(key, value):
    name = _make_c_string(key)
    if isinstance(value, float):
        return "\x01" + name + struct.pack("<d", value)
    if isinstance(value, Binary):
        subtype = value.subtype()
        if subtype == 2:
            value = struct.pack("<i", len(value)) + value
        return "\x05" + name + struct.pack("<i", len(value)) + chr(subtype) + value
    if isinstance(value, Code):
        cstring = _make_c_string(value, False)
        length = struct.pack("<i", len(cstring))
        return "\x0D" + name + length + cstring
    if isinstance(value, str):
        cstring = _make_c_string(value, False)
        length = struct.pack("<i", len(cstring))
        return "\x02" + name + length + cstring
    if isinstance(value, unicode):
        cstring = _make_c_string(value)
        length = struct.pack("<i", len(cstring))
        return "\x02" + name + length + cstring
    if isinstance(value, (dict, SON)):
        return "\x03" + name + _dict_to_bson(value)
    if isinstance(value, (list, tuple)):
        as_dict = SON(zip([str(i) for i in range(len(value))], value))
        return "\x04" + name + _dict_to_bson(as_dict)
    if isinstance(value, ObjectId):
        return "\x07" + name + _shuffle_oid(str(value))
    if value is True:
        return "\x08" + name + "\x01"
    if value is False:
        return "\x08" + name + "\x00"
    if isinstance(value, (int, long)):
        return "\x10" + name + struct.pack("<i", value)
    if isinstance(value, datetime.datetime):
        millis = int(calendar.timegm(value.timetuple()) * 1000 + value.microsecond / 1000)
        return "\x09" + name + struct.pack("<q", millis)
    if value is None:
        return "\x0A" + name
    if isinstance(value, _RE_TYPE):
        pattern = value.pattern
        flags = ""
        if value.flags & re.IGNORECASE:
            flags += "i"
        if value.flags & re.LOCALE:
            flags += "l"
        if value.flags & re.MULTILINE:
            flags += "m"
        if value.flags & re.DOTALL:
            flags += "s"
        if value.flags & re.UNICODE:
            flags += "u"
        if value.flags & re.VERBOSE:
            flags += "x"
        return "\x0B" + name + _make_c_string(pattern) + _make_c_string(flags, False)
    if isinstance(value, DBRef):
        ns = _make_c_string(value.collection())
        return "\x0C" + name + struct.pack("<i", len(ns)) + ns + _shuffle_oid(str(value.id()))
    raise InvalidDocument("cannot convert value of type %s to bson" % type(value))
if _use_c:
    _py_element_to_bson = _element_to_bson
    _c_element_to_bson = _cbson._element_to_bson
    def _element_to_bson(name, value):
        try:
            return _c_element_to_bson(name, value)
        except _cbson.error:
            return _py_element_to_bson(name, value)

def _dict_to_bson(dict):
    try:
        elements = "".join([_element_to_bson(key, value) for (key, value) in dict.iteritems()])
    except AttributeError:
        raise TypeError("argument to from_dict must be a mapping type")

    length = len(elements) + 5
    return struct.pack("<i", length) + elements + "\x00"

def is_valid(bson):
    """Validate that the given string represents valid BSON data.

    Raises TypeError if the data is not an instance of a subclass of str.
    Returns True if the data represents a valid BSON object, False otherwise.

    :Parameters:
      - `bson`: the data to be validated
    """
    if not isinstance(bson, types.StringType):
        raise TypeError("BSON data must be an instance of a subclass of str")

    try:
        remainder = _validate_document(bson)
        return remainder == ""
    except (AssertionError, InvalidBSON):
        return False

def to_dicts(data):
    """Convert binary data to sequence of SON objects.

    Data must be concatenated strings of valid BSON data.

    :Parameters:
      - `data`: bson data
    """
    dicts = []
    while len(data):
        (son, data) = _document_to_dict(data)
        dicts.append(son)
    return dicts

class BSON(str):
    """BSON data.

    Represents binary data storable in and retrievable from Mongo.
    """
    def __new__(cls, bson):
        """Initialize a new BSON object with some data.

        Raises TypeError if `bson` is not an instance of str.

        :Parameters:
          - `bson`: the initial data
        """
        return str.__new__(cls, bson)

    @classmethod
    def from_dict(cls, dict):
        """Create a new BSON object from a python mapping type (like dict).

        Raises TypeError if the argument is not a mapping type, or contains keys
        that are not instance of (str, unicode). Raises InvalidDocument if the
        dictionary cannot be converted to BSON.

        :Parameters:
          - `dict`: mapping type representing a Mongo document
        """
        return cls(_dict_to_bson(dict))

    def to_dict(self):
        """Get the dictionary representation of this data."""
        (son, _) = _document_to_dict(self)
        return son
<commit_msg>just check for instance of dict<commit_after># Copyright 2009 10gen, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tools for dealing with Mongo's BSON data representation.

Generally not needed to be used by application developers."""

import types
import struct
import random
import re
import datetime
import calendar
import logging

from binary import Binary
from code import Code
from objectid import ObjectId
from dbref import DBRef
from son import SON
from errors import InvalidBSON, InvalidDocument, UnsupportedTag

try:
    import _cbson
    _use_c = True
except ImportError:
    _use_c = False

_logger = logging.getLogger("pymongo.bson")
# _logger.setLevel(logging.DEBUG)
# _logger.addHandler(logging.StreamHandler())

def _get_int(data):
    try:
        value = struct.unpack("<i", data[:4])[0]
    except struct.error:
        raise InvalidBSON()

    return (value, data[4:])

def _get_c_string(data):
    try:
        end = data.index("\x00")
    except ValueError:
        raise InvalidBSON()

    return (unicode(data[:end], "utf-8"), data[end + 1:])

def _make_c_string(string, encode=True):
    if encode and isinstance(string, unicode):
        return string.encode("utf-8") + "\x00"
    return string + "\x00"

def _validate_number(data):
    assert len(data) >= 8
    return data[8:]

def _validate_string(data):
    (length, data) = _get_int(data)
    assert len(data) >= length
    assert data[length - 1] == "\x00"
    return data[length:]

def _validate_object(data):
    return _validate_document(data, None)

_valid_array_name = re.compile("^\d+$")
def _validate_array(data):
    return _validate_document(data, _valid_array_name)

def _validate_binary(data):
    (length, data) = _get_int(data)
    # + 1 for the subtype byte
    assert len(data) >= length + 1
    return data[length + 1:]

def _validate_undefined(data):
    return data

_OID_SIZE = 12
def _validate_oid(data):
    assert len(data) >= _OID_SIZE
    return data[_OID_SIZE:]

def _validate_boolean(data):
    assert len(data) >= 1
    return data[1:]

_DATE_SIZE = 8
def _validate_date(data):
    assert len(data) >= _DATE_SIZE
    return data[_DATE_SIZE:]

_validate_null = _validate_undefined

def _validate_regex(data):
    (regex, data) = _get_c_string(data)
    (options, data) = _get_c_string(data)
    return data

def _validate_ref(data):
    data = _validate_string(data)
    return _validate_oid(data)

_validate_code = _validate_string

# still not sure what is actually stored here, but i know how big it is...
def _validate_code_w_scope(data):
    (length, data) = _get_int(data)
    assert len(data) >= length + 1
    return data[length + 1:]

_validate_symbol = _validate_string

def _validate_number_int(data):
    assert len(data) >= 4
    return data[4:]

_element_validator = {
    "\x01": _validate_number,
    "\x02": _validate_string,
    "\x03": _validate_object,
    "\x04": _validate_array,
    "\x05": _validate_binary,
    "\x06": _validate_undefined,
    "\x07": _validate_oid,
    "\x08": _validate_boolean,
    "\x09": _validate_date,
    "\x0A": _validate_null,
    "\x0B": _validate_regex,
    "\x0C": _validate_ref,
    "\x0D": _validate_code,
    "\x0E": _validate_symbol,
    "\x0F": _validate_code_w_scope,
    "\x10": _validate_number_int,
}

def _validate_element_data(type, data):
    try:
        return _element_validator[type](data)
    except KeyError:
        raise InvalidBSON("unrecognized type: %s" % type)

def _validate_element(data, valid_name):
    element_type = data[0]
    (element_name, data) = _get_c_string(data[1:])
    if valid_name:
        assert valid_name.match(element_name), "%r doesn't match %s" % (element_name, valid_name.pattern)
    return _validate_element_data(element_type, data)

def _validate_elements(data, valid_name):
    while data:
        data = _validate_element(data, valid_name)

def _validate_document(data, valid_name=None):
    try:
        obj_size = struct.unpack("<i", data[:4])[0]
    except struct.error:
        raise InvalidBSON()

    assert obj_size <= len(data)
    obj = data[4:obj_size]
    assert len(obj)

    eoo = obj[-1]
    assert eoo == "\x00"

    elements = obj[:-1]
    _validate_elements(elements, valid_name)

    return data[obj_size:]

def _get_number(data):
    return (struct.unpack("<d", data[:8])[0], data[8:])

def _get_string(data):
    return _get_c_string(data[4:])

def _get_object(data):
    return _document_to_dict(data)

def _get_array(data):
    (obj, data) = _get_object(data)
    result = []
    i = 0
    while True:
        try:
            result.append(obj[str(i)])
            i += 1
        except KeyError:
            break
    return (result, data)

def _get_binary(data):
    (length, data) = _get_int(data)
    subtype = ord(data[0])
    data = data[1:]
    if subtype == 2:
        (length2, data) = _get_int(data)
        if length2 != length - 4:
            raise InvalidBSON("invalid binary (st 2) - lengths don't match!")
        length = length2
    return (Binary(data[:length], subtype), data[length:])

def _get_oid(data):
    oid = _shuffle_oid(data[:12])
    return (ObjectId(oid), data[12:])

def _get_boolean(data):
    return (data[0] == "\x01", data[1:])

def _get_date(data):
    seconds = float(struct.unpack("<q", data[:8])[0]) / 1000.0
    return (datetime.datetime.utcfromtimestamp(seconds), data[8:])

def _get_null(data):
    return (None, data)

def _get_regex(data):
    (pattern, data) = _get_c_string(data)
    (bson_flags, data) = _get_c_string(data)
    flags = 0
    if "i" in bson_flags:
        flags |= re.IGNORECASE
    if "l" in bson_flags:
        flags |= re.LOCALE
    if "m" in bson_flags:
        flags |= re.MULTILINE
    if "s" in bson_flags:
        flags |= re.DOTALL
    if "u" in bson_flags:
        flags |= re.UNICODE
    if "x" in bson_flags:
        flags |= re.VERBOSE
    return (re.compile(pattern, flags), data)

def _get_ref(data):
    (collection, data) = _get_c_string(data[4:])
    (oid, data) = _get_oid(data)
    return (DBRef(collection, oid), data)

_element_getter = {
    "\x01": _get_number,
    "\x02": _get_string,
    "\x03": _get_object,
    "\x04": _get_array,
    "\x05": _get_binary,
    "\x06": _get_null, # undefined
    "\x07": _get_oid,
    "\x08": _get_boolean,
    "\x09": _get_date,
    "\x0A": _get_null,
    "\x0B": _get_regex,
    "\x0C": _get_ref,
    "\x0D": _get_string, # code
    "\x0E": _get_string, # symbol
#    "\x0F": _get_code_w_scope
    "\x10": _get_int, # number_int
}

def _element_to_dict(data):
    element_type = data[0]
    (element_name, data) = _get_c_string(data[1:])
    (value, data) = _element_getter[element_type](data)
    return (element_name, value, data)

def _elements_to_dict(data):
    result = SON()
    while data:
        (key, value, data) = _element_to_dict(data)
        result[key] = value
    return result

def _document_to_dict(data):
    obj_size = struct.unpack("<i", data[:4])[0]
    elements = data[4:obj_size - 1]
    return (_elements_to_dict(elements), data[obj_size:])

def _shuffle_oid(data):
    return data[7::-1] + data[:7:-1]
if _use_c:
    _shuffle_oid = _cbson._shuffle_oid

_RE_TYPE = type(_valid_array_name)
def _element_to_bson(key, value):
    name = _make_c_string(key)
    if isinstance(value, float):
        return "\x01" + name + struct.pack("<d", value)
    if isinstance(value, Binary):
        subtype = value.subtype()
        if subtype == 2:
            value = struct.pack("<i", len(value)) + value
        return "\x05" + name + struct.pack("<i", len(value)) + chr(subtype) + value
    if isinstance(value, Code):
        cstring = _make_c_string(value, False)
        length = struct.pack("<i", len(cstring))
        return "\x0D" + name + length + cstring
    if isinstance(value, str):
        cstring = _make_c_string(value, False)
        length = struct.pack("<i", len(cstring))
        return "\x02" + name + length + cstring
    if isinstance(value, unicode):
        cstring = _make_c_string(value)
        length = struct.pack("<i", len(cstring))
        return "\x02" + name + length + cstring
    if isinstance(value, dict):
        return "\x03" + name + _dict_to_bson(value)
    if isinstance(value, (list, tuple)):
        as_dict = SON(zip([str(i) for i in range(len(value))], value))
        return "\x04" + name + _dict_to_bson(as_dict)
    if isinstance(value, ObjectId):
        return "\x07" + name + _shuffle_oid(str(value))
    if value is True:
        return "\x08" + name + "\x01"
    if value is False:
        return "\x08" + name + "\x00"
    if isinstance(value, (int, long)):
        return "\x10" + name + struct.pack("<i", value)
    if isinstance(value, datetime.datetime):
        millis = int(calendar.timegm(value.timetuple()) * 1000 + value.microsecond / 1000)
        return "\x09" + name + struct.pack("<q", millis)
    if value is None:
        return "\x0A" + name
    if isinstance(value, _RE_TYPE):
        pattern = value.pattern
        flags = ""
        if value.flags & re.IGNORECASE:
            flags += "i"
        if value.flags & re.LOCALE:
            flags += "l"
        if value.flags & re.MULTILINE:
            flags += "m"
        if value.flags & re.DOTALL:
            flags += "s"
        if value.flags & re.UNICODE:
            flags += "u"
        if value.flags & re.VERBOSE:
            flags += "x"
        return "\x0B" + name + _make_c_string(pattern) + _make_c_string(flags, False)
    if isinstance(value, DBRef):
        ns = _make_c_string(value.collection())
        return "\x0C" + name + struct.pack("<i", len(ns)) + ns + _shuffle_oid(str(value.id()))
    raise InvalidDocument("cannot convert value of type %s to bson" % type(value))
if _use_c:
    _py_element_to_bson = _element_to_bson
    _c_element_to_bson = _cbson._element_to_bson
    def _element_to_bson(name, value):
        try:
            return _c_element_to_bson(name, value)
        except _cbson.error:
            return _py_element_to_bson(name, value)

def _dict_to_bson(dict):
    try:
        elements = "".join([_element_to_bson(key, value) for (key, value) in dict.iteritems()])
    except AttributeError:
        raise TypeError("argument to from_dict must be a mapping type")

    length = len(elements) + 5
    return struct.pack("<i", length) + elements + "\x00"

def is_valid(bson):
    """Validate that the given string represents valid BSON data.

    Raises TypeError if the data is not an instance of a subclass of str.
    Returns True if the data represents a valid BSON object, False otherwise.

    :Parameters:
      - `bson`: the data to be validated
    """
    if not isinstance(bson, types.StringType):
        raise TypeError("BSON data must be an instance of a subclass of str")

    try:
        remainder = _validate_document(bson)
        return remainder == ""
    except (AssertionError, InvalidBSON):
        return False

def to_dicts(data):
    """Convert binary data to sequence of SON objects.

    Data must be concatenated strings of valid BSON data.

    :Parameters:
      - `data`: bson data
    """
    dicts = []
    while len(data):
        (son, data) = _document_to_dict(data)
        dicts.append(son)
    return dicts

class BSON(str):
    """BSON data.

    Represents binary data storable in and retrievable from Mongo.
    """
    def __new__(cls, bson):
        """Initialize a new BSON object with some data.

        Raises TypeError if `bson` is not an instance of str.

        :Parameters:
          - `bson`: the initial data
        """
        return str.__new__(cls, bson)

    @classmethod
    def from_dict(cls, dict):
        """Create a new BSON object from a python mapping type (like dict).

        Raises TypeError if the argument is not a mapping type, or contains keys
        that are not instance of (str, unicode). Raises InvalidDocument if the
        dictionary cannot be converted to BSON.

        :Parameters:
          - `dict`: mapping type representing a Mongo document
        """
        return cls(_dict_to_bson(dict))

    def to_dict(self):
        """Get the dictionary representation of this data."""
        (son, _) = _document_to_dict(self)
        return son
<|endoftext|>


<commit_before>"""
Improvements on Pyte.
"""
from __future__ import unicode_literals
from pyte.streams import Stream
from pyte.escape import NEL
from pyte import control as ctrl
from collections import defaultdict

from .log import logger

import re

__all__ = (
    'BetterStream',
)

class BetterStream(Stream):
    """
    Extension to the Pyte `Stream` class that also handles "Esc]<num>...BEL"
    sequences. This is used by xterm to set the terminal title.
    """
    escape = Stream.escape.copy()
    escape.update({
        # Call next_line instead of line_feed. We always want to go to the left
        # margin if we receive this, unlike \n, which goes one row down.
        # (Except when LNM has been set.)
        NEL: "next_line",
    })

    def __init__(self, screen):
        super(BetterStream, self).__init__()
        self.listener = screen

        self._validate_screen()

        # Create a regular expression pattern that matches everything what can
        # be considered plain text. This can be used as a very simple lexer
        # that can feed the "plain text part" as one token into the screen.
        special = set([ctrl.ESC, ctrl.CSI, ctrl.NUL, ctrl.DEL]) | set(self.basic)
        self._text_search = re.compile(
            '[^%s]+' % ''.join(re.escape(c) for c in special)).match

        # Start parser.
        self._parser = self._parser_generator()
        self._taking_plain_text = self._parser.send(None)
        self._send = self._parser.send

    def _validate_screen(self):
        """
        Check whether our Screen class has all the required callbacks.
        (We want to verify this statically, before feeding content to the
        screen.)
        """
        for d in [self.basic, self.escape, self.sharp, self.percent, self.csi]:
            for name in d.values():
                assert hasattr(self.listener, name), 'Screen is missing %r' % name

    def feed(self, chars):
        """
        Custom, much more efficient 'feed' function.
        Feed a string of characters to the parser.
        """
        # The original implementation of this function looked like this::
        #
        #     for c in chars:
        #         self._send(c)
        #
        # However, the implementation below does a big optimization. If the
        # parser is possibly expecting a chunk of text (when it is not inside a
        # ESC or CSI escape sequence), then we send that fragment directly to
        # the 'draw' method of the screen.

        # Local copy of functions. (For faster lookups.)
        send = self._send
        taking_plain_text = self._taking_plain_text
        text_search = self._text_search
        draw = self.listener.draw

        # Loop through the chars.
        i = 0
        count = len(chars)

        while i < count:
            # Reading plain text? Don't send characters one by one in the
            # generator, but optimize and send the whole chunk without
            # escapes directly to the listener.
            if taking_plain_text:
                match = text_search(chars, i)
                if match:
                    start, i = match.span()
                    draw(chars[start:i])
                else:
                    taking_plain_text = False

            # The parser expects just one character now. Just send the next one.
            else:
                taking_plain_text = send(chars[i])
                i += 1

        # Remember state for the next 'feed()'.
        self._taking_plain_text = taking_plain_text

    def _parser_generator(self):
        """
        Coroutine that processes VT100 output.

        It's actually a state machine, implemented as a coroutine. So all the
        'state' that we have is stored in local variables.

        This generator is not the most beautiful, but it is as performant as
        possible. When a process generates a lot of output (That is often much
        more than a person would give as input), then this will be the
        bottleneck, because it processes just one character at a time.

        We did many manual optimizations to this function in order to make it
        as efficient as possible. Don't change anything without profiling
        first.
        """
        listener = self.listener

        basic = self.basic
        escape = self.escape
        sharp = self.sharp
        percent = self.percent
        csi = self.csi

        ESC = ctrl.ESC
        CSI = ctrl.CSI
        CTRL_SEQUENCES_ALLOWED_IN_CSI = set([
            ctrl.BEL, ctrl.BS, ctrl.HT, ctrl.LF, ctrl.VT, ctrl.FF, ctrl.CR])

        def create_dispatch_dictionary(source_dict):
            # In order to avoid getting KeyError exceptions below, we make sure
            # that these dictionaries have a dummy handler.
            def dummy(*a, **kw):
                pass

            return defaultdict(
                lambda: dummy,
                {event: getattr(listener, attr) for event, attr in source_dict.items()})

        basic_dispatch = create_dispatch_dictionary(basic)
        sharp_dispatch = create_dispatch_dictionary(sharp)
        percent_dispatch = create_dispatch_dictionary(percent)
        escape_dispatch = create_dispatch_dictionary(escape)
        csi_dispatch = create_dispatch_dictionary(csi)

        while True:
            char = yield True  # (`True` tells the 'send()' function that it
                               # is allowed to send chunks of plain text
                               # directly to the listener, instead of this generator.)

            if char == ESC:  # \x1b
                char = yield

                if char == '[':
                    char = CSI  # Go to CSI.
                else:
                    if char == '#':
                        sharp_dispatch[(yield)]()
                    elif char == '%':
                        percent_dispatch[(yield)]()
                    elif char in '()':
                        listener.set_charset((yield), mode=char)
                    elif char == ']':
                        data = []
                        while True:
                            c = yield
                            if c == '\07':
                                break
                            else:
                                data.append(c)
                        listener.square_close(''.join(data))
                    else:
                        escape_dispatch[char]()
                    continue  # Do not go to CSI.

            if char in basic:  # 'if', not 'elif', because we need to be
                               # able to jump here from Esc[ above in the CSI
                               # section below.
                basic_dispatch[char]()

            elif char == CSI:  # \x9b
                current = ''
                params = []
                private = False

                while True:
                    char = yield
                    if char == '?':
                        private = True
                    elif char in CTRL_SEQUENCES_ALLOWED_IN_CSI:
                        basic_dispatch[char]()
                    elif char in (ctrl.SP, '>'):
                        # Ignore '>' because of 'ESC[>c' (Send device attributes.)
                        pass
                    elif char.isdigit():
                        current += char
                    else:
                        params.append(min(int(current or 0), 9999))

                        if char == ';':
                            current = ''
                        else:
                            try:
                                if private:
                                    csi_dispatch[char](*params, private=True)
                                else:
                                    csi_dispatch[char](*params)
                            except TypeError:
                                # Handler doesn't take params or private attribute.
                                # (Not the cleanest way to handle this, but
                                # it's safe and performant enough.)
                                logger.warning('Dispatch %s failed. params=%s, private=%s',
                                               params, private)
                            break  # Break outside CSI loop.
<commit_msg>Updated syntax to work with Python 2.6<commit_after>"""
Improvements on Pyte.
"""
from __future__ import unicode_literals
from pyte.streams import Stream
from pyte.escape import NEL
from pyte import control as ctrl
from collections import defaultdict

from .log import logger

import re

__all__ = (
    'BetterStream',
)

class BetterStream(Stream):
    """
    Extension to the Pyte `Stream` class that also handles "Esc]<num>...BEL"
    sequences. This is used by xterm to set the terminal title.
    """
    escape = Stream.escape.copy()
    escape.update({
        # Call next_line instead of line_feed. We always want to go to the left
        # margin if we receive this, unlike \n, which goes one row down.
        # (Except when LNM has been set.)
        NEL: "next_line",
    })

    def __init__(self, screen):
        super(BetterStream, self).__init__()
        self.listener = screen

        self._validate_screen()

        # Create a regular expression pattern that matches everything what can
        # be considered plain text. This can be used as a very simple lexer
        # that can feed the "plain text part" as one token into the screen.
        special = set([ctrl.ESC, ctrl.CSI, ctrl.NUL, ctrl.DEL]) | set(self.basic)
        self._text_search = re.compile(
            '[^%s]+' % ''.join(re.escape(c) for c in special)).match

        # Start parser.
        self._parser = self._parser_generator()
        self._taking_plain_text = self._parser.send(None)
        self._send = self._parser.send

    def _validate_screen(self):
        """
        Check whether our Screen class has all the required callbacks.
        (We want to verify this statically, before feeding content to the
        screen.)
        """
        for d in [self.basic, self.escape, self.sharp, self.percent, self.csi]:
            for name in d.values():
                assert hasattr(self.listener, name), 'Screen is missing %r' % name

    def feed(self, chars):
        """
        Custom, much more efficient 'feed' function.
        Feed a string of characters to the parser.
        """
        # The original implementation of this function looked like this::
        #
        #     for c in chars:
        #         self._send(c)
        #
        # However, the implementation below does a big optimization. If the
        # parser is possibly expecting a chunk of text (when it is not inside a
        # ESC or CSI escape sequence), then we send that fragment directly to
        # the 'draw' method of the screen.

        # Local copy of functions. (For faster lookups.)
        send = self._send
        taking_plain_text = self._taking_plain_text
        text_search = self._text_search
        draw = self.listener.draw

        # Loop through the chars.
        i = 0
        count = len(chars)

        while i < count:
            # Reading plain text? Don't send characters one by one in the
            # generator, but optimize and send the whole chunk without
            # escapes directly to the listener.
            if taking_plain_text:
                match = text_search(chars, i)
                if match:
                    start, i = match.span()
                    draw(chars[start:i])
                else:
                    taking_plain_text = False

            # The parser expects just one character now. Just send the next one.
            else:
                taking_plain_text = send(chars[i])
                i += 1

        # Remember state for the next 'feed()'.
        self._taking_plain_text = taking_plain_text

    def _parser_generator(self):
        """
        Coroutine that processes VT100 output.

        It's actually a state machine, implemented as a coroutine. So all the
        'state' that we have is stored in local variables.

        This generator is not the most beautiful, but it is as performant as
        possible. When a process generates a lot of output (That is often much
        more than a person would give as input), then this will be the
        bottleneck, because it processes just one character at a time.

        We did many manual optimizations to this function in order to make it
        as efficient as possible. Don't change anything without profiling
        first.
        """
        listener = self.listener

        basic = self.basic
        escape = self.escape
        sharp = self.sharp
        percent = self.percent
        csi = self.csi

        ESC = ctrl.ESC
        CSI = ctrl.CSI
        CTRL_SEQUENCES_ALLOWED_IN_CSI = set([
            ctrl.BEL, ctrl.BS, ctrl.HT, ctrl.LF, ctrl.VT, ctrl.FF, ctrl.CR])

        def create_dispatch_dictionary(source_dict):
            # In order to avoid getting KeyError exceptions below, we make sure
            # that these dictionaries have a dummy handler.
            def dummy(*a, **kw):
                pass

            return defaultdict(
                lambda: dummy,
                dict((event, getattr(listener, attr)) for event, attr in source_dict.items()))

        basic_dispatch = create_dispatch_dictionary(basic)
        sharp_dispatch = create_dispatch_dictionary(sharp)
        percent_dispatch = create_dispatch_dictionary(percent)
        escape_dispatch = create_dispatch_dictionary(escape)
        csi_dispatch = create_dispatch_dictionary(csi)

        while True:
            char = yield True  # (`True` tells the 'send()' function that it
                               # is allowed to send chunks of plain text
                               # directly to the listener, instead of this generator.)

            if char == ESC:  # \x1b
                char = yield

                if char == '[':
                    char = CSI  # Go to CSI.
                else:
                    if char == '#':
                        sharp_dispatch[(yield)]()
                    elif char == '%':
                        percent_dispatch[(yield)]()
                    elif char in '()':
                        listener.set_charset((yield), mode=char)
                    elif char == ']':
                        data = []
                        while True:
                            c = yield
                            if c == '\07':
                                break
                            else:
                                data.append(c)
                        listener.square_close(''.join(data))
                    else:
                        escape_dispatch[char]()
                    continue  # Do not go to CSI.

            if char in basic:  # 'if', not 'elif', because we need to be
                               # able to jump here from Esc[ above in the CSI
                               # section below.
                basic_dispatch[char]()

            elif char == CSI:  # \x9b
                current = ''
                params = []
                private = False

                while True:
                    char = yield
                    if char == '?':
                        private = True
                    elif char in CTRL_SEQUENCES_ALLOWED_IN_CSI:
                        basic_dispatch[char]()
                    elif char in (ctrl.SP, '>'):
                        # Ignore '>' because of 'ESC[>c' (Send device attributes.)
                        pass
                    elif char.isdigit():
                        current += char
                    else:
                        params.append(min(int(current or 0), 9999))

                        if char == ';':
                            current = ''
                        else:
                            try:
                                if private:
                                    csi_dispatch[char](*params, private=True)
                                else:
                                    csi_dispatch[char](*params)
                            except TypeError:
                                # Handler doesn't take params or private attribute.
                                # (Not the cleanest way to handle this, but
                                # it's safe and performant enough.)
                                logger.warning('Dispatch %s failed. params=%s, private=%s',
                                               params, private)
                            break  # Break outside CSI loop.
<|endoftext|>


<commit_before>"""
Improvements on Pyte.
"""
from __future__ import unicode_literals
from pyte.streams import Stream
from pyte.escape import NEL
from pyte import ctrl
from collections import defaultdict

__all__ = (
    'BetterStream',
)

class BetterStream(Stream):
    """
    Extension to the Pyte `Stream` class that also handles "Esc]<num>...BEL"
    sequences. This is used by xterm to set the terminal title.
    """
    csi = {
        'n': 'cpr',  # Cursor position request.
        'c': 'send_device_attributes',  # csi > Ps c
    }
    csi.update(Stream.csi)

    escape = Stream.escape.copy()
    escape.update({
        # Call next_line instead of line_feed. We always want to go to the left
        # margin if we receive this, unlike \n, which goes one row down.
        # (Except when LNM has been set.)
        NEL: "next_line",
    })

    def __init__(self, screen):
        super(BetterStream, self).__init__()
        self.listener = screen

        # Start parser.
        self._parser = self._parser_generator()
        self._parser.send(None)

    def feed(self, chars):  # TODO: Handle exceptions.
        """
        Custom, much more efficient 'feed' function.
        """
        # Send all input to the parser coroutine.
        send = self._parser.send

        for c in chars:
            send(c)

        # Call the '__after__' function, which is used to update the screen
        # height.
        self.listener.__after__(self)

    def _parser_generator(self):
        """
        Coroutine that processes VT100 output.

        It's actually a state machine, implemented as a coroutine. So all the
        'state' that we have is stored in local variables.
        """
        listener = self.listener
        draw = listener.draw

        # In order to avoid getting KeyError exceptions below, we make sure that
        # these dictionaries resolve to 'Screen.dummy'
        basic = defaultdict(lambda: 'dummy', self.basic)
        escape = defaultdict(lambda: 'dummy', self.escape)
        sharp = defaultdict(lambda: 'dummy', self.sharp)
        percent = defaultdict(lambda: 'dummy', self.percent)
        csi = defaultdict(lambda: 'dummy', self.csi)

        ESC = ctrl.ESC
        CSI = ctrl.CSI
        NUL_OR_DEL = (ctrl.NUL, ctrl.DEL)
        CTRL_SEQUENCES_ALLOWED_IN_CSI = (
            ctrl.BEL, ctrl.BS, ctrl.HT, ctrl.LF, ctrl.VT, ctrl.FF, ctrl.CR)

        def dispatch(event, *args, **flags):
            getattr(listener, event)(*args, **flags)

        while True:
            char = yield

            if char == ESC:  # \x1b
                char = yield

                if char == '[':
                    char = CSI  # Go to CSI.
                else:
                    if char == '#':
                        dispatch(sharp[(yield)])
                    elif char == '%':
                        dispatch(percent[(yield)])
                    elif char in '()':
                        listener.set_charset((yield), mode=char)
                    elif char == ']':
                        data = []
                        while True:
                            c = yield
                            if c == '\07':
                                break
                            else:
                                data.append(c)
                        listener.square_close(''.join(data))
                    else:
                        dispatch(escape[char])
                    continue  # Do not go to CSI.

            if char in basic:  # 'if', not 'elif', because we need to be
                               # able to jump here from Esc[ above in the CSI
                               # section below.
                dispatch(basic[char])

            elif char == CSI:  # \x9b
                current = ''
                params = []
                private = False

                while True:
                    char = yield
                    if char == '?':
                        private = True
                    elif char in CTRL_SEQUENCES_ALLOWED_IN_CSI:
                        dispatch(basic[char])
                    elif char in (ctrl.SP, '>'):
                        # Ignore '>' because of 'ESC[>c' (Send device attributes.)
                        pass
                    elif char.isdigit():
                        current += char
                    else:
                        params.append(min(int(current or 0), 9999))

                        if char == ';':
                            current = ''
                        else:
                            if private:
                                dispatch(csi[char], *params, private=True)
                            else:
                                dispatch(csi[char], *params)
                            break  # Break outside CSI loop.

            elif char not in NUL_OR_DEL:
                draw(char)
<commit_msg>Improved the performance of the Stream class.<commit_after>"""
Improvements on Pyte.
"""
from __future__ import unicode_literals
from pyte.streams import Stream
from pyte.escape import NEL
from pyte import ctrl
from collections import defaultdict

__all__ = (
    'BetterStream',
)

class BetterStream(Stream):
    """
    Extension to the Pyte `Stream` class that also handles "Esc]<num>...BEL"
    sequences. This is used by xterm to set the terminal title.
    """
    csi = {
        'n': 'cpr',  # Cursor position request.
        'c': 'send_device_attributes',  # csi > Ps c
    }
    csi.update(Stream.csi)

    escape = Stream.escape.copy()
    escape.update({
        # Call next_line instead of line_feed. We always want to go to the left
        # margin if we receive this, unlike \n, which goes one row down.
        # (Except when LNM has been set.)
        NEL: "next_line",
    })

    def __init__(self, screen):
        super(BetterStream, self).__init__()
        self.listener = screen

        # Start parser.
        self._parser = self._parser_generator()
        self._parser.send(None)
        self._send = self._parser.send

    def feed(self, chars):
        """
        Custom, much more efficient 'feed' function.
        """
        # Send all input to the parser coroutine.
        send = self._send

        for c in chars:
            send(c)

        # Call the '__after__' function, which is used to update the screen
        # height.
        self.listener.__after__(self)

    def _parser_generator(self):
        """
        Coroutine that processes VT100 output.

        It's actually a state machine, implemented as a coroutine. So all the
        'state' that we have is stored in local variables.

        This generator is not the most beautiful, but it is as performant as
        possible. When a process generates a lot of output (That is often much
        more than a person would give as input), then this will be the
        bottleneck, because it processes just one character at a time.

        We did many manual optimizations to this function in order to make it
        as efficient as possible. Don't change anything without profiling
        first.
        """
        listener = self.listener
        draw = listener.draw

        # In order to avoid getting KeyError exceptions below, we make sure that
        # these dictionaries resolve to 'Screen.dummy'
        basic = defaultdict(lambda: 'dummy', self.basic)
        escape = defaultdict(lambda: 'dummy', self.escape)
        sharp = defaultdict(lambda: 'dummy', self.sharp)
        percent = defaultdict(lambda: 'dummy', self.percent)
        csi = defaultdict(lambda: 'dummy', self.csi)

        ESC = ctrl.ESC
        CSI = ctrl.CSI
        CTRL_SEQUENCES_ALLOWED_IN_CSI = set([
            ctrl.BEL, ctrl.BS, ctrl.HT, ctrl.LF, ctrl.VT, ctrl.FF, ctrl.CR])
        NOT_DRAW = set([ESC, CSI, ctrl.NUL, ctrl.DEL]) | set(basic)

        def dispatch(event, *args, **flags):
            getattr(listener, event)(*args, **flags)

        while True:
            char = yield

            # Handle normal draw operations first. (All overhead here is the
            # most expensive.)
            while char not in NOT_DRAW:
                draw(char)
                char = yield

            if char == ESC:  # \x1b
                char = yield

                if char == '[':
                    char = CSI  # Go to CSI.
                else:
                    if char == '#':
                        dispatch(sharp[(yield)])
                    elif char == '%':
                        dispatch(percent[(yield)])
                    elif char in '()':
                        listener.set_charset((yield), mode=char)
                    elif char == ']':
                        data = []
                        while True:
                            c = yield
                            if c == '\07':
                                break
                            else:
                                data.append(c)
                        listener.square_close(''.join(data))
                    else:
                        dispatch(escape[char])
                    continue  # Do not go to CSI.

            if char in basic:  # 'if', not 'elif', because we need to be
                               # able to jump here from Esc[ above in the CSI
                               # section below.
                dispatch(basic[char])

            elif char == CSI:  # \x9b
                current = ''
                params = []
                private = False

                while True:
                    char = yield
                    if char == '?':
                        private = True
                    elif char in CTRL_SEQUENCES_ALLOWED_IN_CSI:
                        dispatch(basic[char])
                    elif char in (ctrl.SP, '>'):
                        # Ignore '>' because of 'ESC[>c' (Send device attributes.)
                        pass
                    elif char.isdigit():
                        current += char
                    else:
                        params.append(min(int(current or 0), 9999))

                        if char == ';':
                            current = ''
                        else:
                            if private:
                                dispatch(csi[char], *params, private=True)
                            else:
                                dispatch(csi[char], *params)
                            break  # Break outside CSI loop.
<|endoftext|>


<commit_before>__author__ = 'nmearl'

import numpy as np
import time
import pylab
import modeler
import minimizer
import hammer
import utilfuncs
import argparse

def read_data(data_file):
    """Reads in photometric data. File must have at least two columns, error is optional.

    @param data_file: Path of file.
    @return: Numpy arrays of the times, data, and data error.
    """
    data_times, data_fluxes, data_yerr = [], [], []

    with open('{0}'.format(data_file), 'r') as f:
        for line in f.readlines():
            time, flux, err = line.strip().split()[:3]
            data_times.append(float(time))
            data_fluxes.append(float(flux))
            data_yerr.append(float(err))

            # if len(data_yerr) == 0:
            # data_yerr = np.ones(len(data_fluxes))
            # x, y, data_yerr = kepler.get_fits_data('5897826', use_pdc=True, use_slc=False)

    return np.array(data_times), np.array(data_fluxes), np.array(data_yerr)

def read_input(input_file):
    """Reads in the formatted input file.

    @param input_file: Path of file.
    @return: Tuple containing all 15 parameter data.
    """
    with open('{0}'.format(input_file), 'r') as f:
        input_pars = [line.strip() for line in f.readlines() if line[0] is not '#']

    N = [int(i.strip()) for i in input_pars[0].split('#')][0]
    t0 = [float(i.strip()) for i in input_pars[1].split('#')][0]
    maxh = [float(i.strip()) for i in input_pars[2].split('#')][0]
    orbit_error = [float(i.strip()) for i in input_pars[3].split('#')][0]

    masses = [float(i.strip()) for i in input_pars[4].split()]
    radii = [float(i.strip()) for i in input_pars[5].split()]
    fluxes = [float(i.strip()) for i in input_pars[6].split()]
    u1 = [float(i.strip()) for i in input_pars[7].split()]
    u2 = [float(i.strip()) for i in input_pars[8].split()]

    a = [float(i.strip()) for i in input_pars[9].split()]
    e = [float(i.strip()) for i in input_pars[10].split()]
    inc = [float(i.strip()) for i in input_pars[11].split()]
    om = [float(i.strip()) for i in input_pars[12].split()]
    ln = [float(i.strip()) for i in input_pars[13].split()]
    ma = [float(i.strip()) for i in input_pars[14].split()]

    return N, t0, maxh, orbit_error, masses, radii, fluxes, u1, u2, a, e, inc, om, ln, ma

def get_random_pos(N, t0, maxh, orbit_error):
    masses, radii, fluxes, u1, u2, a, e, inc, om, ln, ma = utilfuncs.random_pos(N)
    return N, t0, maxh, orbit_error, masses, radii, fluxes, u1, u2, a, e, inc, om, ln, ma

def plot_model(params, x, y, yerr):
    """Simple plot routine to visualize the modelled input parameters.

    @param params: Tuple of the parameters.
    @param x: Time data.
    @param y: Flux data.
    @param yerr: Flux error.
    """
    N, t0, maxh, orbit_error, masses, radii, fluxes, u1, u2, a, e, inc, om, ln, ma = params

    y_mod = modeler.generate(
        N, t0,
        maxh, orbit_error,
        x,
        masses, radii, fluxes, u1, u2,
        a, e, inc, om, ln, ma
    )

    pylab.plot(x, y, 'k+')
    pylab.plot(x, y_mod, 'r')

    pylab.show()

def main(data_file, fit_method, input_file, nwalkers, niterations, ncores, syspars):
    """The main function mediates the reading of the data and input parameters, and starts the optimization using the
    specified method.

    @param data_file: Path to data file.
    @param fit_method: Currently: 'min' is least squares optimization, and 'mcmc' is the emcee hammer optimization. The
    'plot' option is used to simply plot the returned model based on the supplied parameters.
    @param input_file: Path to input file.
    """

    if input_file:
        params = read_input(input_file)
    elif syspars:
        print('No initial parameter set specified, randomly sampling parameter space...')
        N, t0, maxh, orbit_error = syspars
        params = get_random_pos(int(N), t0, maxh, orbit_error)
        rpars = True
    else:
        print('No initial parameter set specified, randomly sampling parameter space...')
        N = int(raw_input('\tEnter the number of bodies: '))
        t0 = float(raw_input('\tEnter the epoch of coordinates: '))
        maxh = float(raw_input('\tEnter the maximum time step (default: 0.01): '))
        orbit_error = float(raw_input('\tEnter the orbit error tolerance (default: 1e-20): '))

        params = get_random_pos(N, t0, maxh, orbit_error)
        rpars = True

    if not fit_method:
        print('You have not specified a fit method, defaulting to least squares minimization.')

    if '.' in data_file:
        x, y, yerr = read_data(data_file)

    n = int(len(x))

    time_start = time.time()

    if fit_method == 'mcmc':
        hammer.generate(
            params, x[:n], y[:n], yerr[:n],
            nwalkers, niterations, ncores, rpars,
            data_file.split('/')[-1].split('.')[0]
        )

    elif fit_method == 'plot':
        plot_model(params, x[:n], y[:n], yerr[:n])

    else:
        minimizer.generate(
            params, x[:n], y[:n], yerr[:n], fit_method,
            data_file.split('/')[-1].split('.')[0]
        )

    print "Total time:", time.time() - time_start

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Photometric dynamical modeling code.')
    parser.add_argument('data', help='data file containing detrended light curve')
    parser.add_argument('-f', '--fit', help='fit method used to minimize',
                        choices=['mcmc', 'leastsq', 'nelder', 'lbfgsb', 'anneal', 'powell',
                                 'cg', 'newton', 'cobyla', 'slsqp', 'plot'], default='leastsq')
    parser.add_argument('-i', '--input', help='input file containing initial parameters (overrides --system)')
    parser.add_argument('-w', '--walkers', type=int, default=250,
                        help='number of walkers if using mcmc fit method ')
    parser.add_argument('-t', '--iterations', type=int, default=500,
                        help='number of iterations to perform if using mcmc fit method')
    parser.add_argument('-c', '--cores', type=int, default=1,
                        help='number of cores to utilize')
    parser.add_argument('-s', '--system', nargs=4,
                        help='four initial system parameters: N, t0, maxh, orbit_error')

    args = parser.parse_args()
    data_file = args.data
    fit_method = args.fit
    input_file = args.input
    nwalkers = args.walkers
    niterations = args.iterations
    ncores = args.cores
    syspars = map(float, args.system)

    main(data_file, fit_method, input_file, nwalkers, niterations, ncores, syspars)
<commit_msg>Quick fix.<commit_after>__author__ = 'nmearl'

import numpy as np
import time
import pylab
import modeler
import minimizer
import hammer
import utilfuncs
import argparse

def read_data(data_file):
    """Reads in photometric data. File must have at least two columns, error is optional.

    @param data_file: Path of file.
    @return: Numpy arrays of the times, data, and data error.
    """
    data_times, data_fluxes, data_yerr = [], [], []

    with open('{0}'.format(data_file), 'r') as f:
        for line in f.readlines():
            time, flux, err = line.strip().split()[:3]
            data_times.append(float(time))
            data_fluxes.append(float(flux))
            data_yerr.append(float(err))

            # if len(data_yerr) == 0:
            # data_yerr = np.ones(len(data_fluxes))
            # x, y, data_yerr = kepler.get_fits_data('5897826', use_pdc=True, use_slc=False)

    return np.array(data_times), np.array(data_fluxes), np.array(data_yerr)

def read_input(input_file):
    """Reads in the formatted input file.

    @param input_file: Path of file.
    @return: Tuple containing all 15 parameter data.
    """
    with open('{0}'.format(input_file), 'r') as f:
        input_pars = [line.strip() for line in f.readlines() if line[0] is not '#']

    N = [int(i.strip()) for i in input_pars[0].split('#')][0]
    t0 = [float(i.strip()) for i in input_pars[1].split('#')][0]
    maxh = [float(i.strip()) for i in input_pars[2].split('#')][0]
    orbit_error = [float(i.strip()) for i in input_pars[3].split('#')][0]

    masses = [float(i.strip()) for i in input_pars[4].split()]
    radii = [float(i.strip()) for i in input_pars[5].split()]
    fluxes = [float(i.strip()) for i in input_pars[6].split()]
    u1 = [float(i.strip()) for i in input_pars[7].split()]
    u2 = [float(i.strip()) for i in input_pars[8].split()]

    a = [float(i.strip()) for i in input_pars[9].split()]
    e = [float(i.strip()) for i in input_pars[10].split()]
    inc = [float(i.strip()) for i in input_pars[11].split()]
    om = [float(i.strip()) for i in input_pars[12].split()]
    ln = [float(i.strip()) for i in input_pars[13].split()]
    ma = [float(i.strip()) for i in input_pars[14].split()]

    return N, t0, maxh, orbit_error, masses, radii, fluxes, u1, u2, a, e, inc, om, ln, ma

def get_random_pos(N, t0, maxh, orbit_error):
    masses, radii, fluxes, u1, u2, a, e, inc, om, ln, ma = utilfuncs.random_pos(N)
    return N, t0, maxh, orbit_error, masses, radii, fluxes, u1, u2, a, e, inc, om, ln, ma

def plot_model(params, x, y, yerr):
    """Simple plot routine to visualize the modelled input parameters.

    @param params: Tuple of the parameters.
    @param x: Time data.
    @param y: Flux data.
    @param yerr: Flux error.
    """
    N, t0, maxh, orbit_error, masses, radii, fluxes, u1, u2, a, e, inc, om, ln, ma = params

    y_mod = modeler.generate(
        N, t0,
        maxh, orbit_error,
        x,
        masses, radii, fluxes, u1, u2,
        a, e, inc, om, ln, ma
    )

    pylab.plot(x, y, 'k+')
    pylab.plot(x, y_mod, 'r')

    pylab.show()

def main(data_file, fit_method, input_file, nwalkers, niterations, ncores, syspars):
    """The main function mediates the reading of the data and input parameters, and starts the optimization using the
    specified method.

    @param data_file: Path to data file.
    @param fit_method: Currently: 'min' is least squares optimization, and 'mcmc' is the emcee hammer optimization. The
    'plot' option is used to simply plot the returned model based on the supplied parameters.
    @param input_file: Path to input file.
    """

    if input_file:
        params = read_input(input_file)
    elif syspars:
        print('No initial parameter set specified, randomly sampling parameter space...')
        N, t0, maxh, orbit_error = map(float, syspars)
        params = get_random_pos(int(N), t0, maxh, orbit_error)
        rpars = True
    else:
        print('No initial parameter set specified, randomly sampling parameter space...')
        N = int(raw_input('\tEnter the number of bodies: '))
        t0 = float(raw_input('\tEnter the epoch of coordinates: '))
        maxh = float(raw_input('\tEnter the maximum time step (default: 0.01): '))
        orbit_error = float(raw_input('\tEnter the orbit error tolerance (default: 1e-20): '))

        params = get_random_pos(N, t0, maxh, orbit_error)
        rpars = True

    if not fit_method:
        print('You have not specified a fit method, defaulting to least squares minimization.')

    if '.' in data_file:
        x, y, yerr = read_data(data_file)

    n = int(len(x))

    time_start = time.time()

    if fit_method == 'mcmc':
        hammer.generate(
            params, x[:n], y[:n], yerr[:n],
            nwalkers, niterations, ncores, rpars,
            data_file.split('/')[-1].split('.')[0]
        )

    elif fit_method == 'plot':
        plot_model(params, x[:n], y[:n], yerr[:n])

    else:
        minimizer.generate(
            params, x[:n], y[:n], yerr[:n], fit_method,
            data_file.split('/')[-1].split('.')[0]
        )

    print "Total time:", time.time() - time_start

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Photometric dynamical modeling code.')
    parser.add_argument('data', help='data file containing detrended light curve')
    parser.add_argument('-f', '--fit', help='fit method used to minimize',
                        choices=['mcmc', 'leastsq', 'nelder', 'lbfgsb', 'anneal', 'powell',
                                 'cg', 'newton', 'cobyla', 'slsqp', 'plot'], default='leastsq')
    parser.add_argument('-i', '--input', help='input file containing initial parameters (overrides --system)')
    parser.add_argument('-w', '--walkers', type=int, default=250,
                        help='number of walkers if using mcmc fit method ')
    parser.add_argument('-t', '--iterations', type=int, default=500,
                        help='number of iterations to perform if using mcmc fit method')
    parser.add_argument('-c', '--cores', type=int, default=1,
                        help='number of cores to utilize')
    parser.add_argument('-s', '--system', nargs=4,
                        help='four initial system parameters: N, t0, maxh, orbit_error')

    args = parser.parse_args()
    data_file = args.data
    fit_method = args.fit
    input_file = args.input
    nwalkers = args.walkers
    niterations = args.iterations
    ncores = args.cores
    syspars = args.system

    main(data_file, fit_method, input_file, nwalkers, niterations, ncores, syspars)
<|endoftext|>


<commit_before>from django.db import models

class Investment(models.Model):
    symbol          = models.CharField(max_length=5, unique=True)
    name            = models.CharField(max_length=200)

    def __unicode__(self):
        return u'%s' % self.symbol

class Portfolio(models.Model):
    name            = models.CharField(max_length=200)

    def __unicode__(self):
        return u'%s' % self.name

class Lot(models.Model):
    investment      = models.ForeignKey(Investment)
    portfolio       = models.ForeignKey(Portfolio)

    def __unicode__(self):
        return u'%s %s' % (self.investment, self.portfolio)

class Transaction(models.Model):
    lot             = models.ForeignKey(Lot)
    trade_date      = models.DateField()
    price           = models.DecimalField(max_digits=12, decimal_places=4)
    shares          = models.DecimalField(max_digits=15, decimal_places=4)

    def __unicode__(self):
        return u'%s %s' % (self.lot, self.trade_date)

    @property
    def investment(self):
        return self.lot.investment
<commit_msg>Extended models slightly<commit_after>from django.db import models

class Investment(models.Model):
    symbol          = models.CharField(max_length=5, unique=True)
    name            = models.CharField(max_length=200)

    def __unicode__(self):
        return u'%s' % self.symbol

class Portfolio(models.Model):
    name            = models.CharField(max_length=200)

    def __unicode__(self):
        return u'%s' % self.name

class Lot(models.Model):
    investment      = models.ForeignKey(Investment)
    portfolio       = models.ForeignKey(Portfolio)

    def __unicode__(self):
        return u'%s %s' % (self.investment, self.portfolio)

    def purchase_transaction(self):
        return self.transaction_set.order_by('trade_date')[0]

    def gains(self):
        return -sum(t.amount() for t in self.transaction_set.all()) + self.shares() * self.purchase_transaction().price

    def shares(self):
        return sum(t.shares for t in self.transaction_set.all())

class Transaction(models.Model):
    lot             = models.ForeignKey(Lot)
    trade_date      = models.DateField()
    price           = models.DecimalField(max_digits=12, decimal_places=4)
    shares          = models.DecimalField(max_digits=15, decimal_places=4)

    def __unicode__(self):
        return u'%s %s %s' % (self.investment, self.trade_date, self.amount())

    @property
    def investment(self):
        return self.lot.investment

    def amount(self):
        return self.price * self.shares
<|endoftext|>


<commit_before>
# external imports

try:
    import numpy as np
    import numpy.linalg as la
except:
    print 'numpy 1.3 is required'
    raise
try:
    import scipy as sp
    import scipy.stats as stats
    from cg.kdtree import KDTree
    from scipy.spatial.distance import pdist, cdist
except:
    print 'scipy 0.7+ is required'
    raise

class ROD(dict):
    """
    Read Only Dictionary
    """
    def __setitem__(self, *args):
        #raise TypeError, "'Read Only Dictionary (ROD)' object does not support item assignment"
        print args

    def __copy__(self):
        return ROD(self.copy())

import copy
import math
import random
import sys
import time
import unittest
<commit_msg>implemented deepcopy for ROD, see #237<commit_after>
# external imports

try:
    import numpy as np
    import numpy.linalg as la
except:
    print 'numpy 1.3 is required'
    raise
try:
    import scipy as sp
    import scipy.stats as stats
    from cg.kdtree import KDTree
    from scipy.spatial.distance import pdist, cdist
except:
    print 'scipy 0.7+ is required'
    raise

class ROD(dict):
    """
    Read Only Dictionary
    """
    def __setitem__(self, *args):
        raise TypeError, "'Read Only Dictionary (ROD)' object does not support item assignment"

    def __copy__(self):
        """
        Example:
        >>> import copy
        >>> d = ROD({'key':[1,2,3]})
        >>> d2 = copy.copy(d)
        >>> d2 == d
        True
        >>> d2 is d
        False
        >>> d2['key'] is d['key']
        True
        """
        return ROD(self.copy())
    def __deepcopy__(self,memo):
        """
        Example:
        >>> import copy
        >>> d = ROD({'key':[1,2,3]})
        >>> d2 = copy.deepcopy(d)
        >>> d2 == d
        True
        >>> d2 is d
        False
        >>> d2['key'] is d['key']
        False
        """
        import copy
        return ROD(copy.deepcopy(self.copy()))

import copy
import math
import random
import sys
import time
import unittest
<|endoftext|>