<commit_before># Copyright 2014 Algolia
#
# 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.

"""Receives documents from the oplog worker threads and indexes them
    into the backend.

    This file is a document manager for the Algolia search engine.
    """
import base64
import logging
import json
import re
import copy
from datetime import datetime

from threading import Timer, RLock

import bson.json_util
import bson.json_util as bsjson

from algoliasearch import algoliasearch
import urllib3

from mongo_connector import errors
from mongo_connector.compat import u
from mongo_connector.constants import (DEFAULT_COMMIT_INTERVAL,
                                       DEFAULT_MAX_BULK)
from mongo_connector.util import exception_wrapper, retry_until_ok
from mongo_connector.doc_managers.doc_manager_base import DocManagerBase
from mongo_connector.doc_managers.formatters import DefaultDocumentFormatter

decoder = json.JSONDecoder()

def clean_path(dirty):
    """Convert a string of python subscript notation or mongo dot-notation to a
        list of strings.
        """
    # handle python dictionary subscription style, e.g. `"['key1']['key2']"`:
    if re.match(r'^\[', dirty):
        return re.split(r'\'\]\[\'', re.sub(r'^\[\'|\'\]$', '', dirty))
    # handle mongo op dot-notation style, e.g. `"key1.key2"`:
    return dirty.split('.')

def get_at(doc, path, create_anyway=False):
    """Get the value, if any, of the document at the given path, optionally
        mutating the document to create nested dictionaries as necessary.
        """
    node = doc
    last = len(path) - 1
    if last == 0:
        return doc.get(path[0])
    for index, edge in enumerate(path):
        if edge in node:
            node = node[edge]
        elif index == last or not create_anyway:
            # the key doesn't exist, and this is the end of the path:
            return None
        else:
            # create anyway will create any missing nodes:
            node = node[edge] = {}
    return node

def set_at(doc, path, value):
    """Set the value of the document at the given path."""
    node = get_at(doc, path[:-1], create_anyway=True)
    node[path[-1]] = value

def put_at(doc, path, value, append=False):
    """Set or append the given value to the document at the given path"""
    if append:
        get_at(doc, path).append(value)
    else:
        set_at(doc, path, value)

def unix_time(dt=datetime.now()):
    epoch = datetime.utcfromtimestamp(0)
    delta = dt - epoch
    return delta.total_seconds()

def unix_time_millis(dt=datetime.now()):
    return int(round(unix_time(dt) * 1000.0))

def serialize(value):
    """If the value is an BSON ObjectId, cast it to a string."""
    if isinstance(value, bson.objectid.ObjectId):
        return str(value)
    else:
        return value

def filter_value(value, expr):
    """Evaluate the given expression in the context of the given value."""
    if expr == "":
        return True
    try:
        return eval(re.sub(r'\$_', 'value', expr))
    except Exception as e:
        logging.warn("""
            Error raised from expression: {expr} with value {value}
            """.format(**locals()))
        logging.warn(e)
        # return false to prevent potentially sensitive data from being synced:
        return False

class DocManager(DocManagerBase):
    """The DocManager class creates a connection to the Algolia engine and
        adds/removes documents, and in the case of rollback, searches for them.

        Algolia's native 'objectID' field is used to store the unique_key.
        """

    BATCH_SIZE = 1000
    AUTO_COMMIT_DELAY_S = 10

    def __init__(self, url, unique_key='_id', **kwargs):
        """Establish a connection to Algolia using target url
            'APPLICATION_ID:API_KEY:INDEX_NAME'
        """
        application_id, api_key, index = url.split(':')
        self.algolia = algoliasearch.Client(application_id, api_key)
        self.index = self.algolia.initIndex(index)
        self.unique_key = unique_key
        self.last_object_id = None
        self.batch = []
        self.mutex = RLock()
        self.auto_commit = kwargs.pop('auto_commit', True)
        self.run_auto_commit()
        try:
            json = open("algolia_fields_" + index + ".json", 'r')
            self.attributes_filter = decoder.decode(json.read())
            logging.info("Algolia Connector: Start with filter.")
        except IOError:  # No "fields" filter file
            self.attributes_filter = None
            logging.info("Algolia Connector: Start without filter.")
        try:
            json = open("algolia_remap_" + index + ".json", 'r')
            self.attributes_remap = decoder.decode(json.read())
            logging.info("Algolia Connector: Start with remapper.")
        except IOError:  # No "remap" filter file
            self.attributes_remap = None
            logging.info("Algolia Connector: Start without remapper.")
        try:
            f = open("algolia_postproc_" + index + ".py", 'r')
            self.postproc = f.read()
            logging.info("Algolia Connector: Start with post processing.")
        except IOError:  # No "postproc" filter file
            self.postproc = None
            logging.info("Algolia Connector: Start without post processing.")

    def stop(self):
        """ Stops the instance
        """
        self.auto_commit = False

    def apply_remap(self, doc):
        """Copy the values of user-defined fields from the source document to
            user-defined fields in a new target document, then return the
            targetdocument.
            """
        if not self.attributes_remap:
            return doc
        remapped_doc = {}
        for raw_source_key, raw_target_key in self.attributes_remap.items():
            # clean the keys, making a list from possible notations:
            source_key = clean_path(raw_source_key)
            target_key = clean_path(raw_target_key)

            # get the value from the source doc:
            value = get_at(doc, source_key)

            # special case for "_ts" field:
            if source_key == ['_ts'] and target_key == ["*ts*"]:
                value = value if value else str(unix_time_millis())

            set_at(remapped_doc, target_key, value)
        return remapped_doc

    def apply_filter(self, doc, filter):
        """Recursively copy the values of user-defined fields from the source
            document to a new target document by testing each value against a
            corresponding user-defined expression. If the expression returns
            true for a given value, copy that value to the corresponding field
            in the target document. If the special `*all*` filter is used for
            a given document and an adjacent field's expression returns false
            for a given value, remove the document containing that field from
            its parent in the tree of the target document.
            """
        if not filter:
            # alway return a new object:
            return (copy.deepcopy(doc), True)
        filtered_doc = {}
        all_or_nothing = '*all*' in filter
        try:
            items = filter.iteritems()
        except AttributeError:
            items = filter.items();
        for raw_key, expr in items:
            if raw_key == '*all*':
                continue
            key = clean_path(raw_key)
            values = get_at(doc, key)
            state = True
            if type(values) == list:
                append = True
                set_at(filtered_doc, key, [])
            else:
                append = False
                values = [values]
            for value in values:
                if isinstance(value, dict):
                    sub, sub_state = self.apply_filter(value, filter[raw_key])
                    if sub_state:
                        put_at(filtered_doc, key, serialize(sub), append)
                    elif all_or_nothing:
                        node = get_at(filtered_doc, key[:-1])
                        del node[key[-1]]
                        return filtered_doc, False
                elif filter_value(value, filter[raw_key]):
                    put_at(filtered_doc, key, serialize(value), append)
                elif all_or_nothing:
                    return filtered_doc, False
                else:
                    state = False
        return (filtered_doc, state)

    def update_can_be_ignored(self, update_spec):

        if not self.attributes_filter:
            return False

        if "$set" not in update_spec and "$unset" not in update_spec:
            return False

        def attr_can_be_ignored(attr):
            path = clean_path(re.sub(r'\.(?:\$|\d+)\.', '.', attr))
            if get_at(self.attributes_filter, path, False) is not None:
                return False
            return True

        if "$set" in update_spec:
            specAttrs = self.apply_remap(update_spec["$set"])
            for attr in specAttrs:
                if not attr_can_be_ignored(attr):
                    logging.debug("Algolia Connector: Update needed ($set " + attr + ")")
                    return False
        if "$unset" in update_spec:
            specAttrs = self.apply_remap(update_spec["$unset"])
            for attr in specAttrs:
                if not attr_can_be_ignored(attr):
                    logging.debug("Algolia Connector: Update needed ($unset " + attr + ")")
                    return False

        logging.debug("Algolia Connector: Update can be ignored")

        if "$set" in update_spec:
            specAttrs = self.apply_remap(update_spec["$set"])
            for attr in specAttrs:
                if attr_can_be_ignored(attr):
                    logging.debug("Ignoring: $set " + attr)
        if "$unset" in update_spec:
            specAttrs = self.apply_remap(update_spec["$unset"])
            for attr in specAttrs:
                if attr_can_be_ignored(attr):
                    logging.debug("Ignoring: $unset " + attr)

        return True

    def _db_and_collection(self, namespace):
        return namespace.split('.', 1)

    def get_source_doc(self, document_id, namespace):
        db, coll = self._db_and_collection(namespace)
        if not hasattr(self, 'connector_mongo_client'):
            raise Exception('connector mongo client not found')
        doc = self.connector_mongo_client[db][coll].find(
            {'_id': document_id},
        )[0]
        return doc

    def convert_update_spec_to_partial_doc(self, update_spec):
        if "$set" in update_spec:
            specAttrs = update_spec["$set"]
            for attr in specAttrs:
                set_at(update_spec, attr, specAttrs.get(attr))
            del update_spec["$set"]
        if "$unset" in update_spec:
            specAttrs = update_spec["$unset"]
            for attr in specAttrs:
                set_at(update_spec, attr, None)
            del update_spec["$unset"]

    def update(self, document_id, update_spec, namespace = None, timestamp = None):

        if self.update_can_be_ignored(update_spec):
            logging.info("Algolia Connector: Update Skipped (no changes passed the filter)")
            return

        if "$set" not in update_spec and "$unset" not in update_spec:
            logging.info("Algolia Connector: Document full replace")
            self.upsert(update_spec, False, namespace, timestamp)
        elif (self.postproc is not None):
            # if postproc is used, mongo update_spec cannot be applied to a processed doc
            # the full source document from mongo is needed
            logging.info("Algolia Connector: Reprocessing source document")
            doc = self.get_source_doc(document_id, namespace)
            self.upsert(doc, False, namespace, timestamp)
        else:
            # convert mongo update_spec to algolia partial update body
            logging.info("Algolia Connector: Partial update")
            partial_doc = convert_update_spec_to_partial_doc(update_spec)
            self.upsert(partial_doc, True, namespace, timestamp)

    def upsert(self, doc, update = False, namespace = None, timestamp = None):
        """ Update or insert a document into Algolia
        """
        with self.mutex:
            self.last_object_id = serialize(doc.get(self.unique_key) or doc['objectID'])
            filtered_doc, state = self.apply_filter(self.apply_remap(doc),
                                                    self.attributes_filter)
            filtered_doc['objectID'] = self.last_object_id

            #if not state:  # delete in case of update
            #    self.batch.append({'action': 'deleteObject',
            #                       'body': {'objectID': last_object_id}})
            #    return

            if self.postproc is not None:
                exec(re.sub(r"_\$", "filtered_doc", self.postproc))

            self.batch.append({'action': 'partialUpdateObject' if update else 'addObject', 'body': filtered_doc})
            if len(self.batch) >= DocManager.BATCH_SIZE:
                self.commit()

    def remove(self, document_id, namespace = None, timestamp = None):
        """ Removes documents from Algolia
        """
        with self.mutex:
            self.batch.append(
                {'action': 'deleteObject',
                 'body': {'objectID': str(document_id)}})
            if len(self.batch) >= DocManager.BATCH_SIZE:
                self.commit()

    def search(self, start_ts, end_ts):
        """ Called to query Algolia for documents in a time range.
        """
        try:
            params = {
                'numericFilters': '_ts>=%d,_ts<=%d' % (start_ts, end_ts),
                'exhaustive': True,
                'hitsPerPage': 100000000
            }
            return self.index.search('', params)['hits']
        except algoliasearch.AlgoliaException as e:
            raise errors.ConnectionFailed(
                "Could not connect to Algolia Search: %s" % e)

    def commit(self, synchronous=False):
        """ Send the current batch of updates
        """
        try:
            request = {}
            with self.mutex:
                if len(self.batch) == 0:
                    return
                self.index.batch({'requests': self.batch})
                res = self.index.setSettings({'userData': {'lastObjectID': self.last_object_id}})
                self.batch = []
                if synchronous:
                    self.index.waitTask(res['taskID'])
        except (algoliasearch.AlgoliaException, urllib3.exceptions.MaxRetryError) as e:
            raise errors.OperationFailed(
                "Could not connect to Algolia Search: %s" % e)

    def run_auto_commit(self):
        """ Periodically commits to Algolia.
        """
        try:
            self.commit(True)
        except Exception as e:
            logging.warning(e)
        if self.auto_commit:
            Timer(DocManager.AUTO_COMMIT_DELAY_S, self.run_auto_commit).start()

    def get_last_doc(self):
        """ Returns the last document stored in Algolia.
        """
        last_object_id = self.get_last_object_id()
        if last_object_id is None:
            return None
        try:
            return self.index.getObject(str(last_object_id))
        except algoliasearch.AlgoliaException as e:
            raise errors.ConnectionFailed(
                "Could not connect to Algolia Search: %s" % e)

    def get_last_object_id(self):
        try:
            return (self.index.getSettings().get('userData', {})).get('lastObjectID', None)
        except algoliasearch.AlgoliaException as e:
            raise errors.ConnectionFailed(
                "Could not connect to Algolia Search: %s" % e)
