<commit_before># -*- coding: utf-8 -*-
##############################################################################
#
#    Copyright (C) 2014-15 Agile Business Group sagl
#    (<http://www.agilebg.com>)
#
#    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/>.
#
##############################################################################

from openerp import models, api
from openerp.tools.translate import _
from openerp.tools import float_round
import logging
from openerp.exceptions import except_orm

_logger = logging.getLogger(__name__)

class SaleOrderLine(models.Model):
    _inherit = 'sale.order.line'

    @api.onchange('property_ids', 'product_uos_qty')
    def quantity_property_ids_changed(self):
        prop_ctx = self.env.context.copy()
        if 'lang' in prop_ctx:
            del prop_ctx['lang']
        if self.product_id and self.property_ids and self.product_uos_qty:
            if self.product_id.quantity_formula_id:
                prop_dict = {}
                uom_precision = self.pool['decimal.precision'].precision_get(
                    self._cr, self._uid, 'Product UoS')
                for prop in self.env['mrp.property'].with_context(
                    prop_ctx
                ).browse([p.id for p in self.property_ids]):
                    if prop.group_id.name in prop_dict:
                        raise except_orm(
                            _('Error'),
                            _('Property of group %s already present')
                            % prop.group_id.name)
                    prop_dict[prop.group_id.name] = prop.value

                localdict = {
                    'self': self,
                    'cr': self._cr,
                    'uid': self._uid,
                    'properties': prop_dict,
                    'qty_uos': self.product_uos_qty,
                    'product_id': self.product_id.id,
                }
                try:
                    exec self.product_id.quantity_formula_id.\
                        formula_text in localdict
                    try:
                        amount = localdict['result']
                    except KeyError:
                        raise except_orm(
                            _('Error'),
                            _("Formula must contain 'result' variable"))
                    # rounding because decimal values, different from what
                    # displayed on interface, could generate infinite
                    # on_change loop
                    self.product_uom_qty = float_round(
                        amount, precision_digits=uom_precision)
                except KeyError:
                    _logger.warning(
                        "KeyError for formula '%s' and prop_dict '%s'"
                        % (self.product_id.quantity_formula_id.formula_text,
                            prop_dict))

    def product_id_change(
        self, cr, uid, ids, pricelist, product_id, qty=0,
        uom=False, qty_uos=0, uos=False, name='', partner_id=False,
        lang=False, update_tax=True, date_order=False, packaging=False,
        fiscal_position=False, flag=False, context=None
    ):
        res = super(SaleOrderLine, self).product_id_change(
            cr, uid, ids, pricelist, product_id, qty=qty,
            uom=uom, qty_uos=qty_uos, uos=uos,
            name=name, partner_id=partner_id,
            lang=lang, update_tax=update_tax,
            date_order=date_order, packaging=packaging,
            fiscal_position=fiscal_position, flag=flag, context=context)
        if 'value' in res:
            # get empty properties
            res['value']['property_ids'] = []
            # get empty properties dynamic fields
            property_group_pool = self.pool['mrp.property.group']
            group_to_empty_ids = property_group_pool.search(
                cr, uid, [('draw_dynamically', '=', True)], context=context)
            for group in property_group_pool.browse(
                    cr, uid, group_to_empty_ids, context=context
            ):
                res['value'][group.field_id.name] = None
        # Removing product_uos_qty is needed because it can now be used to
        # compute the real quantity. Otherwise, it would be recomputed after
        # the quantity changed. See the automated test for the use case.
        if 'value' in res and 'product_uos_qty' in res['value']:
            del res['value']['product_uos_qty']
        return res
<commit_msg>Quantity recalculation has to be done when property_ids or product_uos_qty<commit_after># -*- coding: utf-8 -*-
##############################################################################
#
#    Copyright (C) 2014-15 Agile Business Group sagl
#    (<http://www.agilebg.com>)
#
#    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/>.
#
##############################################################################

from openerp import models, api
from openerp.tools.translate import _
from openerp.tools import float_round
import logging
from openerp.exceptions import except_orm

_logger = logging.getLogger(__name__)

class SaleOrderLine(models.Model):
    _inherit = 'sale.order.line'

    @api.onchange('property_ids', 'product_uos_qty')
    def quantity_property_ids_changed(self):
        prop_ctx = self.env.context.copy()
        if 'lang' in prop_ctx:
            del prop_ctx['lang']
        if self.product_id and (self.property_ids or self.product_uos_qty):
            if self.product_id.quantity_formula_id:
                prop_dict = {}
                uom_precision = self.pool['decimal.precision'].precision_get(
                    self._cr, self._uid, 'Product UoS')
                for prop in self.env['mrp.property'].with_context(
                    prop_ctx
                ).browse([p.id for p in self.property_ids]):
                    if prop.group_id.name in prop_dict:
                        raise except_orm(
                            _('Error'),
                            _('Property of group %s already present')
                            % prop.group_id.name)
                    prop_dict[prop.group_id.name] = prop.value

                localdict = {
                    'self': self,
                    'cr': self._cr,
                    'uid': self._uid,
                    'properties': prop_dict,
                    'qty_uos': self.product_uos_qty,
                    'product_id': self.product_id.id,
                }
                try:
                    exec self.product_id.quantity_formula_id.\
                        formula_text in localdict
                    try:
                        amount = localdict['result']
                    except KeyError:
                        raise except_orm(
                            _('Error'),
                            _("Formula must contain 'result' variable"))
                    # rounding because decimal values, different from what
                    # displayed on interface, could generate infinite
                    # on_change loop
                    self.product_uom_qty = float_round(
                        amount, precision_digits=uom_precision)
                except KeyError:
                    _logger.warning(
                        "KeyError for formula '%s' and prop_dict '%s'"
                        % (self.product_id.quantity_formula_id.formula_text,
                            prop_dict))

    def product_id_change(
        self, cr, uid, ids, pricelist, product_id, qty=0,
        uom=False, qty_uos=0, uos=False, name='', partner_id=False,
        lang=False, update_tax=True, date_order=False, packaging=False,
        fiscal_position=False, flag=False, context=None
    ):
        res = super(SaleOrderLine, self).product_id_change(
            cr, uid, ids, pricelist, product_id, qty=qty,
            uom=uom, qty_uos=qty_uos, uos=uos,
            name=name, partner_id=partner_id,
            lang=lang, update_tax=update_tax,
            date_order=date_order, packaging=packaging,
            fiscal_position=fiscal_position, flag=flag, context=context)
        if 'value' in res:
            # get empty properties
            res['value']['property_ids'] = []
            # get empty properties dynamic fields
            property_group_pool = self.pool['mrp.property.group']
            group_to_empty_ids = property_group_pool.search(
                cr, uid, [('draw_dynamically', '=', True)], context=context)
            for group in property_group_pool.browse(
                    cr, uid, group_to_empty_ids, context=context
            ):
                res['value'][group.field_id.name] = None
        # Removing product_uos_qty is needed because it can now be used to
        # compute the real quantity. Otherwise, it would be recomputed after
        # the quantity changed. See the automated test for the use case.
        if 'value' in res and 'product_uos_qty' in res['value']:
            del res['value']['product_uos_qty']
        return res
<|endoftext|>


<commit_before>from __future__ import absolute_import

import logging

from functools import partial
from concurrent.futures import ThreadPoolExecutor

import celery
import tornado.web

from tornado import ioloop
from tornado.netutil import bind_unix_socket
from tornado.httpserver import HTTPServer

from .api import control
from .urls import handlers
from .events import Events
from .options import default_options

logger = logging.getLogger(__name__)

class Flower(tornado.web.Application):
    pool_executor_cls = ThreadPoolExecutor
    max_workers = 4

    def __init__(self, options=None, capp=None, events=None,
                 io_loop=None, **kwargs):
        kwargs.update(handlers=handlers)
        super(Flower, self).__init__(**kwargs)
        self.options = options or default_options
        self.io_loop = io_loop or ioloop.IOLoop.instance()
        self.ssl_options = kwargs.get('ssl_options', None)

        self.capp = capp or celery.Celery()
        self.events = events or Events(
            self.capp, db=self.options.db,
            persistent=self.options.persistent,
            enable_events=self.options.enable_events,
            io_loop=self.io_loop,
            max_tasks_in_memory=self.options.max_tasks)
        self.started = False

    def start(self):
        self.pool = self.pool_executor_cls(max_workers=self.max_workers)
        self.events.start()

        if not self.options.unix_socket:
            self.listen(self.options.port, address=self.options.address,
                        ssl_options=self.ssl_options,
                        xheaders=self.options.xheaders)
        else:
            server = HTTPServer(self)
            socket = bind_unix_socket(self.options.unix_socket)
            server.add_socket(socket)

        self.io_loop.add_future(
            control.ControlHandler.update_workers(app=self),
            callback=lambda x: logger.debug('Successfully updated worker cache'))
        self.started = True
        self.io_loop.start()

    def stop(self):
        if self.started:
            self.events.stop()
            self.pool.shutdown(wait=False)
            self.started = False

    def delay(self, method, *args, **kwargs):
        return self.pool.submit(partial(method, *args, **kwargs))

    @property
    def transport(self):
        return getattr(self.capp.connection().transport,
                       'driver_type', None)
<commit_msg>Import bind_unix_socket only if --unix_socket is enabled #519<commit_after>from __future__ import absolute_import

import logging

from functools import partial
from concurrent.futures import ThreadPoolExecutor

import celery
import tornado.web

from tornado import ioloop
from tornado.httpserver import HTTPServer

from .api import control
from .urls import handlers
from .events import Events
from .options import default_options

logger = logging.getLogger(__name__)

class Flower(tornado.web.Application):
    pool_executor_cls = ThreadPoolExecutor
    max_workers = 4

    def __init__(self, options=None, capp=None, events=None,
                 io_loop=None, **kwargs):
        kwargs.update(handlers=handlers)
        super(Flower, self).__init__(**kwargs)
        self.options = options or default_options
        self.io_loop = io_loop or ioloop.IOLoop.instance()
        self.ssl_options = kwargs.get('ssl_options', None)

        self.capp = capp or celery.Celery()
        self.events = events or Events(
            self.capp, db=self.options.db,
            persistent=self.options.persistent,
            enable_events=self.options.enable_events,
            io_loop=self.io_loop,
            max_tasks_in_memory=self.options.max_tasks)
        self.started = False

    def start(self):
        self.pool = self.pool_executor_cls(max_workers=self.max_workers)
        self.events.start()

        if not self.options.unix_socket:
            self.listen(self.options.port, address=self.options.address,
                        ssl_options=self.ssl_options,
                        xheaders=self.options.xheaders)
        else:
            from tornado.netutil import bind_unix_socket
            server = HTTPServer(self)
            socket = bind_unix_socket(self.options.unix_socket)
            server.add_socket(socket)

        self.io_loop.add_future(
            control.ControlHandler.update_workers(app=self),
            callback=lambda x: logger.debug('Successfully updated worker cache'))
        self.started = True
        self.io_loop.start()

    def stop(self):
        if self.started:
            self.events.stop()
            self.pool.shutdown(wait=False)
            self.started = False

    def delay(self, method, *args, **kwargs):
        return self.pool.submit(partial(method, *args, **kwargs))

    @property
    def transport(self):
        return getattr(self.capp.connection().transport,
                       'driver_type', None)
