<commit_before># -*- coding: utf-8 -*-
##############################################################################
#
#    Author: Guewen Baconnier
#    Copyright 2013 Camptocamp SA
#
#    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/>.
#
##############################################################################

# TODO multi-records:
# 2 external records may be converted in 1 openerp record and
# conversely. A processor must accept in input a dict where the keys are
# the name of the external / openerp model and the values are the
# records. And the same for the output records.

from ..connector import ConnectorUnit, MetaConnectorUnit
from ..exception import MappingError
from .binder import Binder

def mapping(func):
    """ Decorator declarating a mapping for a field """
    func.is_mapping = True
    return func

def changed_by(*args):
    """ Decorator for the mappings. When fields are modified, we want to modify
    only the modified fields. Using this decorator, we can specify which fields
    updates should trigger which mapping.

    If ``changed_by`` is empty, the mapping is always active.
    As far as possible, it should be used, thus, when we do an update on
    only a small number of fields on a record, the size of the output
    record will be limited to only the fields really having to be
    modified.

    :param *args: field names which trigger the mapping when modified
    """
    def register_mapping(func):
        func.changed_by = args
        return func
    return register_mapping

class MetaMapper(MetaConnectorUnit):
    """ Metaclass for Mapper """

    def __new__(meta, name, bases, attrs):
        if attrs.get('_map_methods') is None:
            attrs['_map_methods'] = {}

        cls = super(MetaMapper, meta).__new__(meta, name, bases, attrs)

        for base in bases:
            for attr_name, changed_by in getattr(base, '_map_methods', {}).iteritems():
                cls._map_methods.setdefault(attr_name, set()).update(changed_by)

        for attr_name, attr in attrs.iteritems():
            mapping = getattr(attr, 'is_mapping', None)
            if mapping:
                changed_by = getattr(attr, 'changed_by', set())
                cls._map_methods.setdefault(attr_name, set()).update(changed_by)
        return cls

class Mapper(ConnectorUnit):
    """ Transform a record to a defined output """

    __metaclass__ = MetaMapper

    # name of the OpenERP model, to be defined in concrete classes
    _model_name = None

    direct = []  # direct conversion of a field to another (from_attr, to_attr)
    children = []  # conversion of sub-records (from_attr, to_attr, model)

    _map_methods = None

    def _map_direct(self, record, from_attr, to_attr):
        raise NotImplementedError

    def _map_children(self, record, attr, model):
        raise NotImplementedError

    @property
    def map_methods(self):
        for meth, changed_by in self._map_methods.iteritems():
            yield getattr(self, meth), changed_by

    def convert(self, record, fields=None, parent_values=None):
        """ Transform an external record to an OpenERP record or the opposite

        :param record: record to transform
        :param parent_values: openerp record of the containing object
            (e.g. sale_order for a sale_order_line)
        """
        if fields is None:
            fields = {}

        result = {}

        for from_attr, to_attr in self.direct:
            if (not fields or from_attr in fields):
                # XXX not compatible with all
                # record type (wrap
                # records in a standard class representation?)
                value = self._map_direct(record,
                                         from_attr,
                                         to_attr)
                result[to_attr] = value

        for meth, changed_by in self.map_methods:

            if (not fields or
                    not changed_by or
                    set(fields).intersection(changed_by)):
                values = meth(record)
                if not values:
                    continue
                if isinstance(values, dict):
                    result.update(values)
                else:
                    raise ValueError('%s: invalid return value for the '
                                     'mapping method %s' % (values, meth))

        for from_attr, to_attr, model in self.children:
            if (not fields or from_attr in fields):
                values = self._map_children(record, from_attr, model)
                result[to_attr] = values

        return result

    def _sub_convert(self, records, processor, parent_values=None):
        """ return values of a one2many to put in the main record
        do not create the records!
        """
        raise NotImplementedError

class ImportMapper(Mapper):
    """ Transform a record from a backend to an OpenERP record """

    def _get_o2m_binder(self, model_name):
        binder_cls = self.backend.get_class(Binder, model_name)
        return self.binder_cls(self.backend, self.session)

    def _get_o2m_external_identifier(self, record, attr, model):
        # from a record for a model
        return record[attr]

    def _map_direct(self, record, from_attr, to_attr):
        value = record[from_attr]
        attr_type = self.model._columns[to_attr]._type
        if attr_type == 'many2one':
            model = self.model._columns[to_attr]._obj
            ext_id = self._get_o2m_external_identifier(
                    record, from_attr, model)
            binder = self._get_o2m_binder(model)
            value = binder.to_openerp(self.backend, ext_id)

            if not value:
                raise MappingError("Can not find an existing %s for external "
                                   "record %s" % (model, ext_id))
        return value

    def _map_children(self, record, attr, model):
        mapper_cls = self.backend.get_class(ImportMapper, self.model_name)
        mapper = mapper_cls(self.backend, self.session)
        child_records = record[attr]  # XXX not compatible with
                                     # all record types
        return self._sub_convert(child_records, mapper,
                                 parent_values=record)

    def _sub_convert(self, records, mapper, parent_values=None):
        """ return values of a one2many to put in the main record
        do not create the records!
        """
        result = []
        for record in records:
            vals = mapper.convert(record,
                                  parent_values=parent_values)
            result.append((0, 0, vals))
        return result

class ExportMapper(Mapper):
    """ Transform a record from OpenERP to a backend record """

    def _map_children(self, record, attr, model):
        mapper_cls = self.backend.get_class(ExportMapper, self.model_name)
        mapper = mapper_cls(self.backend, self.session)
        child_records = record[attr]  # XXX not compatible with
                                      # all record types
        return self._sub_convert(child_records, mapper,
                                 parent_values=record)

    def _sub_convert(self, records, mapper, parent_values=None):
        """ return values of a one2many to put in the main record
        do not create the records!
        """
        result = []
        for record in records:
            vals = mapper.convert(record,
                                  parent_values=parent_values)
            result.append(vals)
        return result
<commit_msg>[FIX] mapper should read in _all_columns to see the inherits columns<commit_after># -*- coding: utf-8 -*-
##############################################################################
#
#    Author: Guewen Baconnier
#    Copyright 2013 Camptocamp SA
#
#    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/>.
#
##############################################################################

# TODO multi-records:
# 2 external records may be converted in 1 openerp record and
# conversely. A processor must accept in input a dict where the keys are
# the name of the external / openerp model and the values are the
# records. And the same for the output records.

from ..connector import ConnectorUnit, MetaConnectorUnit
from ..exception import MappingError
from .binder import Binder

def mapping(func):
    """ Decorator declarating a mapping for a field """
    func.is_mapping = True
    return func

def changed_by(*args):
    """ Decorator for the mappings. When fields are modified, we want to modify
    only the modified fields. Using this decorator, we can specify which fields
    updates should trigger which mapping.

    If ``changed_by`` is empty, the mapping is always active.
    As far as possible, it should be used, thus, when we do an update on
    only a small number of fields on a record, the size of the output
    record will be limited to only the fields really having to be
    modified.

    :param *args: field names which trigger the mapping when modified
    """
    def register_mapping(func):
        func.changed_by = args
        return func
    return register_mapping

class MetaMapper(MetaConnectorUnit):
    """ Metaclass for Mapper """

    def __new__(meta, name, bases, attrs):
        if attrs.get('_map_methods') is None:
            attrs['_map_methods'] = {}

        cls = super(MetaMapper, meta).__new__(meta, name, bases, attrs)

        for base in bases:
            for attr_name, changed_by in getattr(base, '_map_methods', {}).iteritems():
                cls._map_methods.setdefault(attr_name, set()).update(changed_by)

        for attr_name, attr in attrs.iteritems():
            mapping = getattr(attr, 'is_mapping', None)
            if mapping:
                changed_by = getattr(attr, 'changed_by', set())
                cls._map_methods.setdefault(attr_name, set()).update(changed_by)
        return cls

class Mapper(ConnectorUnit):
    """ Transform a record to a defined output """

    __metaclass__ = MetaMapper

    # name of the OpenERP model, to be defined in concrete classes
    _model_name = None

    direct = []  # direct conversion of a field to another (from_attr, to_attr)
    children = []  # conversion of sub-records (from_attr, to_attr, model)

    _map_methods = None

    def _map_direct(self, record, from_attr, to_attr):
        raise NotImplementedError

    def _map_children(self, record, attr, model):
        raise NotImplementedError

    @property
    def map_methods(self):
        for meth, changed_by in self._map_methods.iteritems():
            yield getattr(self, meth), changed_by

    def convert(self, record, fields=None, parent_values=None):
        """ Transform an external record to an OpenERP record or the opposite

        :param record: record to transform
        :param parent_values: openerp record of the containing object
            (e.g. sale_order for a sale_order_line)
        """
        if fields is None:
            fields = {}

        result = {}

        for from_attr, to_attr in self.direct:
            if (not fields or from_attr in fields):
                # XXX not compatible with all
                # record type (wrap
                # records in a standard class representation?)
                value = self._map_direct(record,
                                         from_attr,
                                         to_attr)
                result[to_attr] = value

        for meth, changed_by in self.map_methods:

            if (not fields or
                    not changed_by or
                    set(fields).intersection(changed_by)):
                values = meth(record)
                if not values:
                    continue
                if isinstance(values, dict):
                    result.update(values)
                else:
                    raise ValueError('%s: invalid return value for the '
                                     'mapping method %s' % (values, meth))

        for from_attr, to_attr, model in self.children:
            if (not fields or from_attr in fields):
                values = self._map_children(record, from_attr, model)
                result[to_attr] = values

        return result

    def _sub_convert(self, records, processor, parent_values=None):
        """ return values of a one2many to put in the main record
        do not create the records!
        """
        raise NotImplementedError