<commit_msg>fixing error a value is None<commit_after># Copyright 2014 Algolia
#
# 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.

"""Receives documents from the oplog worker threads and indexes them
    into the backend.

    This file is a document manager for the Algolia search engine.
    """
import base64
import logging
import json
import re
import copy
from datetime import datetime

from threading import Timer, RLock

import bson.json_util
import bson.json_util as bsjson

from algoliasearch import algoliasearch
import urllib3

from mongo_connector import errors
from mongo_connector.compat import u
from mongo_connector.constants import (DEFAULT_COMMIT_INTERVAL,
                                       DEFAULT_MAX_BULK)
from mongo_connector.util import exception_wrapper, retry_until_ok
from mongo_connector.doc_managers.doc_manager_base import DocManagerBase
from mongo_connector.doc_managers.formatters import DefaultDocumentFormatter

decoder = json.JSONDecoder()

def clean_path(dirty):
    """Convert a string of python subscript notation or mongo dot-notation to a
        list of strings.
        """
    # handle python dictionary subscription style, e.g. `"['key1']['key2']"`:
    if re.match(r'^\[', dirty):
        return re.split(r'\'\]\[\'', re.sub(r'^\[\'|\'\]$', '', dirty))
    # handle mongo op dot-notation style, e.g. `"key1.key2"`:
    return dirty.split('.')

def get_at(doc, path, create_anyway=False):
    """Get the value, if any, of the document at the given path, optionally
        mutating the document to create nested dictionaries as necessary.
        """
    node = doc
    last = len(path) - 1
    if last == 0:
        return doc.get(path[0])
    for index, edge in enumerate(path):
        if edge in node:
            node = node[edge]
        elif index == last or not create_anyway:
            # the key doesn't exist, and this is the end of the path:
            return None
        else:
            # create anyway will create any missing nodes:
            node = node[edge] = {}
    return node

def set_at(doc, path, value):
    """Set the value of the document at the given path."""
    node = get_at(doc, path[:-1], create_anyway=True)
    node[path[-1]] = value

def put_at(doc, path, value, append=False):
    """Set or append the given value to the document at the given path"""
    if append:
        get_at(doc, path).append(value)
    else:
        set_at(doc, path, value)

def unix_time(dt=datetime.now()):
    epoch = datetime.utcfromtimestamp(0)
    delta = dt - epoch
    return delta.total_seconds()

def unix_time_millis(dt=datetime.now()):
    return int(round(unix_time(dt) * 1000.0))

def serialize(value):
    """If the value is an BSON ObjectId, cast it to a string."""
    if isinstance(value, bson.objectid.ObjectId):
        return str(value)
    else:
        return value

def filter_value(value, expr):
    """Evaluate the given expression in the context of the given value."""
    if expr == "":
        return True
    if value is None:
        return False
    try:
        return eval(re.sub(r'\$_', 'value', expr))
    except Exception as e:
        logging.warn("""
            Error raised from expression: {expr} with value {value}
            """.format(**locals()))
        logging.warn(e)
        # return false to prevent potentially sensitive data from being synced:
        return False

class DocManager(DocManagerBase):
    """The DocManager class creates a connection to the Algolia engine and
        adds/removes documents, and in the case of rollback, searches for them.

        Algolia's native 'objectID' field is used to store the unique_key.
        """

    BATCH_SIZE = 1000
    AUTO_COMMIT_DELAY_S = 10

    def __init__(self, url, unique_key='_id', **kwargs):
        """Establish a connection to Algolia using target url
            'APPLICATION_ID:API_KEY:INDEX_NAME'
        """
        application_id, api_key, index = url.split(':')
        self.algolia = algoliasearch.Client(application_id, api_key)
        self.index = self.algolia.initIndex(index)
        self.unique_key = unique_key
        self.last_object_id = None
        self.batch = []
        self.mutex = RLock()
        self.auto_commit = kwargs.pop('auto_commit', True)
        self.run_auto_commit()
        try:
            json = open("algolia_fields_" + index + ".json", 'r')
            self.attributes_filter = decoder.decode(json.read())
            logging.info("Algolia Connector: Start with filter.")
        except IOError:  # No "fields" filter file
            self.attributes_filter = None
            logging.info("Algolia Connector: Start without filter.")
        try:
            json = open("algolia_remap_" + index + ".json", 'r')
            self.attributes_remap = decoder.decode(json.read())
            logging.info("Algolia Connector: Start with remapper.")
        except IOError:  # No "remap" filter file
            self.attributes_remap = None
            logging.info("Algolia Connector: Start without remapper.")
        try:
            f = open("algolia_postproc_" + index + ".py", 'r')
            self.postproc = f.read()
            logging.info("Algolia Connector: Start with post processing.")
        except IOError:  # No "postproc" filter file
            self.postproc = None
            logging.info("Algolia Connector: Start without post processing.")

    def stop(self):
        """ Stops the instance
        """
        self.auto_commit = False

    def apply_remap(self, doc):
        """Copy the values of user-defined fields from the source document to
            user-defined fields in a new target document, then return the
            targetdocument.
            """
        if not self.attributes_remap:
            return doc
        remapped_doc = {}
        for raw_source_key, raw_target_key in self.attributes_remap.items():
            # clean the keys, making a list from possible notations:
            source_key = clean_path(raw_source_key)
            target_key = clean_path(raw_target_key)

            # get the value from the source doc:
            value = get_at(doc, source_key)

            # special case for "_ts" field:
            if source_key == ['_ts'] and target_key == ["*ts*"]:
                value = value if value else str(unix_time_millis())

            set_at(remapped_doc, target_key, value)
        return remapped_doc

    def apply_filter(self, doc, filter):
        """Recursively copy the values of user-defined fields from the source
            document to a new target document by testing each value against a
            corresponding user-defined expression. If the expression returns
            true for a given value, copy that value to the corresponding field
            in the target document. If the special `*all*` filter is used for
            a given document and an adjacent field's expression returns false
            for a given value, remove the document containing that field from
            its parent in the tree of the target document.
            """
        if not filter:
            # alway return a new object:
            return (copy.deepcopy(doc), True)
        filtered_doc = {}
        all_or_nothing = '*all*' in filter
        try:
            items = filter.iteritems()
        except AttributeError:
            items = filter.items();
        for raw_key, expr in items:
            if raw_key == '*all*':
                continue
            key = clean_path(raw_key)
            values = get_at(doc, key)
            state = True
            if type(values) == list:
                append = True
                set_at(filtered_doc, key, [])
            else:
                append = False
                values = [values]
            for value in values:
                if isinstance(value, dict):
                    sub, sub_state = self.apply_filter(value, filter[raw_key])
                    if sub_state:
                        put_at(filtered_doc, key, serialize(sub), append)
                    elif all_or_nothing:
                        node = get_at(filtered_doc, key[:-1])
                        del node[key[-1]]
                        return filtered_doc, False
                elif filter_value(value, filter[raw_key]):
                    put_at(filtered_doc, key, serialize(value), append)
                elif all_or_nothing:
                    return filtered_doc, False
                else:
                    state = False
        return (filtered_doc, state)

    def update_can_be_ignored(self, update_spec):

        if not self.attributes_filter:
            return False

        if "$set" not in update_spec and "$unset" not in update_spec:
            return False

        def attr_can_be_ignored(attr):
            path = clean_path(re.sub(r'\.(?:\$|\d+)\.', '.', attr))
            if get_at(self.attributes_filter, path, False) is not None:
                return False
            return True

        if "$set" in update_spec:
            specAttrs = self.apply_remap(update_spec["$set"])
            for attr in specAttrs:
                if not attr_can_be_ignored(attr):
                    logging.debug("Algolia Connector: Update needed ($set " + attr + ")")
                    return False
        if "$unset" in update_spec:
            specAttrs = self.apply_remap(update_spec["$unset"])
            for attr in specAttrs:
                if not attr_can_be_ignored(attr):
                    logging.debug("Algolia Connector: Update needed ($unset " + attr + ")")
                    return False

        logging.debug("Algolia Connector: Update can be ignored")

        if "$set" in update_spec:
            specAttrs = self.apply_remap(update_spec["$set"])
            for attr in specAttrs:
                if attr_can_be_ignored(attr):
                    logging.debug("Ignoring: $set " + attr)
        if "$unset" in update_spec:
            specAttrs = self.apply_remap(update_spec["$unset"])
            for attr in specAttrs:
                if attr_can_be_ignored(attr):
                    logging.debug("Ignoring: $unset " + attr)

        return True

    def _db_and_collection(self, namespace):
        return namespace.split('.', 1)

    def get_source_doc(self, document_id, namespace):
        db, coll = self._db_and_collection(namespace)
        if not hasattr(self, 'connector_mongo_client'):
            raise Exception('connector mongo client not found')
        doc = self.connector_mongo_client[db][coll].find(
            {'_id': document_id},
        )[0]
        return doc

    def convert_update_spec_to_partial_doc(self, update_spec):
        if "$set" in update_spec:
            specAttrs = update_spec["$set"]
            for attr in specAttrs:
                set_at(update_spec, attr, specAttrs.get(attr))
            del update_spec["$set"]
        if "$unset" in update_spec:
            specAttrs = update_spec["$unset"]
            for attr in specAttrs:
                set_at(update_spec, attr, None)
            del update_spec["$unset"]

    def update(self, document_id, update_spec, namespace = None, timestamp = None):

        if self.update_can_be_ignored(update_spec):
            logging.info("Algolia Connector: Update Skipped (no changes passed the filter)")
            return

        if "$set" not in update_spec and "$unset" not in update_spec:
            logging.info("Algolia Connector: Document full replace")
            self.upsert(update_spec, False, namespace, timestamp)
        elif (self.postproc is not None):
            # if postproc is used, mongo update_spec cannot be applied to a processed doc
            # the full source document from mongo is needed
            logging.info("Algolia Connector: Reprocessing source document")
            doc = self.get_source_doc(document_id, namespace)
            self.upsert(doc, False, namespace, timestamp)
        else:
            # convert mongo update_spec to algolia partial update body
            logging.info("Algolia Connector: Partial update")
            partial_doc = convert_update_spec_to_partial_doc(update_spec)
            self.upsert(partial_doc, True, namespace, timestamp)

    def upsert(self, doc, update = False, namespace = None, timestamp = None):
        """ Update or insert a document into Algolia
        """
        with self.mutex:
            self.last_object_id = serialize(doc.get(self.unique_key) or doc['objectID'])
            filtered_doc, state = self.apply_filter(self.apply_remap(doc),
                                                    self.attributes_filter)
            filtered_doc['objectID'] = self.last_object_id

            #if not state:  # delete in case of update
            #    self.batch.append({'action': 'deleteObject',
            #                       'body': {'objectID': last_object_id}})
            #    return

            if self.postproc is not None:
                exec(re.sub(r"_\$", "filtered_doc", self.postproc))

            self.batch.append({'action': 'partialUpdateObject' if update else 'addObject', 'body': filtered_doc})
            if len(self.batch) >= DocManager.BATCH_SIZE:
                self.commit()

    def remove(self, document_id, namespace = None, timestamp = None):
        """ Removes documents from Algolia
        """
        with self.mutex:
            self.batch.append(
                {'action': 'deleteObject',
                 'body': {'objectID': str(document_id)}})
            if len(self.batch) >= DocManager.BATCH_SIZE:
                self.commit()

    def search(self, start_ts, end_ts):
        """ Called to query Algolia for documents in a time range.
        """
        try:
            params = {
                'numericFilters': '_ts>=%d,_ts<=%d' % (start_ts, end_ts),
                'exhaustive': True,
                'hitsPerPage': 100000000
            }
            return self.index.search('', params)['hits']
        except algoliasearch.AlgoliaException as e:
            raise errors.ConnectionFailed(
                "Could not connect to Algolia Search: %s" % e)

    def commit(self, synchronous=False):
        """ Send the current batch of updates
        """
        try:
            request = {}
            with self.mutex:
                if len(self.batch) == 0:
                    return
                self.index.batch({'requests': self.batch})
                res = self.index.setSettings({'userData': {'lastObjectID': self.last_object_id}})
                self.batch = []
                if synchronous:
                    self.index.waitTask(res['taskID'])
        except (algoliasearch.AlgoliaException, urllib3.exceptions.MaxRetryError) as e:
            raise errors.OperationFailed(
                "Could not connect to Algolia Search: %s" % e)

    def run_auto_commit(self):
        """ Periodically commits to Algolia.
        """
        try:
            self.commit(True)
        except Exception as e:
            logging.warning(e)
        if self.auto_commit:
            Timer(DocManager.AUTO_COMMIT_DELAY_S, self.run_auto_commit).start()

    def get_last_doc(self):
        """ Returns the last document stored in Algolia.
        """
        last_object_id = self.get_last_object_id()
        if last_object_id is None:
            return None
        try:
            return self.index.getObject(str(last_object_id))
        except algoliasearch.AlgoliaException as e:
            raise errors.ConnectionFailed(
                "Could not connect to Algolia Search: %s" % e)

    def get_last_object_id(self):
        try:
            return (self.index.getSettings().get('userData', {})).get('lastObjectID', None)
        except algoliasearch.AlgoliaException as e:
            raise errors.ConnectionFailed(
                "Could not connect to Algolia Search: %s" % e)