<|endoftext|>


<commit_before>from __future__ import absolute_import
import json
from django.db.models import Q
from corehq.apps.accounting.models import Feature, SoftwareProduct, BillingAccount, SoftwarePlanVersion, \
    Subscription, Subscriber, BillingContactInfo, SoftwarePlan
from corehq.apps.accounting.utils import fmt_feature_rate_dict, fmt_product_rate_dict
from corehq.apps.domain.models import Domain
from corehq.apps.hqwebapp.async_handler import BaseAsyncHandler, AsyncHandlerError
from corehq.apps.hqwebapp.encoders import LazyEncoder

class BaseRateAsyncHandler(BaseAsyncHandler):
    """
    Subclass this for interacting with RatesManager.
    """
    allowed_actions = [
        'apply',
        'create',
    ]

    @property
    def name(self):
        return self.data.get('name')

    @property
    def rate_type(self):
        return self.data.get('rate_type')

    @property
    def rate_id(self):
        return self.data.get('rate_id')

    @property
    def create_response(self):
        raise NotImplementedError("create_response is required")

    @property
    def apply_response(self):
        raise NotImplementedError("apply_response is required")

class FeatureRateAsyncHandler(BaseRateAsyncHandler):
    slug = 'features_handler'

    @property
    def create_response(self):
        if Feature.objects.filter(name=self.name).count() > 0:
            raise AsyncHandlerError("Feature '%s' already exists, and likely already "
                                    "in this Software Plan Version." % self.name)
        new_feature, _ = Feature.objects.get_or_create(
            name=self.name,
            feature_type=self.rate_type,
        )
        return fmt_feature_rate_dict(new_feature)

    @property
    def apply_response(self):
        try:
            feature = Feature.objects.get(id=self.rate_id)
            return fmt_feature_rate_dict(feature)
        except Feature.DoesNotExist:
            raise AsyncHandlerError("could not find an existing feature")

class SoftwareProductRateAsyncHandler(BaseRateAsyncHandler):
    slug = 'products_handler'

    @property
    def create_response(self):
        if SoftwareProduct.objects.filter(name=self.name).count() > 0:
            raise AsyncHandlerError("Product '%s' already exists, and likely already "
                                    "in this Software Plan Version." % self.name)
        new_product, _ = SoftwareProduct.objects.get_or_create(
            name=self.name,
            product_type=self.rate_type
        )
        return fmt_product_rate_dict(new_product)

    @property
    def apply_response(self):
        try:
            product = SoftwareProduct.objects.get(id=self.rate_id)
            return fmt_product_rate_dict(product)
        except SoftwareProduct.DoesNotExist:
            raise AsyncHandlerError("could not find an existing product")

class BaseSelect2AsyncHandler(BaseAsyncHandler):

    @property
    def search_string(self):
        return self.data.get('searchString')

    @property
    def existing(self):
        return self.data.getlist('existing[]')

    def _fmt_success(self, response):
        success = json.dumps({
            'results': [{
                'id': r[0],
                'text': r[1],
            } for r in response]
        }, cls=LazyEncoder)
        return success

class Select2RateAsyncHandler(BaseSelect2AsyncHandler):
    """
    Handles the async responses for the select2 widget in the Features & Rates portion
    of the SoftwarePlanVersion form.
    """
    slug = 'select2_rate'
    allowed_actions = [
        'feature_id',
        'product_id',
    ]

    @property
    def feature_id_response(self):
        features = Feature.objects
        if self.existing:
            features = features.exclude(name__in=self.existing)
        if self.search_string:
            features = features.filter(name__istartswith=self.search_string)
        return [(f.id, f.name, f.feature_type) for f in features.all()]

    @property
    def product_id_response(self):
        products = SoftwareProduct.objects
        if self.existing:
            products = products.exclude(name__in=self.existing)
        if self.search_string:
            products = products.filter(name__istartswith=self.search_string)
        return [(p.id, p.name, p.product_type) for p in products.all()]

    def _fmt_success(self, response):
        return json.dumps({
            'results': [{
                'id': r[0],
                'name': r[1],
                'rate_type': r[2],
                'text': '%s (%s)' % (r[1], r[2]),
                'isExisting': True,
            } for r in response]
        })

class Select2BillingInfoHandler(BaseSelect2AsyncHandler):
    slug = 'select2_billing'
    allowed_actions = [
        'country',
        'active_accounts',
        'domain',
        'account',
        'plan_version',
        'new_plan_version',
    ]

    @property
    def country_response(self):
        from django_countries.data import COUNTRIES
        countries = sorted(list(COUNTRIES.items()), key=lambda x: x[1].encode('utf-8'))
        if self.search_string:
            search_string = self.search_string.lower()
            return [x for x in countries if x[1].lower().startswith(search_string)]
        return countries

    @property
    def active_accounts_response(self):
        accounts = BillingAccount.objects.filter(is_active=True)
        if self.search_string:
            accounts = accounts.filter(name__contains=self.search_string)
        return [(a.id, a.name) for a in accounts]

    @property
    def domain_response(self):
        domain_names = [domain['key'] for domain in Domain.get_all(include_docs=False)]
        if self.search_string:
            search_string = self.search_string.lower()
            domain_names = [x for x in domain_names if x.lower().startswith(search_string)]
        return [(name, name) for name in domain_names]

    @property
    def account_response(self):
        accounts = BillingAccount.objects
        if self.search_string:
            accounts = accounts.filter(name__contains=self.search_string)
        return [(a.id, a.name) for a in accounts.order_by('name')]

    @property
    def plan_version_response(self):
        edition = self.data.get('additionalData[edition]')
        product = self.data.get('additionalData[product]')
        plan_versions = SoftwarePlanVersion.objects.filter(
            plan__edition=edition
        ).filter(product_rate__product__product_type=product)
        if self.search_string:
            plan_versions = plan_versions.filter(
                plan__name__contains=self.search_string)
        return [(p.id, p.__str__()) for p in plan_versions.order_by('plan__name')]

    @property
    def new_plan_version_response(self):
        current_version = int(self.data.get('additionalData[current_version]'))
        plan_versions = [x for x in self.plan_version_response if x[0] != current_version]
        return plan_versions

class Select2InvoiceTriggerHandler(BaseSelect2AsyncHandler):
    slug = 'select2_billing'
    allowed_actions = [
        'domain',
    ]

    @property
    def domain_response(self):
        domain_names = [domain['key'] for domain in Domain.get_all(include_docs=False)]
        if self.search_string:
            search_string = self.search_string.lower()
            domain_names = [x for x in domain_names if x.lower().startswith(search_string)]
        return [(d, d) for d in domain_names]

class BaseSingleOptionFilterAsyncHandler(BaseAsyncHandler):

    @property
    def query(self):
        raise NotImplementedError("must return a queryset")

    @property
    def search_string(self):
        return self.data.get('q', None)

    @property
    def page(self):
        return int(self.data.get('page', 1))

    @property
    def paginated_data(self):
        start = (self.page - 1) * self.limit
        end = self.page * self.limit
        return self.query.all()[start:end]

    @property
    def limit(self):
        return self.data.get('limit', 10)

    @property
    def total(self):
        return self.query.count()

    @staticmethod
    def _fmt_select2_data(data_id, data_text):
        return {
            'id': data_id,
            'text': data_text,
        }

    def _fmt_success(self, data):
        return json.dumps({
            'success': True,
            'limit': self.limit,
            'page': self.page,
            'total': self.total,
            'items': data,
        })

class SubscriberFilterAsyncHandler(BaseSingleOptionFilterAsyncHandler):
    slug = 'subscriber_filter'
    allowed_actions = [
        'subscriber',
    ]

    @property
    def query(self):
        query = Subscriber.objects.exclude(domain=None).order_by('domain')
        if self.search_string:
            query = query.filter(domain__istartswith=self.search_string)
        return query

    @property
    def subscriber_response(self):
        return [self._fmt_select2_data(s.domain, s.domain)
                for s in self.paginated_data]

class SubscriptionFilterAsyncHandler(BaseSingleOptionFilterAsyncHandler):
    slug = 'subscription_filter'
    allowed_actions = [
        'contract_id',
    ]

    @property
    def query(self):
        query = Subscription.objects
        if self.action == 'contract_id':
            query = query.exclude(
                salesforce_contract_id=None
            ).exclude(salesforce_contract_id='').order_by('salesforce_contract_id')
            if self.search_string:
                query = query.filter(
                    salesforce_contract_id__istartswith=self.search_string
                )
        return query

    @property
    def contract_id_response(self):
        return [self._fmt_select2_data(
                s.salesforce_contract_id, s.salesforce_contract_id)
                for s in self.paginated_data]

class AccountFilterAsyncHandler(BaseSingleOptionFilterAsyncHandler):
    slug = 'account_filter'
    allowed_actions = [
        'account_name',
        'account_id',
        'dimagi_contact',
    ]

    @property
    def query(self):
        query = BillingAccount.objects.order_by('name')

        if self.action == 'account_name' and self.search_string:
            query = query.filter(name__icontains=self.search_string)

        if self.action == 'account_id':
            query = query.exclude(
                salesforce_account_id=None
            ).exclude(
                salesforce_account_id=''
            ).order_by('salesforce_account_id')
            if self.search_string:
                query = query.filter(
                    salesforce_account_id__istartswith=self.search_string)

        if self.action == 'dimagi_contact':
            query = query.exclude(
                dimagi_contact=''
            ).order_by('dimagi_contact')
            if self.search_string:
                query = query.filter(
                    dimagi_contact__icontains=self.search_string)

        return query

    @property
    def account_name_response(self):
        return [self._fmt_select2_data(a.name, a.name)
                for a in self.paginated_data]

    @property
    def account_id_response(self):
        return [self._fmt_select2_data(a.salesforce_account_id,
                                       a.salesforce_account_id)
                for a in self.paginated_data]

    @property
    def dimagi_contact_response(self):
        return [self._fmt_select2_data(a.dimagi_contact, a.dimagi_contact)
                for a in self.paginated_data]

