<commit_before>import boto.ec2
import boto.exception
import boto.ec2.regioninfo
import boto.ec2.instance
from unittest import TestCase
from mock import patch, Mock
from monocyte.handler import ec2
from monocyte.handler import Resource

class EC2HandlerTest(TestCase):

    def setUp(self):
        self.boto_mock = patch("monocyte.handler.ec2.boto").start()
        self.positive_fake_region = Mock(boto.ec2.regioninfo.EC2RegionInfo)
        self.positive_fake_region.name = "allowed_region"
        self.negative_fake_region = Mock(boto.ec2.regioninfo.EC2RegionInfo)
        self.negative_fake_region.name = "forbbiden_region"

        self.boto_mock.ec2.regions.return_value = [self.positive_fake_region, self.negative_fake_region]
        self.ec2_handler_filter = ec2.Instance(lambda region_name: region_name == self.positive_fake_region.name)

        self.instance_mock = self._given_instance_mock()

    def tearDown(self):
        patch.stopall()

    def test_fetch_unwanted_resources_filtered(self):
        only_resource = list(self.ec2_handler_filter.fetch_unwanted_resources())[0]
        self.assertEquals(only_resource.wrapped, self.instance_mock)

    def test_to_string(self):
        only_resource = list(self.ec2_handler_filter.fetch_unwanted_resources())[0]
        resource_string = self.ec2_handler_filter.to_string(only_resource)

        self.assertTrue(self.instance_mock.id in resource_string)
        self.assertTrue(self.instance_mock.image_id in resource_string)
        self.assertTrue(self.instance_mock.instance_type in resource_string)
        self.assertTrue(self.instance_mock.launch_time in resource_string)
        self.assertTrue(self.instance_mock.public_dns_name in resource_string)
        self.assertTrue(self.instance_mock.key_name in resource_string)
        self.assertTrue(self.instance_mock.region.name in resource_string)

    @patch("monocyte.handler.ec2.print", create=True)
    def test_delete(self, print_mock):
        resource = Resource(self.instance_mock, self.negative_fake_region.name)
        connection = self.boto_mock.ec2.connect_to_region.return_value

        e = boto.exception.EC2ResponseError(412, 'boom')
        e.message = "test"
        connection.terminate_instances.side_effect = e

        deleted_resource = self.ec2_handler_filter.delete(resource)[0]
        self.assertEquals(self.instance_mock, deleted_resource)
        print_mock.assert_called_with("\tTermination test")

    def _given_instance_mock(self):
        instance_mock = Mock(boto.ec2.instance, image_id="ami-1112")
        instance_mock.id = "id-12345"
        instance_mock.instance_type = "m1.small"
        instance_mock.launch_time = "01.01.2015"
        instance_mock.public_dns_name = "test.aws.com"
        instance_mock.key_name = "test-ssh-key"
        instance_mock.region = self.positive_fake_region
        instance_mock._state = boto.ec2.instance.InstanceState(16, "running")
        instance_mock.state = instance_mock._state.name

        self.boto_mock.ec2.connect_to_region.return_value.get_only_instances.return_value = [instance_mock]
        return instance_mock
<commit_msg>fix typo.<commit_after>import boto.ec2
import boto.exception
import boto.ec2.regioninfo
import boto.ec2.instance
from unittest import TestCase
from mock import patch, Mock
from monocyte.handler import ec2
from monocyte.handler import Resource

class EC2HandlerTest(TestCase):

    def setUp(self):
        self.boto_mock = patch("monocyte.handler.ec2.boto").start()
        self.positive_fake_region = Mock(boto.ec2.regioninfo.EC2RegionInfo)
        self.positive_fake_region.name = "allowed_region"
        self.negative_fake_region = Mock(boto.ec2.regioninfo.EC2RegionInfo)
        self.negative_fake_region.name = "forbidden_region"

        self.boto_mock.ec2.regions.return_value = [self.positive_fake_region, self.negative_fake_region]
        self.ec2_handler_filter = ec2.Instance(lambda region_name: region_name == self.positive_fake_region.name)

        self.instance_mock = self._given_instance_mock()

    def tearDown(self):
        patch.stopall()

    def test_fetch_unwanted_resources_filtered(self):
        only_resource = list(self.ec2_handler_filter.fetch_unwanted_resources())[0]
        self.assertEquals(only_resource.wrapped, self.instance_mock)

    def test_to_string(self):
        only_resource = list(self.ec2_handler_filter.fetch_unwanted_resources())[0]
        resource_string = self.ec2_handler_filter.to_string(only_resource)

        self.assertTrue(self.instance_mock.id in resource_string)
        self.assertTrue(self.instance_mock.image_id in resource_string)
        self.assertTrue(self.instance_mock.instance_type in resource_string)
        self.assertTrue(self.instance_mock.launch_time in resource_string)
        self.assertTrue(self.instance_mock.public_dns_name in resource_string)
        self.assertTrue(self.instance_mock.key_name in resource_string)
        self.assertTrue(self.instance_mock.region.name in resource_string)

    @patch("monocyte.handler.ec2.print", create=True)
    def test_delete(self, print_mock):
        resource = Resource(self.instance_mock, self.negative_fake_region.name)
        connection = self.boto_mock.ec2.connect_to_region.return_value

        e = boto.exception.EC2ResponseError(412, 'boom')
        e.message = "test"
        connection.terminate_instances.side_effect = e

        deleted_resource = self.ec2_handler_filter.delete(resource)[0]
        self.assertEquals(self.instance_mock, deleted_resource)
        print_mock.assert_called_with("\tTermination test")

    def _given_instance_mock(self):
        instance_mock = Mock(boto.ec2.instance, image_id="ami-1112")
        instance_mock.id = "id-12345"
        instance_mock.instance_type = "m1.small"
        instance_mock.launch_time = "01.01.2015"
        instance_mock.public_dns_name = "test.aws.com"
        instance_mock.key_name = "test-ssh-key"
        instance_mock.region = self.positive_fake_region
        instance_mock._state = boto.ec2.instance.InstanceState(16, "running")
        instance_mock.state = instance_mock._state.name

        self.boto_mock.ec2.connect_to_region.return_value.get_only_instances.return_value = [instance_mock]
        return instance_mock
<|endoftext|>


<commit_before># Licensed to the StackStorm, Inc ('StackStorm') under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import copy
import functools

import eventlet

from st2common import log as logging
from st2common.models.db.liveaction import LiveActionDB
import st2common.services.action as action_services
from st2common.constants.action import LIVEACTION_STATUS_FAILED
from st2common.constants.action import LIVEACTION_STATUS_TIMED_OUT
from st2common.util.enum import Enum
from st2common.policies.base import ResourcePolicyApplicator

__all__ = [
    'RetryOnPolicy',
    'ExecutionRetryPolicyApplicator'
]

LOG = logging.getLogger(__name__)

VALID_RETRY_STATUSES = [
    LIVEACTION_STATUS_FAILED,
    LIVEACTION_STATUS_TIMED_OUT
]

class RetryOnPolicy(Enum):
    FAILURE = 'failure'  # Retry on execution failure
    TIMEOUT = 'timeout'  # Retry on execution timeout

class ExecutionRetryPolicyApplicator(ResourcePolicyApplicator):
    def __init__(self, policy_ref, policy_type, retry_on, max_retry_count=2, delay=0):
        """
        :param retry_on: Condition to retry the execution on (failure, timeout).
        :type retry_on: ``str``

        :param max_retry_count: Maximum number of times to try to retry an action.
        :type max_retry_count: ``int``

        :param delay: How long to wait before retrying an execution.
        :type delay: ``float``
        """
        super(ExecutionRetryPolicyApplicator, self).__init__(policy_ref=policy_ref,
                                                             policy_type=policy_type)

        self.retry_on = retry_on
        self.max_retry_count = max_retry_count
        self.delay = delay or 0

    def apply_after(self, target):
        target = super(ExecutionRetryPolicyApplicator, self).apply_after(target=target)

        live_action_db = target

        if self._is_live_action_part_of_workflow_action(live_action_db):
            LOG.warning(
                'Retry cannot be applied to this liveaction because it is executed under a '
                'workflow. Use workflow specific retry functionality where applicable. %s',
                live_action_db
            )

            return target

        retry_count = self._get_live_action_retry_count(live_action_db=live_action_db)

        extra = {'live_action_db': live_action_db, 'policy_ref': self._policy_ref,
                 'retry_on': self.retry_on, 'max_retry_count': self.max_retry_count,
                 'current_retry_count': retry_count}

        if live_action_db.status not in VALID_RETRY_STATUSES:
            # Currently we only support retrying on failed action
            LOG.debug('Liveaction not in a valid retry state, not checking retry policy',
                      extra=extra)
            return target

        if (retry_count + 1) > self.max_retry_count:
            LOG.info('Maximum retry count has been reached, not retrying', extra=extra)
            return target

        has_failed = live_action_db.status == LIVEACTION_STATUS_FAILED
        has_timed_out = live_action_db.status == LIVEACTION_STATUS_TIMED_OUT

        # TODO: This is not crash and restart safe, switch to using "DELAYED"
        # status
        if self.delay > 0:
            re_run_live_action = functools.partial(eventlet.spawn_after, self.delay,
                                                   self._re_run_live_action,
                                                   live_action_db=live_action_db)
        else:
            # Even if delay is 0, use a small delay (0.1 seconds) to prevent busy wait
            re_run_live_action = functools.partial(eventlet.spawn_after, 0.1,
                                                   self._re_run_live_action,
                                                   live_action_db=live_action_db)

            re_run_live_action = functools.partial(self._re_run_live_action,
                                                   live_action_db=live_action_db)

        if has_failed and self.retry_on == RetryOnPolicy.FAILURE:
            extra['failure'] = True
            LOG.info('Policy matched (failure), retrying action execution in %s seconds...' %
                     (self.delay), extra=extra)
            re_run_live_action()
            return target

        if has_timed_out and self.retry_on == RetryOnPolicy.TIMEOUT:
            extra['timeout'] = True
            LOG.info('Policy matched (timeout), retrying action execution in %s seconds...' %
                     (self.delay), extra=extra)
            re_run_live_action()
            return target

        return target

    def _is_live_action_part_of_workflow_action(self, live_action_db):
        """
        Retrieve parent info from context of the live action.

        :rtype: ``dict``
        """
        context = getattr(live_action_db, 'context', {})
        parent = context.get('parent', {})
        is_wf_action = (parent is not None and parent != {})

        return is_wf_action

    def _get_live_action_retry_count(self, live_action_db):
        """
        Retrieve current retry count for the provided live action.

        :rtype: ``int``
        """
        # TODO: Ideally we would store retry_count in zookeeper or similar and use locking so we
        # can run multiple instances of st2notififer
        context = getattr(live_action_db, 'context', {})
        retry_count = context.get('policies', {}).get('retry', {}).get('retry_count', 0)

        return retry_count

    def _re_run_live_action(self, live_action_db):
        retry_count = self._get_live_action_retry_count(live_action_db=live_action_db)

        # Add additional policy specific info to the context
        context = getattr(live_action_db, 'context', {})
        new_context = copy.deepcopy(context)
        new_context['policies'] = {}
        new_context['policies']['retry'] = {
            'applied_policy': self._policy_ref,
            'retry_count': (retry_count + 1),
            'retried_liveaction_id': str(live_action_db.id)
        }
        action_ref = live_action_db.action
        parameters = live_action_db.parameters
        new_live_action_db = LiveActionDB(action=action_ref, parameters=parameters,
                                          context=new_context)
        _, action_execution_db = action_services.request(new_live_action_db)
        return action_execution_db