<|endoftext|>


<commit_before>import json
import types
import time
from django.core.cache import cache
from django.db.models.fields.related import ReverseSingleRelatedObjectDescriptor
from django.db.models.fields.related import ForeignKey
from django.http import HttpResponse
from django.views.generic.base import View

from curious import model_registry
from .query import Query
from .utils import report_time

class JSONView(View):

  def _return(self, code, result):
    res = {'result': result}
    return HttpResponse(json.dumps(res), status=code, content_type='application/json')

  def _error(self, code, message):
    res = {'error': {'message': message}}
    return HttpResponse(json.dumps(res), status=code, content_type='application/json')

class ModelListView(JSONView):

  def get(self, request):
    return self._return(200, model_registry.model_names)

class ModelView(JSONView):

  @staticmethod
  def model_to_dict(cls):
    d = dict(model=cls.__name__, relationships=[])
    for f in dir(cls):
      if model_registry.is_rel_allowed(cls, f):
        d['relationships'].append(f)
    return d

  @staticmethod
  def objects_to_dict(objects):
    if len(objects) == 0:
      return dict(fields=[], objects=[])

    fields = []
    fk = []
    obj = objects[0]
    for f in obj._meta.fields:
      fields.append(f.column)
      fk.append(f.name if type(f) == ForeignKey else None)

    packed = []
    for obj in objects:
      values = []
      for column, fk_name in zip(fields, fk):
        value = getattr(obj, column)
        if not type(value) in (long, int, float, bool, types.NoneType):
          value = unicode(value)
        if fk_name is not None:
          v = getattr(obj, fk_name)
          if v is not None:
            model_name = model_registry.getname(v.__class__)
            try:
              url = model_registry.geturl(model_name, v)
            except:
              url = None
            value = (model_name, v.pk, str(v), url)
        values.append(value)
      packed.append(values)

    return dict(fields=fields, objects=packed)

  def get(self, request, model_name):
    try:
      cls = model_registry.getclass(model_name)
    except Exception as e:
      return self._error(404, "Unknown model '%s': %s" % (model_name, str(e)))
    return self._return(200, ModelView.model_to_dict(cls))

  @report_time
  def post(self, request, model_name):
    """
    Fetch objects in batch.
    """

    try:
      data = json.loads(request.body)
    except:
      return self._error(400, "Bad data")

    if 'ids' not in data:
      return self._error(400, "Missing ids array")

    try:
      cls = model_registry.getclass(model_name)
    except:
      return self._error(404, "Unknown model '%s'" % model_name)

    t = time.time()
    r = cls.objects.filter(id__in=data['ids'])
    r = list(r)
    print 'fetch %.2f' % (time.time()-t,)
    t = time.time()
    r = ModelView.objects_to_dict(r)
    print ' pack %.2f' % (time.time()-t,)
    return self._return(200, r)

class ObjectView(JSONView):

  @staticmethod
  def object_to_dict(obj):
    d = {}
    for f in obj._meta.fields:
      d[f.column] = getattr(obj, f.column)
      if not type(d[f.column]) in (long, int, float, bool, types.NoneType):
        d[f.column] = unicode(d[f.column]);
      if type(f) == ForeignKey:
        v = getattr(obj, f.name)
        if v is not None:
          model_name = model_registry.getname(v.__class__)
          d[f.column] = {
            'model': model_name,
            'id': v.pk,
            '__str__': str(v)
          }
          try:
            d[f.column]['url'] = model_registry.geturl(model_name, v)
          except:
            pass
    return d

  def get(self, request, model_name, id):
    try:
      cls = model_registry.getclass(model_name)
    except:
      return self._error(404, "Unknown model '%s'" % model_name)
    try:
      obj = cls.objects.get(pk=id)
    except:
      return self._error(404, "Cannot find instance '%s' on '%s'" % (id, model_name))
    return self._return(200, ObjectView.object_to_dict(obj))

class QueryView(JSONView):

  def get_query_results(self, query, force):
    k = hash(query.query_string)
    v = cache.get(k)
    if v is None or force:
      v = self.run_query(query)
      cache.set(k, v, None)
    return v

  def run_query(self, query):
    res, last_model = query()
    results = []
    for obj_src in res:
      model = type(obj_src[0][0])
      if model._deferred:
        model = model.__base__
      model_name = model_registry.getname(model)
      d = {'model': model_name,
           'objects': [(obj.pk, model_registry.geturl(model_name, obj), src) for obj, src in obj_src]
          }
      results.append(d)
    if last_model is not None:
      return dict(last_model=model_registry.getname(last_model), results=results)
    else:
      return dict(last_model=None, results=results)

  @report_time
  def get(self, request):
    if 'q' not in request.GET:
      return self._error(400, 'Missing query')

    q = request.GET['q']

    try:
      query = Query(q)
    except Exception as e:
      import traceback
      traceback.print_exc()
      return self._error(400, str(e))

    # check query mode
    if 'c' in request.GET:
      return self._return(200, dict(query=q))

    try:
      results = self.get_query_results(query, 'r' in request.GET)
    except Exception as e:
      import traceback
      traceback.print_exc()
      return self._error(400, str(e))

    # print results
    return self._return(200, results)
<commit_msg>Use select_related to get FK objects in one query<commit_after>import json
import types
from django.core.cache import cache
from django.db.models.fields.related import ReverseSingleRelatedObjectDescriptor
from django.db.models.fields.related import ForeignKey
from django.http import HttpResponse
from django.views.generic.base import View

from curious import model_registry
from .query import Query
from .utils import report_time

class JSONView(View):

  def _return(self, code, result):
    res = {'result': result}
    return HttpResponse(json.dumps(res), status=code, content_type='application/json')

  def _error(self, code, message):
    res = {'error': {'message': message}}
    return HttpResponse(json.dumps(res), status=code, content_type='application/json')

class ModelListView(JSONView):

  def get(self, request):
    return self._return(200, model_registry.model_names)