class BillingContactInfoAsyncHandler(BaseSingleOptionFilterAsyncHandler):
    slug = 'billing_contact_filter'
    allowed_actions = [
        'contact_name'
    ]

    @property
    def query(self):
        query = BillingContactInfo.objects.exclude(
            first_name='', last_name='').order_by('first_name', 'last_name')
        if self.search_string:
            query = query.filter(
                Q(first_name__istartswith=self.search_string) |
                Q(last_name__istartswith=self.search_string)
            )
        return query

    @property
    def contact_name_response(self):
        return [self._fmt_select2_data(c.full_name, c.full_name)
                for c in self.paginated_data]

class SoftwarePlanAsyncHandler(BaseSingleOptionFilterAsyncHandler):
    slug = 'software_plan_filter'
    allowed_actions = [
        'name',
    ]

    @property
    def query(self):
        query = SoftwarePlan.objects.order_by('name')
        if self.search_string:
            query = query.filter(name__icontains=self.search_string)
        return query

    @property
    def name_response(self):
        return [self._fmt_select2_data(p.name, p.name)
                for p in self.paginated_data]

class DomainFilterAsyncHandler(BaseSingleOptionFilterAsyncHandler):
    slug = 'domain_filter'
    allowed_actions = [
        'domain_name',
    ]

    @property
    def query(self):
        db = Domain.get_db()
        startkey = self.search_string
        endkey = "{}Z".format(self.search_string) if startkey else ''
        query = db.view(
            'domain/domains',
            reduce=False,
            startkey=startkey,
            endkey=endkey,
            limit=20,
        )
        return query

    @property
    def domain_name_response(self):
        return [self._fmt_select2_data(p['key'], p['key']) for p in self.paginated_data]
<commit_msg>use case insensitive search<commit_after>from __future__ import absolute_import
import json
from django.db.models import Q
from corehq.apps.accounting.models import Feature, SoftwareProduct, BillingAccount, SoftwarePlanVersion, \
    Subscription, Subscriber, BillingContactInfo, SoftwarePlan
from corehq.apps.accounting.utils import fmt_feature_rate_dict, fmt_product_rate_dict
from corehq.apps.domain.models import Domain
from corehq.apps.hqwebapp.async_handler import BaseAsyncHandler, AsyncHandlerError
from corehq.apps.hqwebapp.encoders import LazyEncoder

class BaseRateAsyncHandler(BaseAsyncHandler):
    """
    Subclass this for interacting with RatesManager.
    """
    allowed_actions = [
        'apply',
        'create',
    ]

    @property
    def name(self):
        return self.data.get('name')

    @property
    def rate_type(self):
        return self.data.get('rate_type')

    @property
    def rate_id(self):
        return self.data.get('rate_id')

    @property
    def create_response(self):
        raise NotImplementedError("create_response is required")

    @property
    def apply_response(self):
        raise NotImplementedError("apply_response is required")

class FeatureRateAsyncHandler(BaseRateAsyncHandler):
    slug = 'features_handler'

    @property
    def create_response(self):
        if Feature.objects.filter(name=self.name).count() > 0:
            raise AsyncHandlerError("Feature '%s' already exists, and likely already "
                                    "in this Software Plan Version." % self.name)
        new_feature, _ = Feature.objects.get_or_create(
            name=self.name,
            feature_type=self.rate_type,
        )
        return fmt_feature_rate_dict(new_feature)

    @property
    def apply_response(self):
        try:
            feature = Feature.objects.get(id=self.rate_id)
            return fmt_feature_rate_dict(feature)
        except Feature.DoesNotExist:
            raise AsyncHandlerError("could not find an existing feature")

class SoftwareProductRateAsyncHandler(BaseRateAsyncHandler):
    slug = 'products_handler'

    @property
    def create_response(self):
        if SoftwareProduct.objects.filter(name=self.name).count() > 0:
            raise AsyncHandlerError("Product '%s' already exists, and likely already "
                                    "in this Software Plan Version." % self.name)
        new_product, _ = SoftwareProduct.objects.get_or_create(
            name=self.name,
            product_type=self.rate_type
        )
        return fmt_product_rate_dict(new_product)

    @property
    def apply_response(self):
        try:
            product = SoftwareProduct.objects.get(id=self.rate_id)
            return fmt_product_rate_dict(product)
        except SoftwareProduct.DoesNotExist:
            raise AsyncHandlerError("could not find an existing product")

class BaseSelect2AsyncHandler(BaseAsyncHandler):

    @property
    def search_string(self):
        return self.data.get('searchString')

    @property
    def existing(self):
        return self.data.getlist('existing[]')

    def _fmt_success(self, response):
        success = json.dumps({
            'results': [{
                'id': r[0],
                'text': r[1],
            } for r in response]
        }, cls=LazyEncoder)
        return success

class Select2RateAsyncHandler(BaseSelect2AsyncHandler):
    """
    Handles the async responses for the select2 widget in the Features & Rates portion
    of the SoftwarePlanVersion form.
    """
    slug = 'select2_rate'
    allowed_actions = [
        'feature_id',
        'product_id',
    ]

    @property
    def feature_id_response(self):
        features = Feature.objects
        if self.existing:
            features = features.exclude(name__in=self.existing)
        if self.search_string:
            features = features.filter(name__istartswith=self.search_string)
        return [(f.id, f.name, f.feature_type) for f in features.all()]

    @property
    def product_id_response(self):
        products = SoftwareProduct.objects
        if self.existing:
            products = products.exclude(name__in=self.existing)
        if self.search_string:
            products = products.filter(name__istartswith=self.search_string)
        return [(p.id, p.name, p.product_type) for p in products.all()]

    def _fmt_success(self, response):
        return json.dumps({
            'results': [{
                'id': r[0],
                'name': r[1],
                'rate_type': r[2],
                'text': '%s (%s)' % (r[1], r[2]),
                'isExisting': True,
            } for r in response]
        })

class Select2BillingInfoHandler(BaseSelect2AsyncHandler):
    slug = 'select2_billing'
    allowed_actions = [
        'country',
        'active_accounts',
        'domain',
        'account',
        'plan_version',
        'new_plan_version',
    ]

    @property
    def country_response(self):
        from django_countries.data import COUNTRIES
        countries = sorted(list(COUNTRIES.items()), key=lambda x: x[1].encode('utf-8'))
        if self.search_string:
            search_string = self.search_string.lower()
            return [x for x in countries if x[1].lower().startswith(search_string)]
        return countries

    @property
    def active_accounts_response(self):
        accounts = BillingAccount.objects.filter(is_active=True)
        if self.search_string:
            accounts = accounts.filter(name__icontains=self.search_string)
        return [(a.id, a.name) for a in accounts]

    @property
    def domain_response(self):
        domain_names = [domain['key'] for domain in Domain.get_all(include_docs=False)]
        if self.search_string:
            search_string = self.search_string.lower()
            domain_names = [x for x in domain_names if x.lower().startswith(search_string)]
        return [(name, name) for name in domain_names]

    @property
    def account_response(self):
        accounts = BillingAccount.objects
        if self.search_string:
            accounts = accounts.filter(name__icontains=self.search_string)
        return [(a.id, a.name) for a in accounts.order_by('name')]

    @property
    def plan_version_response(self):
        edition = self.data.get('additionalData[edition]')
        product = self.data.get('additionalData[product]')
        plan_versions = SoftwarePlanVersion.objects.filter(
            plan__edition=edition
        ).filter(product_rate__product__product_type=product)
        if self.search_string:
            plan_versions = plan_versions.filter(
                plan__name__icontains=self.search_string)
        return [(p.id, p.__str__()) for p in plan_versions.order_by('plan__name')]

    @property
    def new_plan_version_response(self):
        current_version = int(self.data.get('additionalData[current_version]'))
        plan_versions = [x for x in self.plan_version_response if x[0] != current_version]
        return plan_versions

class Select2InvoiceTriggerHandler(BaseSelect2AsyncHandler):
    slug = 'select2_billing'
    allowed_actions = [
        'domain',
    ]

    @property
    def domain_response(self):
        domain_names = [domain['key'] for domain in Domain.get_all(include_docs=False)]
        if self.search_string:
            search_string = self.search_string.lower()
            domain_names = [x for x in domain_names if x.lower().startswith(search_string)]
        return [(d, d) for d in domain_names]

class BaseSingleOptionFilterAsyncHandler(BaseAsyncHandler):

    @property
    def query(self):
        raise NotImplementedError("must return a queryset")

    @property
    def search_string(self):
        return self.data.get('q', None)

    @property
    def page(self):
        return int(self.data.get('page', 1))

    @property
    def paginated_data(self):
        start = (self.page - 1) * self.limit
        end = self.page * self.limit
        return self.query.all()[start:end]

    @property
    def limit(self):
        return self.data.get('limit', 10)

    @property
    def total(self):
        return self.query.count()

    @staticmethod
    def _fmt_select2_data(data_id, data_text):
        return {
            'id': data_id,
            'text': data_text,
        }

    def _fmt_success(self, data):
        return json.dumps({
            'success': True,
            'limit': self.limit,
            'page': self.page,
            'total': self.total,
            'items': data,
        })

class SubscriberFilterAsyncHandler(BaseSingleOptionFilterAsyncHandler):
    slug = 'subscriber_filter'
    allowed_actions = [
        'subscriber',
    ]

    @property
    def query(self):
        query = Subscriber.objects.exclude(domain=None).order_by('domain')
        if self.search_string:
            query = query.filter(domain__istartswith=self.search_string)
        return query

    @property
    def subscriber_response(self):
        return [self._fmt_select2_data(s.domain, s.domain)
                for s in self.paginated_data]

class SubscriptionFilterAsyncHandler(BaseSingleOptionFilterAsyncHandler):
    slug = 'subscription_filter'
    allowed_actions = [
        'contract_id',
    ]

    @property
    def query(self):
        query = Subscription.objects
        if self.action == 'contract_id':
            query = query.exclude(
                salesforce_contract_id=None
            ).exclude(salesforce_contract_id='').order_by('salesforce_contract_id')
            if self.search_string:
                query = query.filter(
                    salesforce_contract_id__istartswith=self.search_string
                )
        return query

    @property
    def contract_id_response(self):
        return [self._fmt_select2_data(
                s.salesforce_contract_id, s.salesforce_contract_id)
                for s in self.paginated_data]