<commit_msg>Log a message on invalid retry status in retry policy.<commit_after># Licensed to the StackStorm, Inc ('StackStorm') under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import copy
import functools

import eventlet

from st2common import log as logging
from st2common.models.db.liveaction import LiveActionDB
import st2common.services.action as action_services
from st2common.constants.action import LIVEACTION_STATUS_FAILED
from st2common.constants.action import LIVEACTION_STATUS_TIMED_OUT
from st2common.util.enum import Enum
from st2common.policies.base import ResourcePolicyApplicator

__all__ = [
    'RetryOnPolicy',
    'ExecutionRetryPolicyApplicator'
]

LOG = logging.getLogger(__name__)

VALID_RETRY_STATUSES = [
    LIVEACTION_STATUS_FAILED,
    LIVEACTION_STATUS_TIMED_OUT
]

class RetryOnPolicy(Enum):
    FAILURE = 'failure'  # Retry on execution failure
    TIMEOUT = 'timeout'  # Retry on execution timeout

class ExecutionRetryPolicyApplicator(ResourcePolicyApplicator):
    def __init__(self, policy_ref, policy_type, retry_on, max_retry_count=2, delay=0):
        """
        :param retry_on: Condition to retry the execution on (failure, timeout).
        :type retry_on: ``str``

        :param max_retry_count: Maximum number of times to try to retry an action.
        :type max_retry_count: ``int``

        :param delay: How long to wait before retrying an execution.
        :type delay: ``float``
        """
        super(ExecutionRetryPolicyApplicator, self).__init__(policy_ref=policy_ref,
                                                             policy_type=policy_type)

        self.retry_on = retry_on
        self.max_retry_count = max_retry_count
        self.delay = delay or 0

    def apply_after(self, target):
        target = super(ExecutionRetryPolicyApplicator, self).apply_after(target=target)

        live_action_db = target

        if self._is_live_action_part_of_workflow_action(live_action_db):
            LOG.warning(
                'Retry cannot be applied to this liveaction because it is executed under a '
                'workflow. Use workflow specific retry functionality where applicable. %s',
                live_action_db
            )

            return target

        retry_count = self._get_live_action_retry_count(live_action_db=live_action_db)

        extra = {'live_action_db': live_action_db, 'policy_ref': self._policy_ref,
                 'retry_on': self.retry_on, 'max_retry_count': self.max_retry_count,
                 'current_retry_count': retry_count}

        if live_action_db.status not in VALID_RETRY_STATUSES:
            # Currently we only support retrying on failed action
            LOG.debug('Liveaction not in a valid retry state, not checking retry policy',
                      extra=extra)
            return target

        if (retry_count + 1) > self.max_retry_count:
            LOG.info('Maximum retry count has been reached, not retrying', extra=extra)
            return target

        has_failed = live_action_db.status == LIVEACTION_STATUS_FAILED
        has_timed_out = live_action_db.status == LIVEACTION_STATUS_TIMED_OUT

        # TODO: This is not crash and restart safe, switch to using "DELAYED"
        # status
        if self.delay > 0:
            re_run_live_action = functools.partial(eventlet.spawn_after, self.delay,
                                                   self._re_run_live_action,
                                                   live_action_db=live_action_db)
        else:
            # Even if delay is 0, use a small delay (0.1 seconds) to prevent busy wait
            re_run_live_action = functools.partial(eventlet.spawn_after, 0.1,
                                                   self._re_run_live_action,
                                                   live_action_db=live_action_db)

            re_run_live_action = functools.partial(self._re_run_live_action,
                                                   live_action_db=live_action_db)

        if has_failed and self.retry_on == RetryOnPolicy.FAILURE:
            extra['failure'] = True
            LOG.info('Policy matched (failure), retrying action execution in %s seconds...' %
                     (self.delay), extra=extra)
            re_run_live_action()
            return target

        if has_timed_out and self.retry_on == RetryOnPolicy.TIMEOUT:
            extra['timeout'] = True
            LOG.info('Policy matched (timeout), retrying action execution in %s seconds...' %
                     (self.delay), extra=extra)
            re_run_live_action()
            return target

        LOG.info('Invalid status "%s" for live action "%s", wont retry' %
                (live_action_db.status, str(live_action_db.id)), extra=extra)

        return target

    def _is_live_action_part_of_workflow_action(self, live_action_db):
        """
        Retrieve parent info from context of the live action.

        :rtype: ``dict``
        """
        context = getattr(live_action_db, 'context', {})
        parent = context.get('parent', {})
        is_wf_action = (parent is not None and parent != {})

        return is_wf_action

    def _get_live_action_retry_count(self, live_action_db):
        """
        Retrieve current retry count for the provided live action.

        :rtype: ``int``
        """
        # TODO: Ideally we would store retry_count in zookeeper or similar and use locking so we
        # can run multiple instances of st2notififer
        context = getattr(live_action_db, 'context', {})
        retry_count = context.get('policies', {}).get('retry', {}).get('retry_count', 0)

        return retry_count

    def _re_run_live_action(self, live_action_db):
        retry_count = self._get_live_action_retry_count(live_action_db=live_action_db)

        # Add additional policy specific info to the context
        context = getattr(live_action_db, 'context', {})
        new_context = copy.deepcopy(context)
        new_context['policies'] = {}
        new_context['policies']['retry'] = {
            'applied_policy': self._policy_ref,
            'retry_count': (retry_count + 1),
            'retried_liveaction_id': str(live_action_db.id)
        }
        action_ref = live_action_db.action
        parameters = live_action_db.parameters
        new_live_action_db = LiveActionDB(action=action_ref, parameters=parameters,
                                          context=new_context)
        _, action_execution_db = action_services.request(new_live_action_db)
        return action_execution_db
<|endoftext|>


<commit_before># Licensed to the StackStorm, Inc ('StackStorm') under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys

from oslo_config import cfg

from st2common import config
from st2common import log as logging
from st2common.script_setup import setup as common_setup
from st2common.util.pack_management import download_pack
from st2common.util.virtualenvs import setup_pack_virtualenv

__all__ = [
    'main'
]

LOG = logging.getLogger(__name__)

def _do_register_cli_opts(opts, ignore_errors=False):
    for opt in opts:
        try:
            cfg.CONF.register_cli_opt(opt)
        except:
            if not ignore_errors:
                raise

def _register_cli_opts():
    cli_opts = [
        cfg.StrOpt('pack', default=None, required=True,
                   help='Name of the pack to install.'),
        cfg.BoolOpt('verify-ssl', default=True,
                   help=('Verify SSL certificate of the Git repo from which the pack is '
                         'downloaded.')),
        cfg.BoolOpt('force', default=False,
                    help='True to force pack installation and ignore install '
                         'lock file if it exists.'),
    ]
    _do_register_cli_opts(cli_opts)

def get_and_set_proxy_config():
    https_proxy = os.environ.get('https_proxy', None)
    http_proxy = os.environ.get('http_proxy', None)
    proxy_ca_bundle_path = os.environ.get('proxy_ca_bundle_path', None)
    no_proxy = os.environ.get('no_proxy', None)

    proxy_config = {}

    if http_proxy or https_proxy:
        LOG.debug('Using proxy %s', http_proxy if http_proxy else https_proxy)

        proxy_config = {
            'https_proxy': https_proxy,
            'http_proxy': http_proxy,
            'proxy_ca_bundle_path': proxy_ca_bundle_path,
            'no_proxy': no_proxy
        }

    if https_proxy and not os.environ.get('https_proxy', None):
        os.environ['https_proxy'] = https_proxy

    if http_proxy and not os.environ.get('http_proxy', None):
        os.environ['http_proxy'] = http_proxy

    if no_proxy and not os.environ.get('no_proxy', None):
        os.environ['no_proxy'] = no_proxy

    if proxy_ca_bundle_path and not os.environ.get('proxy_ca_bundle_path', None):
        os.environ['no_proxy'] = no_proxy

    return proxy_config