class ImportMapper(Mapper):
    """ Transform a record from a backend to an OpenERP record """

    def _get_o2m_binder(self, model_name):
        binder_cls = self.backend.get_class(Binder, model_name)
        return self.binder_cls(self.backend, self.session)

    def _get_o2m_external_identifier(self, record, attr, model):
        # from a record for a model
        return record[attr]

    def _map_direct(self, record, from_attr, to_attr):
        value = record[from_attr]
        attr_type = self.model._all_columns[to_attr].column._type
        if attr_type == 'many2one':
            model = self.model._all_columns[to_attr].column._obj
            ext_id = self._get_o2m_external_identifier(
                    record, from_attr, model)
            binder = self._get_o2m_binder(model)
            value = binder.to_openerp(self.backend, ext_id)

            if not value:
                raise MappingError("Can not find an existing %s for external "
                                   "record %s" % (model, ext_id))
        return value

    def _map_children(self, record, attr, model):
        mapper_cls = self.backend.get_class(ImportMapper, self.model_name)
        mapper = mapper_cls(self.backend, self.session)
        child_records = record[attr]  # XXX not compatible with
                                     # all record types
        return self._sub_convert(child_records, mapper,
                                 parent_values=record)

    def _sub_convert(self, records, mapper, parent_values=None):
        """ return values of a one2many to put in the main record
        do not create the records!
        """
        result = []
        for record in records:
            vals = mapper.convert(record,
                                  parent_values=parent_values)
            result.append((0, 0, vals))
        return result

class ExportMapper(Mapper):
    """ Transform a record from OpenERP to a backend record """

    def _map_children(self, record, attr, model):
        mapper_cls = self.backend.get_class(ExportMapper, self.model_name)
        mapper = mapper_cls(self.backend, self.session)
        child_records = record[attr]  # XXX not compatible with
                                      # all record types
        return self._sub_convert(child_records, mapper,
                                 parent_values=record)

    def _sub_convert(self, records, mapper, parent_values=None):
        """ return values of a one2many to put in the main record
        do not create the records!
        """
        result = []
        for record in records:
            vals = mapper.convert(record,
                                  parent_values=parent_values)
            result.append(vals)
        return result
<|endoftext|>


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

"""Script to run packman via command line

Usage:
    pkm get [--components=<list> --components-file=<path> --exclude=<list>]
    pkm pack [--components=<list> --components-file=<path> --exclude=<list>]
    pkm make [--components=<list> --components-file=<path> --exclude=<list>]
    pkm --version

Arguments:
    pack     Packs component configured in components file
    get      Gets component configured in components file
    make     Gets AND (yeah!) Packs.. don't ya kno!

Options:
    -h --help                   Show this screen.
    -c --components=<list>      Comma Separated list of component names
    -x --exclude=<list>         Comma Separated list of excluded components
    -f --components-file=<path> Components file path
    --verbose                   a LOT of output (NOT YET IMPLEMENTED)
    -v --version                Display current version of sandman and exit

"""

from __future__ import absolute_import
from docopt import docopt

from packman.packman import packman_runner

def main(test_options=None):
    """Main entry point for script."""
    import pkg_resources
    version = None
    try:
        version = pkg_resources.get_distribution('packman').version
    except Exception as e:
        print(e)
    finally:
        del pkg_resources

    # TODO: implement verbose output
    options = test_options or docopt(__doc__, version=version)
    print(options)
    if options['pack']:
        packman_runner('pack',
                       options.get('--components_file'),
                       options.get('--components'),
                       options.get('--exclude'))
    elif options['get']:
        packman_runner('get',
                       options.get('--components_file'),
                       options.get('--components'),
                       options.get('--exclude'))
    elif options['make']:
        packman_runner('get',
                       options.get('--components_file'),
                       options.get('--components'),
                       options.get('--exclude'))
        packman_runner('pack',
                       options.get('--components_file'),
                       options.get('--components'),
                       options.get('--exclude'))

if __name__ == '__main__':
    main()
<commit_msg>now checking for distribution support<commit_after>#!/usr/bin/env python
########
# Copyright (c) 2014 GigaSpaces Technologies Ltd. All rights reserved
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
#    * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    * See the License for the specific language governing permissions and
#    * limitations under the License.

"""Script to run packman via command line

Usage:
    pkm get [--components=<list> --components-file=<path> --exclude=<list>]
    pkm pack [--components=<list> --components-file=<path> --exclude=<list>]
    pkm make [--components=<list> --components-file=<path> --exclude=<list>]
    pkm --version

Arguments:
    pack     Packs component configured in components file
    get      Gets component configured in components file
    make     Gets AND (yeah!) Packs.. don't ya kno!

Options:
    -h --help                   Show this screen.
    -c --components=<list>      Comma Separated list of component names
    -x --exclude=<list>         Comma Separated list of excluded components
    -f --components-file=<path> Components file path
    --verbose                   a LOT of output (NOT YET IMPLEMENTED)
    -v --version                Display current version of sandman and exit

"""

from __future__ import absolute_import
from docopt import docopt
from platform import dist
from sys import exit

from packman.packman import packman_runner

dist_list = ('Ubuntu', 'debian', 'CentOS')

def main(test_options=None):
    """Main entry point for script."""
    distro = dist()[0]
    print('Distribution Identified: ', distro)
    if not distro in (dist_list):
      print('Your distribution is not supported.'
            'Supported Disributions are:')
      for dist in dist_list:
        print('    ', dist)
      exit(1)
    import pkg_resources
    version = None
    try:
        version = pkg_resources.get_distribution('packman').version
    except Exception as e:
        print(e)
    finally:
        del pkg_resources

    # TODO: implement verbose output
    options = test_options or docopt(__doc__, version=version)
    print(options)
    if options['pack']:
        packman_runner('pack',
                       options.get('--components_file'),
                       options.get('--components'),
                       options.get('--exclude'))
    elif options['get']:
        packman_runner('get',
                       options.get('--components_file'),
                       options.get('--components'),
                       options.get('--exclude'))
    elif options['make']:
        packman_runner('get',
                       options.get('--components_file'),
                       options.get('--components'),
                       options.get('--exclude'))
        packman_runner('pack',
                       options.get('--components_file'),
                       options.get('--components'),
                       options.get('--exclude'))

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


<commit_before># Copyright (c) 2012 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import re
import traceback
import sys

from buildbot.process.properties import Properties
from buildbot.schedulers.trysched import TryBase
from buildbot.schedulers.trysched import BadJobfile
from twisted.internet import defer
from twisted.python import log
from twisted.web import client

from master.factory.commands import DEFAULT_TESTS

def text_to_dict(text):
  """Converts a series of key=value lines in a string into a multi-value dict.
  """
  parsed = {}
  regexp = r'([A-Za-z\-\_]+)\=(.+)\w*$'
  for match in filter(None, (re.match(regexp, i) for i in text.splitlines())):
    parsed.setdefault(match.group(1), []).append(match.group(2))
  return parsed

def flatten(options, key, default):
  """Expects only one value."""
  if options.setdefault(key, []) == []:
    options[key] = default
  elif len(options[key]) == 1:
    options[key] = options[key][0]
  else:
    raise ValueError('Expecting one value, got many', key, options[key])

def try_int(options, key, default):
  """Expects a int."""
  flatten(options, key, default)
  if options[key] != default:
    options[key] = int(options[key])

def try_bool(options, key, default):
  """Expects a bool."""
  flatten(options, key, default)
  if options[key] != default:
    if isinstance(options[key], bool):
      pass
    elif options[key].lower() in ('true', '1'):
      options[key] = True
    elif options[key].lower() in ('false', '0'):
      options[key] = False
    else:
      raise ValueError('Expecting a bool, got non-bool', key, options[key])

def comma_separated(options, key):
  """Splits comma separated strings into multiple values."""
  for i in options.setdefault(key, [])[:]:
    if ',' in i:
      options[key].remove(i)
      options[key].extend(x for x in i.split(',') if x)

def dict_comma(values, valid_keys, default):
  """Splits comma separated strings with : to create a dict."""
  out = {}
  # We accept two formats:
  # 1. Builder(s) only, will run default tests.  Eg. win,mac
  # 2. Builder(s) with test(s): Eg. win,mac:test1:filter1,test2:filter2,test3
  #    In this example, win will run default tests, and mac will be ran with
  #    test1:filter1, test2:filter2, and test3.
  # In no cases are builders allowed to be specified after tests are specified.

  for value in values:
    sections_match = re.match(r'^([^\:]+):?(.*)$', value)
    if sections_match:
      builders_raw, tests_raw = sections_match.groups()
      builders = builders_raw.split(',')
      tests = [test for test in tests_raw.split(',') if test]
      for builder in builders:
        if not builder:
          # Empty builder means trailing comma.
          raise BadJobfile('Bad input - trailing comma: %s' % value)
        if builder not in valid_keys:
          # Drop unrecognized builders.
          log.msg('%s not in builders, dropping.' % builder)
          continue
        if not tests or builder != builders[-1]:
          # We can only add filters to the last builder.
          out.setdefault(builder, set()).add(default)
        else:
          for test in tests:
            if len(test.split(':')) > 2:
              # Don't support test:filter:foo.
              raise BadJobfile('Failed to process test %s' % value)
            out.setdefault(builder, set()).add(test)
    else:
      raise BadJobfile('Failed to process value %s' % value)

  return out