class AccountFilterAsyncHandler(BaseSingleOptionFilterAsyncHandler):
    slug = 'account_filter'
    allowed_actions = [
        'account_name',
        'account_id',
        'dimagi_contact',
    ]

    @property
    def query(self):
        query = BillingAccount.objects.order_by('name')

        if self.action == 'account_name' and self.search_string:
            query = query.filter(name__icontains=self.search_string)

        if self.action == 'account_id':
            query = query.exclude(
                salesforce_account_id=None
            ).exclude(
                salesforce_account_id=''
            ).order_by('salesforce_account_id')
            if self.search_string:
                query = query.filter(
                    salesforce_account_id__istartswith=self.search_string)

        if self.action == 'dimagi_contact':
            query = query.exclude(
                dimagi_contact=''
            ).order_by('dimagi_contact')
            if self.search_string:
                query = query.filter(
                    dimagi_contact__icontains=self.search_string)

        return query

    @property
    def account_name_response(self):
        return [self._fmt_select2_data(a.name, a.name)
                for a in self.paginated_data]

    @property
    def account_id_response(self):
        return [self._fmt_select2_data(a.salesforce_account_id,
                                       a.salesforce_account_id)
                for a in self.paginated_data]

    @property
    def dimagi_contact_response(self):
        return [self._fmt_select2_data(a.dimagi_contact, a.dimagi_contact)
                for a in self.paginated_data]

class BillingContactInfoAsyncHandler(BaseSingleOptionFilterAsyncHandler):
    slug = 'billing_contact_filter'
    allowed_actions = [
        'contact_name'
    ]

    @property
    def query(self):
        query = BillingContactInfo.objects.exclude(
            first_name='', last_name='').order_by('first_name', 'last_name')
        if self.search_string:
            query = query.filter(
                Q(first_name__istartswith=self.search_string) |
                Q(last_name__istartswith=self.search_string)
            )
        return query

    @property
    def contact_name_response(self):
        return [self._fmt_select2_data(c.full_name, c.full_name)
                for c in self.paginated_data]

class SoftwarePlanAsyncHandler(BaseSingleOptionFilterAsyncHandler):
    slug = 'software_plan_filter'
    allowed_actions = [
        'name',
    ]

    @property
    def query(self):
        query = SoftwarePlan.objects.order_by('name')
        if self.search_string:
            query = query.filter(name__icontains=self.search_string)
        return query

    @property
    def name_response(self):
        return [self._fmt_select2_data(p.name, p.name)
                for p in self.paginated_data]

class DomainFilterAsyncHandler(BaseSingleOptionFilterAsyncHandler):
    slug = 'domain_filter'
    allowed_actions = [
        'domain_name',
    ]

    @property
    def query(self):
        db = Domain.get_db()
        startkey = self.search_string
        endkey = "{}Z".format(self.search_string) if startkey else ''
        query = db.view(
            'domain/domains',
            reduce=False,
            startkey=startkey,
            endkey=endkey,
            limit=20,
        )
        return query

    @property
    def domain_name_response(self):
        return [self._fmt_select2_data(p['key'], p['key']) for p in self.paginated_data]
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

"""Copyright 2014 Roger R Labbe Jr.

filterpy library.
http://github.com/rlabbe/filterpy

Documentation at:
https://filterpy.readthedocs.org

Supporting book at:
https://github.com/rlabbe/Kalman-and-Bayesian-Filters-in-Python

This is licensed under an MIT license. See the readme.MD file
for more information.
"""

from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import math
import numpy as np
import numpy.linalg as linalg
import matplotlib.pyplot as plt
import scipy.sparse as sp
import scipy.sparse.linalg as spln
import scipy.stats
from scipy.stats import norm
from matplotlib.patches import Ellipse
import random

def gaussian(x, mean, var):
    """returns normal distribution (pdf) for x given a Gaussian with the
    specified mean and variance. All must be scalars.

    gaussian (1,2,3) is equivalent to scipy.stats.norm(2,math.sqrt(3)).pdf(1)
    It is quite a bit faster albeit much less flexible than the latter.

    **Parameters**

    x : scalar
        The value for which we compute the probability

    mean : scalar
        Mean of the Gaussian

    var : scalar
        Variance of the Gaussian

    **Returns**

    probability : float
        probability of x for the Gaussian (mean, var). E.g. 0.101 denotes
        10.1%.
    """

    return math.exp((-0.5*(x-mean)**2)/var) / math.sqrt(2*math.pi*var)

def mul (mean1, var1, mean2, var2):
    """ multiply Gaussians (mean1, var1) with (mean2, var2) and return the
    results as a tuple (mean,var).

    var1 and var2 are variances - sigma squared in the usual parlance.
    """

    mean = (var1*mean2 + var2*mean1) / (var1 + var2)
    var = 1 / (1/var1 + 1/var2)
    return (mean, var)

def add (mean1, var1, mean2, var2):
    """ add the Gaussians (mean1, var1) with (mean2, var2) and return the
    results as a tuple (mean,var).

    var1 and var2 are variances - sigma squared in the usual parlance.
    """

    return (mean1+mean2, var1+var2)

def multivariate_gaussian(x, mu, cov):
    """ This is designed to replace scipy.stats.multivariate_normal
    which is not available before version 0.14. You may either pass in a
    multivariate set of data:

       multivariate_gaussian (array([1,1]), array([3,4]), eye(2)*1.4)
       multivariate_gaussian (array([1,1,1]), array([3,4,5]), 1.4)

    or unidimensional data:
       multivariate_gaussian(1, 3, 1.4)

    In the multivariate case if cov is a scalar it is interpreted as eye(n)*cov

    The function gaussian() implements the 1D (univariate)case, and is much
    faster than this function.

    equivalent calls:
       multivariate_gaussian(1, 2, 3)
       scipy.stats.multivariate_normal(2,3).pdf(1)

    **Parameters**

    x : float, or np.array-like
       Value to compute the probability for. May be a scalar if univariate,
       or any type that can be converted to an np.array (list, tuple, etc).
       np.array is best for speed.

    mu :  float, or np.array-like
       mean for the Gaussian . May be a scalar if univariate,  or any type
       that can be converted to an np.array (list, tuple, etc).np.array is
       best for speed.

    cov :  float, or np.array-like
       Covariance for the Gaussian . May be a scalar if univariate,  or any
       type that can be converted to an np.array (list, tuple, etc).np.array is
       best for speed.

    **Returns**

    probability : float
        probability for x for the Gaussian (mu,cov)
    """

    scipy.stats.multivariate_normal
    # force all to numpy.array type
    x   = np.array(x, copy=False, ndmin=1)
    mu  = np.array(mu,copy=False, ndmin=1)

    nx = len(mu)
    cov = _to_cov(cov, nx)

    norm_coeff = nx*math.log(2*math.pi) + np.linalg.slogdet(cov)[1]

    err = x - mu
    if (sp.issparse(cov)):
        numerator = spln.spsolve(cov, err).T.dot(err)
    else:
        numerator = np.linalg.solve(cov, err).T.dot(err)

    return math.exp(-0.5*(norm_coeff + numerator))

def multivariate_multiply(m1, c1, m2, c2):
    """ Multiplies the two multivariate Gaussians together and returns the
    results as the tuple (mean, covariance).

    **example**

    m, c = multivariate_multiply([7.0, 2], [[1.0, 2.0], [2.0, 1.0]],
                                 [3.2, 0], [[8.0, 1.1], [1.1,8.0]])

    **Parameters**

    m1 : array-like
        Mean of first Gaussian. Must be convertable to an 1D array via
        numpy.asarray(), For example 6, [6], [6, 5], np.array([3, 4, 5, 6])
        are all valid.

    c1 : matrix-like
        Mean of first Gaussian. Must be convertable to an 2D array via
        numpy.asarray().

     m2 : array-like
        Mean of second Gaussian. Must be convertable to an 1D array via
        numpy.asarray(), For example 6, [6], [6, 5], np.array([3, 4, 5, 6])
        are all valid.

    c2 : matrix-like
        Mean of second Gaussian. Must be convertable to an 2D array via
        numpy.asarray().

    **Returns**

    m : ndarray
        mean of the result

    c : ndarray
        covariance of the result
    """

    C1 = np.asarray(c1)
    C2 = np.asarray(c2)
    M1 = np.asarray(m1)
    M2 = np.asarray(m2)

    sum_inv = np.linalg.inv(C1+C2)
    C3 = np.dot(C1, sum_inv).dot(C2)

    M3 = (np.dot(C2, sum_inv).dot(M1) +
          np.dot(C1, sum_inv).dot(M2))

    return M3, C3

def plot_gaussian(mean, variance,
                  mean_line=False,
                  xlim=None,
                  xlabel=None,
                  ylabel=None):
    """ plots the normal distribution with the given mean and variance. x-axis
    contains the mean, the y-axis shows the probability.

    **parameters**

    mean_line : boolean
        draws a line at x=mean

    xlim: (float,float), optional
        specify the limits for the x axis as tuple (low,high).
        If not specified, limits will be automatically chosen to be 'nice'

    xlabel : str,optional
        label for the x-axis

    ylabel : str, optional
        label for the y-axis
    """

    sigma = math.sqrt(variance)
    n = scipy.stats.norm(mean, sigma)

    if xlim is None:
        min_x = n.ppf(0.001)
        max_x = n.ppf(0.999)
    else:
        min_x = xlim[0]
        max_x = xlim[1]
    xs = np.arange(min_x, max_x, (max_x - min_x) / 1000)
    plt.plot(xs,n.pdf(xs))
    plt.xlim((min_x, max_x))

    if mean_line:
        plt.axvline(mean)
    if xlabel:
       plt.xlabel(xlabel)
    if ylabel:
       plt.ylabel(ylabel)

def covariance_ellipse(P, deviations=1):
    """ returns a tuple defining the ellipse representing the 2 dimensional
    covariance matrix P.

    **Parameters**

    P : nd.array shape (2,2)
       covariance matrix

    deviations : int (optional, default = 1)
       # of standard deviations. Default is 1.

    Returns (angle_radians, width_radius, height_radius)
    """

    U,s,v = linalg.svd(P)
    orientation = math.atan2(U[1,0],U[0,0])
    width  = deviations*math.sqrt(s[0])
    height = deviations*math.sqrt(s[1])

    assert width >= height
    return (orientation, width, height)