def main(argv):
    _register_cli_opts()

    # Parse CLI args, set up logging
    common_setup(config=config, setup_db=False, register_mq_exchanges=False,
                 register_internal_trigger_types=False)

    pack = cfg.CONF.pack
    verify_ssl = cfg.CONF.verify_ssl
    force = cfg.CONF.force

    proxy_config = get_and_set_proxy_config()

    # 1. Download the pac
    LOG.info('Installing pack "%s"' % (pack))
    result = download_pack(pack=pack, verify_ssl=verify_ssl, force=force,
                           proxy_config=proxy_config, force_permissions=True)

    success = result[2][0]

    if success:
        LOG.info('Successfuly installed pack "%s"' % (pack))
    else:
        error = result[2][1]
        LOG.error('Failed to install pack "%s": %s' % (pack, error))
        sys.exit(2)

    # 2. Setup pack virtual environment
    LOG.info('Setting up virtualenv for pack "%s"' % (pack))
    setup_pack_virtualenv(pack_name=pack, update=False, logger=LOG,
                          proxy_config=proxy_config, use_python3=False,
                          no_download=True)
    LOG.info('Successfuly set up virtualenv for pack "%s"' % (pack))
<commit_msg>Add support for installing multiple packs at once to "st2-pack-install" CLI script.<commit_after># Licensed to the StackStorm, Inc ('StackStorm') under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import sys

from oslo_config import cfg

from st2common import config
from st2common import log as logging
from st2common.script_setup import setup as common_setup
from st2common.util.pack_management import download_pack
from st2common.util.virtualenvs import setup_pack_virtualenv

__all__ = [
    'main'
]

LOG = logging.getLogger(__name__)

def _do_register_cli_opts(opts, ignore_errors=False):
    for opt in opts:
        try:
            cfg.CONF.register_cli_opt(opt)
        except:
            if not ignore_errors:
                raise

def _register_cli_opts():
    cli_opts = [
        cfg.MultiStrOpt('pack', default=None, required=True, positional=True,
                        help='Name of the pack to install.'),
        cfg.BoolOpt('verify-ssl', default=True,
                   help=('Verify SSL certificate of the Git repo from which the pack is '
                         'downloaded.')),
        cfg.BoolOpt('force', default=False,
                    help='True to force pack installation and ignore install '
                         'lock file if it exists.'),
    ]
    _do_register_cli_opts(cli_opts)

def get_and_set_proxy_config():
    https_proxy = os.environ.get('https_proxy', None)
    http_proxy = os.environ.get('http_proxy', None)
    proxy_ca_bundle_path = os.environ.get('proxy_ca_bundle_path', None)
    no_proxy = os.environ.get('no_proxy', None)

    proxy_config = {}

    if http_proxy or https_proxy:
        LOG.debug('Using proxy %s', http_proxy if http_proxy else https_proxy)

        proxy_config = {
            'https_proxy': https_proxy,
            'http_proxy': http_proxy,
            'proxy_ca_bundle_path': proxy_ca_bundle_path,
            'no_proxy': no_proxy
        }

    if https_proxy and not os.environ.get('https_proxy', None):
        os.environ['https_proxy'] = https_proxy

    if http_proxy and not os.environ.get('http_proxy', None):
        os.environ['http_proxy'] = http_proxy

    if no_proxy and not os.environ.get('no_proxy', None):
        os.environ['no_proxy'] = no_proxy

    if proxy_ca_bundle_path and not os.environ.get('proxy_ca_bundle_path', None):
        os.environ['no_proxy'] = no_proxy

    return proxy_config

def main(argv):
    _register_cli_opts()

    # Parse CLI args, set up logging
    common_setup(config=config, setup_db=False, register_mq_exchanges=False,
                 register_internal_trigger_types=False)

    packs = cfg.CONF.pack
    verify_ssl = cfg.CONF.verify_ssl
    force = cfg.CONF.force

    proxy_config = get_and_set_proxy_config()

    for pack in packs:
        # 1. Download the pack
        LOG.info('Installing pack "%s"' % (pack))
        result = download_pack(pack=pack, verify_ssl=verify_ssl, force=force,
                               proxy_config=proxy_config, force_permissions=True)

        success = result[2][0]

        if success:
            LOG.info('Successfuly installed pack "%s"' % (pack))
        else:
            error = result[2][1]
            LOG.error('Failed to install pack "%s": %s' % (pack, error))
            sys.exit(2)

        # 2. Setup pack virtual environment
        LOG.info('Setting up virtualenv for pack "%s"' % (pack))
        setup_pack_virtualenv(pack_name=pack, update=False, logger=LOG,
                              proxy_config=proxy_config, use_python3=False,
                              no_download=True)
        LOG.info('Successfuly set up virtualenv for pack "%s"' % (pack))

    return 0
<|endoftext|>


<commit_before>import json
import os
import logging
from grapheneexchange import GrapheneExchange
from stakemachine.storage import Storage
log = logging.getLogger(__name__)

class MissingSettingsException(Exception):
    pass