def parse_options(options, builders, pools):
  """Converts try job settings into a dict.

  The dict is received as dict(str, list(str)).
  """
  try:
    # Flush empty values.
    options = dict((k, v) for k, v in options.iteritems() if v)

    flatten(options, 'name', 'Unnamed')
    flatten(options, 'user', 'John Doe')
    flatten(options, 'requester', None)

    comma_separated(options, 'email')
    for email in options['email']:
      # Access to a protected member XXX of a client class
      # pylint: disable=W0212
      if not TryJobBase._EMAIL_VALIDATOR.match(email):
        raise BadJobfile("'%s' is an invalid email address!" % email)

    flatten(options, 'patch', None)
    flatten(options, 'root', None)
    try_int(options, 'patchlevel', 0)
    flatten(options, 'branch', None)

    flatten(options, 'revision', None)
    options.setdefault('orig_revision', options['revision'])

    flatten(options, 'reason', '%s: %s' % (options['user'], options['name']))
    try_bool(options, 'clobber', False)

    flatten(options, 'project', pools.default_pool_name if pools else None)
    flatten(options, 'repository', None)

    # Code review info. Enforce numbers.
    try_int(options, 'patchset', None)
    try_int(options, 'issue', None)

    # Manages bot selection and test filtering. It is preferable to use
    # multiple bot=unit_tests:Gtest.Filter lines. DEFAULT_TESTS is a marker to
    # specify that the default tests should be run for this builder.
    if not isinstance(options.get('bot'), dict):
      options['bot'] = dict_comma(
          options.get('bot', []), builders, DEFAULT_TESTS)
    if not options['bot'] and pools:
      options['bot'] = pools.Select(None, options['project'])
    if not options['bot']:
      raise BadJobfile('No builder selected')
    comma_separated(options, 'testfilter')
    if options['testfilter']:
      for k in options['bot']:
        options['bot'][k].update(options['testfilter'])
    # Convert the set back to list.
    for bot in options['bot']:
      options['bot'][bot] = list(options['bot'][bot])

    log.msg(
        'Chose %s for job %s' %
        (','.join(options['bot']), options['reason']))
    return options
  except (TypeError, ValueError), e:
    lines = [
        (i[0].rsplit('/', 1)[-1], i[1])
        for i in traceback.extract_tb(sys.exc_info()[2])
    ]
    raise BadJobfile('Failed to parse the metadata: %r' % options, e, lines)

class TryJobBase(TryBase):
  compare_attrs = TryBase.compare_attrs + (
      'pools', 'last_good-urls', 'code_review_sites')

  # Simplistic email matching regexp.
  _EMAIL_VALIDATOR = re.compile(
      r'[a-zA-Z0-9][a-zA-Z0-9\.\+\-\_]*@[a-zA-Z0-9\.\-]+\.[a-zA-Z]{2,3}$')

  _PROPERTY_SOURCE = 'Try job'

  def __init__(self, name, pools, properties,
               last_good_urls, code_review_sites):
    TryBase.__init__(self, name, pools.ListBuilderNames(), properties or {})
    self.pools = pools
    pools.SetParent(self)
    self.last_good_urls = last_good_urls
    self.code_review_sites = code_review_sites
    self._last_lkgr = {}
    self.valid_builders = []

  def setServiceParent(self, parent):
    TryBase.setServiceParent(self, parent)
    self.valid_builders = self.master.botmaster.builders.keys()

  def gotChange(self, change, important):  # pylint: disable=R0201
    log.msg('ERROR: gotChange was unexpectedly called.')

  def parse_options(self, options):
    return parse_options(options, self.valid_builders, self.pools)

  def get_props(self, builder, options):
    """Current job extra properties that are not related to the source stamp.
    Initialize with the Scheduler's base properties.
    """
    keys = (
      'clobber',
      'issue',
      'orig_revision',
      'patch_url',
      'patchset',
      'requester',
      'rietveld',
      'root',
      'try_job_key',
    )
    # All these settings have no meaning when False or not set, so don't set
    # them in that case.
    properties = dict((i, options[i]) for i in keys if options.get(i))
    properties['testfilter'] = options['bot'].get(builder, None)
    props = Properties()
    props.updateFromProperties(self.properties)
    props.update(properties, self._PROPERTY_SOURCE)
    return props

  def create_buildset(self, ssid, parsed_job):
    log.msg('Creating try job(s) %s' % ssid)
    result = None
    for builder in parsed_job['bot']:
      result = self.addBuildsetForSourceStamp(ssid=ssid,
          reason=parsed_job['name'],
          external_idstring=parsed_job['name'],
          builderNames=[builder],
          properties=self.get_props(builder, parsed_job))
    return result

  def SubmitJob(self, parsed_job, changeids):
    if not parsed_job['bot']:
      raise BadJobfile(
          'incoming Try job did not specify any allowed builder names')

    # Verify the try job patch is not more than 20MB.
    if parsed_job.get('patch'):
      patchsize = len(parsed_job['patch'])
      if patchsize > 20*1024*1024:  # 20MB
        raise BadJobfile('incoming Try job patch is %s bytes, '
                        'must be less than 20MB' % (patchsize))

    d = self.master.db.sourcestamps.addSourceStamp(
        branch=parsed_job['branch'],
        revision=parsed_job['revision'],
        patch_body=parsed_job['patch'],
        patch_level=parsed_job['patchlevel'],
        patch_subdir=parsed_job['root'],
        project=parsed_job['project'],
        repository=parsed_job['repository'] or '',
        changeids=changeids)

    d.addCallback(self.create_buildset, parsed_job)
    d.addErrback(log.err, "Failed to queue a try job!")
    return d

  def get_lkgr(self, options):
    """Grabs last known good revision number if necessary."""

    # TODO: crbug.com/233418 - Only use the blink url for jobs with the blink
    # project set. This will require us to always set the project in blink
    # configs.
    project = options['project'] or 'chrome'
    if project == 'chrome' and all('_layout' in bot for bot in options['bot']):
      project = 'blink'

    options['rietveld'] = (self.code_review_sites or {}).get(project)
    self._last_lkgr.setdefault(project, None)
    last_good_url = (self.last_good_urls or {}).get(project)

    if options['revision'] or not last_good_url:
      return defer.succeed(0)

    def Success(result):
      try:
        new_value = int(result.strip())
      except (TypeError, ValueError):
        new_value = None
      if new_value and (not self._last_lkgr[project] or
                        new_value > self._last_lkgr[project]):
        self._last_lkgr[project] = new_value
      options['revision'] = self._last_lkgr[project] or 'HEAD'

    def Failure(result):
      options['revision'] = self._last_lkgr[project] or 'HEAD'

    connection = client.getPage(last_good_url, agent='buildbot')
    connection.addCallbacks(Success, Failure)
    return connection
<commit_msg>Always set orig_revision build property<commit_after># Copyright (c) 2012 The Chromium Authors. All rights reserved.
# Use of this source code is governed by a BSD-style license that can be
# found in the LICENSE file.

import re
import traceback
import sys

from buildbot.process.properties import Properties
from buildbot.schedulers.trysched import TryBase
from buildbot.schedulers.trysched import BadJobfile
from twisted.internet import defer
from twisted.python import log
from twisted.web import client

from master.factory.commands import DEFAULT_TESTS

def text_to_dict(text):
  """Converts a series of key=value lines in a string into a multi-value dict.
  """
  parsed = {}
  regexp = r'([A-Za-z\-\_]+)\=(.+)\w*$'
  for match in filter(None, (re.match(regexp, i) for i in text.splitlines())):
    parsed.setdefault(match.group(1), []).append(match.group(2))
  return parsed

def flatten(options, key, default):
  """Expects only one value."""
  if options.setdefault(key, []) == []:
    options[key] = default
  elif len(options[key]) == 1:
    options[key] = options[key][0]
  else:
    raise ValueError('Expecting one value, got many', key, options[key])

def try_int(options, key, default):
  """Expects a int."""
  flatten(options, key, default)
  if options[key] != default:
    options[key] = int(options[key])

def try_bool(options, key, default):
  """Expects a bool."""
  flatten(options, key, default)
  if options[key] != default:
    if isinstance(options[key], bool):
      pass
    elif options[key].lower() in ('true', '1'):
      options[key] = True
    elif options[key].lower() in ('false', '0'):
      options[key] = False
    else:
      raise ValueError('Expecting a bool, got non-bool', key, options[key])

def comma_separated(options, key):
  """Splits comma separated strings into multiple values."""
  for i in options.setdefault(key, [])[:]:
    if ',' in i:
      options[key].remove(i)
      options[key].extend(x for x in i.split(',') if x)