def plot_covariance_ellipse(mean, cov=None,
            variance = 1.0, std=None,
             ellipse=None, title=None, axis_equal=True,
             facecolor='none', edgecolor='#004080',
             alpha=1.0, xlim=None, ylim=None):
    """ plots the covariance ellipse where

    mean is a (x,y) tuple for the mean of the covariance (center of ellipse)

    cov is a 2x2 covariance matrix.

    variance is the normal sigma^2 that we want to plot. If list-like,
    ellipses for all ellipses will be ploted. E.g. [1,2] will plot the
    sigma^2 = 1 and sigma^2 = 2 ellipses.

    ellipse is a (angle,width,height) tuple containing the angle in radians,
    and width and height radii.

    You may provide either cov or ellipse, but not both.

    plt.show() is not called, allowing you to plot multiple things on the
    same figure.
    """

    assert cov is None or ellipse is None
    assert not (cov is None and ellipse is None)

    if cov is not None:
        ellipse = covariance_ellipse(cov)

    if axis_equal:
        plt.axis('equal')

    if title is not None:
        plt.title (title)

    compute_std = False
    if std is None:
        std = variance
        compute_std = True

    if np.isscalar(std):
            std = [std]

    if compute_std:
        std = np.sqrt(np.asarray(std))

    ax = plt.gca()

    angle = np.degrees(ellipse[0])
    width = ellipse[1] * 2.
    height = ellipse[2] * 2.

    for sd in std:
        e = Ellipse(xy=mean, width=sd*width, height=sd*height, angle=angle,
                    facecolor=facecolor,
                    edgecolor=edgecolor,
                    alpha=alpha,
                    lw=2)
        ax.add_patch(e)
    plt.scatter(mean[0], mean[1], marker='+') # mark the center
    if xlim is not None:
        ax.set_xlim(xlim)

    if ylim is not None:
        ax.set_ylim(ylim)

def norm_cdf (x_range, mu, var=1, std=None):
    """ computes the probability that a Gaussian distribution lies
    within a range of values.

    **Paramateters**

    x_range : (float, float)
        tuple of range to compute probability for

    mu : float
        mean of the Gaussian

    var : float, optional
        variance of the Gaussian. Ignored if std is provided

    std : float, optional
       standard deviation of the Gaussian. This overrides the var parameter

    **Returns**

    probability : float
        probability that Gaussian is within x_range. E.g. .1 means 10%.
    """

    if std is None:
        std = math.sqrt(var)
    return abs(norm.cdf(x_range[0], loc=mu, scale=std) -
               norm.cdf(x_range[1], loc=mu, scale=std))

def _is_inside_ellipse(x, y, ex, ey, orientation, width, height):

    co = np.cos(orientation)
    so = np.sin(orientation)

    xx = x*co + y*so
    yy = y*co - x*so

    return (xx / width)**2 + (yy / height)**2 <= 1.

    return ((x-ex)*co - (y-ey)*so)**2/width**2 + \
           ((x-ex)*so + (y-ey)*co)**2/height**2 <= 1

def _to_cov(x,n):
    """ If x is a scalar, returns a covariance matrix generated from it
    as the identity matrix multiplied by x. The dimension will be nxn.
    If x is already a numpy array then it is returned unchanged.
    """
    try:
        x.shape
        if type(x) != np.ndarray:
            x = np.asarray(x)[0]
        return x
    except:
        return np.eye(n) * x

def _do_plot_test():

    from numpy.random import multivariate_normal
    p = np.array([[32, 15],[15., 40.]])

    x,y = multivariate_normal(mean=(0,0), cov=p, size=5000).T
    sd = 2
    a,w,h = covariance_ellipse(p,sd)
    print (np.degrees(a), w, h)

    count = 0
    color=[]
    for i in range(len(x)):
        if _is_inside_ellipse(x[i], y[i], 0, 0, a, w, h):
            color.append('b')
            count += 1
        else:
            color.append('r')
    plt.scatter(x,y,alpha=0.2, c=color)

    plt.axis('equal')

    plot_covariance_ellipse(mean=(0., 0.),
                            cov = p,
                            std=sd,
                            facecolor='none')

    print (count / len(x))

def plot_std_vs_var():
    plt.figure()
    x = (0,0)
    P = np.array([[3,1],[1,3]])
    plot_covariance_ellipse(x, P, std=[1,2,3], facecolor='g', alpha=.2)
    plot_covariance_ellipse(x, P, variance=[1,2,3], facecolor='r', alpha=.5)

def rand_student_t(df, mu=0, std=1):
    """return random number distributed by student's t distribution with
    `df` degrees of freedom with the specified mean and standard deviation.
    """
    x = random.gauss(0, std)
    y = 2.0*random.gammavariate(0.5*df, 2.0)
    return x / (math.sqrt(y/df)) + mu

if __name__ == '__main__':
    plot_std_vs_var()
    plt.figure()

    _do_plot_test()

    #test_gaussian()

    # test conversion of scalar to covariance matrix
    x  = multivariate_gaussian(np.array([1,1]), np.array([3,4]), np.eye(2)*1.4)
    x2 = multivariate_gaussian(np.array([1,1]), np.array([3,4]), 1.4)
    assert x == x2

    # test univarate case
    rv = norm(loc = 1., scale = np.sqrt(2.3))
    x2 = multivariate_gaussian(1.2, 1., 2.3)
    x3 = gaussian(1.2, 1., 2.3)

    assert rv.pdf(1.2) == x2
    assert abs(x2- x3) < 0.00000001

    cov = np.array([[1.0, 1.0],
                    [1.0, 1.1]])

    plt.figure()
    P = np.array([[2,0],[0,2]])
    plot_covariance_ellipse((2,7), cov=cov, variance=[1,2], facecolor='g', title='my title', alpha=.2)
    plt.show()

    print("all tests passed")
<commit_msg>Added array support to gaussian()<commit_after># -*- coding: utf-8 -*-

"""Copyright 2014 Roger R Labbe Jr.

filterpy library.
http://github.com/rlabbe/filterpy

Documentation at:
https://filterpy.readthedocs.org

Supporting book at:
https://github.com/rlabbe/Kalman-and-Bayesian-Filters-in-Python

This is licensed under an MIT license. See the readme.MD file
for more information.
"""

from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import math
import numpy as np
import numpy.linalg as linalg
import matplotlib.pyplot as plt
import scipy.sparse as sp
import scipy.sparse.linalg as spln
import scipy.stats
from scipy.stats import norm
from matplotlib.patches import Ellipse
import random

def gaussian(x, mean, var):
    """returns normal distribution (pdf) for x given a Gaussian with the
    specified mean and variance. All must be scalars.

    gaussian (1,2,3) is equivalent to scipy.stats.norm(2,math.sqrt(3)).pdf(1)
    It is quite a bit faster albeit much less flexible than the latter.

    **Parameters**

    x : scalar or array-like
        The value for which we compute the probability

    mean : scalar
        Mean of the Gaussian

    var : scalar
        Variance of the Gaussian

    **Returns**

    probability : float
        probability of x for the Gaussian (mean, var). E.g. 0.101 denotes
        10.1%.

    **Examples**

    gaussian(8, 1, 2)
    gaussian([8, 7, 9], 1, 2)
    """

    return (np.exp((-0.5*(np.asarray(x)-mean)**2)/var) /
            math.sqrt(2*math.pi*var))

def mul (mean1, var1, mean2, var2):
    """ multiply Gaussians (mean1, var1) with (mean2, var2) and return the
    results as a tuple (mean,var).

    var1 and var2 are variances - sigma squared in the usual parlance.
    """

    mean = (var1*mean2 + var2*mean1) / (var1 + var2)
    var = 1 / (1/var1 + 1/var2)
    return (mean, var)

def add (mean1, var1, mean2, var2):
    """ add the Gaussians (mean1, var1) with (mean2, var2) and return the
    results as a tuple (mean,var).

    var1 and var2 are variances - sigma squared in the usual parlance.
    """

    return (mean1+mean2, var1+var2)

def multivariate_gaussian(x, mu, cov):
    """ This is designed to replace scipy.stats.multivariate_normal
    which is not available before version 0.14. You may either pass in a
    multivariate set of data:

       multivariate_gaussian (array([1,1]), array([3,4]), eye(2)*1.4)
       multivariate_gaussian (array([1,1,1]), array([3,4,5]), 1.4)

    or unidimensional data:
       multivariate_gaussian(1, 3, 1.4)

    In the multivariate case if cov is a scalar it is interpreted as eye(n)*cov

    The function gaussian() implements the 1D (univariate)case, and is much
    faster than this function.

    equivalent calls:
       multivariate_gaussian(1, 2, 3)
       scipy.stats.multivariate_normal(2,3).pdf(1)

    **Parameters**

    x : float, or np.array-like
       Value to compute the probability for. May be a scalar if univariate,
       or any type that can be converted to an np.array (list, tuple, etc).
       np.array is best for speed.

    mu :  float, or np.array-like
       mean for the Gaussian . May be a scalar if univariate,  or any type
       that can be converted to an np.array (list, tuple, etc).np.array is
       best for speed.

    cov :  float, or np.array-like
       Covariance for the Gaussian . May be a scalar if univariate,  or any
       type that can be converted to an np.array (list, tuple, etc).np.array is
       best for speed.

    **Returns**

    probability : float
        probability for x for the Gaussian (mu,cov)
    """

    scipy.stats.multivariate_normal
    # force all to numpy.array type
    x   = np.array(x, copy=False, ndmin=1)
    mu  = np.array(mu,copy=False, ndmin=1)

    nx = len(mu)
    cov = _to_cov(cov, nx)

    norm_coeff = nx*math.log(2*math.pi) + np.linalg.slogdet(cov)[1]

    err = x - mu
    if (sp.issparse(cov)):
        numerator = spln.spsolve(cov, err).T.dot(err)
    else:
        numerator = np.linalg.solve(cov, err).T.dot(err)

    return math.exp(-0.5*(norm_coeff + numerator))

def multivariate_multiply(m1, c1, m2, c2):
    """ Multiplies the two multivariate Gaussians together and returns the
    results as the tuple (mean, covariance).

    **example**

    m, c = multivariate_multiply([7.0, 2], [[1.0, 2.0], [2.0, 1.0]],
                                 [3.2, 0], [[8.0, 1.1], [1.1,8.0]])

    **Parameters**

    m1 : array-like
        Mean of first Gaussian. Must be convertable to an 1D array via
        numpy.asarray(), For example 6, [6], [6, 5], np.array([3, 4, 5, 6])
        are all valid.

    c1 : matrix-like
        Mean of first Gaussian. Must be convertable to an 2D array via
        numpy.asarray().

     m2 : array-like
        Mean of second Gaussian. Must be convertable to an 1D array via
        numpy.asarray(), For example 6, [6], [6, 5], np.array([3, 4, 5, 6])
        are all valid.

    c2 : matrix-like
        Mean of second Gaussian. Must be convertable to an 2D array via
        numpy.asarray().

    **Returns**

    m : ndarray
        mean of the result

    c : ndarray
        covariance of the result
    """

    C1 = np.asarray(c1)
    C2 = np.asarray(c2)
    M1 = np.asarray(m1)
    M2 = np.asarray(m2)

    sum_inv = np.linalg.inv(C1+C2)
    C3 = np.dot(C1, sum_inv).dot(C2)

    M3 = (np.dot(C2, sum_inv).dot(M1) +
          np.dot(C1, sum_inv).dot(M2))

    return M3, C3