class BaseStrategy():
    """ This is the base strategy that allows to share commonly used
        methods, such as sell, buy, borrow, cancel, and many more!

        The purpose of this class is to inherit or even overwrite some
        of this calls in your individual bots it required.

        .. note:: Each bot has it's own markets that it serves, hence,
                 we distinguish between `all` orders of a market that
                 are owned by an account and `mine` which are merely the
                 orders that have been created in a market by this bot
                 specifically. For this reasons, every bot stores it's
                 orders in a `json` file on the disk to be able to
                 distinguish its own orders from others!
    """

    def __init__(self, *args, **kwargs):
        self.state = {"orders" : {}}

        for arg in args :
            if isinstance(arg, GrapheneExchange):
                self.dex = arg
        for key in kwargs:
            setattr(self, key, kwargs[key])

        if "name" not in kwargs:
            raise MissingSettingsException("Missing parameter 'name'!")

        self.settings = self.config.bots[self.name]
        self.opened_orders = []
        self.storage = Storage(self.name, self.config)
        self.restore()

        if "markets" not in self.settings:
            raise MissingSettingsException("markets")

    def _cancel_set(self, toCancel):
        numCanceled = 0
        for orderId in toCancel:
            log.info("Canceling %s" % orderId)
            self.dex.cancel(orderId)
            self.orderCanceled(orderId)
            for m in self.state["orders"]:
                if orderId in self.state["orders"][m]:
                    self.state["orders"][m].remove(orderId)
            numCanceled += 1
        return numCanceled

    def cancel_all(self, markets=None, side="both") :
        """ Cancel all the account's orders **of all market** including
            those orders of other bot instances

            :param str side: cancel only "sell", "buy", "both" side(s) (default: "both")
            :return: number of canceld orders
            :rtype: number

        """
        if not markets:
            markets = self.settings["markets"]
        curOrders = self.dex.returnOpenOrders()
        toCancel = set()
        for m in markets:
            if m in curOrders:
                for o in curOrders[m]:
                    if o["type"] is side or side is "both":
                        toCancel.add(o["orderNumber"])
        return self._cancel_set(toCancel)

    def cancel_mine(self, markets=None, side="both") :
        """ Cancel only the orders of this particular bot in all markets
            the bot serves.

            :param str side: cancel only "sell", "buy", "both" side(s) (default: "both")
            :return: number of canceld orders
            :rtype: number
        """
        if not markets:
            markets = self.settings["markets"]
        curOrders = self.dex.returnOpenOrders()
        state = self.getState()
        toCancel = set()
        for m in markets:
            for currentOrderStates in curOrders[m]:
                stateOrderId = currentOrderStates["orderNumber"]
                if m not in state["orders"]:
                    continue
                if stateOrderId in state["orders"][m]:
                    if currentOrderStates["type"] is side or side is "both":
                        toCancel.add(currentOrderStates["orderNumber"])
        return self._cancel_set(toCancel)

    def cancel_this_markets(self, markets=None, side="both") :
        """ Cancel all orders in all markets of that are served by this
            bot.

            :param str side: cancel only "sell", "buy", "both" side(s) (default: "both")
            :return: number of canceld orders
            :rtype: number
        """
        if not markets:
            markets = self.settings["markets"]
        orders = self.dex.returnOpenOrders()
        toCancel = set()
        for m in markets:
            for o in orders[m]:
                if o["type"] is side or side is "both":
                    toCancel.add(o["orderNumber"])
        return self._cancel_set(toCancel)

    def cancel_all_sell_orders(self):
        """ alias for ``self.cancel_all("sell")``
        """
        return self.cancel_all("sell")

    def cancel_all_buy_orders(self):
        """ alias for ``self.cancel_all("buy")``
        """
        return self.cancel_all("buy")

    def cancel_my_sell_orders(self):
        """ alias for ``self.cancel_mine("sell")``
        """
        return self.cancel_mine("sell")

    def cancel_my_buy_orders(self):
        """ alias for ``self.cancel_mine("buy")``
        """
        return self.cancel_mine("buy")

    def cancel_all_bid_orders(self):
        """ alias for ``self.cancel_all("buy")``
        """
        return self.cancel_all("buy")

    def cancel_all_ask_orders(self):
        """ alias for ``self.cancel_all("sell")``
        """
        return self.cancel_all("sell")

    def cancel_my_bid_orders(self):
        """ alias for ``self.cancel_my_buys()``
        """
        return self.cancel_my_buys()

    def cancel_my_ask_orders(self):
        """ alias for ``self.cancel_my_sells()``
        """
        return self.cancel_my_sells()

    def cancel(self, orderid):
        """ Cancel the order with id ``orderid``
        """
        log.info("Canceling %s" % orderid)
        return self.dex.cancel(orderid)

    def getState(self):
        """ Return the stored state of the bot. This includes the
            ``orders`` that have been placed by this bot
        """
        return self.state

    def setState(self, key, value):
        """ Set the full state

            :param str key: Key
            :param Object value: Value
        """
        self.state[key] = value

    def store(self):
        """ Evaluate the changes (orders) made by the bot and store the
            state on disk.
        """
        state = self.getState()
        myorders = state["orders"]
        curOrders = self.dex.returnOpenOrdersIds()
        for market in self.settings["markets"] :
            if market not in myorders:
                myorders[market] = []
            if market in curOrders:
                for orderid in curOrders[market] :
                    if market not in self.opened_orders or \
                            orderid not in self.opened_orders[market] :
                        myorders[market].append(orderid)
                        self.orderPlaced(orderid)

        state["orders"] = myorders
        self.storage.store(state)

    def restore(self):
        """ Restore the data stored on the disk
        """
        self.state = self.storage.restore()

    def loadMarket(self, notify=True):
        """ Load the markets and compare the stored orders with the
            still open orders. Calls ``orderFilled(orderid)`` for orders no
            longer open (i.e. fully filled)
        """
        #: Load Open Orders for the markets and store them for later
        self.opened_orders = self.dex.returnOpenOrdersIds()

        #: Have orders been matched?
        old_orders = self.getState()["orders"]
        cur_orders = self.dex.returnOpenOrdersIds()
        for market in self.settings["markets"] :
            if market in old_orders:
                for orderid in old_orders[market] :
                    if orderid not in cur_orders[market] :
                        # Remove it from the state
                        if orderid in self.state["orders"][market]:
                            self.state["orders"][market].remove(orderid)
                        # Execute orderFilled call
                        if notify :
                            self.orderFilled(orderid)

    def getMyOrders(self):
        """ Return open orders for this bot
        """
        myOrders = {}
        for market in self.settings["markets"] :
            if market in self.state["orders"]:
                myOrders[market] = self.state["orders"][market]
            else:
                myOrders[market] = []
        return myOrders

    def returnBalances(self):
        """ This is a wrapper for self.dex.returnBalances()
            that limits the amounts such that the reserves are always
            available in the account. The reserves are defined in the
            configuration with:

            ```
            reserves:
             - BTS: 1000
             - USD: 10
            ```
        """
        balances = self.dex.returnBalances()
        if not hasattr(self.config, "reserves"):
            return balances

        reserves = getattr(self.config, "reserves")
        for a in balances:
            balances[a] -= reserves.get(a, 0)
            if balances[a] < 0:
                balances[a] = 0.0
        return balances

    def sell(self, market, price, amount, expiration=60 * 60 * 24, **kwargs):
        """ Places a sell order in a given market (sell ``quote``, buy
            ``base`` in market ``quote_base``). Required POST parameters
            are "currencyPair", "rate", and "amount". If successful, the
            method will return the order creating (signed) transaction.

            :param str currencyPair: Return results for a particular market only (default: "all")
            :param float price: price denoted in ``base``/``quote``
            :param number amount: Amount of ``quote`` to sell

            Prices/Rates are denoted in 'base', i.e. the USD_BTS market
            is priced in BTS per USD.

            **Example:** in the USD_BTS market, a price of 300 means
            a USD is worth 300 BTS

            .. note::

                All prices returned are in the **reveresed** orientation as the
                market. I.e. in the BTC/BTS market, prices are BTS per BTS.
                That way you can multiply prices with `1.05` to get a +5%.
        """
        quote, base = market.split(self.config.market_separator)
        log.warning(" - Selling %f %s for %s @%f %s/%s" % (amount, quote, base, price, base, quote))
        self.dex.sell(market, price, amount, expiration, **kwargs)

    def buy(self, market, price, amount, expiration=60 * 60 * 24, **kwargs):
        """ Places a buy order in a given market (buy ``quote``, sell
            ``base`` in market ``quote_base``). Required POST parameters
            are "currencyPair", "rate", and "amount". If successful, the
            method will return the order creating (signed) transaction.

            :param str currencyPair: Return results for a particular market only (default: "all")
            :param float price: price denoted in ``base``/``quote``
            :param number amount: Amount of ``quote`` to buy

            Prices/Rates are denoted in 'base', i.e. the USD_BTS market
            is priced in BTS per USD.

            **Example:** in the USD_BTS market, a price of 300 means
            a USD is worth 300 BTS

            .. note::

                All prices returned are in the **reveresed** orientation as the
                market. I.e. in the BTC/BTS market, prices are BTS per BTS.
                That way you can multiply prices with `1.05` to get a +5%.
        """
        quote, base = market.split(self.config.market_separator)
        log.warning(" - Buying %f %s with %s @%f %s/%s" % (amount, quote, base, price, base, quote))
        self.dex.buy(market, price, amount, expiration, **kwargs)

    def init(self) :
        """ Initialize the bot's individual settings
        """
        log.debug("Init. Please define `%s.init()`" % self.name)

    def place(self) :
        """ Place orders
        """
        log.debug("Place order. Please define `%s.place()`" % self.name)

    def tick(self) :
        """ Tick every block
        """
        log.debug("New block. Please define `%s.tick()`" % self.name)

    def asset_tick(self) :
        """ Tick every block
        """
        log.debug("Asset Updated. Please define `%s.asset_tick()`" % self.name)

    def orderFilled(self, oid):
        """ An order has been fully filled

            :param str oid: The order object id
        """
        log.debug("Order Filled. Please define `%s.orderFilled(%s)`" % (self.name, oid))

#    def orderMatched(self, oid):
#        """ An order has been machted / partially filled
#
#            :param str oid: The order object id
#        """
#        log.debug("An order has been matched: %s" % oid)

    def orderPlaced(self, oid):
        """ An order has been placed

            :param str oid: The order object id
        """
        log.debug("New Order. Please define `%s.orderPlaced(%s)`" % (self.name, oid))

    def orderCanceled(self, oid):
        """ An order has been canceld

            :param str oid: The order object id
        """
        log.debug("Order Canceld. Please define `%s.orderCanceled(%s)`" % (self.name, oid))
<commit_msg>[base] do not store state in safe_mode<commit_after>import json
import os
import logging
from grapheneexchange import GrapheneExchange
from stakemachine.storage import Storage
log = logging.getLogger(__name__)

class MissingSettingsException(Exception):
    pass