def dict_comma(values, valid_keys, default):
  """Splits comma separated strings with : to create a dict."""
  out = {}
  # We accept two formats:
  # 1. Builder(s) only, will run default tests.  Eg. win,mac
  # 2. Builder(s) with test(s): Eg. win,mac:test1:filter1,test2:filter2,test3
  #    In this example, win will run default tests, and mac will be ran with
  #    test1:filter1, test2:filter2, and test3.
  # In no cases are builders allowed to be specified after tests are specified.

  for value in values:
    sections_match = re.match(r'^([^\:]+):?(.*)$', value)
    if sections_match:
      builders_raw, tests_raw = sections_match.groups()
      builders = builders_raw.split(',')
      tests = [test for test in tests_raw.split(',') if test]
      for builder in builders:
        if not builder:
          # Empty builder means trailing comma.
          raise BadJobfile('Bad input - trailing comma: %s' % value)
        if builder not in valid_keys:
          # Drop unrecognized builders.
          log.msg('%s not in builders, dropping.' % builder)
          continue
        if not tests or builder != builders[-1]:
          # We can only add filters to the last builder.
          out.setdefault(builder, set()).add(default)
        else:
          for test in tests:
            if len(test.split(':')) > 2:
              # Don't support test:filter:foo.
              raise BadJobfile('Failed to process test %s' % value)
            out.setdefault(builder, set()).add(test)
    else:
      raise BadJobfile('Failed to process value %s' % value)

  return out

def parse_options(options, builders, pools):
  """Converts try job settings into a dict.

  The dict is received as dict(str, list(str)).
  """
  try:
    # Flush empty values.
    options = dict((k, v) for k, v in options.iteritems() if v)

    flatten(options, 'name', 'Unnamed')
    flatten(options, 'user', 'John Doe')
    flatten(options, 'requester', None)

    comma_separated(options, 'email')
    for email in options['email']:
      # Access to a protected member XXX of a client class
      # pylint: disable=W0212
      if not TryJobBase._EMAIL_VALIDATOR.match(email):
        raise BadJobfile("'%s' is an invalid email address!" % email)

    flatten(options, 'patch', None)
    flatten(options, 'root', None)
    try_int(options, 'patchlevel', 0)
    flatten(options, 'branch', None)

    flatten(options, 'revision', None)
    options.setdefault('orig_revision', options['revision'])

    flatten(options, 'reason', '%s: %s' % (options['user'], options['name']))
    try_bool(options, 'clobber', False)

    flatten(options, 'project', pools.default_pool_name if pools else None)
    flatten(options, 'repository', None)

    # Code review info. Enforce numbers.
    try_int(options, 'patchset', None)
    try_int(options, 'issue', None)

    # Manages bot selection and test filtering. It is preferable to use
    # multiple bot=unit_tests:Gtest.Filter lines. DEFAULT_TESTS is a marker to
    # specify that the default tests should be run for this builder.
    if not isinstance(options.get('bot'), dict):
      options['bot'] = dict_comma(
          options.get('bot', []), builders, DEFAULT_TESTS)
    if not options['bot'] and pools:
      options['bot'] = pools.Select(None, options['project'])
    if not options['bot']:
      raise BadJobfile('No builder selected')
    comma_separated(options, 'testfilter')
    if options['testfilter']:
      for k in options['bot']:
        options['bot'][k].update(options['testfilter'])
    # Convert the set back to list.
    for bot in options['bot']:
      options['bot'][bot] = list(options['bot'][bot])

    log.msg(
        'Chose %s for job %s' %
        (','.join(options['bot']), options['reason']))
    return options
  except (TypeError, ValueError), e:
    lines = [
        (i[0].rsplit('/', 1)[-1], i[1])
        for i in traceback.extract_tb(sys.exc_info()[2])
    ]
    raise BadJobfile('Failed to parse the metadata: %r' % options, e, lines)

class TryJobBase(TryBase):
  compare_attrs = TryBase.compare_attrs + (
      'pools', 'last_good-urls', 'code_review_sites')

  # Simplistic email matching regexp.
  _EMAIL_VALIDATOR = re.compile(
      r'[a-zA-Z0-9][a-zA-Z0-9\.\+\-\_]*@[a-zA-Z0-9\.\-]+\.[a-zA-Z]{2,3}$')

  _PROPERTY_SOURCE = 'Try job'

  def __init__(self, name, pools, properties,
               last_good_urls, code_review_sites):
    TryBase.__init__(self, name, pools.ListBuilderNames(), properties or {})
    self.pools = pools
    pools.SetParent(self)
    self.last_good_urls = last_good_urls
    self.code_review_sites = code_review_sites
    self._last_lkgr = {}
    self.valid_builders = []

  def setServiceParent(self, parent):
    TryBase.setServiceParent(self, parent)
    self.valid_builders = self.master.botmaster.builders.keys()

  def gotChange(self, change, important):  # pylint: disable=R0201
    log.msg('ERROR: gotChange was unexpectedly called.')

  def parse_options(self, options):
    return parse_options(options, self.valid_builders, self.pools)

  def get_props(self, builder, options):
    """Current job extra properties that are not related to the source stamp.
    Initialize with the Scheduler's base properties.
    """
    always_included_keys = (
      'orig_revision',
      'testfilter',
    )

    optional_keys = (
      'clobber',
      'issue',
      'patch_url',
      'patchset',
      'requester',
      'rietveld',
      'root',
      'try_job_key',
    )

    # All these settings have no meaning when False or not set, so don't set
    # them in that case.
    properties = dict((i, options[i]) for i in optional_keys if options.get(i))

    # These settings are meaningful even if the value evaluates to False
    # or None. Note that when options don't contain given key, it will
    # be set to None.
    properties.update(dict((i, options.get(i)) for i in always_included_keys))

    props = Properties()
    props.updateFromProperties(self.properties)
    props.update(properties, self._PROPERTY_SOURCE)
    return props

  def create_buildset(self, ssid, parsed_job):
    log.msg('Creating try job(s) %s' % ssid)
    result = None
    for builder in parsed_job['bot']:
      result = self.addBuildsetForSourceStamp(ssid=ssid,
          reason=parsed_job['name'],
          external_idstring=parsed_job['name'],
          builderNames=[builder],
          properties=self.get_props(builder, parsed_job))
    return result

  def SubmitJob(self, parsed_job, changeids):
    if not parsed_job['bot']:
      raise BadJobfile(
          'incoming Try job did not specify any allowed builder names')

    # Verify the try job patch is not more than 20MB.
    if parsed_job.get('patch'):
      patchsize = len(parsed_job['patch'])
      if patchsize > 20*1024*1024:  # 20MB
        raise BadJobfile('incoming Try job patch is %s bytes, '
                        'must be less than 20MB' % (patchsize))

    d = self.master.db.sourcestamps.addSourceStamp(
        branch=parsed_job['branch'],
        revision=parsed_job['revision'],
        patch_body=parsed_job['patch'],
        patch_level=parsed_job['patchlevel'],
        patch_subdir=parsed_job['root'],
        project=parsed_job['project'],
        repository=parsed_job['repository'] or '',
        changeids=changeids)

    d.addCallback(self.create_buildset, parsed_job)
    d.addErrback(log.err, "Failed to queue a try job!")
    return d

  def get_lkgr(self, options):
    """Grabs last known good revision number if necessary."""

    # TODO: crbug.com/233418 - Only use the blink url for jobs with the blink
    # project set. This will require us to always set the project in blink
    # configs.
    project = options['project'] or 'chrome'
    if project == 'chrome' and all('_layout' in bot for bot in options['bot']):
      project = 'blink'

    options['rietveld'] = (self.code_review_sites or {}).get(project)
    self._last_lkgr.setdefault(project, None)
    last_good_url = (self.last_good_urls or {}).get(project)

    if options['revision'] or not last_good_url:
      return defer.succeed(0)

    def Success(result):
      try:
        new_value = int(result.strip())
      except (TypeError, ValueError):
        new_value = None
      if new_value and (not self._last_lkgr[project] or
                        new_value > self._last_lkgr[project]):
        self._last_lkgr[project] = new_value
      options['revision'] = self._last_lkgr[project] or 'HEAD'

    def Failure(result):
      options['revision'] = self._last_lkgr[project] or 'HEAD'

    connection = client.getPage(last_good_url, agent='buildbot')
    connection.addCallbacks(Success, Failure)
    return connection
<|endoftext|>


<commit_before><commit_msg>Update msct_multiatlas_seg.py<commit_after><|endoftext|>


<commit_before>#!/hpc/astro/users/amp2217/yt-x86_64/bin/python
# coding: utf-8

""" Script for using the Rewinder to infer the Galactic host potential """

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os, sys
import shutil
import copy
import logging
from datetime import datetime
import multiprocessing

# Third-party
import astropy.units as u
from astropy.io.misc import fnpickle, fnunpickle
from astropy.utils.console import color_print
import emcee
import matplotlib.pyplot as plt
import numpy as np
np.seterr(all="ignore")
import scipy
scipy.seterr(all="ignore")
import triangle
import yaml

try:
    from emcee.utils import MPIPool
except ImportError:
    color_print("Failed to import MPIPool from emcee! MPI functionality "
                "won't work.", "yellow")

# Project
from streams import usys
from streams.coordinates import _gc_to_hel, _hel_to_gc
from streams.inference import (ModelParameter, StreamModel,
                               LogNormalPrior, LogUniformPrior)
from streams.io.sgr import mass_selector
from streams.observation.gaia import RRLyraeErrorModel
import streams.potential as sp
from streams.util import make_path

global pool
pool = None

# Create logger
logger = logging.getLogger(__name__)

def _null(*args, **kwargs):
    return 0.