def plot_gaussian(mean, variance,
                  mean_line=False,
                  xlim=None,
                  xlabel=None,
                  ylabel=None):
    """ plots the normal distribution with the given mean and variance. x-axis
    contains the mean, the y-axis shows the probability.

    **parameters**

    mean_line : boolean
        draws a line at x=mean

    xlim: (float,float), optional
        specify the limits for the x axis as tuple (low,high).
        If not specified, limits will be automatically chosen to be 'nice'

    xlabel : str,optional
        label for the x-axis

    ylabel : str, optional
        label for the y-axis
    """

    sigma = math.sqrt(variance)
    n = scipy.stats.norm(mean, sigma)

    if xlim is None:
        min_x = n.ppf(0.001)
        max_x = n.ppf(0.999)
    else:
        min_x = xlim[0]
        max_x = xlim[1]
    xs = np.arange(min_x, max_x, (max_x - min_x) / 1000)
    plt.plot(xs,n.pdf(xs))
    plt.xlim((min_x, max_x))

    if mean_line:
        plt.axvline(mean)
    if xlabel:
       plt.xlabel(xlabel)
    if ylabel:
       plt.ylabel(ylabel)

def covariance_ellipse(P, deviations=1):
    """ returns a tuple defining the ellipse representing the 2 dimensional
    covariance matrix P.

    **Parameters**

    P : nd.array shape (2,2)
       covariance matrix

    deviations : int (optional, default = 1)
       # of standard deviations. Default is 1.

    Returns (angle_radians, width_radius, height_radius)
    """

    U,s,v = linalg.svd(P)
    orientation = math.atan2(U[1,0],U[0,0])
    width  = deviations*math.sqrt(s[0])
    height = deviations*math.sqrt(s[1])

    assert width >= height
    return (orientation, width, height)

def plot_covariance_ellipse(mean, cov=None,
            variance = 1.0, std=None,
             ellipse=None, title=None, axis_equal=True,
             facecolor='none', edgecolor='#004080',
             alpha=1.0, xlim=None, ylim=None):
    """ plots the covariance ellipse where

    mean is a (x,y) tuple for the mean of the covariance (center of ellipse)

    cov is a 2x2 covariance matrix.

    variance is the normal sigma^2 that we want to plot. If list-like,
    ellipses for all ellipses will be ploted. E.g. [1,2] will plot the
    sigma^2 = 1 and sigma^2 = 2 ellipses.

    ellipse is a (angle,width,height) tuple containing the angle in radians,
    and width and height radii.

    You may provide either cov or ellipse, but not both.

    plt.show() is not called, allowing you to plot multiple things on the
    same figure.
    """

    assert cov is None or ellipse is None
    assert not (cov is None and ellipse is None)

    if cov is not None:
        ellipse = covariance_ellipse(cov)

    if axis_equal:
        plt.axis('equal')

    if title is not None:
        plt.title (title)

    compute_std = False
    if std is None:
        std = variance
        compute_std = True

    if np.isscalar(std):
            std = [std]

    if compute_std:
        std = np.sqrt(np.asarray(std))

    ax = plt.gca()

    angle = np.degrees(ellipse[0])
    width = ellipse[1] * 2.
    height = ellipse[2] * 2.

    for sd in std:
        e = Ellipse(xy=mean, width=sd*width, height=sd*height, angle=angle,
                    facecolor=facecolor,
                    edgecolor=edgecolor,
                    alpha=alpha,
                    lw=2)
        ax.add_patch(e)
    plt.scatter(mean[0], mean[1], marker='+') # mark the center
    if xlim is not None:
        ax.set_xlim(xlim)

    if ylim is not None:
        ax.set_ylim(ylim)

def norm_cdf (x_range, mu, var=1, std=None):
    """ computes the probability that a Gaussian distribution lies
    within a range of values.

    **Paramateters**

    x_range : (float, float)
        tuple of range to compute probability for

    mu : float
        mean of the Gaussian

    var : float, optional
        variance of the Gaussian. Ignored if std is provided

    std : float, optional
       standard deviation of the Gaussian. This overrides the var parameter

    **Returns**

    probability : float
        probability that Gaussian is within x_range. E.g. .1 means 10%.
    """

    if std is None:
        std = math.sqrt(var)
    return abs(norm.cdf(x_range[0], loc=mu, scale=std) -
               norm.cdf(x_range[1], loc=mu, scale=std))

def _is_inside_ellipse(x, y, ex, ey, orientation, width, height):

    co = np.cos(orientation)
    so = np.sin(orientation)

    xx = x*co + y*so
    yy = y*co - x*so

    return (xx / width)**2 + (yy / height)**2 <= 1.

    return ((x-ex)*co - (y-ey)*so)**2/width**2 + \
           ((x-ex)*so + (y-ey)*co)**2/height**2 <= 1

def _to_cov(x,n):
    """ If x is a scalar, returns a covariance matrix generated from it
    as the identity matrix multiplied by x. The dimension will be nxn.
    If x is already a numpy array then it is returned unchanged.
    """
    try:
        x.shape
        if type(x) != np.ndarray:
            x = np.asarray(x)[0]
        return x
    except:
        return np.eye(n) * x

def _do_plot_test():

    from numpy.random import multivariate_normal
    p = np.array([[32, 15],[15., 40.]])

    x,y = multivariate_normal(mean=(0,0), cov=p, size=5000).T
    sd = 2
    a,w,h = covariance_ellipse(p,sd)
    print (np.degrees(a), w, h)

    count = 0
    color=[]
    for i in range(len(x)):
        if _is_inside_ellipse(x[i], y[i], 0, 0, a, w, h):
            color.append('b')
            count += 1
        else:
            color.append('r')
    plt.scatter(x,y,alpha=0.2, c=color)

    plt.axis('equal')

    plot_covariance_ellipse(mean=(0., 0.),
                            cov = p,
                            std=sd,
                            facecolor='none')

    print (count / len(x))

def plot_std_vs_var():
    plt.figure()
    x = (0,0)
    P = np.array([[3,1],[1,3]])
    plot_covariance_ellipse(x, P, std=[1,2,3], facecolor='g', alpha=.2)
    plot_covariance_ellipse(x, P, variance=[1,2,3], facecolor='r', alpha=.5)

def rand_student_t(df, mu=0, std=1):
    """return random number distributed by student's t distribution with
    `df` degrees of freedom with the specified mean and standard deviation.
    """
    x = random.gauss(0, std)
    y = 2.0*random.gammavariate(0.5*df, 2.0)
    return x / (math.sqrt(y/df)) + mu

if __name__ == '__main__':
    plot_std_vs_var()
    plt.figure()

    _do_plot_test()

    #test_gaussian()

    # test conversion of scalar to covariance matrix
    x  = multivariate_gaussian(np.array([1,1]), np.array([3,4]), np.eye(2)*1.4)
    x2 = multivariate_gaussian(np.array([1,1]), np.array([3,4]), 1.4)
    assert x == x2

    # test univarate case
    rv = norm(loc = 1., scale = np.sqrt(2.3))
    x2 = multivariate_gaussian(1.2, 1., 2.3)
    x3 = gaussian(1.2, 1., 2.3)

    assert rv.pdf(1.2) == x2
    assert abs(x2- x3) < 0.00000001

    cov = np.array([[1.0, 1.0],
                    [1.0, 1.1]])

    plt.figure()
    P = np.array([[2,0],[0,2]])
    plot_covariance_ellipse((2,7), cov=cov, variance=[1,2], facecolor='g', title='my title', alpha=.2)
    plt.show()

    print("all tests passed")
<|endoftext|>


<commit_before>"""
wrappers.py

Wrappers for mido messages.

The intent of these is to have a more friendly way of
displaying and working with control and sysex messages.
"""
# I'm not sure if that goal is achieved. May have gone overboard with classes

# Should these be wrappers, subclasses, or entirely different???
import re

class WrappedMessage(object):
    # These are wrapped mido messages.
    # Like normal messages, but with some extra metadata.
    type = "message"

    def __init__(self, message):
        """
        Wrap a mido message.
        """
        # I suppose we could use FrozenMessages here,
        # but we are all responsible adults, right?
        self.message = message

    def __str__(self):
        return str(self.message)

    def __repr__(self):
        return "<{} {!r}>".format(self.type, self.message)

class WrappedSysEx(WrappedMessage):
    type = "sysex"

    def __init__(self, message, match=None):
        super().__init__(self, message)
        self._process(match)

    def _process(self, match):
        pass

class WrappedGMSystemOn(WrappedSysEx):
    # GM System ON, F0 7E 7F 09 01 F7
    type = "gm_on"
    REGEX = re.compile(rb'\x7E\x7F\x09\x01\xF7', re.S)

class WrappedMIDIMasterVolume(WrappedSysEx):
    # MIDI Master Volume, F0 7F 7F 04 01 ** mm F7
    type = "master_vol"
    REGEX = re.compile(rb'\x7F\x7F\x04\x01.(.)', re.S)
    def _process(self, match):
        self.value, = match.group(1)

class WrappedMIDIMasterTuning(WrappedSysEx):
    # MIDI Master Tuning, F0 43 1* 27 30 00 00 *m *l ** F7
    type = "master_tuning"
    REGEX = re.compile(rb'\x43[\x10-\x1F]\x27\x30\x00\x00(..).', re.S)
    def _process(self, match):
        self.msb, self.lsb = match.group(1)
        # we take the least significant nybble of each.
        ml = ((self.msb & 0xF) << 4) | (self.lsb & 0xF)
        # subtract 0x80 and clamp to range -100 +100
        self.value = max(-100, min(ml-0x80, +100))

class WrappedReverbChorus(WrappedSysEx):
    # Reverb Type, F0 43 1n 4C 02 01 00 mm ll F7
    # Chorus Type, F0 43 1n 4C 02 01 20 mm ll F7
    REGEX = re.compile(rb'\x43[\x10-\x1F]\x4C\x02\x01\x00([\x00\x20]..)', re.S)
    def _process(self, match):
        cat, self.msb, self.lsb = match.group(1)
        if cat == 0x00:
            self.type = "reverb"
        elif cat == 0x20:
            self.type = "chorus"

def wrap_sysex(message):

    # Put the data into a bytes object so we can regex it
    data = bytes(message.data)
    # Match one by one
    for sysex_class in (WrappedGMSystemOn, WrappedMIDIMasterVolume,
            WrappedMIDIMasterTuning, WrappedReverbChorus):
        match = sysex_class.REGEX.fullmatch(data)
        if match is not None:
            return sysex_class(message, match)
    # if not, just wrap it generically
    return WrappedSysEx(message)