class BaseStrategy():
    """ This is the base strategy that allows to share commonly used
        methods, such as sell, buy, borrow, cancel, and many more!

        The purpose of this class is to inherit or even overwrite some
        of this calls in your individual bots it required.

        .. note:: Each bot has it's own markets that it serves, hence,
                 we distinguish between `all` orders of a market that
                 are owned by an account and `mine` which are merely the
                 orders that have been created in a market by this bot
                 specifically. For this reasons, every bot stores it's
                 orders in a `json` file on the disk to be able to
                 distinguish its own orders from others!
    """

    def __init__(self, *args, **kwargs):
        self.state = {"orders" : {}}

        for arg in args :
            if isinstance(arg, GrapheneExchange):
                self.dex = arg
        for key in kwargs:
            setattr(self, key, kwargs[key])

        if "name" not in kwargs:
            raise MissingSettingsException("Missing parameter 'name'!")

        self.settings = self.config.bots[self.name]
        self.opened_orders = []
        self.storage = Storage(self.name, self.config)
        self.restore()

        if "markets" not in self.settings:
            raise MissingSettingsException("markets")

    def _cancel_set(self, toCancel):
        numCanceled = 0
        for orderId in toCancel:
            log.info("Canceling %s" % orderId)
            self.dex.cancel(orderId)
            self.orderCanceled(orderId)
            for m in self.state["orders"]:
                if orderId in self.state["orders"][m]:
                    self.state["orders"][m].remove(orderId)
            numCanceled += 1
        return numCanceled

    def cancel_all(self, markets=None, side="both") :
        """ Cancel all the account's orders **of all market** including
            those orders of other bot instances

            :param str side: cancel only "sell", "buy", "both" side(s) (default: "both")
            :return: number of canceld orders
            :rtype: number

        """
        if not markets:
            markets = self.settings["markets"]
        curOrders = self.dex.returnOpenOrders()
        toCancel = set()
        for m in markets:
            if m in curOrders:
                for o in curOrders[m]:
                    if o["type"] is side or side is "both":
                        toCancel.add(o["orderNumber"])
        return self._cancel_set(toCancel)

    def cancel_mine(self, markets=None, side="both") :
        """ Cancel only the orders of this particular bot in all markets
            the bot serves.

            :param str side: cancel only "sell", "buy", "both" side(s) (default: "both")
            :return: number of canceld orders
            :rtype: number
        """
        if not markets:
            markets = self.settings["markets"]
        curOrders = self.dex.returnOpenOrders()
        state = self.getState()
        toCancel = set()
        for m in markets:
            for currentOrderStates in curOrders[m]:
                stateOrderId = currentOrderStates["orderNumber"]
                if m not in state["orders"]:
                    continue
                if stateOrderId in state["orders"][m]:
                    if currentOrderStates["type"] is side or side is "both":
                        toCancel.add(currentOrderStates["orderNumber"])
        return self._cancel_set(toCancel)

    def cancel_this_markets(self, markets=None, side="both") :
        """ Cancel all orders in all markets of that are served by this
            bot.

            :param str side: cancel only "sell", "buy", "both" side(s) (default: "both")
            :return: number of canceld orders
            :rtype: number
        """
        if not markets:
            markets = self.settings["markets"]
        orders = self.dex.returnOpenOrders()
        toCancel = set()
        for m in markets:
            for o in orders[m]:
                if o["type"] is side or side is "both":
                    toCancel.add(o["orderNumber"])
        return self._cancel_set(toCancel)

    def cancel_all_sell_orders(self):
        """ alias for ``self.cancel_all("sell")``
        """
        return self.cancel_all("sell")

    def cancel_all_buy_orders(self):
        """ alias for ``self.cancel_all("buy")``
        """
        return self.cancel_all("buy")

    def cancel_my_sell_orders(self):
        """ alias for ``self.cancel_mine("sell")``
        """
        return self.cancel_mine("sell")

    def cancel_my_buy_orders(self):
        """ alias for ``self.cancel_mine("buy")``
        """
        return self.cancel_mine("buy")

    def cancel_all_bid_orders(self):
        """ alias for ``self.cancel_all("buy")``
        """
        return self.cancel_all("buy")

    def cancel_all_ask_orders(self):
        """ alias for ``self.cancel_all("sell")``
        """
        return self.cancel_all("sell")

    def cancel_my_bid_orders(self):
        """ alias for ``self.cancel_my_buys()``
        """
        return self.cancel_my_buys()

    def cancel_my_ask_orders(self):
        """ alias for ``self.cancel_my_sells()``
        """
        return self.cancel_my_sells()

    def cancel(self, orderid):
        """ Cancel the order with id ``orderid``
        """
        log.info("Canceling %s" % orderid)
        return self.dex.cancel(orderid)

    def getState(self):
        """ Return the stored state of the bot. This includes the
            ``orders`` that have been placed by this bot
        """
        return self.state

    def setState(self, key, value):
        """ Set the full state

            :param str key: Key
            :param Object value: Value
        """
        self.state[key] = value

    def store(self):
        """ Evaluate the changes (orders) made by the bot and store the
            state on disk.
        """
        state = self.getState()
        myorders = state["orders"]
        curOrders = self.dex.returnOpenOrdersIds()
        for market in self.settings["markets"] :
            if market not in myorders:
                myorders[market] = []
            if market in curOrders:
                for orderid in curOrders[market] :
                    if market not in self.opened_orders or \
                            orderid not in self.opened_orders[market] :
                        myorders[market].append(orderid)
                        self.orderPlaced(orderid)

        state["orders"] = myorders
        if not self.config.safe_mode:
            self.storage.store(state)

    def restore(self):
        """ Restore the data stored on the disk
        """
        self.state = self.storage.restore()

    def loadMarket(self, notify=True):
        """ Load the markets and compare the stored orders with the
            still open orders. Calls ``orderFilled(orderid)`` for orders no
            longer open (i.e. fully filled)
        """
        #: Load Open Orders for the markets and store them for later
        self.opened_orders = self.dex.returnOpenOrdersIds()

        #: Have orders been matched?
        old_orders = self.getState()["orders"]
        cur_orders = self.dex.returnOpenOrdersIds()
        for market in self.settings["markets"] :
            if market in old_orders:
                for orderid in old_orders[market] :
                    if orderid not in cur_orders[market] :
                        # Remove it from the state
                        if orderid in self.state["orders"][market]:
                            self.state["orders"][market].remove(orderid)
                        # Execute orderFilled call
                        if notify :
                            self.orderFilled(orderid)

    def getMyOrders(self):
        """ Return open orders for this bot
        """
        myOrders = {}
        for market in self.settings["markets"] :
            if market in self.state["orders"]:
                myOrders[market] = self.state["orders"][market]
            else:
                myOrders[market] = []
        return myOrders

    def returnBalances(self):
        """ This is a wrapper for self.dex.returnBalances()
            that limits the amounts such that the reserves are always
            available in the account. The reserves are defined in the
            configuration with:

            ```
            reserves:
             - BTS: 1000
             - USD: 10
            ```
        """
        balances = self.dex.returnBalances()
        if not hasattr(self.config, "reserves"):
            return balances

        reserves = getattr(self.config, "reserves")
        for a in balances:
            balances[a] -= reserves.get(a, 0)
            if balances[a] < 0:
                balances[a] = 0.0
        return balances

    def sell(self, market, price, amount, expiration=60 * 60 * 24, **kwargs):
        """ Places a sell order in a given market (sell ``quote``, buy
            ``base`` in market ``quote_base``). Required POST parameters
            are "currencyPair", "rate", and "amount". If successful, the
            method will return the order creating (signed) transaction.

            :param str currencyPair: Return results for a particular market only (default: "all")
            :param float price: price denoted in ``base``/``quote``
            :param number amount: Amount of ``quote`` to sell

            Prices/Rates are denoted in 'base', i.e. the USD_BTS market
            is priced in BTS per USD.

            **Example:** in the USD_BTS market, a price of 300 means
            a USD is worth 300 BTS

            .. note::

                All prices returned are in the **reveresed** orientation as the
                market. I.e. in the BTC/BTS market, prices are BTS per BTS.
                That way you can multiply prices with `1.05` to get a +5%.
        """
        quote, base = market.split(self.config.market_separator)
        log.warning(" - Selling %f %s for %s @%f %s/%s" % (amount, quote, base, price, base, quote))
        self.dex.sell(market, price, amount, expiration, **kwargs)

    def buy(self, market, price, amount, expiration=60 * 60 * 24, **kwargs):
        """ Places a buy order in a given market (buy ``quote``, sell
            ``base`` in market ``quote_base``). Required POST parameters
            are "currencyPair", "rate", and "amount". If successful, the
            method will return the order creating (signed) transaction.

            :param str currencyPair: Return results for a particular market only (default: "all")
            :param float price: price denoted in ``base``/``quote``
            :param number amount: Amount of ``quote`` to buy

            Prices/Rates are denoted in 'base', i.e. the USD_BTS market
            is priced in BTS per USD.

            **Example:** in the USD_BTS market, a price of 300 means
            a USD is worth 300 BTS

            .. note::

                All prices returned are in the **reveresed** orientation as the
                market. I.e. in the BTC/BTS market, prices are BTS per BTS.
                That way you can multiply prices with `1.05` to get a +5%.
        """
        quote, base = market.split(self.config.market_separator)
        log.warning(" - Buying %f %s with %s @%f %s/%s" % (amount, quote, base, price, base, quote))
        self.dex.buy(market, price, amount, expiration, **kwargs)

    def init(self) :
        """ Initialize the bot's individual settings
        """
        log.debug("Init. Please define `%s.init()`" % self.name)

    def place(self) :
        """ Place orders
        """
        log.debug("Place order. Please define `%s.place()`" % self.name)

    def tick(self) :
        """ Tick every block
        """
        log.debug("New block. Please define `%s.tick()`" % self.name)

    def asset_tick(self) :
        """ Tick every block
        """
        log.debug("Asset Updated. Please define `%s.asset_tick()`" % self.name)

    def orderFilled(self, oid):
        """ An order has been fully filled

            :param str oid: The order object id
        """
        log.debug("Order Filled. Please define `%s.orderFilled(%s)`" % (self.name, oid))

#    def orderMatched(self, oid):
#        """ An order has been machted / partially filled
#
#            :param str oid: The order object id
#        """
#        log.debug("An order has been matched: %s" % oid)

    def orderPlaced(self, oid):
        """ An order has been placed

            :param str oid: The order object id
        """
        log.debug("New Order. Please define `%s.orderPlaced(%s)`" % (self.name, oid))

    def orderCanceled(self, oid):
        """ An order has been canceld

            :param str oid: The order object id
        """
        log.debug("Order Canceld. Please define `%s.orderCanceled(%s)`" % (self.name, oid))
<|endoftext|>


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

# Copyright 2015 University of Chicago

import sys
import argparse
import logging
import pytz
import os
import pwd
import grp
import datetime

import elasticsearch
import elasticsearch.helpers

ES_NODES = ['uct2-es-head.mwt2.org:9200', 'uct2-es-door.mwt2.org:9200']
TZ_NAME = "US/Central"
TIMEZONE = pytz.timezone(TZ_NAME)

def get_es_client():
    """ Instantiate DB client and pass connection back """
    return elasticsearch.Elasticsearch(hosts=ES_NODES,
                                       retry_on_timeout=True,
                                       max_retries=10,
                                       timeout=300)

def create_record(dirpath, num_files, date, dir_size=None,
                  index=None, ceph_fs=False):
    """
    Query directory at dirpath for information and store information in ES

    :param dirpath: string storing path to file
    :parm num_files: number of files present in dirpath
    :param date: datetime.date with date to use
    :param dir_size: size of directory and contents in bytes
    :param index: ES index for records
    :param ceph_fs: Is the directory on a CephFS filesystem
    :return: dictionary with record if successful, empty dict otherwise
    """
    if not os.path.isdir(dirpath):
        return {}
    if not index:
        index = "stash-space-{0}-{1:0>2}".format(date.year, date.month)
    dir_info = os.stat(dirpath)
    uid = dir_info.st_uid
    gid = dir_info.st_gid
    ctime = datetime.datetime.fromtimestamp(dir_info.st_ctime, TIMEZONE)
    mtime = datetime.datetime.fromtimestamp(dir_info.st_mtime, TIMEZONE)
    record_fields = {'@timestamp': date.isoformat(),
                     'size': dir_size,
                     'num_files': num_files,
                     'ctime': ctime,
                     'mtime': mtime,
                     'user': pwd.getpwuid(uid).pw_name,
                     'group': grp.getgrgid(gid).gr_name,
                     'path': dirpath}
    record = {'_index': index,
              '_source': record_fields,
              '_op_type': 'index',
              '_type': 'space_usage_record'}
    return record