def get_pool(config):
    """ Given a config structure, return an MPIPool, a Python
        multiprocessing.Pool, or None.
    """

    if config.get("mpi", False):
        # Initialize the MPI pool
        pool = MPIPool()

        # Make sure the thread we're running on is the master
        if not pool.is_master():
            pool.wait()
            sys.exit(0)
        logger.info("Running with MPI.")

    elif config.get("threads", 0) > 1:
        logger.info("Running with multiprocessing on {} cores."\
                    .format(config["threads"]))
        pool = multiprocessing.Pool(config["threads"])
    else:
        logger.info("Running serial...parallel panda is sad :(")
        pool = None

    return pool

def read_simulation(config):
    """ TODO: """

    if config["simulation"]["source"] == "sgr":
        m = config["simulation"]["satellite_mass"]
        particles_today, satellite_today, time = mass_selector(m)
    elif config["simulation"]["source"] == "lm10":
        #particles_today, satellite_today, time = mass_selector(config["mass"])
        # TODO
        pass
    else:
        raise ValueError("Invalid particle_source: {particle_source}"\
                         .format(config))

    particles = particles_today(N=config["particles"]["N"],
                                expr=config["particles"]["selection_expr"])
    satellite = satellite_today()
    t1,t2 = time()

    return t1,t2,satellite,particles

def _parse_quantity(q):
    try:
        val,unit = q.split()
    except AttributeError:
        val = q
        unit = u.dimensionless_unscaled

    return u.Quantity(float(val), unit)

def main(config_file, job_name=None):
    """ TODO: """

    # read in configurable parameters
    with open(config_file) as f:
        config = yaml.load(f.read())

    # TODO: write a separate script just for plotting...
    # get a pool object given the configuration parameters
    pool = get_pool(config)

    # determine the output data path
    path = make_path(config)
    sampler_file = os.path.join(path, "sampler_data.pickle")

    logger.info("Will write output to '{0}'...".format(path))

    make_plots = config.get("make_plots", False)

    # get the potential object specified
    Potential = getattr(sp, config["potential"]["class_name"])
    potential = Potential()
    logger.debug("Using {} potential...".format(potential))

    # Actually get simulation data
    np.random.seed(config["seed"])
    t1,t2,_satellite,_particles = read_simulation(config)

    # TODO: right now error specification in yml doesn't propagate
    factor = config["errors"].get("global_factor", 1.)
    particle_errors = dict()
    for k,v in config["errors"].items():
        if k == "global_factor": continue
        particle_errors["{}_err".format(k)] = _parse_quantity(v)

    error_model = RRLyraeErrorModel(units=usys,
                                    factor=factor,
                                    **particle_errors)
    obs_data, obs_error = _particles.observe(error_model)

    # satellite has different errors from individual stars...
    # from: http://iopscience.iop.org/1538-4357/618/1/L25/pdf/18807.web.pdf
    sat_error_model = RRLyraeErrorModel(units=usys,
                                        mul_err=1000.*u.mas/u.yr,
                                        mub_err=1000.*u.mas/u.yr,
                                        D_err=1.,
                                        vr_err=200.*u.km/u.s)
    sat_obs_data, sat_obs_error = _satellite.observe(sat_error_model)

    particles = _particles.copy()
    satellite = _satellite.copy()

    model_parameters = []
    ##########################################################################
    # Potential parameters
    #
    potential_params = config["potential"].get("parameters", dict())
    for name,meta in potential_params.items():
        p = getattr(potential, name)

        if meta.has_key("range"):
            # TODO: fix when astropy fixed...
            #lo,hi = map(u.Quantity, meta["range"])
            lo = _parse_quantity(meta["range"][0])
            hi = _parse_quantity(meta["range"][1])
            prior = LogUniformPrior(lo.decompose(usys).value,
                                    hi.decompose(usys).value)
        else:
            prior = LogPrior()

        model_parameters.append(ModelParameter(target=p,
                                               attr="_value",
                                               ln_prior=prior))

    ##########################################################################
    # Particle parameters
    #
    particle_params = config["particles"].get("parameters", [])

    try:
        Nparticles = config["particles"]["N"]
    except KeyError:
        raise ValueError("Must specify number of partices in config file!")

    # time unbound / escape time (tub)
    if "tub" in particle_params:
        lo = [t2] * len(particles)
        hi = [t1] * len(particles)
        prior = LogUniformPrior(lo, hi)
        model_parameters.append(ModelParameter(target=particles,
                                               attr="tub",
                                               ln_prior=prior))

    # here I monte carlo transform the error distribution from observed
    #   to cartesian, then take np.cov and use that to sample new particle
    #   positions
    O = np.array([np.random.normal(obs_data, obs_error) \
                    for ii in range(1000)])
    X = _hel_to_gc(O)

    obs_error_gc = []
    for ii in range(Nparticles):
        obs_error_gc.append(np.cov(X[:,ii].T))

    obs_error_gc = np.array(obs_error_gc)
    obs_data_gc = _hel_to_gc(obs_data)

    # true positions of particles (flat_X)
    if "_X" in particle_params:
        particles._X = _hel_to_gc(obs_data)
        particles.obs_data = obs_data
        particles.obs_error = obs_error

        prior = LogNormalPrior(obs_data_gc, cov=obs_error_gc)
        p = ModelParameter(target=particles, attr="_X", ln_prior=prior)
        p.ln_prior = _null # THIS IS A HACK?
        model_parameters.append(p)

    ##########################################################################
    # Satellite parameters
    #
    # TODO: infer the damn satellite position!
    satellite_config = config.get("satellite", dict())
    satellite_params = satellite_config.get("parameters", [])

    # here I monte carlo transform the error distribution from observed
    #   to cartesian, then take np.cov and use that to sample new satellite
    #   positions
    O = np.array([np.random.normal(sat_obs_data, sat_obs_error) \
                    for ii in range(1000)])
    X = _hel_to_gc(O)

    sat_obs_error_gc = np.array([np.cov(X[:,0].T)])
    sat_obs_data_gc = _hel_to_gc(sat_obs_data)

    # true position of the satellite
    if "_X" in satellite_params:
        satellite._X = _hel_to_gc(sat_obs_data)
        satellite.obs_data = sat_obs_data
        satellite.obs_error = sat_obs_error

        prior = LogNormalPrior(sat_obs_data_gc, cov=sat_obs_error_gc)
        p = ModelParameter(target=satellite, attr="_X", ln_prior=prior)
        p.ln_prior = _null # THIS IS A HACK?
        model_parameters.append(p)

    # read in the number of walkers to use
    Nwalkers = config.get("walkers", "auto")
    if str(Nwalkers).lower() == "auto":
        Nwalkers = model.ndim*2 + 2

    # now create the model
    model = StreamModel(potential, satellite, particles,
                        parameters=model_parameters)
    p0 = model.sample(size=Nwalkers)

    if not os.path.exists(sampler_file):
        sampler = emcee.EnsembleSampler(Nwalkers, model.ndim, model,
                                        args=(t1, t2, -1.),
                                        pool=pool)

        Nburn_in = config.get("burn_in", 0)
        Nsteps = config["steps"]
        if Nburn_in > 0:
            pos, xx, yy = sampler.run_mcmc(p0, Nburn_in)
            sampler.reset()
        else:
            pos = p0

        pos, prob, state = sampler.run_mcmc(pos, Nsteps)

        # write the sampler to a pickle file
        sampler.lnprobfn = None
        sampler.pool = None
        fnpickle(sampler, sampler_file)

        pool.close()

    else:
        sampler = fnunpickle(sampler_file)

    if make_plots:

        # TODO: plot observed data / true particles

        # Make a corner plot for the potential parameters
        Npp = len(potential_params) # number of potential parameters
        pparams = model.parameters[:Npp]

        # First, just samples from the priors:
        fig = triangle.corner(p0[:,:Npp],
                    truths=[p.target._truth for p in pparams],
                    extents=[(p._ln_prior.a,p._ln_prior.b) for p in pparams],
                    labels=[p.target.latex for p in pparams])
        fig.savefig(os.path.join(path, "potential_corner_prior.png"))

        # Now the actual chains, extents from the priors
        fig = triangle.corner(sampler.flatchain[:,:Npp],
                    truths=[p.target._truth for p in pparams],
                    extents=[(p._ln_prior.a,p._ln_prior.b) for p in pparams],
                    labels=[p.target.latex for p in pparams])
        fig.savefig(os.path.join(path, "potential_corner.png"))

        # ---------
        # Now make 7x7 corner plots for each particle
        # TODO: need these plots to fail if not specified in config...
        true_obs_data = _gc_to_hel(_particles._X)
        for ii in range(Nparticles):
            tub = sampler.flatchain[:,Npp+ii]

            start = Npp+Nparticles + 6*ii
            stop = start + 6
            XX = sampler.flatchain[:,start:stop]
            OO = _gc_to_hel(XX)
            truths = np.append(_particles.tub[ii], true_obs_data[ii])

            extents = zip(obs_data[ii] - 3*obs_error[ii], \
                          obs_data[ii] + 3*obs_error[ii])
            extents = [(t2,t1)] + extents
            fig = triangle.corner(np.hstack((tub[:,np.newaxis], OO)),
                                  labels=['tub','l','b','D',\
                                          r'$\mu_l$', r'$\mu_l$','$v_r$'],
                                  extents=extents,
                                  truths=truths)
            fig.suptitle("Particle {0}".format(ii))
            fig.savefig(os.path.join(path, "particle_{0}_corner.png"\
                                     .format(ii)))
            plt.clf()

        # ---------
        # Now make 6x6 corner plot for satellite
        true_obs_data = _gc_to_hel(_satellite._X)
        start = Npp + Nparticles + 6*Nparticles
        stop = start + 6

        XX = sampler.flatchain[:,start:stop]
        OO = _gc_to_hel(XX)
        truths = np.squeeze(true_obs_data)

        so = np.squeeze(sat_obs_data)
        se = np.squeeze(sat_obs_error)
        extents = zip(so - 3*se, \
                      so + 3*se)

        # First, just samples from the priors:
        fig = triangle.corner(p0[:,start:stop],
                    truths=truths,
                    extents=extents,
                    labels=['l','b','D',\
                            r'$\mu_l$', r'$\mu_l$','$v_r$'],)
        fig.savefig(os.path.join(path, "satellite_corner_prior.png"))

        fig = triangle.corner(OO,
                              labels=['l','b','D',\
                                      r'$\mu_l$', r'$\mu_l$','$v_r$'],
                              extents=extents,
                              truths=truths)
        fig.suptitle("Satellite")
        fig.savefig(os.path.join(path, "satellite_corner.png"))

        # fig = plt.figure(figsize=(6,4))
        # ax = fig.add_subplot(111)
        # for jj in range(Npp) + [10]:
        #     ax.cla()
        #     for ii in range(Nwalkers):
        #         ax.plot(sampler.chain[ii,:,jj], drawstyle='step')

        #     fig.savefig(os.path.join(path, "{0}.png".format(jj)))

    sys.exit(0)