class WrappedControlChange(WrappedMessage):
    type = "control_change"

    def __init__(self, message):
        super().__init__(self, message)
        self.channel = message.channel
        self.control = message.control
        self.control_value = message.value
        self._process(self.control_value)

    def _process(self, value):
        pass

class WrappedLocal(WrappedControlChange):
    type = "local"
    CONTROL = 0x7A

    def _process(self, value):
        # highest bit: 1 for ON, 0 for OFF.
        self.local = value >= 64

def wrap_control(message):
    for control_class in (WrappedLocal,):
        if message.control == control_class.CONTROL:
            return control_class(message)
    return WrappedControlChange(message)

class WrappedProgramChange(WrappedMessage):
    type = "program_change"

    def __init__(self, message):
        super().__init__(self, message)
        self.channel = message.channel
        self.program = message.program

def wrap(message):
    if message.type == "program_change":
        return WrappedProgramChange(message)
    elif message.type == "control_change":
        return wrap_control(message)
    elif message.type == "sysex":
        return wrap_sysex(message)
    else:
        return WrappedMessage(message)

<commit_msg>change wrappers to use enums<commit_after>"""
wrappers.py

Wrappers for mido messages.

The intent of these is to have a more friendly way of
displaying and working with control and sysex messages.
"""
# I'm not sure if that goal is achieved. May have gone overboard with classes

# Should these be wrappers, subclasses, or entirely different???
import re

from ..enums import ReverbType, ChorusType

class WrappedMessage(object):
    # These are wrapped mido messages.
    # Like normal messages, but with some extra metadata.
    type = "message"

    def __init__(self, message):
        """
        Wrap a mido message.
        """
        # I suppose we could use FrozenMessages here,
        # but we are all responsible adults, right?
        self.message = message

    def __str__(self):
        return str(self.message)

    def __repr__(self):
        return "<{} {!r}>".format(self.type, self.message)

class WrappedSysEx(WrappedMessage):
    type = "sysex"

    def __init__(self, message, match=None):
        super().__init__(self, message)
        self._process(match)

    def _process(self, match):
        pass

class WrappedGMSystemOn(WrappedSysEx):
    # GM System ON, F0 7E 7F 09 01 F7
    type = "gm_on"
    REGEX = re.compile(rb'\x7E\x7F\x09\x01\xF7', re.S)

class WrappedMIDIMasterVolume(WrappedSysEx):
    # MIDI Master Volume, F0 7F 7F 04 01 ** mm F7
    type = "master_vol"
    REGEX = re.compile(rb'\x7F\x7F\x04\x01.(.)', re.S)
    def _process(self, match):
        self.value, = match.group(1)

class WrappedMIDIMasterTuning(WrappedSysEx):
    # MIDI Master Tuning, F0 43 1* 27 30 00 00 *m *l ** F7
    type = "master_tuning"
    REGEX = re.compile(rb'\x43[\x10-\x1F]\x27\x30\x00\x00(..).', re.S)
    def _process(self, match):
        self.msb, self.lsb = match.group(1)
        # we take the least significant nybble of each.
        ml = ((self.msb & 0xF) << 4) | (self.lsb & 0xF)
        # subtract 0x80 and clamp to range -100 +100
        self.value = max(-100, min(ml-0x80, +100))

class WrappedReverbChorus(WrappedSysEx):
    # Reverb Type, F0 43 1n 4C 02 01 00 mm ll F7
    # Chorus Type, F0 43 1n 4C 02 01 20 mm ll F7
    TABLE = (None, {})

    def _process(self, match):
        self.msb, self.lsb = match.group(1)
        m_default, m_mapping = self.TABLE
        try:
            l_default, l_mapping = m_mapping[self.msb]
        except KeyError:
            self.value = m_default
        else:
            self.value = l_mapping.get(self.lsb, l_default)

class WrappedReverbType(WrappedReverbChorus):
    # Reverb Type, F0 43 1n 4C 02 01 00 mm ll F7
    type = "reverb"
    REGEX = re.compile(rb'\x43[\x10-\x1F]\x4C\x02\x01\x00\x00(..)', re.S)
    TABLE = (ReverbType.OFF, {
        0x01: (ReverbType.HALL1, {
            0x10: ReverbType.HALL2,
            0x11: ReverbType.HALL3,
        }),
        0x02: (ReverbType.ROOM, {
            0x11: ReverbType.ROOM1,
            0x13: ReverbType.ROOM2,
        }),
        0x03: (ReverbType.STAGE, {
            0x10: ReverbType.STAGE1,
            0x11: ReverbType.STAGE2,
        }),
        0x04: (ReverbType.PLATE, {
            0x10: ReverbType.PLATE1,
            0x11: ReverbType.PLATE2,
        }),
    })

class WrappedChorusType(WrappedReverbChorus):
    # Chorus Type, F0 43 1n 4C 02 01 20 mm ll F7
    type = "chorus"
    REGEX = re.compile(rb'\x43[\x10-\x1F]\x4C\x02\x01\x00\x20(..)', re.S)
    TABLE = (ChorusType.OFF, {
        0x40: (ChorusType.THRU, {}),
        0x41: (ChorusType.CHORUS, {
            0x02: ChorusType.CHORUS2
        }),
        0x42: (ChorusType.CELESTE, {
            0x11: ChorusType.CHORUS1
        }),
        0x43: (ChorusType.FLANGER, {
            0x08: ChorusType.FLANGER1,
            0x11: ChorusType.FLANGER2
        })
    })

def wrap_sysex(message):

    # Put the data into a bytes object so we can regex it
    data = bytes(message.data)
    # Match one by one
    for sysex_class in (WrappedGMSystemOn, WrappedMIDIMasterVolume,
            WrappedMIDIMasterTuning, WrappedReverbType, WrappedChorusType):
        match = sysex_class.REGEX.fullmatch(data)
        if match is not None:
            return sysex_class(message, match)
    # if not, just wrap it generically
    return WrappedSysEx(message)

class WrappedControlChange(WrappedMessage):
    type = "control_change"

    def __init__(self, message):
        super().__init__(self, message)
        self.channel = message.channel
        self.control = message.control
        self.control_value = message.value
        self._process(self.control_value)

    def _process(self, value):
        pass

class WrappedLocal(WrappedControlChange):
    type = "local"
    CONTROL = 0x7A

    def _process(self, value):
        # highest bit: 1 for ON, 0 for OFF.
        self.local = value >= 64

def wrap_control(message):
    for control_class in (WrappedLocal,):
        if message.control == control_class.CONTROL:
            return control_class(message)
    return WrappedControlChange(message)

class WrappedProgramChange(WrappedMessage):
    type = "program_change"

    def __init__(self, message):
        super().__init__(self, message)
        self.channel = message.channel
        self.program = message.program

def wrap(message):
    if message.type == "program_change":
        return WrappedProgramChange(message)
    elif message.type == "control_change":
        return wrap_control(message)
    elif message.type == "sysex":
        return wrap_sysex(message)
    else:
        return WrappedMessage(message)

<|endoftext|>


<commit_before>#!/usr/local/sal/Python.framework/Versions/Current/bin/python3

import pathlib
import plistlib
import re
import subprocess
import sys
from xml.etree import ElementTree

import macsesh
from SystemConfiguration import (
    SCDynamicStoreCreate,
    SCDynamicStoreCopyValue,
    SCDynamicStoreCopyConsoleUser,
)

import sal

MODEL_PATH = pathlib.Path("/usr/local/sal/model_cache")
MEMORY_EXPONENTS = {"KB": 0, "MB": 1, "GB": 2, "TB": 3}
__version__ = "1.1.0"

def main():
    machine_results = {"facts": {"checkin_module_version": __version__}}
    extras = {}
    extras["hostname"] = get_hostname()
    extras["os_family"] = "Darwin"
    extras["console_user"] = get_console_user()[0]
    extras.update(process_system_profile())
    machine_results["extra_data"] = extras
    sal.set_checkin_results("Machine", machine_results)

def process_system_profile():
    machine_results = {}
    system_profile = get_sys_profile()

    if not system_profile:
        # We can't continue if system_profiler dies.
        return machine_results

    machine_results["serial"] = system_profile["SPHardwareDataType"][0].get(
        "serial_number"
    )
    os_version = system_profile["SPSoftwareDataType"][0].get("os_version").split()[1]
    if os_version == "X":
        os_version = (
            system_profile["SPSoftwareDataType"][0].get("os_version").split()[2]
        )
    machine_results["operating_system"] = os_version
    machine_results["machine_model"] = system_profile["SPHardwareDataType"][0].get(
        "machine_model"
    )

    udid = system_profile["SPHardwareDataType"][0].get("provisioning_UDID")
    friendly_model = get_friendly_model(serial=machine_results["serial"], udid=udid)
    if friendly_model:
        machine_results["machine_model_friendly"] = friendly_model
    if system_profile["SPHardwareDataType"][0].get("chip_type", None):
        machine_results["cpu_type"] = system_profile["SPHardwareDataType"][0].get(
            "chip_type", ""
        )
    else:
        machine_results["cpu_type"] = system_profile["SPHardwareDataType"][0].get(
            "cpu_type", ""
        )
    machine_results["cpu_speed"] = system_profile["SPHardwareDataType"][0].get(
        "current_processor_speed", ""
    )
    machine_results["memory"] = system_profile["SPHardwareDataType"][0].get(
        "physical_memory", ""
    )
    machine_results["memory_kb"] = process_memory(machine_results["memory"])

    for device in system_profile["SPStorageDataType"]:
        if device["mount_point"] == "/":
            # div by 1000.0 to
            # a) Convert to Apple base 10 kilobytes
            # b) Cast to python floats
            machine_results["hd_space"] = device["free_space_in_bytes"]
            machine_results["hd_total"] = device["size_in_bytes"]
            # We want the % used, not of free space, so invert.
            machine_results["hd_percent"] = "{:.2f}".format(
                abs(
                    float(machine_results["hd_space"]) / machine_results["hd_total"] - 1
                )
                * 100
            )

    return machine_results

def get_hostname():
    name_type = sal.sal_pref("NameType", default="ComputerName")
    net_config = SCDynamicStoreCreate(None, "net", None, None)
    return get_machine_name(net_config, name_type)

def get_machine_name(net_config, nametype):
    """Return the ComputerName of this Mac."""
    sys_info = SCDynamicStoreCopyValue(net_config, "Setup:/System")
    if sys_info:
        return sys_info.get(nametype)
    return subprocess.check_output(
        ["/usr/sbin/scutil", "--get", "ComputerName"], text=True
    ).strip()