class ModelView(JSONView):

  @staticmethod
  def model_to_dict(cls):
    d = dict(model=cls.__name__, relationships=[])
    for f in dir(cls):
      if model_registry.is_rel_allowed(cls, f):
        d['relationships'].append(f)
    return d

  @staticmethod
  def objects_to_dict(objects):
    if len(objects) == 0:
      return dict(fields=[], objects=[])

    fields = []
    fk = []
    obj = objects[0]
    for f in obj._meta.fields:
      fields.append(f.column)
      fk.append(f.name if type(f) == ForeignKey else None)

    packed = []
    for obj in objects:
      values = []
      for column, fk_name in zip(fields, fk):
        value = getattr(obj, column)
        if not type(value) in (long, int, float, bool, types.NoneType):
          value = unicode(value)
        if fk_name is not None:
          v = getattr(obj, fk_name)
          if v is not None:
            model_name = model_registry.getname(v.__class__)
            try:
              url = model_registry.geturl(model_name, v)
            except:
              url = None
            value = (model_name, v.pk, str(v), url)
        values.append(value)
      packed.append(values)

    return dict(fields=fields, objects=packed)

  def get(self, request, model_name):
    try:
      cls = model_registry.getclass(model_name)
    except Exception as e:
      return self._error(404, "Unknown model '%s': %s" % (model_name, str(e)))
    return self._return(200, ModelView.model_to_dict(cls))

  @report_time
  def post(self, request, model_name):
    """
    Fetch objects in batch.
    """

    try:
      data = json.loads(request.body)
    except:
      return self._error(400, "Bad data")

    if 'ids' not in data:
      return self._error(400, "Missing ids array")

    try:
      cls = model_registry.getclass(model_name)
    except:
      return self._error(404, "Unknown model '%s'" % model_name)

    fks = []
    for f in cls._meta.fields:
      if type(f) == ForeignKey:
        fks.append(f.name)

    q = cls.objects.filter(id__in=data['ids'])
    if len(fks) > 0:
      q = q.select_related(*fks)
    r = ModelView.objects_to_dict(list(q))
    return self._return(200, r)

class ObjectView(JSONView):

  @staticmethod
  def object_to_dict(obj):
    d = {}
    for f in obj._meta.fields:
      d[f.column] = getattr(obj, f.column)
      if not type(d[f.column]) in (long, int, float, bool, types.NoneType):
        d[f.column] = unicode(d[f.column]);
      if type(f) == ForeignKey:
        v = getattr(obj, f.name)
        if v is not None:
          model_name = model_registry.getname(v.__class__)
          d[f.column] = {
            'model': model_name,
            'id': v.pk,
            '__str__': str(v)
          }
          try:
            d[f.column]['url'] = model_registry.geturl(model_name, v)
          except:
            pass
    return d

  def get(self, request, model_name, id):
    try:
      cls = model_registry.getclass(model_name)
    except:
      return self._error(404, "Unknown model '%s'" % model_name)
    try:
      obj = cls.objects.get(pk=id)
    except:
      return self._error(404, "Cannot find instance '%s' on '%s'" % (id, model_name))
    return self._return(200, ObjectView.object_to_dict(obj))

class QueryView(JSONView):

  def get_query_results(self, query, force):
    k = hash(query.query_string)
    v = cache.get(k)
    if v is None or force:
      v = self.run_query(query)
      cache.set(k, v, None)
    return v

  def run_query(self, query):
    res, last_model = query()
    results = []
    for obj_src in res:
      model = type(obj_src[0][0])
      if model._deferred:
        model = model.__base__
      model_name = model_registry.getname(model)
      d = {'model': model_name,
           'objects': [(obj.pk, model_registry.geturl(model_name, obj), src) for obj, src in obj_src]
          }
      results.append(d)
    if last_model is not None:
      return dict(last_model=model_registry.getname(last_model), results=results)
    else:
      return dict(last_model=None, results=results)

  @report_time
  def get(self, request):
    if 'q' not in request.GET:
      return self._error(400, 'Missing query')

    q = request.GET['q']

    try:
      query = Query(q)
    except Exception as e:
      import traceback
      traceback.print_exc()
      return self._error(400, str(e))

    # check query mode
    if 'c' in request.GET:
      return self._return(200, dict(query=q))

    try:
      results = self.get_query_results(query, 'r' in request.GET)
    except Exception as e:
      import traceback
      traceback.print_exc()
      return self._error(400, str(e))

    # print results
    return self._return(200, results)
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- coding: utf-8; py-indent-offset:4 -*-
###############################################################################
#
# Copyright (C) 2015,2016 Daniel Rodriguez
#
# Licensed under the Simplified BSD License
#
###############################################################################
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import collections
from datetime import datetime, timedelta
import time as _time
import json
import threading
import requests
import sseclient
import logging
import backtrader as bt
from backtrader.metabase import MetaParams
from backtrader.utils.py3 import queue, with_metaclass
from backtrader.utils import AutoDict
from backtrader.position import Position
from copy import copy

class MetaSingleton(MetaParams):
    '''Metaclass to make a metaclassed class a singleton'''
    def __init__(cls, name, bases, dct):
        super(MetaSingleton, cls).__init__(name, bases, dct)
        cls._singleton = None

    def __call__(cls, *args, **kwargs):
        if cls._singleton is None:
            cls._singleton = (
                super(MetaSingleton, cls).__call__(*args, **kwargs))

        return cls._singleton

class SharpPointStore(with_metaclass(MetaSingleton, object)):
    '''Singleton class wrapping to control the connections to SharpPoint.

    Params:

      - ``token`` (default:``None``): API access token

      - ``account`` (default: ``None``): account id

      - ``practice`` (default: ``False``): use the test environment

      - ``account_tmout`` (default: ``10.0``): refresh period for account
        value/cash refresh
    '''

    BrokerCls = None  # broker class will autoregister
    DataCls = None  # data class will auto register
    BackTestCls = None

    params = (
        ('gateway', 'http://localhost:5000/'),
        ('token', ''),
        ('account', ''),
        ('login', None),
        ('practice', False),
        ('loglevel', logging.INFO),
        ('account_tmout', 30.0),
    )

    _DTEPOCH = datetime(1970, 1, 1)
    _ENVPRACTICE = 'practice'
    _ENVLIVE = 'live'

    @classmethod
    def getdata(cls, *args, **kwargs):
        '''Returns ``DataCls`` with args, kwargs'''
        return cls.DataCls(*args, **kwargs)

    @classmethod
    def getbroker(cls, *args, **kwargs):
        '''Returns broker with *args, **kwargs from registered ``BrokerCls``
or ``BackTestCls``
        '''
        backtest = kwargs.pop('backtest', False)
        if backtest:
            return cls.BackTestCls(*args, **kwargs)
        return cls.BrokerCls(*args, **kwargs)

    def __init__(self):
        super(SharpPointStore, self).__init__()

        self.positions = collections.defaultdict(Position)
        self.notifs = collections.deque()  # store notifications for cerebro

        self._env = None  # reference to cerebro for general notifications
        self.broker = None  # broker instance
        self.datas = list()  # datas that have registered over start

        self._orders = collections.OrderedDict()  # map order.ref to oid
        self._ordersrev = collections.OrderedDict()  # map oid to order.ref
        self._transpend = collections.defaultdict(collections.deque)

        self._oenv = self._ENVPRACTICE if self.p.practice else self._ENVLIVE

        self._cash = 0.0
        self._value = 0.0
        self.q_account = queue.Queue()
        self.q_ordercreate = queue.Queue()
        self.q_orderclose = queue.Queue()
        self.streaming_events()

    def start(self, data=None, broker=None):
        # Datas require some processing to kickstart data reception
        if data is None and broker is None:
            self.cash = None
            return

        if data is not None:
            self._env = data._env
            # For datas simulate a queue with None to kickstart co
            self.datas.append(data)

            if self.broker is not None:
                self.broker.data_started(data)

        elif broker is not None:
            self.broker = broker
            self.broker_threads()

    def stop(self):
        # signal end of thread
        if self.broker is not None:
            self.q_ordercreate.put(None)
            self.q_orderclose.put(None)
            self.q_account.put(None)

    def put_notification(self, msg, *args, **kwargs):
        self.notifs.append((msg, args, kwargs))

    def get_notifications(self):
        '''Return the pending "store" notifications'''
        self.notifs.append(None)  # put a mark / threads could still append
        return [x for x in iter(self.notifs.popleft, None)]

    def get_positions(self):
        pass

    def get_granularity(self, timeframe, compression):
        pass

    def get_instrument(self, dataname):
        pass

    def _get_request(self, method, **kwargs):
        if self.p.gateway is not None:
            return requests.get(self.p.gateway + method,
                                **kwargs)

    def _post_request(self, method, **kwargs):
        if self.p.gateway is not None:
            return requests.post(self.p.gateway + method,
                                 **kwargs)

    def streaming_events(self, tmout=None):
        q = queue.Queue()
        kwargs = {'q': q, 'tmout': tmout}

        t = threading.Thread(target=self._t_streaming_listener, kwargs=kwargs)
        t.daemon = True
        t.start()

    def updateposition(self, data):
        """Update position from streamer"""
        position = Position(data['Qty'],
                            abs(data['TotalAmt']/data['Qty']))
        self.positions[data['ProdCode']] = position

    def getposition(self, data, clone=False):
        position =  self.positions[data]
        if clone:
            return copy(position)
        else:
            return position

    def _t_streaming_listener(self, q, tmout=None):
        response = self._get_request("log/subscribe", stream=True)
        client = sseclient.SSEClient(response)
        for event in client.events():
            data = json.loads(event.data)
            info = data.get('data', None)
            oref = 0

            if event.event == "AccountPositionPush":
                if self.p.loglevel <= logging.DEBUG:
                    print(event.event, data)
                self.updateposition(data)
                continue
            if self.broker is None:
                continue

            try:
                oref = info['Ref2']
            except:
                if self.p.loglevel <= logging.DEBUG:
                    print("Unhandled event")
                continue

            if self.p.loglevel <= logging.DEBUG:
                print(event)
            if event.event == "OrderBeforeSendReport":
                if self.p.loglevel <= logging.DEBUG:
                    print(data)
                self.broker._submit(oref)
            elif event.event == "OrderRequestFailed":
                if self.p.loglevel <= logging.DEBUG:
                    print(data)
                self.broker._reject(oref)
            elif event.event == "OrderReport":
                if self.p.loglevel <= logging.DEBUG:
                    print(data)
                status = int(info['Status'])
                if status == 4:
                    self.broker._accept(oref)
                elif status == 6:
                    self.broker._cancel(oref)
            elif event.event == "TradeReport":
                if self.p.loglevel <= logging.DEBUG:
                    print(data)
                qty = int(info['Qty'])
                price = float(info['Price'])
                self.broker._fill(oref, qty, price)

    def streaming_prices(self, dataname, tmout=None):
        q = queue.Queue()
        kwargs = {'q': q, 'dataname': dataname, 'tmout': tmout}
        t = threading.Thread(target=self._t_streaming_prices, kwargs=kwargs)
        t.daemon = True
        t.start()
        return q

    def _t_streaming_prices(self, dataname, q, tmout):
        r = self._get_request("ticker/subscribe/" + dataname)

    def get_cash(self):
        return self._cash

    def get_value(self):
        return self._value

    def broker_threads(self):
        self.q_account.put(True)  # force an immediate update
        t = threading.Thread(target=self._t_account)
        t.daemon = True
        t.start()

        t = threading.Thread(target=self._t_order_create)
        t.daemon = True
        t.start()

        t = threading.Thread(target=self._t_order_cancel)
        t.daemon = True
        t.start()

    _ORDEREXECS = {
        bt.Order.Market: 'market',
        bt.Order.Limit: 'limit',
        bt.Order.Stop: 'stop',
        bt.Order.StopLimit: 'stop',
    }

    def isloggedin(self):
        login_info = self._get_request("login-info").json()
        if self.p.loglevel <= logging.DEBUG:
            print("login-info", login_info)
        return int(login_info['status']) != -1

    def setlogin(self, login):
        self.p.login = login
        self.q_account.put(True)  # force an immediate update

    def _t_account(self):
        if self.p.loglevel <= logging.DEBUG:
            print("t_account")
        while True:
            try:
                msg = self.q_account.get(timeout=self.p.account_tmout)
                if msg is None:
                    break  # end of thread
            except queue.Empty:  # tmout -> time to refresh
                pass
            try:
                if self.p.login is not None and not self.isloggedin():
                    if self.p.loglevel <= logging.DEBUG:
                        print("login", self.p.login)
                    r = self._post_request("login", json=self.p.login)
            except Exception as e:
                self.put_notification(e)
                continue

    def order_create(self, order, **kwargs):
        okwargs = {"DecInPrice": 0,
                   "ClOrderId": "test2",
                   "OpenClose": 0,
                   "CondType": 0,
                   "OrderType": 0,
                   "ValidType": 0,
                   "StopType": 0,
                   "OrderOptions": 0}
        if order.isbuy():
            okwargs['BuySell'] = "B"
        elif order.issell():
            okwargs['BuySell'] = "S"
        okwargs['Price'] = order.created.price
        okwargs['Qty'] = abs(order.created.size)
        okwargs['ProdCode'] = order.data._dataname
        okwargs['Ref'] = kwargs.get('Ref', '')
        order.ref = "{:%Y%m%d%H%M%S}".format(datetime.now())
        self._orders[order.ref] = order
        okwargs['Ref2'] = str(order.ref)
        if self.p.loglevel <= logging.DEBUG:
            print(okwargs)
        self.q_ordercreate.put((order.ref, okwargs,))
        return order

    def order_by_ref(self, oref):
        return self._orders.get(oref, None)

    def _t_order_create(self):
        while True:
            msg = self.q_ordercreate.get()
            if msg is None:
                break
            oref, okwargs = msg
            if self.p.loglevel <= logging.DEBUG:
                print(msg)
            try:
                r = self._post_request("order/add", json=okwargs)
            except Exception as e:
                self.put_notification(e)
                self.broker._reject(order.ref)
                return

    def order_cancel(self, order):
        self.q_orderclose.put(order.ref)
        return order

    def _t_order_cancel(self):
        while True:
            order = self.q_orderclose.get()
            if oref is None:
                break

            oid = self._orders.get(oref, None)
            if oid is None:
                continue  # the order is no longer there
            try:
                o = self.oapi.close_order(self.p.account, oid)
            except Exception as e:
                continue  # not cancelled - FIXME: notify
            self.broker._cancel(oref)