if __name__ == "__main__":
    from argparse import ArgumentParser

    parser = ArgumentParser(description="")
    parser.add_argument("-v", "--verbose", action="store_true",
                        dest="verbose", default=False,
                        help="Be chatty! (default = False)")
    parser.add_argument("-q", "--quiet", action="store_true", dest="quiet",
                    default=False, help="Be quiet! (default = False)")

    parser.add_argument("-f", "--file", dest="file", default="streams.cfg",
                    help="Path to the configuration file to run with.")
    parser.add_argument("-n", "--name", dest="job_name", default=None,
                    help="Name of the output.")

    args = parser.parse_args()

    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)
    elif args.quiet:
        logging.basicConfig(level=logging.ERROR)
    else:
        logging.basicConfig(level=logging.INFO)

    try:
        main(args.file, args.job_name)
    except:
        raise
        sys.exit(1)

    sys.exit(0)
<commit_msg>broke out making model from config file<commit_after>#!/hpc/astro/users/amp2217/yt-x86_64/bin/python
# coding: utf-8

""" Script for using the Rewinder to infer the Galactic host potential """

from __future__ import division, print_function

__author__ = "adrn <adrn@astro.columbia.edu>"

# Standard library
import os, sys
import shutil
import copy
import logging
from datetime import datetime
import multiprocessing

# Third-party
import astropy.units as u
from astropy.io.misc import fnpickle, fnunpickle
from astropy.utils.console import color_print
import emcee
import matplotlib.pyplot as plt
import numpy as np
np.seterr(all="ignore")
import scipy
scipy.seterr(all="ignore")
import triangle
import yaml

try:
    from emcee.utils import MPIPool
except ImportError:
    color_print("Failed to import MPIPool from emcee! MPI functionality "
                "won't work.", "yellow")

# Project
from streams import usys
from streams.coordinates import _gc_to_hel, _hel_to_gc
from streams.inference import (ModelParameter, StreamModel,
                               LogNormalPrior, LogUniformPrior)
from streams.io.sgr import mass_selector
from streams.observation.gaia import RRLyraeErrorModel
import streams.potential as sp
from streams.util import make_path

global pool
pool = None

# Create logger
logger = logging.getLogger(__name__)

def _null(*args, **kwargs):
    return 0.

def get_pool(config):
    """ Given a config structure, return an MPIPool, a Python
        multiprocessing.Pool, or None.
    """

    if config.get("mpi", False):
        # Initialize the MPI pool
        pool = MPIPool()

        # Make sure the thread we're running on is the master
        if not pool.is_master():
            pool.wait()
            sys.exit(0)
        logger.info("Running with MPI.")

    elif config.get("threads", 0) > 1:
        logger.info("Running with multiprocessing on {} cores."\
                    .format(config["threads"]))
        pool = multiprocessing.Pool(config["threads"])
    else:
        logger.info("Running serial...parallel panda is sad :(")
        pool = None

    return pool

def read_simulation(config):
    """ TODO: """

    if config["simulation"]["source"] == "sgr":
        m = config["simulation"]["satellite_mass"]
        particles_today, satellite_today, time = mass_selector(m)
    elif config["simulation"]["source"] == "lm10":
        #particles_today, satellite_today, time = mass_selector(config["mass"])
        # TODO
        pass
    else:
        raise ValueError("Invalid particle_source: {particle_source}"\
                         .format(config))

    particles = particles_today(N=config["particles"]["N"],
                                expr=config["particles"]["selection_expr"])
    satellite = satellite_today()
    t1,t2 = time()

    return t1,t2,satellite,particles

def _parse_quantity(q):
    try:
        val,unit = q.split()
    except AttributeError:
        val = q
        unit = u.dimensionless_unscaled

    return u.Quantity(float(val), unit)

def build_stream_model(config):
    """ """

    # get the potential object specified
    Potential = getattr(sp, config["potential"]["class_name"])
    potential = Potential()
    logger.debug("Using {} potential...".format(potential))

    # Actually get simulation data
    np.random.seed(config["seed"])
    t1,t2,_satellite,_particles = read_simulation(config)

    # TODO: right now error specification in yml doesn't propagate
    factor = config["errors"].get("global_factor", 1.)
    particle_errors = dict()
    for k,v in config["errors"].items():
        if k == "global_factor": continue
        particle_errors["{}_err".format(k)] = _parse_quantity(v)

    error_model = RRLyraeErrorModel(units=usys,
                                    factor=factor,
                                    **particle_errors)
    obs_data, obs_error = _particles.observe(error_model)

    # satellite has different errors from individual stars...
    # from: http://iopscience.iop.org/1538-4357/618/1/L25/pdf/18807.web.pdf
    sat_error_model = RRLyraeErrorModel(units=usys,
                                        mul_err=1000.*u.mas/u.yr,
                                        mub_err=1000.*u.mas/u.yr,
                                        D_err=1.,
                                        vr_err=200.*u.km/u.s)
    sat_obs_data, sat_obs_error = _satellite.observe(sat_error_model)

    particles = _particles.copy()
    satellite = _satellite.copy()

    model_parameters = []
    ##########################################################################
    # Potential parameters
    #
    potential_params = config["potential"].get("parameters", dict())
    for name,meta in potential_params.items():
        p = getattr(potential, name)

        if meta.has_key("range"):
            # TODO: fix when astropy fixed...
            #lo,hi = map(u.Quantity, meta["range"])
            lo = _parse_quantity(meta["range"][0])
            hi = _parse_quantity(meta["range"][1])
            prior = LogUniformPrior(lo.decompose(usys).value,
                                    hi.decompose(usys).value)
        else:
            prior = LogPrior()

        model_parameters.append(ModelParameter(target=p,
                                               attr="_value",
                                               ln_prior=prior))

    ##########################################################################
    # Particle parameters
    #
    particle_params = config["particles"].get("parameters", [])

    try:
        Nparticles = config["particles"]["N"]
    except KeyError:
        raise ValueError("Must specify number of partices in config file!")

    # time unbound / escape time (tub)
    if "tub" in particle_params:
        lo = [t2] * len(particles)
        hi = [t1] * len(particles)
        prior = LogUniformPrior(lo, hi)
        model_parameters.append(ModelParameter(target=particles,
                                               attr="tub",
                                               ln_prior=prior))

    # here I monte carlo transform the error distribution from observed
    #   to cartesian, then take np.cov and use that to sample new particle
    #   positions
    O = np.array([np.random.normal(obs_data, obs_error) \
                    for ii in range(1000)])
    X = _hel_to_gc(O)

    obs_error_gc = []
    for ii in range(Nparticles):
        obs_error_gc.append(np.cov(X[:,ii].T))

    obs_error_gc = np.array(obs_error_gc)
    obs_data_gc = _hel_to_gc(obs_data)

    # true positions of particles (flat_X)
    if "_X" in particle_params:
        particles._X = _hel_to_gc(obs_data)
        particles.obs_data = obs_data
        particles.obs_error = obs_error

        prior = LogNormalPrior(obs_data_gc, cov=obs_error_gc)
        p = ModelParameter(target=particles, attr="_X", ln_prior=prior)
        p.ln_prior = _null # THIS IS A HACK?
        model_parameters.append(p)

    ##########################################################################
    # Satellite parameters
    #
    satellite_config = config.get("satellite", dict())
    satellite_params = satellite_config.get("parameters", [])

    # here I monte carlo transform the error distribution from observed
    #   to cartesian, then take np.cov and use that to sample new satellite
    #   positions
    O = np.array([np.random.normal(sat_obs_data, sat_obs_error) \
                    for ii in range(1000)])
    X = _hel_to_gc(O)

    sat_obs_error_gc = np.array([np.cov(X[:,0].T)])
    sat_obs_data_gc = _hel_to_gc(sat_obs_data)

    # true position of the satellite
    if "_X" in satellite_params:
        satellite._X = _hel_to_gc(sat_obs_data)
        satellite.obs_data = sat_obs_data
        satellite.obs_error = sat_obs_error

        prior = LogNormalPrior(sat_obs_data_gc, cov=sat_obs_error_gc)
        p = ModelParameter(target=satellite, attr="_X", ln_prior=prior)
        p.ln_prior = _null # THIS IS A HACK?
        model_parameters.append(p)

    # now create the model
    model = StreamModel(potential, satellite, particles,
                        parameters=model_parameters)

    return model