def get_friendly_model(serial, udid):
    """Return friendly model name"""
    cmd = ["/usr/sbin/ioreg", "-arc", "IOPlatformDevice", "-k", "product-name"]
    try:
        out = subprocess.check_output(cmd)
    except:
        pass
    if out:
        try:
            data = plistlib.loads(out)
            if len(data) != 0:
                return data[0].get("product-name").encode("ascii", "ignore").strip()
        except:
            pass

    # set up cache file for this udid...create dir,
    MODEL_PATH.mkdir(mode=0o755, parents=True, exist_ok=True)

    # name cache for this udid
    UDID_CACHE_PATH = pathlib.Path(MODEL_PATH, "%s.txt" % (udid))
    for cache_file in MODEL_PATH.iterdir():
        # clean up any other files in dir
        if cache_file != UDID_CACHE_PATH:
            try:
                cache_file.unlink()
            except:
                pass

    if not UDID_CACHE_PATH.exists():
        model = cleanup_model(query_apple_support(serial))
        if model:
            UDID_CACHE_PATH.write_text(model)
    else:
        try:
            model = UDID_CACHE_PATH.read_text().strip()
        except:
            model = None

    return model

def get_model_code(serial):
    # Workaround for machines with dummy serial numbers.
    if "serial" in serial.lower():
        return

    if 12 <= len(serial) <= 13:
        if serial.startswith("S"):
            # Remove S prefix from scanned codes.
            serial = serial[1:]
        return serial[8:].upper()

    elif 11 <= len(serial) <= 12:
        # 2010 Mac Pros starting with H or Y are 11 characters
        return serial[8:].upper()

    return

def query_apple_support(serial):
    model_code = get_model_code(serial)
    tree = ElementTree.ElementTree()
    session = macsesh.Session()
    response = session.get(
        f"https://support-sp.apple.com/sp/product?cc={model_code}&lang=en_US"
    )
    try:
        tree = ElementTree.fromstring(response.text)
    except ElementTree.ParseError:
        tree = None
    return tree.findtext("configCode") if tree else None

def cleanup_model(model):
    cleanup_res = [(re.compile(r"inch ? "), "inch, "), (re.compile(r"  "), " ")]
    if model:
        for pattern, replacement in cleanup_res:
            model = pattern.sub(replacement, model)
    return model

def process_memory(amount):
    """Convert the amount of memory like '4 GB' to the size in kb as int"""
    try:
        memkb = int(amount[:-3]) * 1024 ** MEMORY_EXPONENTS[amount[-2:]]
    except ValueError:
        memkb = int(float(amount[:-3])) * 1024 ** MEMORY_EXPONENTS[amount[-2:]]
    return memkb

def get_sys_profile():
    """Get sysprofiler info.

    Returns:
        System Profiler report for networking, drives, and hardware as a
        plist dict, or an empty dict.
    """
    command = [
        "/usr/sbin/system_profiler",
        "-xml",
        "SPHardwareDataType",
        "SPStorageDataType",
        "SPSoftwareDataType",
    ]
    try:
        output = subprocess.check_output(command)
    except subprocess.CalledProcessError:
        return {}

    try:
        system_profile = plistlib.loads(output)
    except plistlib.InvalidException:
        system_profile = {}

    # sytem_profiler gives us back an array; convert to a dict with just
    # the data we care about.
    results = {}
    for data_type in system_profile:
        key = data_type["_dataType"]
        results[key] = data_type["_items"]

    return results

def get_console_user():
    """Get informatino about the console user

    Returns:
        3-Tuple of (str) username, (int) uid, (int) gid
    """
    return SCDynamicStoreCopyConsoleUser(None, None, None)

if __name__ == "__main__":
    main()
<commit_msg>decode<commit_after>#!/usr/local/sal/Python.framework/Versions/Current/bin/python3

import pathlib
import plistlib
import re
import subprocess
import sys
from xml.etree import ElementTree

import macsesh
from SystemConfiguration import (
    SCDynamicStoreCreate,
    SCDynamicStoreCopyValue,
    SCDynamicStoreCopyConsoleUser,
)

import sal

MODEL_PATH = pathlib.Path("/usr/local/sal/model_cache")
MEMORY_EXPONENTS = {"KB": 0, "MB": 1, "GB": 2, "TB": 3}
__version__ = "1.1.0"

def main():
    machine_results = {"facts": {"checkin_module_version": __version__}}
    extras = {}
    extras["hostname"] = get_hostname()
    extras["os_family"] = "Darwin"
    extras["console_user"] = get_console_user()[0]
    extras.update(process_system_profile())
    machine_results["extra_data"] = extras
    sal.set_checkin_results("Machine", machine_results)

def process_system_profile():
    machine_results = {}
    system_profile = get_sys_profile()

    if not system_profile:
        # We can't continue if system_profiler dies.
        return machine_results

    machine_results["serial"] = system_profile["SPHardwareDataType"][0].get(
        "serial_number"
    )
    os_version = system_profile["SPSoftwareDataType"][0].get("os_version").split()[1]
    if os_version == "X":
        os_version = (
            system_profile["SPSoftwareDataType"][0].get("os_version").split()[2]
        )
    machine_results["operating_system"] = os_version
    machine_results["machine_model"] = system_profile["SPHardwareDataType"][0].get(
        "machine_model"
    )

    udid = system_profile["SPHardwareDataType"][0].get("provisioning_UDID")
    friendly_model = get_friendly_model(serial=machine_results["serial"], udid=udid)
    if friendly_model:
        machine_results["machine_model_friendly"] = friendly_model
    if system_profile["SPHardwareDataType"][0].get("chip_type", None):
        machine_results["cpu_type"] = system_profile["SPHardwareDataType"][0].get(
            "chip_type", ""
        )
    else:
        machine_results["cpu_type"] = system_profile["SPHardwareDataType"][0].get(
            "cpu_type", ""
        )
    machine_results["cpu_speed"] = system_profile["SPHardwareDataType"][0].get(
        "current_processor_speed", ""
    )
    machine_results["memory"] = system_profile["SPHardwareDataType"][0].get(
        "physical_memory", ""
    )
    machine_results["memory_kb"] = process_memory(machine_results["memory"])

    for device in system_profile["SPStorageDataType"]:
        if device["mount_point"] == "/":
            # div by 1000.0 to
            # a) Convert to Apple base 10 kilobytes
            # b) Cast to python floats
            machine_results["hd_space"] = device["free_space_in_bytes"]
            machine_results["hd_total"] = device["size_in_bytes"]
            # We want the % used, not of free space, so invert.
            machine_results["hd_percent"] = "{:.2f}".format(
                abs(
                    float(machine_results["hd_space"]) / machine_results["hd_total"] - 1
                )
                * 100
            )

    return machine_results

def get_hostname():
    name_type = sal.sal_pref("NameType", default="ComputerName")
    net_config = SCDynamicStoreCreate(None, "net", None, None)
    return get_machine_name(net_config, name_type)

def get_machine_name(net_config, nametype):
    """Return the ComputerName of this Mac."""
    sys_info = SCDynamicStoreCopyValue(net_config, "Setup:/System")
    if sys_info:
        return sys_info.get(nametype)
    return subprocess.check_output(
        ["/usr/sbin/scutil", "--get", "ComputerName"], text=True
    ).strip()

def get_friendly_model(serial, udid):
    """Return friendly model name"""
    cmd = ["/usr/sbin/ioreg", "-arc", "IOPlatformDevice", "-k", "product-name"]
    try:
        out = subprocess.check_output(cmd)
    except:
        pass
    if out:
        try:
            data = plistlib.loads(out)
            if len(data) != 0:
                return (
                    data[0]
                    .get("product-name")
                    .encode("ascii", "ignore")
                    .decode()
                    .strip()
                )
        except:
            pass

    # set up cache file for this udid...create dir,
    MODEL_PATH.mkdir(mode=0o755, parents=True, exist_ok=True)

    # name cache for this udid
    UDID_CACHE_PATH = pathlib.Path(MODEL_PATH, "%s.txt" % (udid))
    for cache_file in MODEL_PATH.iterdir():
        # clean up any other files in dir
        if cache_file != UDID_CACHE_PATH:
            try:
                cache_file.unlink()
            except:
                pass

    if not UDID_CACHE_PATH.exists():
        model = cleanup_model(query_apple_support(serial))
        if model:
            UDID_CACHE_PATH.write_text(model)
    else:
        try:
            model = UDID_CACHE_PATH.read_text().strip()
        except:
            model = None

    return model

def get_model_code(serial):
    # Workaround for machines with dummy serial numbers.
    if "serial" in serial.lower():
        return

    if 12 <= len(serial) <= 13:
        if serial.startswith("S"):
            # Remove S prefix from scanned codes.
            serial = serial[1:]
        return serial[8:].upper()

    elif 11 <= len(serial) <= 12:
        # 2010 Mac Pros starting with H or Y are 11 characters
        return serial[8:].upper()

    return

def query_apple_support(serial):
    model_code = get_model_code(serial)
    tree = ElementTree.ElementTree()
    session = macsesh.Session()
    response = session.get(
        f"https://support-sp.apple.com/sp/product?cc={model_code}&lang=en_US"
    )
    try:
        tree = ElementTree.fromstring(response.text)
    except ElementTree.ParseError:
        tree = None
    return tree.findtext("configCode") if tree else None

def cleanup_model(model):
    cleanup_res = [(re.compile(r"inch ? "), "inch, "), (re.compile(r"  "), " ")]
    if model:
        for pattern, replacement in cleanup_res:
            model = pattern.sub(replacement, model)
    return model

def process_memory(amount):
    """Convert the amount of memory like '4 GB' to the size in kb as int"""
    try:
        memkb = int(amount[:-3]) * 1024 ** MEMORY_EXPONENTS[amount[-2:]]
    except ValueError:
        memkb = int(float(amount[:-3])) * 1024 ** MEMORY_EXPONENTS[amount[-2:]]
    return memkb

def get_sys_profile():
    """Get sysprofiler info.

    Returns:
        System Profiler report for networking, drives, and hardware as a
        plist dict, or an empty dict.
    """
    command = [
        "/usr/sbin/system_profiler",
        "-xml",
        "SPHardwareDataType",
        "SPStorageDataType",
        "SPSoftwareDataType",
    ]
    try:
        output = subprocess.check_output(command)
    except subprocess.CalledProcessError:
        return {}

    try:
        system_profile = plistlib.loads(output)
    except plistlib.InvalidException:
        system_profile = {}

    # sytem_profiler gives us back an array; convert to a dict with just
    # the data we care about.
    results = {}
    for data_type in system_profile:
        key = data_type["_dataType"]
        results[key] = data_type["_items"]

    return results

def get_console_user():
    """Get informatino about the console user

    Returns:
        3-Tuple of (str) username, (int) uid, (int) gid
    """
    return SCDynamicStoreCopyConsoleUser(None, None, None)

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