def get_ceph_dir_size(root):
    """
    Get the space used by a directory and it's contents on a CephFS
    :param root: path to directory
    :return: number of bytes used by the directory and it's contents
    """
    return os.stat(root, follow_symlinks=False).st_size

def get_dir_size(root, dir_sizes, inodes):
    """
    Get the size of a directory and it's contents

    :param root: path to the directory
    :param dir_sizes: a dictionary of directories with sizes of those
                      directories
    :param inodes: a set with list of inodes that have been visited
    :return: number of bytes used by the directory and it's contents
    """

    total_size = 0
    entries = os.listdir(root)
    for entry in entries:
        entry_name = os.path.join(root, entry)
        inode = os.stat(entry_name).st_ino
        if inode in inodes:
            continue
        if os.path.isfile(entry_name) and not os.path.islink(entry_name):
            try:
                total_size += os.stat(entry_name).st_size
                inodes.add(inode)
            except OSError:
                continue
        elif os.path.isdir(entry_name) and not os.path.islink(entry_name):
            if entry_name in dir_sizes:
                total_size += dir_sizes[entry_name]
                continue
            dir_size = get_dir_size(entry_name, dir_sizes, inodes)
            dir_sizes[entry_name] = dir_size
            inodes.add(inode)
    return total_size

def traverse_directory(dirpath, index=None, ceph_fs=False):
    """
    Traverse subdirectories and create a set of records
    with space usage
    :param dirpath: path to directory to
    :param index: the ES index to store records in
    :param ceph_fs: Whether the dirpath is a CephFS directory
    :return: Nothing
    """
    current_date = TIMEZONE.localize(datetime.datetime.combine(datetime.date.today(),
                                                               datetime.time(0, 0, 0)))

    if not os.path.isdir(dirpath):
        return
    records = []
    if ceph_fs:
        topdown = True
    else:
        topdown = False
        dir_sizes = {}
    for root, dirs, files in os.walk(dirpath, topdown=topdown):
        if not ceph_fs:
            inodes = set()
            size = get_dir_size(root, dir_sizes, inodes)
        else:
            size = get_ceph_dir_size(root)
        record = create_record(root, len(files), current_date, size,
                               index, ceph_fs)
        if record:
            records.append(record)
    save_records(records)

def save_records(records=None):
    """
    Save a job record to ES
    """
    es_client = get_es_client()
    elasticsearch.helpers.bulk(es_client,
                               records,
                               stats_only=False)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Get space usage on a given '
                                                 'filesystem and store in ES')
    parser.add_argument('--index', dest='index', default=None,
                        help='ES index to store records in')
    parser.add_argument('--is-ceph', dest='ceph_fs', default=False,
                        type=bool, action='store_true',
                        help='Is the directory a CephFS directory?')
    parser.add_argument("directory", default=None,
                        help="Directory to examine")
    args = parser.parse_args(sys.argv[1:])
    if not os.path.isdir(args.directory):
        sys.stderr.write("{0} must be a directory\n".format(args.directory))
        sys.exit(1)
    traverse_directory(args.directory, args.index, args.ceph_fs)
<commit_msg>Handle missing entries for passwd and groups, other fixes<commit_after>#!/usr/bin/env python

# Copyright 2015 University of Chicago

import sys
import argparse
import logging
import pytz
import os
import pwd
import grp
import datetime

import elasticsearch
import elasticsearch.helpers

ES_NODES = ['uct2-es-head.mwt2.org:9200', 'uct2-es-door.mwt2.org:9200']
TZ_NAME = "US/Central"
TIMEZONE = pytz.timezone(TZ_NAME)

def get_es_client():
    """ Instantiate DB client and pass connection back """
    return elasticsearch.Elasticsearch(hosts=ES_NODES,
                                       retry_on_timeout=True,
                                       max_retries=10,
                                       timeout=300)

def create_record(dirpath, num_files, date, dir_size=None,
                  index=None, ceph_fs=False):
    """
    Query directory at dirpath for information and store information in ES

    :param dirpath: string storing path to file
    :parm num_files: number of files present in dirpath
    :param date: datetime.date with date to use
    :param dir_size: size of directory and contents in bytes
    :param index: ES index for records
    :param ceph_fs: Is the directory on a CephFS filesystem
    :return: dictionary with record if successful, empty dict otherwise
    """
    if not os.path.isdir(dirpath):
        return {}
    if not index:
        index = "stash-space-{0}-{1:0>2}".format(date.year, date.month)
    dir_info = os.stat(dirpath)
    uid = dir_info.st_uid
    gid = dir_info.st_gid
    ctime = datetime.datetime.fromtimestamp(dir_info.st_ctime, TIMEZONE)
    mtime = datetime.datetime.fromtimestamp(dir_info.st_mtime, TIMEZONE)
    try:
       group = grp.getgrgid(gid).gr_name
    except KeyError:
       group = gid
    try:
       user = pwd.getpwuid(uid).pw_name
    except KeyError:
       user = uid
    record_fields = {'@timestamp': date.isoformat(),
                     'size': dir_size,
                     'num_files': num_files,
                     'ctime': ctime.isoformat(),
                     'mtime': mtime.isoformat(),
                     'user': user,
                     'group': group,
                     'path': dirpath}
    record = {'_index': index,
              '_source': record_fields,
              '_op_type': 'index',
              '_type': 'space_usage_record'}
    return record

def get_ceph_dir_size(root):
    """
    Get the space used by a directory and it's contents on a CephFS
    :param root: path to directory
    :return: number of bytes used by the directory and it's contents
    """
    return os.stat(root, follow_symlinks=False).st_size

def get_dir_size(root, dir_sizes, inodes):
    """
    Get the size of a directory and it's contents

    :param root: path to the directory
    :param dir_sizes: a dictionary of directories with sizes of those
                      directories
    :param inodes: a set with list of inodes that have been visited
    :return: number of bytes used by the directory and it's contents
    """

    total_size = 0
    entries = os.listdir(root)
    for entry in entries:
        entry_name = os.path.join(root, entry)
        try:
            entry_stat = os.stat(entry_name)
        except OSError:
            continue
        if entry_stat.st_ino in inodes:
            continue
        if os.path.isfile(entry_name) and not os.path.islink(entry_name):
            try:
                total_size += entry_stat.st_size
                inodes.add(entry_stat.st_ino)
            except OSError:
                continue
        elif os.path.isdir(entry_name) and not os.path.islink(entry_name):
            if entry_name in dir_sizes:
                total_size += dir_sizes[entry_name]
                continue
            dir_size = get_dir_size(entry_name, dir_sizes, inodes)
            dir_size += entry_stat.st_size
            dir_sizes[entry_name] = dir_size
            total_size += dir_size
            inodes.add(entry_stat.st_ino)
    total_size += os.stat(root).st_size
    dir_sizes[root] = total_size
    return total_size

def traverse_directory(dirpath, index=None, ceph_fs=False):
    """
    Traverse subdirectories and create a set of records
    with space usage
    :param dirpath: path to directory to
    :param index: the ES index to store records in
    :param ceph_fs: Whether the dirpath is a CephFS directory
    :return: Nothing
    """
    current_date = TIMEZONE.localize(datetime.datetime.combine(datetime.date.today(),
                                                               datetime.time(0, 0, 0)))

    if not os.path.isdir(dirpath):
        return
    records = []
    if ceph_fs:
        topdown = True
    else:
        topdown = False
        dir_sizes = {}
    for root, dirs, files in os.walk(dirpath, topdown=topdown):
        if not ceph_fs:
            inodes = set()
            size = get_dir_size(root, dir_sizes, inodes)
        else:
            size = get_ceph_dir_size(root)
        record = create_record(root, len(files), current_date, size,
                               index, ceph_fs)
        if record:
            records.append(record)
    save_records(records)

def save_records(records=None):
    """
    Save a job record to ES
    """
    es_client = get_es_client()
    print elasticsearch.helpers.bulk(es_client,
                               records,
                               stats_only=False)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Get space usage on a given '
                                                 'filesystem and store in ES')
    parser.add_argument('--index', dest='index', default=None,
                        help='ES index to store records in')
    parser.add_argument('--is-ceph', dest='ceph_fs', default=False,
                        action='store_true',
                        help='Is the directory a CephFS directory?')
    parser.add_argument("directory", default=None,
                        help="Directory to examine")
    args = parser.parse_args(sys.argv[1:])
    if not os.path.isdir(args.directory):
        sys.stderr.write("{0} must be a directory\n".format(args.directory))
        sys.exit(1)
    traverse_directory(args.directory, args.index, args.ceph_fs)
<|endoftext|>


<commit_before>import numpy as np
from scipy.stats import scoreatpercentile as sap

from statsmodels.compat.pandas import Substitution
from statsmodels.sandbox.nonparametric import kernels

def _select_sigma(X):
    """
    Returns the smaller of std(X, ddof=1) or normalized IQR(X) over axis 0.

    References
    ----------
    Silverman (1986) p.47
    """