def main(config_file, job_name=None):
    """ TODO: """

    # read in configurable parameters
    with open(config_file) as f:
        config = yaml.load(f.read())

    # TODO: write a separate script just for plotting...
    # get a pool object given the configuration parameters
    pool = get_pool(config)

    # determine the output data path
    path = make_path(config)
    sampler_file = os.path.join(path, "sampler_data.pickle")

    logger.info("Will write output to '{0}'...".format(path))

    make_plots = config.get("make_plots", False)

    model = build_stream_model(config)

    # read in the number of walkers to use
    Nwalkers = config.get("walkers", "auto")
    if str(Nwalkers).lower() == "auto":
        Nwalkers = model.ndim*2 + 2

    p0 = model.sample(size=Nwalkers)

    if not os.path.exists(sampler_file):
        sampler = emcee.EnsembleSampler(Nwalkers, model.ndim, model,
                                        args=(t1, t2, -1.),
                                        pool=pool)

        Nburn_in = config.get("burn_in", 0)
        Nsteps = config["steps"]
        if Nburn_in > 0:
            pos, xx, yy = sampler.run_mcmc(p0, Nburn_in)
            sampler.reset()
        else:
            pos = p0

        pos, prob, state = sampler.run_mcmc(pos, Nsteps)

        # write the sampler to a pickle file
        sampler.lnprobfn = None
        sampler.pool = None
        fnpickle(sampler, sampler_file)

        pool.close()

    else:
        sampler = fnunpickle(sampler_file)

    if make_plots:

        # TODO: plot observed data / true particles

        # Make a corner plot for the potential parameters
        Npp = len(potential_params) # number of potential parameters
        pparams = model.parameters[:Npp]

        # First, just samples from the priors:
        fig = triangle.corner(p0[:,:Npp],
                    truths=[p.target._truth for p in pparams],
                    extents=[(p._ln_prior.a,p._ln_prior.b) for p in pparams],
                    labels=[p.target.latex for p in pparams])
        fig.savefig(os.path.join(path, "potential_corner_prior.png"))

        # Now the actual chains, extents from the priors
        fig = triangle.corner(sampler.flatchain[:,:Npp],
                    truths=[p.target._truth for p in pparams],
                    extents=[(p._ln_prior.a,p._ln_prior.b) for p in pparams],
                    labels=[p.target.latex for p in pparams])
        fig.savefig(os.path.join(path, "potential_corner.png"))

        # ---------
        # Now make 7x7 corner plots for each particle
        # TODO: need these plots to fail if not specified in config...
        true_obs_data = _gc_to_hel(_particles._X)
        for ii in range(Nparticles):
            tub = sampler.flatchain[:,Npp+ii]

            start = Npp+Nparticles + 6*ii
            stop = start + 6
            XX = sampler.flatchain[:,start:stop]
            OO = _gc_to_hel(XX)
            truths = np.append(_particles.tub[ii], true_obs_data[ii])

            extents = zip(obs_data[ii] - 3*obs_error[ii], \
                          obs_data[ii] + 3*obs_error[ii])
            extents = [(t2,t1)] + extents
            fig = triangle.corner(np.hstack((tub[:,np.newaxis], OO)),
                                  labels=['tub','l','b','D',\
                                          r'$\mu_l$', r'$\mu_l$','$v_r$'],
                                  extents=extents,
                                  truths=truths)
            fig.suptitle("Particle {0}".format(ii))
            fig.savefig(os.path.join(path, "particle_{0}_corner.png"\
                                     .format(ii)))
            plt.clf()

        # ---------
        # Now make 6x6 corner plot for satellite
        true_obs_data = _gc_to_hel(_satellite._X)
        start = Npp + Nparticles + 6*Nparticles
        stop = start + 6

        XX = sampler.flatchain[:,start:stop]
        OO = _gc_to_hel(XX)
        truths = np.squeeze(true_obs_data)

        so = np.squeeze(sat_obs_data)
        se = np.squeeze(sat_obs_error)
        extents = zip(so - 3*se, \
                      so + 3*se)

        # First, just samples from the priors:
        fig = triangle.corner(p0[:,start:stop],
                    truths=truths,
                    extents=extents,
                    labels=['l','b','D',\
                            r'$\mu_l$', r'$\mu_l$','$v_r$'],)
        fig.savefig(os.path.join(path, "satellite_corner_prior.png"))

        fig = triangle.corner(OO,
                              labels=['l','b','D',\
                                      r'$\mu_l$', r'$\mu_l$','$v_r$'],
                              extents=extents,
                              truths=truths)
        fig.suptitle("Satellite")
        fig.savefig(os.path.join(path, "satellite_corner.png"))

        # fig = plt.figure(figsize=(6,4))
        # ax = fig.add_subplot(111)
        # for jj in range(Npp) + [10]:
        #     ax.cla()
        #     for ii in range(Nwalkers):
        #         ax.plot(sampler.chain[ii,:,jj], drawstyle='step')

        #     fig.savefig(os.path.join(path, "{0}.png".format(jj)))

    sys.exit(0)

if __name__ == "__main__":
    from argparse import ArgumentParser

    parser = ArgumentParser(description="")
    parser.add_argument("-v", "--verbose", action="store_true",
                        dest="verbose", default=False,
                        help="Be chatty! (default = False)")
    parser.add_argument("-q", "--quiet", action="store_true", dest="quiet",
                    default=False, help="Be quiet! (default = False)")

    parser.add_argument("-f", "--file", dest="file", default="streams.cfg",
                    help="Path to the configuration file to run with.")
    parser.add_argument("-n", "--name", dest="job_name", default=None,
                    help="Name of the output.")

    args = parser.parse_args()

    if args.verbose:
        logging.basicConfig(level=logging.DEBUG)
    elif args.quiet:
        logging.basicConfig(level=logging.ERROR)
    else:
        logging.basicConfig(level=logging.INFO)

    try:
        main(args.file, args.job_name)
    except:
        raise
        sys.exit(1)

    sys.exit(0)
<|endoftext|>


<commit_before># copies data from SGID to the app and makes necessary optimizations

import arcpy
import settings
from settings import fieldnames
import spreadsheet
from os import path
from build_json import parse_fields
from collections import namedtuple

fiveFields = [
              fieldnames.ID,
              fieldnames.NAME,
              fieldnames.ADDRESS,
              fieldnames.CITY,
              fieldnames.TYPE
              ]
logger = None
errors = []
field_type_mappings = {'Integer': 'LONG',
                       'String': 'TEXT',
                       'SmallInteger': 'SHORT'}

def run(logr, test_layer=None):
    global logger, errors
    logger = logr

    logger.logMsg('processing query layers\n')
    update_query_layers(test_layer)

    logger.logMsg('processing related tables\n')
    update_related_tables(test_layer)

    return errors

def update_related_tables(test_layer=None):
    for t in spreadsheet.get_related_tables():
        name = t[fieldnames.sgidName]
        if test_layer and name != test_layer:
            continue
        try:
            if name.startswith('SGID10'):
                logger.logMsg('\nProcessing: {}'.format(name.split('.')[-1]))
                localTbl = path.join(settings.fgd, name.split('.')[-1])
                remoteTbl = path.join(settings.sgid, name)
                update(localTbl, remoteTbl)
                validate_fields([f.name for f in arcpy.ListFields(localTbl)], t[fieldnames.fields], name)
        except:
            errors.append('Execution error trying to update fgdb with {}:\n{}'.format(name, logger.logError()))

def update(local, remote):
    logger.logMsg('updating: {} \n    from: {}'.format(local, remote))
    if not arcpy.Exists(local):
        logger.logMsg('creating new local feature class')
        arcpy.Copy_management(remote, local)
    else:
        arcpy.TruncateTable_management(local)
        try:
            arcpy.Append_management(remote, local, 'NO_TEST')
        except:
            with arcpy.da.Editor(settings.fgd) as edit:
                logger.logMsg('append failed, using insert cursor')
                flds = [f.name for f in arcpy.ListFields(remote)]
                with arcpy.da.SearchCursor(remote, flds) as rcur, arcpy.da.InsertCursor(local, flds) as icur:
                    for row in rcur:
                        icur.insertRow(row)