if __name__ == '__main__':
    s = SharpPointStore()
<commit_msg>put try catch block to catch errors in position<commit_after>#!/usr/bin/env python
# -*- coding: utf-8; py-indent-offset:4 -*-
###############################################################################
#
# Copyright (C) 2015,2016 Daniel Rodriguez
#
# Licensed under the Simplified BSD License
#
###############################################################################
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import collections
from datetime import datetime, timedelta
import time as _time
import json
import threading
import requests
import sseclient
import logging
import backtrader as bt
from backtrader.metabase import MetaParams
from backtrader.utils.py3 import queue, with_metaclass
from backtrader.utils import AutoDict
from backtrader.position import Position
from copy import copy

class MetaSingleton(MetaParams):
    '''Metaclass to make a metaclassed class a singleton'''
    def __init__(cls, name, bases, dct):
        super(MetaSingleton, cls).__init__(name, bases, dct)
        cls._singleton = None

    def __call__(cls, *args, **kwargs):
        if cls._singleton is None:
            cls._singleton = (
                super(MetaSingleton, cls).__call__(*args, **kwargs))

        return cls._singleton

class SharpPointStore(with_metaclass(MetaSingleton, object)):
    '''Singleton class wrapping to control the connections to SharpPoint.

    Params:

      - ``token`` (default:``None``): API access token

      - ``account`` (default: ``None``): account id

      - ``practice`` (default: ``False``): use the test environment

      - ``account_tmout`` (default: ``10.0``): refresh period for account
        value/cash refresh
    '''

    BrokerCls = None  # broker class will autoregister
    DataCls = None  # data class will auto register
    BackTestCls = None

    params = (
        ('gateway', 'http://localhost:5000/'),
        ('token', ''),
        ('account', ''),
        ('login', None),
        ('practice', False),
        ('loglevel', logging.INFO),
        ('account_tmout', 30.0),
    )

    _DTEPOCH = datetime(1970, 1, 1)
    _ENVPRACTICE = 'practice'
    _ENVLIVE = 'live'

    @classmethod
    def getdata(cls, *args, **kwargs):
        '''Returns ``DataCls`` with args, kwargs'''
        return cls.DataCls(*args, **kwargs)

    @classmethod
    def getbroker(cls, *args, **kwargs):
        '''Returns broker with *args, **kwargs from registered ``BrokerCls``
or ``BackTestCls``
        '''
        backtest = kwargs.pop('backtest', False)
        if backtest:
            return cls.BackTestCls(*args, **kwargs)
        return cls.BrokerCls(*args, **kwargs)

    def __init__(self):
        super(SharpPointStore, self).__init__()

        self.positions = collections.defaultdict(Position)
        self.notifs = collections.deque()  # store notifications for cerebro

        self._env = None  # reference to cerebro for general notifications
        self.broker = None  # broker instance
        self.datas = list()  # datas that have registered over start

        self._orders = collections.OrderedDict()  # map order.ref to oid
        self._ordersrev = collections.OrderedDict()  # map oid to order.ref
        self._transpend = collections.defaultdict(collections.deque)

        self._oenv = self._ENVPRACTICE if self.p.practice else self._ENVLIVE

        self._cash = 0.0
        self._value = 0.0
        self.q_account = queue.Queue()
        self.q_ordercreate = queue.Queue()
        self.q_orderclose = queue.Queue()
        self.streaming_events()

    def start(self, data=None, broker=None):
        # Datas require some processing to kickstart data reception
        if data is None and broker is None:
            self.cash = None
            return

        if data is not None:
            self._env = data._env
            # For datas simulate a queue with None to kickstart co
            self.datas.append(data)

            if self.broker is not None:
                self.broker.data_started(data)

        elif broker is not None:
            self.broker = broker
            self.broker_threads()

    def stop(self):
        # signal end of thread
        if self.broker is not None:
            self.q_ordercreate.put(None)
            self.q_orderclose.put(None)
            self.q_account.put(None)

    def put_notification(self, msg, *args, **kwargs):
        self.notifs.append((msg, args, kwargs))

    def get_notifications(self):
        '''Return the pending "store" notifications'''
        self.notifs.append(None)  # put a mark / threads could still append
        return [x for x in iter(self.notifs.popleft, None)]

    def get_positions(self):
        pass

    def get_granularity(self, timeframe, compression):
        pass

    def get_instrument(self, dataname):
        pass

    def _get_request(self, method, **kwargs):
        if self.p.gateway is not None:
            return requests.get(self.p.gateway + method,
                                **kwargs)

    def _post_request(self, method, **kwargs):
        if self.p.gateway is not None:
            return requests.post(self.p.gateway + method,
                                 **kwargs)

    def streaming_events(self, tmout=None):
        q = queue.Queue()
        kwargs = {'q': q, 'tmout': tmout}

        t = threading.Thread(target=self._t_streaming_listener, kwargs=kwargs)
        t.daemon = True
        t.start()

    def updateposition(self, data):
        """Update position from streamer"""
        try:
            position = Position(data['Qty'],
                                abs(data['TotalAmt']/data['Qty']))
            self.positions[data['ProdCode']] = position
        except KeyError:
            print("key-error in updateposition", data)

    def getposition(self, data, clone=False):
        position =  self.positions[data]
        if clone:
            return copy(position)
        else:
            return position

    def _t_streaming_listener(self, q, tmout=None):
        response = self._get_request("log/subscribe", stream=True)
        client = sseclient.SSEClient(response)
        for event in client.events():
            data = json.loads(event.data)
            info = data.get('data', None)
            oref = 0

            if event.event == "AccountPositionPush":
                if self.p.loglevel <= logging.DEBUG:
                    print(event.event, data['data'])
                self.updateposition(data['data'])
                continue
            if self.broker is None:
                continue

            try:
                oref = info['Ref2']
            except:
                if self.p.loglevel <= logging.DEBUG:
                    print("Unhandled event")
                continue

            if self.p.loglevel <= logging.DEBUG:
                print(event)
            if event.event == "OrderBeforeSendReport":
                if self.p.loglevel <= logging.DEBUG:
                    print(data)
                self.broker._submit(oref)
            elif event.event == "OrderRequestFailed":
                if self.p.loglevel <= logging.DEBUG:
                    print(data)
                self.broker._reject(oref)
            elif event.event == "OrderReport":
                if self.p.loglevel <= logging.DEBUG:
                    print(data)
                status = int(info['Status'])
                if status == 4:
                    self.broker._accept(oref)
                elif status == 6:
                    self.broker._cancel(oref)
            elif event.event == "TradeReport":
                if self.p.loglevel <= logging.DEBUG:
                    print(data)
                qty = int(info['Qty'])
                price = float(info['Price'])
                self.broker._fill(oref, qty, price)

    def streaming_prices(self, dataname, tmout=None):
        q = queue.Queue()
        kwargs = {'q': q, 'dataname': dataname, 'tmout': tmout}
        t = threading.Thread(target=self._t_streaming_prices, kwargs=kwargs)
        t.daemon = True
        t.start()
        return q

    def _t_streaming_prices(self, dataname, q, tmout):
        r = self._get_request("ticker/subscribe/" + dataname)

    def get_cash(self):
        return self._cash

    def get_value(self):
        return self._value

    def broker_threads(self):
        self.q_account.put(True)  # force an immediate update
        t = threading.Thread(target=self._t_account)
        t.daemon = True
        t.start()

        t = threading.Thread(target=self._t_order_create)
        t.daemon = True
        t.start()

        t = threading.Thread(target=self._t_order_cancel)
        t.daemon = True
        t.start()

    _ORDEREXECS = {
        bt.Order.Market: 'market',
        bt.Order.Limit: 'limit',
        bt.Order.Stop: 'stop',
        bt.Order.StopLimit: 'stop',
    }

    def isloggedin(self):
        login_info = self._get_request("login-info").json()
        if self.p.loglevel <= logging.DEBUG:
            print("login-info", login_info)
        return int(login_info['status']) != -1

    def setlogin(self, login):
        self.p.login = login
        self.q_account.put(True)  # force an immediate update

    def _t_account(self):
        if self.p.loglevel <= logging.DEBUG:
            print("t_account")
        while True:
            try:
                msg = self.q_account.get(timeout=self.p.account_tmout)
                if msg is None:
                    break  # end of thread
            except queue.Empty:  # tmout -> time to refresh
                pass
            try:
                if self.p.login is not None and not self.isloggedin():
                    if self.p.loglevel <= logging.DEBUG:
                        print("login", self.p.login)
                    r = self._post_request("login", json=self.p.login)
            except Exception as e:
                self.put_notification(e)
                continue

    def order_create(self, order, **kwargs):
        okwargs = {"DecInPrice": 0,
                   "ClOrderId": "test2",
                   "OpenClose": 0,
                   "CondType": 0,
                   "OrderType": 0,
                   "ValidType": 0,
                   "StopType": 0,
                   "OrderOptions": 0}
        if order.isbuy():
            okwargs['BuySell'] = "B"
        elif order.issell():
            okwargs['BuySell'] = "S"
        okwargs['Price'] = order.created.price
        okwargs['Qty'] = abs(order.created.size)
        okwargs['ProdCode'] = order.data._dataname
        okwargs['Ref'] = kwargs.get('Ref', '')
        order.ref = "{:%Y%m%d%H%M%S}".format(datetime.now())
        self._orders[order.ref] = order
        okwargs['Ref2'] = str(order.ref)
        if self.p.loglevel <= logging.DEBUG:
            print(okwargs)
        self.q_ordercreate.put((order.ref, okwargs,))
        return order

    def order_by_ref(self, oref):
        return self._orders.get(oref, None)

    def _t_order_create(self):
        while True:
            msg = self.q_ordercreate.get()
            if msg is None:
                break
            oref, okwargs = msg
            if self.p.loglevel <= logging.DEBUG:
                print(msg)
            try:
                r = self._post_request("order/add", json=okwargs)
            except Exception as e:
                self.put_notification(e)
                self.broker._reject(order.ref)
                return

    def order_cancel(self, order):
        self.q_orderclose.put(order.ref)
        return order

    def _t_order_cancel(self):
        while True:
            order = self.q_orderclose.get()
            if oref is None:
                break

            oid = self._orders.get(oref, None)
            if oid is None:
                continue  # the order is no longer there
            try:
                o = self.oapi.close_order(self.p.account, oid)
            except Exception as e:
                continue  # not cancelled - FIXME: notify
            self.broker._cancel(oref)

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