#    normalize = norm.ppf(.75) - norm.ppf(.25)
    normalize = 1.349
#    IQR = np.subtract.reduce(percentile(X, [75,25],
#                             axis=axis), axis=axis)/normalize
    IQR = (sap(X, 75) - sap(X, 25))/normalize
    return np.minimum(np.std(X, axis=0, ddof=1), IQR)

## Univariate Rule of Thumb Bandwidths ##
def bw_scott(x, kernel=None):
    """
    Scott's Rule of Thumb

    Parameters
    ----------
    x : array_like
        Array for which to get the bandwidth
    kernel : CustomKernel object
        Unused

    Returns
    -------
    bw : float
        The estimate of the bandwidth

    Notes
    -----
    Returns 1.059 * A * n ** (-1/5.) where ::

       A = min(std(x, ddof=1), IQR/1.349)
       IQR = np.subtract.reduce(np.percentile(x, [75,25]))

    References
    ----------

    Scott, D.W. (1992) Multivariate Density Estimation: Theory, Practice, and
        Visualization.
    """
    A = _select_sigma(x)
    n = len(x)
    return 1.059 * A * n ** (-0.2)

def bw_silverman(x, kernel=None):
    """
    Silverman's Rule of Thumb

    Parameters
    ----------
    x : array_like
        Array for which to get the bandwidth
    kernel : CustomKernel object
        Unused

    Returns
    -------
    bw : float
        The estimate of the bandwidth

    Notes
    -----
    Returns .9 * A * n ** (-1/5.) where ::

       A = min(std(x, ddof=1), IQR/1.349)
       IQR = np.subtract.reduce(np.percentile(x, [75,25]))

    References
    ----------

    Silverman, B.W. (1986) `Density Estimation.`
    """
    A = _select_sigma(x)
    n = len(x)
    return .9 * A * n ** (-0.2)

def bw_normal_reference(x, kernel=kernels.Gaussian):
    """
    Plug-in bandwidth with kernel specific constant based on normal reference.

    This bandwidth minimizes the mean integrated square error if the true
    distribution is the normal. This choice is an appropriate bandwidth for
    single peaked distributions that are similar to the normal distribution.

    Parameters
    ----------
    x : array_like
        Array for which to get the bandwidth
    kernel : CustomKernel object
        Used to calculate the constant for the plug-in bandwidth.

    Returns
    -------
    bw : float
        The estimate of the bandwidth

    Notes
    -----
    Returns C * A * n ** (-1/5.) where ::

       A = min(std(x, ddof=1), IQR/1.349)
       IQR = np.subtract.reduce(np.percentile(x, [75,25]))
       C = constant from Hansen (2009)

    When using a Gaussian kernel this is equivalent to the 'scott' bandwidth up
    to two decimal places. This is the accuracy to which the 'scott' constant is
    specified.

    References
    ----------

    Silverman, B.W. (1986) `Density Estimation.`
    Hansen, B.E. (2009) `Lecture Notes on Nonparametrics.`
    """
    C = kernel.normal_reference_constant
    A = _select_sigma(x)
    n = len(x)
    return C * A * n ** (-0.2)

## Plug-In Methods ##

## Least Squares Cross-Validation ##

## Helper Functions ##

bandwidth_funcs = {
    "scott": bw_scott,
    "silverman": bw_silverman,
    "normal_reference": bw_normal_reference,
}

@Substitution(", ".join(sorted(bandwidth_funcs.keys())))
def select_bandwidth(x, bw, kernel):
    """
    Selects bandwidth for a selection rule bw

    this is a wrapper around existing bandwidth selection rules

    Parameters
    ----------
    x : array_like
        Array for which to get the bandwidth
    bw : str
        name of bandwidth selection rule, currently supported are:
        %s
    kernel : not used yet

    Returns
    -------
    bw : float
        The estimate of the bandwidth
    """
    bw = bw.lower()
    if bw not in bandwidth_funcs:
        raise ValueError("Bandwidth %s not understood" % bw)
    bandwidth = bandwidth_funcs[bw](x, kernel)
    if bandwidth == 0:
        # eventually this can fall back on another selection criterion.
        err = "Selected KDE bandwidth is 0. Cannot estiamte density."
        raise RuntimeError(err)
    else:
        return bandwidth
<commit_msg>correcting typo<commit_after>import numpy as np
from scipy.stats import scoreatpercentile as sap

from statsmodels.compat.pandas import Substitution
from statsmodels.sandbox.nonparametric import kernels

def _select_sigma(X):
    """
    Returns the smaller of std(X, ddof=1) or normalized IQR(X) over axis 0.

    References
    ----------
    Silverman (1986) p.47
    """
#    normalize = norm.ppf(.75) - norm.ppf(.25)
    normalize = 1.349
#    IQR = np.subtract.reduce(percentile(X, [75,25],
#                             axis=axis), axis=axis)/normalize
    IQR = (sap(X, 75) - sap(X, 25))/normalize
    return np.minimum(np.std(X, axis=0, ddof=1), IQR)

## Univariate Rule of Thumb Bandwidths ##
def bw_scott(x, kernel=None):
    """
    Scott's Rule of Thumb

    Parameters
    ----------
    x : array_like
        Array for which to get the bandwidth
    kernel : CustomKernel object
        Unused

    Returns
    -------
    bw : float
        The estimate of the bandwidth

    Notes
    -----
    Returns 1.059 * A * n ** (-1/5.) where ::

       A = min(std(x, ddof=1), IQR/1.349)
       IQR = np.subtract.reduce(np.percentile(x, [75,25]))

    References
    ----------

    Scott, D.W. (1992) Multivariate Density Estimation: Theory, Practice, and
        Visualization.
    """
    A = _select_sigma(x)
    n = len(x)
    return 1.059 * A * n ** (-0.2)

def bw_silverman(x, kernel=None):
    """
    Silverman's Rule of Thumb

    Parameters
    ----------
    x : array_like
        Array for which to get the bandwidth
    kernel : CustomKernel object
        Unused

    Returns
    -------
    bw : float
        The estimate of the bandwidth

    Notes
    -----
    Returns .9 * A * n ** (-1/5.) where ::

       A = min(std(x, ddof=1), IQR/1.349)
       IQR = np.subtract.reduce(np.percentile(x, [75,25]))

    References
    ----------

    Silverman, B.W. (1986) `Density Estimation.`
    """
    A = _select_sigma(x)
    n = len(x)
    return .9 * A * n ** (-0.2)

def bw_normal_reference(x, kernel=kernels.Gaussian):
    """
    Plug-in bandwidth with kernel specific constant based on normal reference.

    This bandwidth minimizes the mean integrated square error if the true
    distribution is the normal. This choice is an appropriate bandwidth for
    single peaked distributions that are similar to the normal distribution.

    Parameters
    ----------
    x : array_like
        Array for which to get the bandwidth
    kernel : CustomKernel object
        Used to calculate the constant for the plug-in bandwidth.

    Returns
    -------
    bw : float
        The estimate of the bandwidth

    Notes
    -----
    Returns C * A * n ** (-1/5.) where ::

       A = min(std(x, ddof=1), IQR/1.349)
       IQR = np.subtract.reduce(np.percentile(x, [75,25]))
       C = constant from Hansen (2009)

    When using a Gaussian kernel this is equivalent to the 'scott' bandwidth up
    to two decimal places. This is the accuracy to which the 'scott' constant is
    specified.

    References
    ----------

    Silverman, B.W. (1986) `Density Estimation.`
    Hansen, B.E. (2009) `Lecture Notes on Nonparametrics.`
    """
    C = kernel.normal_reference_constant
    A = _select_sigma(x)
    n = len(x)
    return C * A * n ** (-0.2)

## Plug-In Methods ##

## Least Squares Cross-Validation ##

## Helper Functions ##

bandwidth_funcs = {
    "scott": bw_scott,
    "silverman": bw_silverman,
    "normal_reference": bw_normal_reference,
}

@Substitution(", ".join(sorted(bandwidth_funcs.keys())))
def select_bandwidth(x, bw, kernel):
    """
    Selects bandwidth for a selection rule bw

    this is a wrapper around existing bandwidth selection rules

    Parameters
    ----------
    x : array_like
        Array for which to get the bandwidth
    bw : str
        name of bandwidth selection rule, currently supported are:
        %s
    kernel : not used yet

    Returns
    -------
    bw : float
        The estimate of the bandwidth
    """
    bw = bw.lower()
    if bw not in bandwidth_funcs:
        raise ValueError("Bandwidth %s not understood" % bw)
    bandwidth = bandwidth_funcs[bw](x, kernel)
    if bandwidth == 0:
        # eventually this can fall back on another selection criterion.
        err = "Selected KDE bandwidth is 0. Cannot estimate density."
        raise RuntimeError(err)
    else:
        return bandwidth
<|endoftext|>


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-
###############################################################################
#
# ODOO (ex OpenERP)
# Open Source Management Solution
# Copyright (C) 2001-2015 Micronaet S.r.l. (<https://micronaet.com>)
# Developer: Nicola Riolini @thebrush (<https://it.linkedin.com/in/thebrush>)
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#
###############################################################################

import os
import logging
from openerp.osv import fields, osv, expression
from datetime import datetime, timedelta
from openerp.tools.translate import _
from openerp.tools import (DEFAULT_SERVER_DATE_FORMAT,
    DEFAULT_SERVER_DATETIME_FORMAT,
    DATETIME_FORMATS_MAP,
    float_compare)
import xlrd
import pdb

_logger = logging.getLogger(__name__)