def update_query_layers(test_layer=None):
    for l in spreadsheet.get_query_layers():
        fcname = l[fieldnames.sgidName]
        if test_layer and fcname != test_layer:
            continue
        try:
            # only try to update rows with valid sgid names and data sources
            if fcname.startswith('SGID10'):
                # update fgd from SGID
                logger.logMsg('\nProcessing: {}'.format(fcname.split('.')[-1]))
                localFc = path.join(settings.fgd, fcname.split('.')[-1])
                remoteFc = path.join(settings.sgid, fcname)
                update(localFc, remoteFc)

                # APP-SPECIFIC OPTIMIZATIONS
                # make sure that it has the five main fields for the fdg only
                upper_fields = [x.name.upper() for x in arcpy.ListFields(localFc)]
                for f in fiveFields:
                    if not f in upper_fields:
                        logger.logMsg('{} not found. Adding to {}'.format(f, localFc))

                        # get mapped field properties
                        if not l[f] == 'n/a':
                            mappedFld = arcpy.ListFields(localFc, l[f])[0]
                        else:
                            mappedFld = namedtuple('literal', 'precision scale length')(**{'precision': 0,
                                                                                'scale': 0,
                                                                                'length': 3})

                        arcpy.AddField_management(localFc, f, 'String',
                                                  mappedFld.precision, mappedFld.scale, mappedFld.length)

                    # calc field
                    expression = l[f]
                    if not expression == 'n/a':
                        expression = '(!{}!.encode("utf-8"))'.format(expression)
                    else:
                        expression = '"{}"'.format(expression)
                    arcpy.CalculateField_management(localFc, f, expression, 'PYTHON')

                validate_fields([f.name for f in arcpy.ListFields(localFc)], l[fieldnames.fields], fcname)
        except:
            errors.append('Execution error trying to update fgdb with {}:\n{}'.format(fcname, logger.logError().strip()))

def validate_fields(dataFields, fieldString, datasetName):
    msg = '{}: Could not find matches in the source data for the following fields from the query layers spreadsheet: {}'
    dataFields = set(dataFields)
    spreadsheetFields = set([f[0] for f in parse_fields(fieldString)])

    invalidFields = spreadsheetFields - dataFields

    if len(invalidFields) > 0:
        er = msg.format(datasetName, ', '.join(invalidFields))
        errors.append(er)
        return er
    else:
        return []

if __name__ == '__main__':
    from agrc import logging
    import sys

    logger = logging.Logger()

    # first argument is optionally the SGID feature class or table name
    if len(sys.argv) == 2:
        run(logger, sys.argv[1])
    else:
        run(logger)
    print('done')<commit_msg>Handle numeric id fields.<commit_after># copies data from SGID to the app and makes necessary optimizations

import arcpy
import settings
from settings import fieldnames
import spreadsheet
from os import path
from build_json import parse_fields
from collections import namedtuple

fiveFields = [
              fieldnames.ID,
              fieldnames.NAME,
              fieldnames.ADDRESS,
              fieldnames.CITY,
              fieldnames.TYPE
              ]
logger = None
errors = []
field_type_mappings = {'Integer': 'LONG',
                       'String': 'TEXT',
                       'SmallInteger': 'SHORT'}

def run(logr, test_layer=None):
    global logger, errors
    logger = logr

    logger.logMsg('processing query layers\n')
    update_query_layers(test_layer)

    logger.logMsg('processing related tables\n')
    update_related_tables(test_layer)

    return errors

def update_related_tables(test_layer=None):
    for t in spreadsheet.get_related_tables():
        name = t[fieldnames.sgidName]
        if test_layer and name != test_layer:
            continue
        try:
            if name.startswith('SGID10'):
                logger.logMsg('\nProcessing: {}'.format(name.split('.')[-1]))
                localTbl = path.join(settings.fgd, name.split('.')[-1])
                remoteTbl = path.join(settings.sgid, name)
                update(localTbl, remoteTbl)
                validate_fields([f.name for f in arcpy.ListFields(localTbl)], t[fieldnames.fields], name)
        except:
            errors.append('Execution error trying to update fgdb with {}:\n{}'.format(name, logger.logError()))

def update(local, remote):
    logger.logMsg('updating: {} \n    from: {}'.format(local, remote))
    if not arcpy.Exists(local):
        logger.logMsg('creating new local feature class')
        arcpy.Copy_management(remote, local)
    else:
        arcpy.TruncateTable_management(local)
        try:
            arcpy.Append_management(remote, local, 'NO_TEST')
        except:
            with arcpy.da.Editor(settings.fgd) as edit:
                logger.logMsg('append failed, using insert cursor')
                flds = [f.name for f in arcpy.ListFields(remote)]
                with arcpy.da.SearchCursor(remote, flds) as rcur, arcpy.da.InsertCursor(local, flds) as icur:
                    for row in rcur:
                        icur.insertRow(row)

def update_query_layers(test_layer=None):
    for l in spreadsheet.get_query_layers():
        fcname = l[fieldnames.sgidName]
        if test_layer and fcname != test_layer:
            continue
        try:
            # only try to update rows with valid sgid names and data sources
            if fcname.startswith('SGID10'):
                # update fgd from SGID
                logger.logMsg('\nProcessing: {}'.format(fcname.split('.')[-1]))
                localFc = path.join(settings.fgd, fcname.split('.')[-1])
                remoteFc = path.join(settings.sgid, fcname)
                update(localFc, remoteFc)

                # APP-SPECIFIC OPTIMIZATIONS
                # make sure that it has the five main fields for the fdg only
                upper_fields = [x.name.upper() for x in arcpy.ListFields(localFc)]
                for f in fiveFields:
                    if not f in upper_fields:
                        logger.logMsg('{} not found. Adding to {}'.format(f, localFc))

                        # get mapped field properties
                        if not l[f] == 'n/a':
                            mappedFld = arcpy.ListFields(localFc, l[f])[0]
                        else:
                            mappedFld = namedtuple('literal', 'precision scale length')(**{'precision': 0,
                                                                                'scale': 0,
                                                                                'length': 3})

                        arcpy.AddField_management(localFc, f, 'String',
                                                  mappedFld.precision, mappedFld.scale, mappedFld.length)

                    # calc field
                    expression = l[f]
                    if not expression == 'n/a':
                        if arcpy.ListFields(localFc, l[f])[0].type != 'String':
                            expression = 'str(!{}!)'.format(expression)
                        else:
                            expression = '(!{}!.encode("utf-8"))'.format(expression)
                    else:
                        expression = '"{}"'.format(expression)
                    arcpy.CalculateField_management(localFc, f, expression, 'PYTHON')

                validate_fields([f.name for f in arcpy.ListFields(localFc)], l[fieldnames.fields], fcname)
        except:
            errors.append('Execution error trying to update fgdb with {}:\n{}'.format(fcname, logger.logError().strip()))

def validate_fields(dataFields, fieldString, datasetName):
    msg = '{}: Could not find matches in the source data for the following fields from the query layers spreadsheet: {}'
    dataFields = set(dataFields)
    spreadsheetFields = set([f[0] for f in parse_fields(fieldString)])

    invalidFields = spreadsheetFields - dataFields

    if len(invalidFields) > 0:
        er = msg.format(datasetName, ', '.join(invalidFields))
        errors.append(er)
        return er
    else:
        return []

if __name__ == '__main__':
    from agrc import logging
    import sys

    logger = logging.Logger()

    # first argument is optionally the SGID feature class or table name
    if len(sys.argv) == 2:
        run(logger, sys.argv[1])
    else:
        run(logger)
    print('done')<|endoftext|>


<commit_before>'''
update_sgid.py

handles updating data in SGID from the app database
'''
import logging
from os import path

import arcpy
import settings

utm = arcpy.SpatialReference(26912)
period_replacement = '___'
logger = logging.getLogger('forklift')

def update_sgid_data(source, destination):
    logger.debug('updating: %s', destination)

    arcpy.DeleteRows_management(destination)
    arcpy.Append_management(source, destination, 'NO_TEST')

def update_sgid_for_crates(crates_from_slip):
    updated_crates = [crate for crate in crates_from_slip if crate['was_updated'] and not crate['name'].startswith('DirectFrom')]

    for crate_slip in updated_crates:
        sgid_name = path.basename(crate_slip['destination']).replace(period_replacement, '.')
        destination = path.join(settings.sgid[sgid_name.split('.')[1]], sgid_name)

        if sgid_name.startswith('SGID10'):
            update_sgid_data(crate_slip['destination'], destination)
<commit_msg>fix: make SGID update work on unqualified feature class names<commit_after>'''
update_sgid.py

handles updating data in SGID from the app database
'''
import logging
from os import path

import arcpy
import settings

utm = arcpy.SpatialReference(26912)
period_replacement = '___'
logger = logging.getLogger('forklift')

def update_sgid_data(source, destination):
    logger.debug('updating: %s', destination)

    arcpy.DeleteRows_management(destination)
    arcpy.Append_management(source, destination, 'NO_TEST')

def update_sgid_for_crates(crates_from_slip):
    logger.info('getting SGID names lookup')
    sgid_lookup = {}
    arcpy.env.workspace = path.join(settings.dbConnects, 'SGID10.sde')
    for dataset in arcpy.ListTables() + arcpy.ListFeatureClasses():
        sgid_lookup[dataset.split('.')[-1]] = dataset
    arcpy.env.workspace = None

    updated_crates = [crate for crate in crates_from_slip if crate['was_updated'] and not crate['name'].startswith('DirectFrom')]

    for crate_slip in updated_crates:
        sgid_name = sgid_lookup[crate_slip['name']]
        destination = path.join(settings.sgid[sgid_name.split('.')[1]], sgid_name)

        if sgid_name.startswith('SGID10'):
            update_sgid_data(crate_slip['destination'], destination)
<|endoftext|>