<commit_before>from datetime import datetime
from grow.common import utils
from grow.pods import locales as locales_lib
from grow.pods.documents import collection
import collections
import csv as csv_lib
import itertools
import jinja2
import json as json_lib
import markdown
import re

@utils.memoize_tag
def categories(collection=None, collections=None, reverse=None, order_by=None,
               _pod=None):
  if isinstance(collection, collection.Collection):
    collection = collection
  elif isinstance(collection, basestring):
    collection = _pod.get_collection(collection)
  else:
    text = '{} must be a Collection instance or a collection path, found: {}.'
    raise ValueError(text.format(collection, type(collection)))

  category_list = collection.list_categories()

  def order_func(doc):
    return category_list.index(doc.category)

  docs = [doc for doc in collection.list_docs(reverse=reverse)]
  docs = sorted(docs, key=order_func)
  items = itertools.groupby(docs, key=order_func)
  return ((category_list[index], pages) for index, pages in items)

def LocaleIterator(iterator, locale):
  locale = str(locale)
  for i, line in enumerate(iterator):
    if i == 0 or line.startswith(locale):
      yield line

_no_locale = '__no_locale'

@utils.memoize_tag
def csv(path, locale=_no_locale, _pod=None):
  fp = _pod.open_file(path)
  if locale is not _no_locale:
    fp = LocaleIterator(fp, locale=locale)
  rows = []
  for row in csv_lib.DictReader(fp):
    data = {}
    for header, cell in row.iteritems():
      if cell is None:
        cell = ''
      data[header] = cell.decode('utf-8')
    rows.append(data)
  return rows

@utils.memoize_tag
def docs(collection, locale=None, order_by=None, _pod=None):
  collection = _pod.get_collection(collection)
  return collection.list_docs(locale=locale, order_by=order_by)

@utils.memoize_tag
def statics(pod_path, locale=None, _pod=None):
  return _pod.list_statics(pod_path, locale=locale)

def markdown_filter(value):
  try:
    if isinstance(value, unicode):
      value = value.decode('utf-8')
    return markdown.markdown(value)
  except UnicodeEncodeError:
    return markdown.markdown(value)

_slug_regex = re.compile(r'[\t !"#$%&\'()*\-/<=>?@\[\\\]^_`{|},.]+')

def slug_filter(value):
  result = []
  for word in _slug_regex.split(value.lower()):
    if word:
      result.append(word)
  return unicode(u'-'.join(result))

@utils.memoize_tag
def static(path, locale=None, _pod=None):
  return _pod.get_static(path, locale=locale)

class Menu(object):

  def __init__(self):
    self.items = collections.OrderedDict()

  def build(self, nodes):
    self._recursive_build(self.items, None, nodes)

  def iteritems(self):
    return self.items.iteritems()

  def _recursive_build(self, tree, parent, nodes):
    children = [n for n in nodes if n.parent == parent]
    for child in children:
      tree[child] = collections.OrderedDict()
      self._recursive_build(tree[child], child, nodes)

@utils.memoize_tag
def nav(collection=None, locale=None, _pod=None):
  collection_obj = _pod.get_collection('/content/' + collection)
  results = collection_obj.list_docs(order_by='order', locale=locale)
  menu = Menu()
  menu.build(results)
  return menu

@utils.memoize_tag
def breadcrumb(doc, _pod=None):
  pass

@utils.memoize_tag
def url(pod_path, locale=None, _pod=None):
  doc = _pod.get_doc(pod_path, locale=locale)
  return doc.url

@utils.memoize_tag
def get_doc(pod_path, locale=None, _pod=None):
  return _pod.get_doc(pod_path, locale=locale)

@jinja2.contextfilter
def render_filter(ctx, template):
  if isinstance(template, basestring):
    template = ctx.environment.from_string(template)
  return template.render(ctx)

@jinja2.contextfilter
def parsedatetime_filter(ctx, date_string, string_format):
  return datetime.strptime(date_string, string_format)

@jinja2.contextfilter
def deeptrans(ctx, obj):
  return _deep_gettext(ctx, obj)

@jinja2.contextfilter
def jsonify(ctx, obj, *args, **kwargs):
  return json_lib.dumps(obj, *args, **kwargs)

def _deep_gettext(ctx, fields):
  if isinstance(fields, dict):
    new_dct = {}
    for key, val in fields.iteritems():
      if isinstance(val, (dict, list, set)):
        new_dct[key] = _deep_gettext(ctx, val)
      elif isinstance(val, basestring):
        new_dct[key] = _gettext_alias(ctx, val)
      else:
        new_dct[key] = val
    return new_dct
  elif isinstance(fields, (list, set)):
    for i, val in enumerate(fields):
      if isinstance(val, (dict, list, set)):
        fields[i] = _deep_gettext(ctx, val)
      elif isinstance(val, basestring):
        fields[i] = _gettext_alias(ctx, val)
      else:
        fields[i] = val
    return fields

def _gettext_alias(__context, *args, **kwargs):
  return __context.call(__context.resolve('gettext'), *args, **kwargs)

@utils.memoize_tag
def yaml(path, _pod):
  fields = utils.parse_yaml(_pod.read_file(path), pod=_pod)
  return utils.untag_fields(fields)

@utils.memoize_tag
def json(path, _pod):
  fp = _pod.open_file(path)
  return json_lib.load(fp)

def date(datetime_obj=None, _pod=None, **kwargs):
  _from = kwargs.get('from', None)
  if datetime_obj is None:
    datetime_obj = datetime.now()
  elif isinstance(datetime_obj, basestring) and _from is not None:
    datetime_obj = datetime.strptime(datetime_obj, _from)
  return datetime_obj

@utils.memoize_tag
def locales(codes, _pod=None):
  return locales_lib.Locale.parse_codes(codes)
<commit_msg>Fix caching issue with collections.<commit_after>from datetime import datetime
from grow.common import utils
from grow.pods import locales as locales_lib
from grow.pods.documents import collection as collection_lib
import collections
import csv as csv_lib
import itertools
import jinja2
import json as json_lib
import markdown
import re

def categories(collection=None, collections=None, reverse=None, order_by=None,
               _pod=None, use_cache=False):
  if isinstance(collection, collection_lib.Collection):
    collection = collection
  elif isinstance(collection, basestring):
    collection = _pod.get_collection(collection)
  else:
    text = '{} must be a Collection instance or a collection path, found: {}.'
    raise ValueError(text.format(collection, type(collection)))

  category_list = collection.list_categories()

  def order_func(doc):
    return category_list.index(doc.category)

  docs = [doc for doc in collection.list_docs(reverse=reverse)]
  docs = sorted(docs, key=order_func)
  items = itertools.groupby(docs, key=order_func)
  return ((category_list[index], pages) for index, pages in items)

def LocaleIterator(iterator, locale):
  locale = str(locale)
  for i, line in enumerate(iterator):
    if i == 0 or line.startswith(locale):
      yield line

_no_locale = '__no_locale'

@utils.memoize_tag
def csv(path, locale=_no_locale, _pod=None):
  fp = _pod.open_file(path)
  if locale is not _no_locale:
    fp = LocaleIterator(fp, locale=locale)
  rows = []
  for row in csv_lib.DictReader(fp):
    data = {}
    for header, cell in row.iteritems():
      if cell is None:
        cell = ''
      data[header] = cell.decode('utf-8')
    rows.append(data)
  return rows

@utils.memoize_tag
def docs(collection, locale=None, order_by=None, _pod=None):
  collection = _pod.get_collection(collection)
  return collection.list_docs(locale=locale, order_by=order_by)

@utils.memoize_tag
def collections(path, _pod=None):
  return _pod.get_collection(collection)

@utils.memoize_tag
def statics(pod_path, locale=None, _pod=None):
  return _pod.list_statics(pod_path, locale=locale)

def markdown_filter(value):
  try:
    if isinstance(value, unicode):
      value = value.decode('utf-8')
    return markdown.markdown(value)
  except UnicodeEncodeError:
    return markdown.markdown(value)

_slug_regex = re.compile(r'[\t !"#$%&\'()*\-/<=>?@\[\\\]^_`{|},.]+')

def slug_filter(value):
  result = []
  for word in _slug_regex.split(value.lower()):
    if word:
      result.append(word)
  return unicode(u'-'.join(result))

@utils.memoize_tag
def static(path, locale=None, _pod=None):
  return _pod.get_static(path, locale=locale)

class Menu(object):

  def __init__(self):
    self.items = collections.OrderedDict()

  def build(self, nodes):
    self._recursive_build(self.items, None, nodes)

  def iteritems(self):
    return self.items.iteritems()

  def _recursive_build(self, tree, parent, nodes):
    children = [n for n in nodes if n.parent == parent]
    for child in children:
      tree[child] = collections.OrderedDict()
      self._recursive_build(tree[child], child, nodes)

@utils.memoize_tag
def nav(collection=None, locale=None, _pod=None):
  collection_obj = _pod.get_collection('/content/' + collection)
  results = collection_obj.list_docs(order_by='order', locale=locale)
  menu = Menu()
  menu.build(results)
  return menu

@utils.memoize_tag
def breadcrumb(doc, _pod=None):
  pass

@utils.memoize_tag
def url(pod_path, locale=None, _pod=None):
  doc = _pod.get_doc(pod_path, locale=locale)
  return doc.url