class MrpProductionWorkcenterLine(osv.osv):
    """ Model name: Lavoration
    """

    _inherit = 'mrp.production.workcenter.line'

    # Override so update after update MRP:
    def update_product_level_from_production(self, cr, uid, ids, context=None):
        """ Update product level from production
        """
        _logger.info('Update marketed product medium')
        product_pool = self.pool.get('product.product')
        company_pool = self.pool.get('res.company')

        # ---------------------------------------------------------------------
        # A. Update marketed product:
        # ---------------------------------------------------------------------
        # Fixed parameters:
        columns_position = {
            'date': 1,
            # 'invoice': 7,
            'default_code': 20,
            # 'name': 21,
            'qty': 22,
            # 'uom': 23,
            # 'price': 24,
        }
        start_test = 'INVOICE DATE'
        sheet_name = 'BASE'

        # Dynamic parameters:
        company_ids = company_pool.search(cr, uid, [], context=context)
        company = company_pool.browse(
            cr, uid, company_ids, context=context)[0]

        stock_level_external_excel = os.path.expanduser(
            company.stock_level_external_excel or '~/VTA PCA 2011A.xlsx')
        stock_level_days = company.stock_level_days
        if not stock_level_external_excel:
            raise osv.except_osv(
                _('Error stock management'),
                _('Setup the parameter in company form (input file)'),
                )

        now = datetime.now()
        from_date = now - timedelta(days=stock_level_days)
        now_text = '%s 00:00:00' % now.strftime(
             DEFAULT_SERVER_DATE_FORMAT)
        from_text = '%s 00:00:00' % from_date.strftime(
             DEFAULT_SERVER_DATE_FORMAT)

        # A1. Search product marketed:
        product_ids = product_pool.search(cr, uid, [
            ('default_code', '=ilike', 'R%'),
            ('default_code', '=ilike', 'S%'),
            # ('default_code', 'not =ilike', '%X'),
            ], context=context)

        _logger.warning('Imported product #%s [%s - %s]' % (
            len(product_ids),
            from_text,
            now_text,
            ))
        pdb.set_trace()

        # A2. Prepare dict for medium
        product_medium = {}
        for product in product_pool.browse(
                cr, uid, product_ids, context=context):
            default_code = product.default_code
            if default_code.endswith('X'):
                continue
            product_medium[default_code] = 0.0

        # A3. Load data from Excel:
        try:
            wb = xlrd.open_workbook(stock_level_external_excel)
        except:
            _logger.error(
                '[ERROR] Cannot read XLS file: %s' % stock_level_external_excel
            )
            return False

        ws = wb.sheet_by_name(sheet_name)
        _logger.info('Read XLS file: %s' % stock_level_external_excel)
        start = False
        for row in range(ws.nrows):
            date = ws.cell(row, columns_position['date']).value
            if not start and date == start_test:
                start = True
                continue

            default_code = ws.cell(row, columns_position['default_code']).value
            if not(start and date and default_code in product_medium):
                continue

            # Load data for medium
            qty = ws.cell(row, columns_position['qty']).value
            product_medium[default_code] += qty

        # A4. Update product medium
        _logger.warning('Product found: %s' % len(product_medium))
        for product in product_medium:
            total = product_medium[product]
            if product.manual_stock_level:
                continue

            if not total:
                medium_stock_qty = 0.0
            else:
                medium_stock_qty = total / stock_level_days

            product_pool.write(cr, uid, [product.id], {
                'medium_stock_qty': medium_stock_qty,
                'min_stock_level':
                    product.day_min_level * medium_stock_qty,
                'max_stock_level':
                    product.day_max_level * medium_stock_qty,
                'ready_stock_level':
                    product.day_max_ready_level * medium_stock_qty,
                }, context=context)

        # ---------------------------------------------------------------------
        # B. Update original procedure from MRP:
        # ---------------------------------------------------------------------
        return super(MrpProductionWorkcenterLine, self).\
            update_product_level_from_production(cr, uid, ids, context=context)

class ResCompany(osv.osv):
    """ Model name: Parameters
    """

    _inherit = 'res.company'

    _columns = {
        'stock_level_external_excel': fields.char(
            'Stock movement file', size=100,
            help='Path for external XLSX file for stock movement',
        )
    }
<commit_msg>add log<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-
###############################################################################
#
# ODOO (ex OpenERP)
# Open Source Management Solution
# Copyright (C) 2001-2015 Micronaet S.r.l. (<https://micronaet.com>)
# Developer: Nicola Riolini @thebrush (<https://it.linkedin.com/in/thebrush>)
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
# See the GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#
###############################################################################

import os
import logging
from openerp.osv import fields, osv, expression
from datetime import datetime, timedelta
from openerp.tools.translate import _
from openerp.tools import (DEFAULT_SERVER_DATE_FORMAT,
    DEFAULT_SERVER_DATETIME_FORMAT,
    DATETIME_FORMATS_MAP,
    float_compare)
import xlrd
import pdb

_logger = logging.getLogger(__name__)

class MrpProductionWorkcenterLine(osv.osv):
    """ Model name: Lavoration
    """

    _inherit = 'mrp.production.workcenter.line'

    # Override so update after update MRP:
    def update_product_level_from_production(self, cr, uid, ids, context=None):
        """ Update product level from production
        """
        _logger.info('Update marketed product medium')
        product_pool = self.pool.get('product.product')
        company_pool = self.pool.get('res.company')

        # ---------------------------------------------------------------------
        # A. Update marketed product:
        # ---------------------------------------------------------------------
        # Fixed parameters:
        columns_position = {
            'date': 1,
            # 'invoice': 7,
            'default_code': 20,
            # 'name': 21,
            'qty': 22,
            # 'uom': 23,
            # 'price': 24,
        }
        start_test = 'INVOICE DATE'
        sheet_name = 'BASE'

        # Dynamic parameters:
        company_ids = company_pool.search(cr, uid, [], context=context)
        company = company_pool.browse(
            cr, uid, company_ids, context=context)[0]

        stock_level_external_excel = os.path.expanduser(
            company.stock_level_external_excel or '~/VTA PCA 2011A.xlsx')
        stock_level_days = company.stock_level_days
        if not stock_level_external_excel:
            raise osv.except_osv(
                _('Error stock management'),
                _('Setup the parameter in company form (input file)'),
                )

        now = datetime.now()
        from_date = now - timedelta(days=stock_level_days)
        now_text = '%s 00:00:00' % now.strftime(
             DEFAULT_SERVER_DATE_FORMAT)
        from_text = '%s 00:00:00' % from_date.strftime(
             DEFAULT_SERVER_DATE_FORMAT)

        # A1. Search product marketed:
        product_ids = product_pool.search(cr, uid, [
            '|',
            ('default_code', '=ilike', 'R%'),
            ('default_code', '=ilike', 'S%'),
            # ('default_code', 'not =ilike', '%X'),
            ], context=context)

        _logger.warning('Imported product #%s [%s - %s]' % (
            len(product_ids),
            from_text,
            now_text,
            ))

        # A2. Prepare dict for medium
        product_medium = {}
        for product in product_pool.browse(
                cr, uid, product_ids, context=context):
            default_code = product.default_code
            if default_code.endswith('X'):
                continue
            product_medium[default_code] = 0.0

        # A3. Load data from Excel:
        try:
            wb = xlrd.open_workbook(stock_level_external_excel)
        except:
            _logger.error(
                '[ERROR] Cannot read XLS file: %s' % stock_level_external_excel
            )
            return False

        ws = wb.sheet_by_name(sheet_name)
        _logger.info('Read XLS file: %s' % stock_level_external_excel)
        start = False
        pdb.set_trace()
        for row in range(ws.nrows):
            date = ws.cell(row, columns_position['date']).value
            if not start and date == start_test:
                _logger.info('%s. Line not used pre start' % (row + 1))
                start = True
                continue

            if from_text > date > now_text:  # Out of period range:
                _logger.info('%s. Line not used out of range %s' % (
                    row + 1, date))
                continue

            default_code = ws.cell(row, columns_position['default_code']).value
            if not(start and date and default_code in product_medium):
                _logger.info('%s. Line not used no imported product: %s' % (
                    row + 1,
                    default_code,
                ))
                continue

            # Load data for medium
            qty = ws.cell(row, columns_position['qty']).value
            product_medium[default_code] += qty
            _logger.info('%s. Line used %s - %s' % (
                row + 1,
                default_code,
                qty,
            ))

        # A4. Update product medium
        _logger.warning('Product found: %s' % len(product_medium))
        for product in product_medium:
            total = product_medium[product]
            if product.manual_stock_level:
                continue

            if not total:
                medium_stock_qty = 0.0
            else:
                medium_stock_qty = total / stock_level_days

            product_pool.write(cr, uid, [product.id], {
                'medium_stock_qty': medium_stock_qty,
                'min_stock_level':
                    product.day_min_level * medium_stock_qty,
                'max_stock_level':
                    product.day_max_level * medium_stock_qty,
                'ready_stock_level':
                    product.day_max_ready_level * medium_stock_qty,
                }, context=context)

        # ---------------------------------------------------------------------
        # B. Update original procedure from MRP:
        # ---------------------------------------------------------------------
        return super(MrpProductionWorkcenterLine, self).\
            update_product_level_from_production(cr, uid, ids, context=context)

class ResCompany(osv.osv):
    """ Model name: Parameters
    """

    _inherit = 'res.company'

    _columns = {
        'stock_level_external_excel': fields.char(
            'Stock movement file', size=100,
            help='Path for external XLSX file for stock movement',
        )
    }
<|endoftext|>