@utils.memoize_tag
def get_doc(pod_path, locale=None, _pod=None):
  return _pod.get_doc(pod_path, locale=locale)

@jinja2.contextfilter
def render_filter(ctx, template):
  if isinstance(template, basestring):
    template = ctx.environment.from_string(template)
  return template.render(ctx)

@jinja2.contextfilter
def parsedatetime_filter(ctx, date_string, string_format):
  return datetime.strptime(date_string, string_format)

@jinja2.contextfilter
def deeptrans(ctx, obj):
  return _deep_gettext(ctx, obj)

@jinja2.contextfilter
def jsonify(ctx, obj, *args, **kwargs):
  return json_lib.dumps(obj, *args, **kwargs)

def _deep_gettext(ctx, fields):
  if isinstance(fields, dict):
    new_dct = {}
    for key, val in fields.iteritems():
      if isinstance(val, (dict, list, set)):
        new_dct[key] = _deep_gettext(ctx, val)
      elif isinstance(val, basestring):
        new_dct[key] = _gettext_alias(ctx, val)
      else:
        new_dct[key] = val
    return new_dct
  elif isinstance(fields, (list, set)):
    for i, val in enumerate(fields):
      if isinstance(val, (dict, list, set)):
        fields[i] = _deep_gettext(ctx, val)
      elif isinstance(val, basestring):
        fields[i] = _gettext_alias(ctx, val)
      else:
        fields[i] = val
    return fields

def _gettext_alias(__context, *args, **kwargs):
  return __context.call(__context.resolve('gettext'), *args, **kwargs)

@utils.memoize_tag
def yaml(path, _pod):
  fields = utils.parse_yaml(_pod.read_file(path), pod=_pod)
  return utils.untag_fields(fields)

@utils.memoize_tag
def json(path, _pod):
  fp = _pod.open_file(path)
  return json_lib.load(fp)

def date(datetime_obj=None, _pod=None, **kwargs):
  _from = kwargs.get('from', None)
  if datetime_obj is None:
    datetime_obj = datetime.now()
  elif isinstance(datetime_obj, basestring) and _from is not None:
    datetime_obj = datetime.strptime(datetime_obj, _from)
  return datetime_obj

@utils.memoize_tag
def locales(codes, _pod=None):
  return locales_lib.Locale.parse_codes(codes)
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# vim: set ts=4 sw=4 expandtab :

# Copyright (c) 2017 Timothy Savannah - All Rights Reserved
#   This code is licensed under the terms of the APACHE license version 2.0
#
#  extractMtree.py - Extracts the mtree from all available packages,
#    attempting short-reads where possible to limit bandwidth,
#    and assemble a compressed json database for use by whatprovides_upstream

import gzip
import os
import json
import re
import subprocess
import sys
import tarfile
import tempfile
import time

from io import BytesIO

import func_timeout
from nonblock import bgwrite

try:
    PermissionError
except NameError:
    PermissionError = IOError

PROVIDES_DB_LOCATION = "/var/lib/pacman/.providesDB"

global isVerbose
isVerbose = False

def decompressZlib(data):

    devnull = open(os.devnull, 'w')

    pipe = subprocess.Popen(['/usr/bin/zcat'], shell=False, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=devnull)
    pipe.stdin.write(data)
    pipe.stdin.close()

    results = pipe.stdout.read()
    pipe.wait()

    devnull.close()

    return results

def decompressXz(data):

    devnull = open(os.devnull, 'w')
    pipe = subprocess.Popen(['/usr/bin/xzcat'], shell=False, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=devnull)

    bgwrite(pipe.stdin, data, ioPrio=1, closeWhenFinished=True)

    results = pipe.stdout.read()
    pipe.wait()

    devnull.close()

    return results

SIZE_IDX_START = 124
SIZE_IDX_END = 124 + 12

def getSize(header):

    trySection = header[SIZE_IDX_START : SIZE_IDX_END]
    return int(trySection, 8)

    return int(trySection, 8)

def getFiles(dataStr):
    lines = dataStr.split('\n')

    rePat = re.compile('^\.(?P<filename>.+) time')
    ret = []

    for line in lines:
        if not line or line[0] != '.':
            continue
        matchObj = rePat.match(line)
        if not matchObj:
            # Uh oh..
            continue
        ret.append(matchObj.groupdict()['filename'])

    return ret

def fetchFromUrl(url, numBytes):

    devnull = open(os.devnull, 'w')

    if numBytes:
        limitBytes = '| head -c %d' %(numBytes, )
    else:
        limitBytes = ''

    if not isVerbose:
        extraArgs = '--silent'
    else:
        extraArgs = ""

    pipe = subprocess.Popen("/usr/bin/curl '-k' %s '%s' %s" %(extraArgs, url, limitBytes), shell=True, stdout=subprocess.PIPE, stderr=devnull)

    urlContents = pipe.stdout.read()
    ret = pipe.wait()

    devnull.close()

    if b'404 Not Found' in urlContents and '-x86_64' in url:
        return fetchFromUrl(url.replace('-x86_64', '-any'), numBytes)

    return urlContents

def getAllPackages():

    devnull = open(os.devnull, 'w')

    pipe = subprocess.Popen(["/usr/bin/pacman", "-Sl"], shell=False, stdout=subprocess.PIPE, stderr=devnull)

    contents = pipe.stdout.read()

    pipe.wait()

    devnull.close()

    return [tuple(x.split(' ')[0:3]) for x in contents.decode('utf-8').split('\n') if x and ' ' in x]

# USE_ARCH - Package arch to use. TODO: Allow others
USE_ARCH = "x86_64"

def getRepoUrl():
    nextLine = True

    repoRE = re.compile('^[ \t]*Server[ \t]*=[ \t]*(?P<server_url>.+)$')

    with open('/etc/pacman.d/mirrorlist', 'rt') as f:
        nextLine = f.readline()
        while nextLine != '':
            matchObj = repoRE.match(nextLine.strip())
            if matchObj:
                groupDict = matchObj.groupdict()
                if groupDict['server_url']:
                    ret = groupDict['server_url'].replace('$repo', '%s').replace('$arch', USE_ARCH)
                    while ret.endswith('/'):
                        ret = ret[:-1]

                    ret += '/%s'
                    return ret

            nextLine = f.readline()

    raise Exception('Failed to find repo URL from /etc/pacman.d/mirrorlist. Are any not commented?')

#REPO_URL = "http://mirrors.acm.wpi.edu/archlinux/%s/os/x86_64/%s"

def getFileData(filename, decodeWith=None):
    with open(filename, 'rb') as f:
        contents = f.read()

    if decodeWith:
        contents = contents.decode(decodeWith)
    return contents

if __name__ == '__main__':

    args = sys.argv[1:]

    if '-v' in args:
        isVerbose = True
        args.remove('-v')
    if '--verbose' in args:
        isVerbose = True
        args.remove('--verbose')

    if len(args) == 0:
        allPackages = getAllPackages()
    else:
        sys.stderr.write('Too many arguments\n')
        sys.exit(1)

    if not os.access(PROVIDES_DB_LOCATION, os.W_OK):
        sys.stdout.write('Cannot write to "%s". Will create temp file.\n' %((PROVIDES_DB_LOCATION, )) )
        result = False
        while result not in ('y', 'n'):
            sys.stdout.write('Continue? (y/n): ')
            sys.stdout.flush()
            result = sys.stdin.readline().strip().lower()

        if result == 'n':
            sys.exit(2)

    REPO_URL = getRepoUrl()
    print ( "Using first repo from /etc/pacman.d/mirrorlist:\n\t%s\n" %(REPO_URL, ))

    results = {}

    def doOne(repoName, packageName, versionInfo, results, fetchedData=None, useTarMod=False):
        global isVerbose

        if fetchedData is None:
            finalUrl = REPO_URL %( repoName, packageName + "-" + versionInfo + "-x86_64.pkg.tar.xz" )
            if isVerbose:
                print ( "Fetching url: " + finalUrl )

            if useTarMod is False:
                maxSize = 1024 * 80
            else:
                maxSize = None
            tarContents = fetchFromUrl(finalUrl, maxSize)
        else:
            tarContents = fetchedData

        if useTarMod is False:
            data = decompressXz(tarContents[:1024 * 80])

            mtreeIdx = data.index(b'.MTREE')

            headerStart = data[mtreeIdx:]

            try:
                mtreeSize = getSize(headerStart)
                compressedData = headerStart[512 : 512 + mtreeSize]
            except:
                return doOne(repoName, packageName, versionInfo, results, fetchedData, useTarMod=True)
        else:
            if isVerbose:
                print ( "Fallback..." )
            data = decompressXz(tarContents)

            bio = BytesIO()
            bio.write(data)
            bio.seek(0)

            tf = tarfile.open(fileobj=bio)

            extractedMtreeFile = tf.extractfile('.MTREE')
            compressedData = extractedMtreeFile.read()
            try:
                extractedMtreeFile.close()
            except:
                pass

        mtreeData = decompressZlib(compressedData).decode('utf-8')

        files = getFiles(mtreeData)

        results[packageName] = files
        if isVerbose:
            sys.stdout.write("Got %d files.\n\n" %(len(files), ))

#    if len(args) != 1:
#        fileData = getFileData(args)
#        doOne('core', 'binutils', '2.28.0-2', results, fileData)
#        sys.exit(0)

    failedPackages = []

    def runThroughPackages(allPackages, results, failedPackages, timeout=5):
        global isVerbose

        for repoName, packageName, versionInfo in allPackages:
            time.sleep(1.5)
            if isVerbose:
                sys.stdout.write("Processing %s - %s: " %(repoName, packageName) )
                sys.stdout.flush()
            try:
                func_timeout.func_timeout(5, doOne, (repoName, packageName, versionInfo, results))
            except func_timeout.FunctionTimedOut as fte:
                try:
                    failedPackages.append ( (repoName, packageName, versionInfo) )
                    errStr = 'Error TIMEOUT processing %s - %s : FunctionTimedOut\n\n' %(repoName, packageName )
                    sys.stderr.write(errStr)
                    results[packageName] = errStr
                except:
                    pass
            except KeyboardInterrupt as ke:
                raise ke
            except Exception as e:
                try:
                    failedPackages.append ( (repoName, packageName, versionInfo) )
                    errStr = 'Error processing %s - %s : < %s >: %s\n\n' %(repoName, packageName, e.__class__.__name__, str(e))
                    sys.stderr.write(errStr)
                    results[packageName] = errStr
                except:
                    pass

    try:
        runThroughPackages(allPackages, results, failedPackages, timeout=5)
    except KeyboardInterrupt as ke:
        pass

    if failedPackages:
        sys.stderr.write("Need to retry %d packages. Resting for a minute....\n\n" %(len(failedPackages), ))
        time.sleep(60)

        # Now give up to 8 minutes for each package to complete. This covers if a PAX_FORMAT on a huge file,
        #   slow mirror, etc.
        newFailedPackages = []
        try:
            runThroughPackages( failedPackages, results, newFailedPackages, timeout=(60 * 8))
        except KeyboardInterrupt as ke:
            pass

        if newFailedPackages:
            sys.stderr.write('After completing, still %d failed packages.\nFailed after retry: %s\n\n' %(len(newFailedPackages), '\n'.join(['\t[%s] %s-%s  \t%s' %(failedP[0], failedP[1], failedP[2], results[failedP[1]]  ) for failedP in newFailedPackages]) ) )

    compressed = gzip.compress( json.dumps(results).encode('utf-8') )

    try:
        with open(PROVIDES_DB_LOCATION, 'wb') as f:
            f.write( compressed )
    except PermissionError:
        tempFile = tempfile.NamedTemporaryFile(mode='wb', delete=False)
        sys.stderr.write('Failed to open "%s" for writing. Dumping to tempfile: "%s"\n' %(PROVIDES_DB_LOCATION, tempFile.name, ))
        tempFile.write( compressed )
        tempFile.close()

# vim: set ts=4 sw=4 expandtab :
<commit_msg>Refresh pacman database if able at start of extractMtree.py<commit_after>#!/usr/bin/env python
#
# vim: set ts=4 sw=4 expandtab :

# Copyright (c) 2017 Timothy Savannah - All Rights Reserved
#   This code is licensed under the terms of the APACHE license version 2.0
#
#  extractMtree.py - Extracts the mtree from all available packages,
#    attempting short-reads where possible to limit bandwidth,
#    and assemble a compressed json database for use by whatprovides_upstream

import gzip
import os
import json
import re
import subprocess
import sys
import tarfile
import tempfile
import time

from io import BytesIO

import func_timeout
from nonblock import bgwrite

try:
    PermissionError
except NameError:
    PermissionError = IOError

PROVIDES_DB_LOCATION = "/var/lib/pacman/.providesDB"

global isVerbose
isVerbose = False

def decompressZlib(data):

    devnull = open(os.devnull, 'w')

    pipe = subprocess.Popen(['/usr/bin/zcat'], shell=False, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=devnull)
    pipe.stdin.write(data)
    pipe.stdin.close()

    results = pipe.stdout.read()
    pipe.wait()

    devnull.close()

    return results

def decompressXz(data):

    devnull = open(os.devnull, 'w')
    pipe = subprocess.Popen(['/usr/bin/xzcat'], shell=False, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=devnull)

    bgwrite(pipe.stdin, data, ioPrio=1, closeWhenFinished=True)

    results = pipe.stdout.read()
    pipe.wait()

    devnull.close()

    return results

SIZE_IDX_START = 124
SIZE_IDX_END = 124 + 12

def getSize(header):

    trySection = header[SIZE_IDX_START : SIZE_IDX_END]
    return int(trySection, 8)

    return int(trySection, 8)

def getFiles(dataStr):
    lines = dataStr.split('\n')

    rePat = re.compile('^\.(?P<filename>.+) time')
    ret = []

    for line in lines:
        if not line or line[0] != '.':
            continue
        matchObj = rePat.match(line)
        if not matchObj:
            # Uh oh..
            continue
        ret.append(matchObj.groupdict()['filename'])

    return ret

def fetchFromUrl(url, numBytes):

    devnull = open(os.devnull, 'w')

    if numBytes:
        limitBytes = '| head -c %d' %(numBytes, )
    else:
        limitBytes = ''

    if not isVerbose:
        extraArgs = '--silent'
    else:
        extraArgs = ""

    pipe = subprocess.Popen("/usr/bin/curl '-k' %s '%s' %s" %(extraArgs, url, limitBytes), shell=True, stdout=subprocess.PIPE, stderr=devnull)

    urlContents = pipe.stdout.read()
    ret = pipe.wait()

    devnull.close()

    if b'404 Not Found' in urlContents and '-x86_64' in url:
        return fetchFromUrl(url.replace('-x86_64', '-any'), numBytes)

    return urlContents

def getAllPackages():

    devnull = open(os.devnull, 'w')

    pipe = subprocess.Popen(["/usr/bin/pacman", "-Sl"], shell=False, stdout=subprocess.PIPE, stderr=devnull)

    contents = pipe.stdout.read()

    pipe.wait()

    devnull.close()

    return [tuple(x.split(' ')[0:3]) for x in contents.decode('utf-8').split('\n') if x and ' ' in x]

# USE_ARCH - Package arch to use. TODO: Allow others
USE_ARCH = "x86_64"

def getRepoUrl():
    nextLine = True

    repoRE = re.compile('^[ \t]*Server[ \t]*=[ \t]*(?P<server_url>.+)$')

    with open('/etc/pacman.d/mirrorlist', 'rt') as f:
        nextLine = f.readline()
        while nextLine != '':
            matchObj = repoRE.match(nextLine.strip())
            if matchObj:
                groupDict = matchObj.groupdict()
                if groupDict['server_url']:
                    ret = groupDict['server_url'].replace('$repo', '%s').replace('$arch', USE_ARCH)
                    while ret.endswith('/'):
                        ret = ret[:-1]

                    ret += '/%s'
                    return ret

            nextLine = f.readline()

    raise Exception('Failed to find repo URL from /etc/pacman.d/mirrorlist. Are any not commented?')

#REPO_URL = "http://mirrors.acm.wpi.edu/archlinux/%s/os/x86_64/%s"

def getFileData(filename, decodeWith=None):
    with open(filename, 'rb') as f:
        contents = f.read()

    if decodeWith:
        contents = contents.decode(decodeWith)
    return contents

if __name__ == '__main__':

    args = sys.argv[1:]

    if '-v' in args:
        isVerbose = True
        args.remove('-v')
    if '--verbose' in args:
        isVerbose = True
        args.remove('--verbose')

    if len(args) == 0:
        allPackages = getAllPackages()
    else:
        sys.stderr.write('Too many arguments\n')
        sys.exit(1)

    if not os.access(PROVIDES_DB_LOCATION, os.W_OK):
        sys.stdout.write('Cannot write to "%s". Will create temp file.\n' %((PROVIDES_DB_LOCATION, )) )
        result = False
        while result not in ('y', 'n'):
            sys.stdout.write('Continue? (y/n): ')
            sys.stdout.flush()
            result = sys.stdin.readline().strip().lower()

        if result == 'n':
            sys.exit(2)

    if os.getuid() != 0:
        sys.stderr.write('WARNING: Cannot refresh pacman database.\n')
    else:
        ret = subprocess.Popen(['/usr/bin/pacman', '-Sy'], shell=False).wait()
        if ret != 0:
            sys.stderr.write('WARNING: pacman -Sy returned non-zero: %d\n' %(ret,))

    REPO_URL = getRepoUrl()
    print ( "Using first repo from /etc/pacman.d/mirrorlist:\n\t%s\n" %(REPO_URL, ))

    results = {}

    def doOne(repoName, packageName, versionInfo, results, fetchedData=None, useTarMod=False):
        global isVerbose

        if fetchedData is None:
            finalUrl = REPO_URL %( repoName, packageName + "-" + versionInfo + "-x86_64.pkg.tar.xz" )
            if isVerbose:
                print ( "Fetching url: " + finalUrl )

            if useTarMod is False:
                maxSize = 1024 * 80
            else:
                maxSize = None
            tarContents = fetchFromUrl(finalUrl, maxSize)
        else:
            tarContents = fetchedData

        if useTarMod is False:
            data = decompressXz(tarContents[:1024 * 80])

            mtreeIdx = data.index(b'.MTREE')

            headerStart = data[mtreeIdx:]

            try:
                mtreeSize = getSize(headerStart)
                compressedData = headerStart[512 : 512 + mtreeSize]
            except:
                return doOne(repoName, packageName, versionInfo, results, fetchedData, useTarMod=True)
        else:
            if isVerbose:
                print ( "Fallback..." )
            data = decompressXz(tarContents)

            bio = BytesIO()
            bio.write(data)
            bio.seek(0)

            tf = tarfile.open(fileobj=bio)

            extractedMtreeFile = tf.extractfile('.MTREE')
            compressedData = extractedMtreeFile.read()
            try:
                extractedMtreeFile.close()
            except:
                pass

        mtreeData = decompressZlib(compressedData).decode('utf-8')

        files = getFiles(mtreeData)

        results[packageName] = files
        if isVerbose:
            sys.stdout.write("Got %d files.\n\n" %(len(files), ))

#    if len(args) != 1:
#        fileData = getFileData(args)
#        doOne('core', 'binutils', '2.28.0-2', results, fileData)
#        sys.exit(0)

    failedPackages = []

    def runThroughPackages(allPackages, results, failedPackages, timeout=5):
        global isVerbose

        for repoName, packageName, versionInfo in allPackages:
            time.sleep(1.5)
            if isVerbose:
                sys.stdout.write("Processing %s - %s: " %(repoName, packageName) )
                sys.stdout.flush()
            try:
                func_timeout.func_timeout(5, doOne, (repoName, packageName, versionInfo, results))
            except func_timeout.FunctionTimedOut as fte:
                try:
                    failedPackages.append ( (repoName, packageName, versionInfo) )
                    errStr = 'Error TIMEOUT processing %s - %s : FunctionTimedOut\n\n' %(repoName, packageName )
                    sys.stderr.write(errStr)
                    results[packageName] = errStr
                except:
                    pass
            except KeyboardInterrupt as ke:
                raise ke
            except Exception as e:
                try:
                    failedPackages.append ( (repoName, packageName, versionInfo) )
                    errStr = 'Error processing %s - %s : < %s >: %s\n\n' %(repoName, packageName, e.__class__.__name__, str(e))
                    sys.stderr.write(errStr)
                    results[packageName] = errStr
                except:
                    pass

    try:
        runThroughPackages(allPackages, results, failedPackages, timeout=5)
    except KeyboardInterrupt as ke:
        pass

    if failedPackages:
        sys.stderr.write("Need to retry %d packages. Resting for a minute....\n\n" %(len(failedPackages), ))
        time.sleep(60)

        # Now give up to 8 minutes for each package to complete. This covers if a PAX_FORMAT on a huge file,
        #   slow mirror, etc.
        newFailedPackages = []
        try:
            runThroughPackages( failedPackages, results, newFailedPackages, timeout=(60 * 8))
        except KeyboardInterrupt as ke:
            pass

        if newFailedPackages:
            sys.stderr.write('After completing, still %d failed packages.\nFailed after retry: %s\n\n' %(len(newFailedPackages), '\n'.join(['\t[%s] %s-%s  \t%s' %(failedP[0], failedP[1], failedP[2], results[failedP[1]]  ) for failedP in newFailedPackages]) ) )

    compressed = gzip.compress( json.dumps(results).encode('utf-8') )

    try:
        with open(PROVIDES_DB_LOCATION, 'wb') as f:
            f.write( compressed )
    except PermissionError:
        tempFile = tempfile.NamedTemporaryFile(mode='wb', delete=False)
        sys.stderr.write('Failed to open "%s" for writing. Dumping to tempfile: "%s"\n' %(PROVIDES_DB_LOCATION, tempFile.name, ))
        tempFile.write( compressed )
        tempFile.close()

# vim: set ts=4 sw=4 expandtab :
<|endoftext|>