<commit_before># Core modules
import os
import re
import logging
import platform
import subprocess
import sys
import time
import datetime
import socket

# Needed to identify server uniquely
import uuid
try:
    from hashlib import md5
except ImportError: # Python < 2.5
    from md5 import new as md5

import modules

from util import getOS
from config import get_version
from checks import gethostname

import checks.system.unix as u
import checks.system.win32 as w32

from checks.nagios import Nagios
from checks.build import Hudson

from checks.db.mysql import MySql
from checks.db.mongo import MongoDb
from checks.db.couch import CouchDb
from checks.db.mcache import Memcache

from checks.queue import RabbitMq
from checks.web import Apache, Nginx
from checks.ganglia import Ganglia
from checks.cassandra import Cassandra
from checks.datadog import Dogstreams, DdForwarder

from checks.jmx import Jvm, Tomcat, ActiveMQ, Solr
from checks.cacti import Cacti
from checks.varnish import Varnish

from checks.db.elastic import ElasticSearch, ElasticSearchClusterStatus

from checks.wmi_check import WMICheck

from checks.ec2 import EC2

from resources.processes import Processes as ResProcesses

def getUuid():
    # Generate a unique name that will stay constant between
    # invocations, such as platform.node() + uuid.getnode()
    # Use uuid5, which does not depend on the clock and is
    # recommended over uuid3.
    # This is important to be able to identify a server even if
    # its drives have been wiped clean.
    # Note that this is not foolproof but we can reconcile servers
    # on the back-end if need be, based on mac addresses.
    return uuid.uuid5(uuid.NAMESPACE_DNS, platform.node() + str(uuid.getnode())).hex

class checks(object):
    def __init__(self, agentConfig, emitters):
        self.agentConfig = agentConfig
        self.os = getOS()
        self.plugins = None
        self.emitters = emitters
        self.checksLogger = logging.getLogger('checks')
        self.metadata_interval = int(agentConfig.get('metadata_interval', 5 * 60 * 60))
        self.metadata_start = time.time()
        socket.setdefaulttimeout(15)

        # Unix System Checks
        self._unix_system_checks = {
            'disk': u.Disk(self.checksLogger),
            'io': u.IO(),
            'load': u.Load(self.checksLogger),
            'memory': u.Memory(self.checksLogger),
            'network': u.Network(self.checksLogger),
            'processes': u.Processes(),
            'cpu': u.Cpu()
        }

        # Win32 System Checks
        self._win32_system_checks = {
            'disk': w32.Disk(self.checksLogger),
            'io': w32.IO(self.checksLogger),
            'proc': w32.Processes(self.checksLogger),
            'memory': w32.Memory(self.checksLogger),
            'network': w32.Network(self.checksLogger),
            'cpu': w32.Cpu(self.checksLogger)
        }

        # Old-style metric checks
        self._apache = Apache(self.checksLogger)
        self._couchdb = CouchDb(self.checksLogger)
        self._mongodb = MongoDb(self.checksLogger)
        self._mysql = MySql(self.checksLogger)
        self._rabbitmq = RabbitMq()
        self._ganglia = Ganglia(self.checksLogger)
        self._cassandra = Cassandra()
        self._dogstream = Dogstreams.init(self.checksLogger, self.agentConfig)
        self._ddforwarder = DdForwarder(self.checksLogger, self.agentConfig)
        self._ec2 = EC2(self.checksLogger)

        # Metric Checks
        self._metrics_checks = [
            Cacti(self.checksLogger),
            Varnish(self.checksLogger),
            ElasticSearch(self.checksLogger),
            Jvm(self.checksLogger),
            Tomcat(self.checksLogger),
            ActiveMQ(self.checksLogger),
            Solr(self.checksLogger),
            WMICheck(self.checksLogger),
            Nginx(self.checksLogger),
            Memcache(self.checksLogger),
        ]

        # Custom metric checks
        for module_spec in [s.strip() for s in self.agentConfig.get('custom_checks', '').split(',')]:
            if len(module_spec) == 0: continue
            try:
                self._metrics_checks.append(modules.load(module_spec, 'Check')(self.checksLogger))
                self.checksLogger.info("Registered custom check %s" % module_spec)
            except Exception, e:
                self.checksLogger.exception('Unable to load custom check module %s' % module_spec)

        # Event Checks
        self._event_checks = [
            ElasticSearchClusterStatus(self.checksLogger),
            Nagios(socket.gethostname())
        ]

        # Resource Checks
        self._resources_checks = [
            ResProcesses(self.checksLogger,self.agentConfig)
        ]

    def get_metadata(self):
        metadata = self._ec2.get_metadata()
        if metadata.get('hostname'):
            metadata['ec2-hostname'] = metadata.get('hostname')

        # if hostname is set in the configuration file
        # use that instead of gethostname
        # gethostname is vulnerable to 2 hosts: x.domain1, x.domain2
        # will cause both to be aliased (see #157)
        if self.agentConfig.get('hostname'):
            metadata['agent-hostname'] = self.agentConfig.get('hostname')
            metadata['hostname'] = metadata['agent-hostname']
        else:
            try:
                metadata["hostname"] = socket.gethostname()
            except:
                pass
        try:
            metadata["fqdn"] = socket.getfqdn()
        except:
            pass

        return metadata

    def should_send_metadata(self):
        # If the interval has passed, send the metadata again
        now = time.time()
        if now - self.metadata_start >= self.metadata_interval:
            self.checksLogger.debug('Metadata interval has passed. Sending metadata.')
            self.metadata_start = now
            return True

        return False

    def doChecks(self, firstRun=False, systemStats=False, checksd=None):
        """Actual work
        """
        self.checksLogger.info("Starting checks")
        checksData = {
            'collection_timestamp': time.time(),
            'os' : self.os,
            'python': sys.version,
            'agentVersion' : self.agentConfig['version'],
            'apiKey': self.agentConfig['api_key'],
            'events': {},
            'resources': {}
        }
        metrics = []
        events = {}

        # Run the system checks. Checks will depend on the OS
        if self.os == 'windows':
            # Win32 system checks
            metrics.extend(self._win32_system_checks['disk'].check(self.agentConfig))
            metrics.extend(self._win32_system_checks['memory'].check(self.agentConfig))
            metrics.extend(self._win32_system_checks['cpu'].check(self.agentConfig))
            metrics.extend(self._win32_system_checks['network'].check(self.agentConfig))
            metrics.extend(self._win32_system_checks['io'].check(self.agentConfig))
            metrics.extend(self._win32_system_checks['proc'].check(self.agentConfig))
        else:
            # Unix system checks
            sys_checks = self._unix_system_checks

            diskUsage = sys_checks['disk'].check(self.agentConfig)
            if diskUsage is not False and len(diskUsage) == 2:
                checksData["diskUsage"] = diskUsage[0]
                checksData["inodes"] = diskUsage[1]

            loadAvrgs = sys_checks['load'].check(self.agentConfig)
            checksData.update({
                'loadAvrg1': loadAvrgs['1'],
                'loadAvrg5': loadAvrgs['5'],
                'loadAvrg15': loadAvrgs['15']
            })

            memory = sys_checks['memory'].check(self.agentConfig)
            checksData.update({
                'memPhysUsed' : memory.get('physUsed'),
                'memPhysFree' : memory.get('physFree'),
                'memPhysTotal' : memory.get('physTotal'),
                'memPhysUsable' : memory.get('physUsable'),
                'memSwapUsed' : memory.get('swapUsed'),
                'memSwapFree' : memory.get('swapFree'),
                'memSwapTotal' : memory.get('swapTotal'),
                'memCached' : memory.get('physCached'),
                'memBuffers': memory.get('physBuffers'),
                'memShared': memory.get('physShared')
            })

            ioStats = sys_checks['io'].check(self.checksLogger, self.agentConfig)
            if ioStats:
                checksData['ioStats'] = ioStats

            processes = sys_checks['processes'].check(self.checksLogger, self.agentConfig)
            checksData.update({'processes': processes})

            networkTraffic = sys_checks['network'].check(self.agentConfig)
            checksData.update({'networkTraffic': networkTraffic})

            cpuStats = sys_checks['cpu'].check(self.checksLogger, self.agentConfig)
            if cpuStats is not False and cpuStats is not None:
                checksData.update(cpuStats)

        # Run old-style checks
        apacheStatus = self._apache.check(self.agentConfig)
        mysqlStatus = self._mysql.check(self.agentConfig)
        rabbitmq = self._rabbitmq.check(self.checksLogger, self.agentConfig)
        mongodb = self._mongodb.check(self.agentConfig)
        couchdb = self._couchdb.check(self.agentConfig)
        gangliaData = self._ganglia.check(self.agentConfig)
        cassandraData = self._cassandra.check(self.checksLogger, self.agentConfig)
        dogstreamData = self._dogstream.check(self.agentConfig)
        ddforwarderData = self._ddforwarder.check(self.agentConfig)

        if gangliaData is not False and gangliaData is not None:
            checksData['ganglia'] = gangliaData

        if cassandraData is not False and cassandraData is not None:
            checksData['cassandra'] = cassandraData

        # Apache Status
        if apacheStatus:
            checksData.update(apacheStatus)

        # MySQL Status
        if mysqlStatus:
            checksData.update(mysqlStatus)

        # RabbitMQ
        if rabbitmq:
            checksData['rabbitMQ'] = rabbitmq

        # MongoDB
        if mongodb:
            if mongodb.has_key('events'):
                events['Mongo'] = mongodb['events']['Mongo']
                del mongodb['events']
            checksData['mongoDB'] = mongodb

        # CouchDB
        if couchdb:
            checksData['couchDB'] = couchdb

        if dogstreamData:
            dogstreamEvents = dogstreamData.get('dogstreamEvents', None)
            if dogstreamEvents:
                if 'dogstream' in checksData['events']:
                    events['dogstream'].extend(dogstreamEvents)
                else:
                    events['dogstream'] = dogstreamEvents
                del dogstreamData['dogstreamEvents']

            checksData.update(dogstreamData)

        if ddforwarderData:
            checksData['datadog'] = ddforwarderData

        # Include server indentifiers
        checksData['internalHostname'] = gethostname(self.agentConfig)
        checksData['uuid'] = getUuid()
        self.checksLogger.debug('doChecks: added uuid %s' % checksData['uuid'])

        # Process the event checks.
        for event_check in self._event_checks:
            event_data = event_check.check(self.checksLogger, self.agentConfig)
            if event_data:
                events[event_check.key] = event_data

        # Include system stats on first postback
        if firstRun:
            checksData['systemStats'] = systemStats
            # Also post an event in the newsfeed
            events['System'] = [{'api_key': self.agentConfig['api_key'],
                                               'host': checksData['internalHostname'],
                                               'timestamp': int(time.mktime(datetime.datetime.now().timetuple())),
                                               'event_type':'Agent Startup',
                                               'msg_text': 'Version %s' % get_version()
                                            }]

        if firstRun or self.should_send_metadata():
            # Collect metadata
            checksData['meta'] = self.get_metadata()
            # Add static tags from the configuration file
            if self.agentConfig['tags'] is not None:
                checksData['tags'] = self.agentConfig['tags']

        # Resources checks
        if self.os != 'windows':
            has_resource = False
            for resources_check in self._resources_checks:
                resources_check.check()
                snaps = resources_check.pop_snapshots()
                if snaps:
                    has_resource = True
                    res_value = { 'snaps': snaps,
                                  'format_version': resources_check.get_format_version() }
                    res_format = resources_check.describe_format_if_needed()
                    if res_format is not None:
                        res_value['format_description'] = res_format
                    checksData['resources'][resources_check.RESOURCE_KEY] = res_value

            if has_resource:
                checksData['resources']['meta'] = {
                            'api_key': self.agentConfig['api_key'],
                            'host': checksData['internalHostname'],
                        }

        for metrics_check in self._metrics_checks:
            res = metrics_check.check(self.agentConfig)
            if res:
                metrics.extend(res)

        # checks.d checks
        if checksd:
            for check in checksd:
                check_cls = check['class']
                for instance in check['instances']:
                    try:
                        # Run the check for each configuration
                        check_cls.check(instance)
                        metrics.extend(check_cls.get_metrics())
                        if check_cls.has_events():
                            if check['name'] not in events:
                                events[check['name']] = []
                            for ev in check_cls.get_events():
                                events[check['name']].append(ev)
                    except Exception:
                        self.checksLogger.exception("Check %s failed" % check_cls.name)

        # Store the metrics in the payload
        checksData['metrics'] = metrics

        # Store the events in the payload
        checksData['events'] = events

        # Send back data
        self.checksLogger.debug("checksData: %s" % checksData)
        for emitter in self.emitters:
            emitter(checksData, self.checksLogger, self.agentConfig)
        self.checksLogger.info("Checks done")<commit_msg>remove old cacti imports<commit_after># Core modules
import os
import re
import logging
import platform
import subprocess
import sys
import time
import datetime
import socket

# Needed to identify server uniquely
import uuid
try:
    from hashlib import md5
except ImportError: # Python < 2.5
    from md5 import new as md5

import modules

from util import getOS
from config import get_version
from checks import gethostname

import checks.system.unix as u
import checks.system.win32 as w32

from checks.nagios import Nagios
from checks.build import Hudson

from checks.db.mysql import MySql
from checks.db.mongo import MongoDb
from checks.db.couch import CouchDb
from checks.db.mcache import Memcache

from checks.queue import RabbitMq
from checks.web import Apache, Nginx
from checks.ganglia import Ganglia
from checks.cassandra import Cassandra
from checks.datadog import Dogstreams, DdForwarder

from checks.jmx import Jvm, Tomcat, ActiveMQ, Solr
from checks.varnish import Varnish

from checks.db.elastic import ElasticSearch, ElasticSearchClusterStatus

from checks.wmi_check import WMICheck

from checks.ec2 import EC2

from resources.processes import Processes as ResProcesses

def getUuid():
    # Generate a unique name that will stay constant between
    # invocations, such as platform.node() + uuid.getnode()
    # Use uuid5, which does not depend on the clock and is
    # recommended over uuid3.
    # This is important to be able to identify a server even if
    # its drives have been wiped clean.
    # Note that this is not foolproof but we can reconcile servers
    # on the back-end if need be, based on mac addresses.
    return uuid.uuid5(uuid.NAMESPACE_DNS, platform.node() + str(uuid.getnode())).hex

class checks(object):
    def __init__(self, agentConfig, emitters):
        self.agentConfig = agentConfig
        self.os = getOS()
        self.plugins = None
        self.emitters = emitters
        self.checksLogger = logging.getLogger('checks')
        self.metadata_interval = int(agentConfig.get('metadata_interval', 5 * 60 * 60))
        self.metadata_start = time.time()
        socket.setdefaulttimeout(15)

        # Unix System Checks
        self._unix_system_checks = {
            'disk': u.Disk(self.checksLogger),
            'io': u.IO(),
            'load': u.Load(self.checksLogger),
            'memory': u.Memory(self.checksLogger),
            'network': u.Network(self.checksLogger),
            'processes': u.Processes(),
            'cpu': u.Cpu()
        }

        # Win32 System Checks
        self._win32_system_checks = {
            'disk': w32.Disk(self.checksLogger),
            'io': w32.IO(self.checksLogger),
            'proc': w32.Processes(self.checksLogger),
            'memory': w32.Memory(self.checksLogger),
            'network': w32.Network(self.checksLogger),
            'cpu': w32.Cpu(self.checksLogger)
        }

        # Old-style metric checks
        self._apache = Apache(self.checksLogger)
        self._couchdb = CouchDb(self.checksLogger)
        self._mongodb = MongoDb(self.checksLogger)
        self._mysql = MySql(self.checksLogger)
        self._rabbitmq = RabbitMq()
        self._ganglia = Ganglia(self.checksLogger)
        self._cassandra = Cassandra()
        self._dogstream = Dogstreams.init(self.checksLogger, self.agentConfig)
        self._ddforwarder = DdForwarder(self.checksLogger, self.agentConfig)
        self._ec2 = EC2(self.checksLogger)

        # Metric Checks
        self._metrics_checks = [
            Varnish(self.checksLogger),
            ElasticSearch(self.checksLogger),
            Jvm(self.checksLogger),
            Tomcat(self.checksLogger),
            ActiveMQ(self.checksLogger),
            Solr(self.checksLogger),
            WMICheck(self.checksLogger),
            Nginx(self.checksLogger),
            Memcache(self.checksLogger),
        ]

        # Custom metric checks
        for module_spec in [s.strip() for s in self.agentConfig.get('custom_checks', '').split(',')]:
            if len(module_spec) == 0: continue
            try:
                self._metrics_checks.append(modules.load(module_spec, 'Check')(self.checksLogger))
                self.checksLogger.info("Registered custom check %s" % module_spec)
            except Exception, e:
                self.checksLogger.exception('Unable to load custom check module %s' % module_spec)

        # Event Checks
        self._event_checks = [
            ElasticSearchClusterStatus(self.checksLogger),
            Nagios(socket.gethostname())
        ]

        # Resource Checks
        self._resources_checks = [
            ResProcesses(self.checksLogger,self.agentConfig)
        ]

    def get_metadata(self):
        metadata = self._ec2.get_metadata()
        if metadata.get('hostname'):
            metadata['ec2-hostname'] = metadata.get('hostname')

        # if hostname is set in the configuration file
        # use that instead of gethostname
        # gethostname is vulnerable to 2 hosts: x.domain1, x.domain2
        # will cause both to be aliased (see #157)
        if self.agentConfig.get('hostname'):
            metadata['agent-hostname'] = self.agentConfig.get('hostname')
            metadata['hostname'] = metadata['agent-hostname']
        else:
            try:
                metadata["hostname"] = socket.gethostname()
            except:
                pass
        try:
            metadata["fqdn"] = socket.getfqdn()
        except:
            pass

        return metadata

    def should_send_metadata(self):
        # If the interval has passed, send the metadata again
        now = time.time()
        if now - self.metadata_start >= self.metadata_interval:
            self.checksLogger.debug('Metadata interval has passed. Sending metadata.')
            self.metadata_start = now
            return True

        return False

    def doChecks(self, firstRun=False, systemStats=False, checksd=None):
        """Actual work
        """
        self.checksLogger.info("Starting checks")
        checksData = {
            'collection_timestamp': time.time(),
            'os' : self.os,
            'python': sys.version,
            'agentVersion' : self.agentConfig['version'],
            'apiKey': self.agentConfig['api_key'],
            'events': {},
            'resources': {}
        }
        metrics = []
        events = {}

        # Run the system checks. Checks will depend on the OS
        if self.os == 'windows':
            # Win32 system checks
            metrics.extend(self._win32_system_checks['disk'].check(self.agentConfig))
            metrics.extend(self._win32_system_checks['memory'].check(self.agentConfig))
            metrics.extend(self._win32_system_checks['cpu'].check(self.agentConfig))
            metrics.extend(self._win32_system_checks['network'].check(self.agentConfig))
            metrics.extend(self._win32_system_checks['io'].check(self.agentConfig))
            metrics.extend(self._win32_system_checks['proc'].check(self.agentConfig))
        else:
            # Unix system checks
            sys_checks = self._unix_system_checks

            diskUsage = sys_checks['disk'].check(self.agentConfig)
            if diskUsage is not False and len(diskUsage) == 2:
                checksData["diskUsage"] = diskUsage[0]
                checksData["inodes"] = diskUsage[1]

            loadAvrgs = sys_checks['load'].check(self.agentConfig)
            checksData.update({
                'loadAvrg1': loadAvrgs['1'],
                'loadAvrg5': loadAvrgs['5'],
                'loadAvrg15': loadAvrgs['15']
            })

            memory = sys_checks['memory'].check(self.agentConfig)
            checksData.update({
                'memPhysUsed' : memory.get('physUsed'),
                'memPhysFree' : memory.get('physFree'),
                'memPhysTotal' : memory.get('physTotal'),
                'memPhysUsable' : memory.get('physUsable'),
                'memSwapUsed' : memory.get('swapUsed'),
                'memSwapFree' : memory.get('swapFree'),
                'memSwapTotal' : memory.get('swapTotal'),
                'memCached' : memory.get('physCached'),
                'memBuffers': memory.get('physBuffers'),
                'memShared': memory.get('physShared')
            })

            ioStats = sys_checks['io'].check(self.checksLogger, self.agentConfig)
            if ioStats:
                checksData['ioStats'] = ioStats

            processes = sys_checks['processes'].check(self.checksLogger, self.agentConfig)
            checksData.update({'processes': processes})

            networkTraffic = sys_checks['network'].check(self.agentConfig)
            checksData.update({'networkTraffic': networkTraffic})

            cpuStats = sys_checks['cpu'].check(self.checksLogger, self.agentConfig)
            if cpuStats is not False and cpuStats is not None:
                checksData.update(cpuStats)

        # Run old-style checks
        apacheStatus = self._apache.check(self.agentConfig)
        mysqlStatus = self._mysql.check(self.agentConfig)
        rabbitmq = self._rabbitmq.check(self.checksLogger, self.agentConfig)
        mongodb = self._mongodb.check(self.agentConfig)
        couchdb = self._couchdb.check(self.agentConfig)
        gangliaData = self._ganglia.check(self.agentConfig)
        cassandraData = self._cassandra.check(self.checksLogger, self.agentConfig)
        dogstreamData = self._dogstream.check(self.agentConfig)
        ddforwarderData = self._ddforwarder.check(self.agentConfig)

        if gangliaData is not False and gangliaData is not None:
            checksData['ganglia'] = gangliaData

        if cassandraData is not False and cassandraData is not None:
            checksData['cassandra'] = cassandraData

        # Apache Status
        if apacheStatus:
            checksData.update(apacheStatus)

        # MySQL Status
        if mysqlStatus:
            checksData.update(mysqlStatus)

        # RabbitMQ
        if rabbitmq:
            checksData['rabbitMQ'] = rabbitmq

        # MongoDB
        if mongodb:
            if mongodb.has_key('events'):
                events['Mongo'] = mongodb['events']['Mongo']
                del mongodb['events']
            checksData['mongoDB'] = mongodb

        # CouchDB
        if couchdb:
            checksData['couchDB'] = couchdb

        if dogstreamData:
            dogstreamEvents = dogstreamData.get('dogstreamEvents', None)
            if dogstreamEvents:
                if 'dogstream' in checksData['events']:
                    events['dogstream'].extend(dogstreamEvents)
                else:
                    events['dogstream'] = dogstreamEvents
                del dogstreamData['dogstreamEvents']

            checksData.update(dogstreamData)

        if ddforwarderData:
            checksData['datadog'] = ddforwarderData

        # Include server indentifiers
        checksData['internalHostname'] = gethostname(self.agentConfig)
        checksData['uuid'] = getUuid()
        self.checksLogger.debug('doChecks: added uuid %s' % checksData['uuid'])

        # Process the event checks.
        for event_check in self._event_checks:
            event_data = event_check.check(self.checksLogger, self.agentConfig)
            if event_data:
                events[event_check.key] = event_data

        # Include system stats on first postback
        if firstRun:
            checksData['systemStats'] = systemStats
            # Also post an event in the newsfeed
            events['System'] = [{'api_key': self.agentConfig['api_key'],
                                               'host': checksData['internalHostname'],
                                               'timestamp': int(time.mktime(datetime.datetime.now().timetuple())),
                                               'event_type':'Agent Startup',
                                               'msg_text': 'Version %s' % get_version()
                                            }]

        if firstRun or self.should_send_metadata():
            # Collect metadata
            checksData['meta'] = self.get_metadata()
            # Add static tags from the configuration file
            if self.agentConfig['tags'] is not None:
                checksData['tags'] = self.agentConfig['tags']

        # Resources checks
        if self.os != 'windows':
            has_resource = False
            for resources_check in self._resources_checks:
                resources_check.check()
                snaps = resources_check.pop_snapshots()
                if snaps:
                    has_resource = True
                    res_value = { 'snaps': snaps,
                                  'format_version': resources_check.get_format_version() }
                    res_format = resources_check.describe_format_if_needed()
                    if res_format is not None:
                        res_value['format_description'] = res_format
                    checksData['resources'][resources_check.RESOURCE_KEY] = res_value

            if has_resource:
                checksData['resources']['meta'] = {
                            'api_key': self.agentConfig['api_key'],
                            'host': checksData['internalHostname'],
                        }

        for metrics_check in self._metrics_checks:
            res = metrics_check.check(self.agentConfig)
            if res:
                metrics.extend(res)

        # checks.d checks
        if checksd:
            for check in checksd:
                check_cls = check['class']
                for instance in check['instances']:
                    try:
                        # Run the check for each configuration
                        check_cls.check(instance)
                        metrics.extend(check_cls.get_metrics())
                        if check_cls.has_events():
                            if check['name'] not in events:
                                events[check['name']] = []
                            for ev in check_cls.get_events():
                                events[check['name']].append(ev)
                    except Exception:
                        self.checksLogger.exception("Check %s failed" % check_cls.name)

        # Store the metrics in the payload
        checksData['metrics'] = metrics

        # Store the events in the payload
        checksData['events'] = events

        # Send back data
        self.checksLogger.debug("checksData: %s" % checksData)
        for emitter in self.emitters:
            emitter(checksData, self.checksLogger, self.agentConfig)
        self.checksLogger.info("Checks done")
<|endoftext|>


<commit_before>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Exercise 25: Check Tic Tac Toe

This exercise is Part 2 of 4 of the Tic Tac Toe exercise series. The
other exercises are: Part 1, Part 3, and Part 4.

As you may have guessed, we are trying to build up to a full tic-tac-toe
board. However, this is significantly more than half an hour of coding,
so we’re doing it in pieces.

Today, we will simply focus on checking whether someone has WON a game
of Tic Tac Toe, not worrying about how the moves were made.

If a game of Tic Tac Toe is represented as a list of lists, like so:

    game = [[1, 2, 0],
            [2, 1, 0],
            [2, 1, 1]]

where a 0 means an empty square, a 1 means that player 1 put their token
in that space, and a 2 means that player 2 put their token in that space.

Your task this week: given a 3 by 3 list of lists that represents a Tic
Tac Toe game board, tell me whether anyone has won, and tell me which
player won, if any. A Tic Tac Toe win is 3 in a row - either in a row, a
column, or a diagonal. Don’t worry about the case where TWO people have
won - assume that in every board there will only be one winner.

Here are some more examples to work with:

winner_is_2 = [[2, 2, 0],
               [2, 1, 0],
               [2, 1, 1]]

winner_is_1 = [[1, 2, 0],
               [2, 1, 0],
               [2, 1, 1]]

winner_is_also_1 = [[0, 1, 0],
                    [2, 1, 0],
                    [2, 1, 1]]

no_winner = [[1, 2, 0],
             [2, 1, 0],
             [2, 1, 2]]

also_no_winner = [[1, 2, 0],
                  [2, 1, 0],
                  [2, 1, 0]]

winner_schema = [[0, 1, 2],
                 [3, 4, 5],
                 [6, 7, 8]]

"""

def wins_positions():
    return (((0, 0), (0, 1), (0, 2)), ((1, 0), (1, 1), (1, 2)),
            ((2, 0), (2, 1), (2, 2)), ((0, 0), (1, 0), (2, 0)),
            ((0, 1), (1, 1), (2, 1)), ((0, 2), (1, 2), (2, 2)),
            ((0, 0), (1, 1), (2, 2)), ((0, 2), (1, 1), (2, 0)))

def get_winner(players_positions):

    busy_position = 0

    for player in (1, 2):
        for positions in wins_positions():
            for position in positions:
                if position in players_positions.get(player):
                    busy_position += 1
                    if busy_position == 3:
                        return "Player %i wins." % player
            busy_position = 0
    else:
        return "This game hasn't winner."

def get_players_positions(game):
    positions = []
    players_positions = dict()

    for player in (1, 2):
        for num_row, row_value in enumerate(game):
            for num_col, col_value in enumerate(row_value):
                if col_value == player:
                    positions.append((num_row, num_col))
        players_positions[player] = positions[:]
        positions.clear()
    return players_positions

if __name__ == "__main__":

    games = [[[2, 2, 0],
              [2, 1, 0],
              [2, 1, 1]],
             [[1, 2, 0],
              [2, 1, 0],
              [2, 1, 1]],
             [[0, 1, 0],
              [2, 1, 0],
              [2, 1, 1]],
             [[1, 2, 0],
              [2, 1, 0],
              [2, 1, 2]]]

    for game in games:
        print(get_winner(get_players_positions(game)))

<commit_msg>Update check_tic_tac_toe.py<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Exercise 26: Check Tic Tac Toe

This exercise is Part 2 of 4 of the Tic Tac Toe exercise series. The
other exercises are: Part 1, Part 3, and Part 4.

As you may have guessed, we are trying to build up to a full tic-tac-toe
board. However, this is significantly more than half an hour of coding,
so we’re doing it in pieces.

Today, we will simply focus on checking whether someone has WON a game
of Tic Tac Toe, not worrying about how the moves were made.

If a game of Tic Tac Toe is represented as a list of lists, like so:

    game = [[1, 2, 0],
            [2, 1, 0],
            [2, 1, 1]]

where a 0 means an empty square, a 1 means that player 1 put their token
in that space, and a 2 means that player 2 put their token in that space.

Your task this week: given a 3 by 3 list of lists that represents a Tic
Tac Toe game board, tell me whether anyone has won, and tell me which
player won, if any. A Tic Tac Toe win is 3 in a row - either in a row, a
column, or a diagonal. Don’t worry about the case where TWO people have
won - assume that in every board there will only be one winner.

Here are some more examples to work with:

winner_is_2 = [[2, 2, 0],
               [2, 1, 0],
               [2, 1, 1]]

winner_is_1 = [[1, 2, 0],
               [2, 1, 0],
               [2, 1, 1]]

winner_is_also_1 = [[0, 1, 0],
                    [2, 1, 0],
                    [2, 1, 1]]

no_winner = [[1, 2, 0],
             [2, 1, 0],
             [2, 1, 2]]

also_no_winner = [[1, 2, 0],
                  [2, 1, 0],
                  [2, 1, 0]]

winner_schema = [[0, 1, 2],
                 [3, 4, 5],
                 [6, 7, 8]]

"""

def wins_positions():
    return (((0, 0), (0, 1), (0, 2)), ((1, 0), (1, 1), (1, 2)),
            ((2, 0), (2, 1), (2, 2)), ((0, 0), (1, 0), (2, 0)),
            ((0, 1), (1, 1), (2, 1)), ((0, 2), (1, 2), (2, 2)),
            ((0, 0), (1, 1), (2, 2)), ((0, 2), (1, 1), (2, 0)))

def get_winner(players_positions):

    busy_position = 0

    for player in (1, 2):
        for positions in wins_positions():
            for position in positions:
                if position in players_positions.get(player):
                    busy_position += 1
                    if busy_position == 3:
                        return "Player %i wins." % player
            busy_position = 0
    else:
        return "This game hasn't winner."

def get_players_positions(game):
    positions = []
    players_positions = dict()

    for player in (1, 2):
        for num_row, row_value in enumerate(game):
            for num_col, col_value in enumerate(row_value):
                if col_value == player:
                    positions.append((num_row, num_col))
        players_positions[player] = positions[:]
        positions.clear()
    return players_positions

if __name__ == "__main__":

    games = [[[2, 2, 0],
              [2, 1, 0],
              [2, 1, 1]],
             [[1, 2, 0],
              [2, 1, 0],
              [2, 1, 1]],
             [[0, 1, 0],
              [2, 1, 0],
              [2, 1, 1]],
             [[1, 2, 0],
              [2, 1, 0],
              [2, 1, 2]]]

    for game in games:
        print(get_winner(get_players_positions(game)))

<|endoftext|>


<commit_before>"""Add the effects of instrumental noise into the simulation.

This is separated out into two tasks. The first, :class:`ReceiverTemperature`
adds in the effects of instrumental noise bias into the data. The second,
:class:`SampleNoise`, takes a timestream which is assumed to be the expected (or
average) value and returns an observed time stream.

Tasks
=====

.. autosummary::
    :toctree:

    ReceiverTemperature
    SampleNoise
"""
# === Start Python 2/3 compatibility
from __future__ import absolute_import, division, print_function, unicode_literals
from future.builtins import *  # noqa  pylint: disable=W0401, W0614
from future.builtins.disabled import *  # noqa  pylint: disable=W0401, W0614

# === End Python 2/3 compatibility

import contextlib

import numpy as np

from caput import config
from caput.time import STELLAR_S
from cora.util import nputil

from ..core import task, containers, io
from ..util import tools

class ReceiverTemperature(task.SingleTask):
    """Add a basic receiver temperature term into the data.

    This class adds in an uncorrelated, frequency and time independent receiver
    noise temperature to the data. As it is uncorrelated this will only affect
    the auto-correlations. Note this only adds in the offset to the visibility,
    to add the corresponding random fluctuations to subsequently use the
    :class:`SampleNoise` task.

    Attributes
    ----------
    recv_temp : float
        The receiver temperature in Kelvin.
    """

    recv_temp = config.Property(proptype=float, default=0.0)

    def process(self, data):

        # Iterate over the products to find the auto-correlations and add the noise into them
        for pi, prod in enumerate(data.prodstack):

            # Great an auto!
            if prod[0] == prod[1]:
                data.vis[:, pi] += self.recv_temp

        return data

class GaussianNoise(task.SingleTask):
    """Add Gaussian distributed noise to a visibility dataset.

    Note that this is an approximation to the actual noise distribution good only
    when T_recv >> T_sky and delta_time * delta_freq >> 1.

    Attributes
    ----------
    ndays : float
        Multiplies the number of samples in each measurement.
    seed : int
        Random seed for the noise generation.
    set_weights : bool
        Set the weights to the appropriate values.
    recv_temp : bool
        The temperature of the noise to add.
    """

    recv_temp = config.Property(proptype=float, default=50.0)
    ndays = config.Property(proptype=float, default=733.0)
    seed = config.Property(proptype=int, default=None)
    set_weights = config.Property(proptype=bool, default=True)

    def setup(self, manager=None):
        """Set the telescope instance if a manager object is given.

        This is used to simulate noise for visibilities that are stacked
        over redundant baselines.

        Parameters
        ----------
        manager : manager.ProductManager, optional
            The telescope/manager used to set the `redundancy`. If not set,
            `redundancy` is derived from the data.
        """
        if manager is not None:
            self.telescope = io.get_telescope(manager)
        else:
            self.telescope = None

    def process(self, data):
        """Generate a noisy dataset.

        Parameters
        ----------
        data : :class:`containers.SiderealStream` or :class:`containers.TimeStream`
            The expected (i.e. noiseless) visibility dataset. Note the modification
            is done in place.

        Returns
        -------
        data_noise : same as :param:`data`
            The sampled (i.e. noisy) visibility dataset.
        """

        data.redistribute("freq")

        visdata = data.vis[:]

        # Get the time interval
        if isinstance(data, containers.SiderealStream):
            dt = 240 * (data.ra[1] - data.ra[0]) * STELLAR_S
            ntime = len(data.ra)
        else:
            dt = data.time[1] - data.time[0]
            ntime = len(data.time)

        # TODO: this assumes uniform channels
        df = data.index_map["freq"]["width"][0] * 1e6
        nfreq = data.vis.local_shape[0]
        nprod = len(data.index_map["prod"])
        ninput = len(data.index_map["input"])

        # Consider if this data is stacked over redundant baselines or not.
        if (self.telescope is not None) and (nprod == self.telescope.nbase):
            redundancy = self.telescope.redundancy
        elif nprod == nfeed * (nfeed + 1) / 2:
            redundancy = np.ones(nprod)
        else:
            raise ValueError("Unexpected number of products")

        # Calculate the number of samples, this is a 1D array for the prod axis.
        nsamp = int(self.ndays * dt * df) * redundancy
        std = self.recv_temp / np.sqrt(nsamp)

        with mpi_random_seed(self.seed):
            noise = std[np.newaxis, :, np.newaxis] * nputil.complex_std_normal(
                (nfreq, nprod, ntime)
            )

        # Iterate over the products to find the auto-correlations and add the noise into them
        for pi, prod in enumerate(data.index_map["prod"]):

            # Auto: multiply by sqrt(2) because auto has twice the variance
            if prod[0] == prod[1]:
                visdata[:, pi].real += np.sqrt(2) * noise[:, pi].real

            else:
                visdata[:, pi] += noise[:, pi]

        # Construct and set the correct weights in place
        if self.set_weights:
            for lfi, fi in visdata.enumerate(0):
                data.weight[fi] = 1.0 / std[:, np.newaxis] ** 2

        return data

class SampleNoise(task.SingleTask):
    """Add properly distributed noise to a visibility dataset.

    This task draws properly (complex Wishart) distributed samples from an input
    visibility dataset which is assumed to represent the expectation.

    See http://link.springer.com/article/10.1007%2Fs10440-010-9599-x for a
    discussion of the Bartlett decomposition for complex Wishart distributed
    quantities.

    Attributes
    ----------
    sample_frac : float
        Multiplies the number of samples in each measurement. For instance this
        could be a duty cycle if the correlator was not keeping up, or could be
        larger than one if multiple measurements were combined.
    seed : int
        Random seed for the noise generation.
    set_weights : bool
        Set the weights to the appropriate values.
    """

    sample_frac = config.Property(proptype=float, default=1.0)
    seed = config.Property(proptype=int, default=None)
    set_weights = config.Property(proptype=bool, default=True)

    def process(self, data_exp):
        """Generate a noisy dataset.

        Parameters
        ----------
        data_exp : :class:`containers.SiderealStream` or :class:`containers.TimeStream`
            The expected (i.e. noiseless) visibility dataset. Must be the full
            triangle. Make sure you have added an instrumental noise bias if you
            want instrumental noise.

        Returns
        -------
        data_samp : same as :param:`data_exp`
            The sampled (i.e. noisy) visibility dataset.
        """

        from caput.time import STELLAR_S
        from ..util import _fast_tools

        data_exp.redistribute("freq")

        nfeed = len(data_exp.index_map["input"])

        # Get a reference to the base MPIArray. Attempting to do this in the
        # loop fails if not all ranks enter the loop (as there is an implied MPI
        # Barrier)
        vis_data = data_exp.vis[:]

        # Get the time interval
        if isinstance(data_exp, containers.SiderealStream):
            dt = 240 * (data_exp.ra[1] - data_exp.ra[0]) * STELLAR_S
        else:
            dt = data_exp.time[1] - data_exp.time[0]

        with mpi_random_seed(self.seed):

            # Iterate over frequencies
            for lfi, fi in vis_data.enumerate(0):

                # Get the frequency interval
                df = data_exp.index_map["freq"]["width"][fi] * 1e6

                # Calculate the number of samples
                nsamp = int(self.sample_frac * dt * df)

                # Iterate over time
                for lti, ti in vis_data.enumerate(2):

                    # Unpack visibilites into full matrix
                    vis_utv = vis_data[lfi, :, lti].view(np.ndarray).copy()
                    vis_mat = np.zeros((nfeed, nfeed), dtype=vis_utv.dtype)
                    _fast_tools._unpack_product_array_fast(
                        vis_utv, vis_mat, np.arange(nfeed), nfeed
                    )

                    vis_samp = draw_complex_wishart(vis_mat, nsamp) / nsamp

                    vis_data[lfi, :, lti] = vis_samp[np.triu_indices(nfeed)]

                # Construct and set the correct weights in place
                if self.set_weights:
                    autos = tools.extract_diagonal(vis_data[lfi], axis=0).real
                    weight_fac = nsamp ** 0.5 / autos
                    tools.apply_gain(
                        data_exp.weight[fi][np.newaxis, ...],
                        weight_fac[np.newaxis, ...],
                        out=data_exp.weight[fi][np.newaxis, ...],
                    )

        return data_exp

def standard_complex_wishart(m, n):
    """Draw a standard Wishart matrix.

    Parameters
    ----------
    m : integer
        Number of variables (i.e. size of matrix).
    n : integer
        Number of measurements the covariance matrix is estimated from.

    Returns
    -------
    B : np.ndarray[m, m]
    """

    from scipy.stats import gamma

    # Fill in normal variables in the lower triangle
    T = np.zeros((m, m), dtype=np.complex128)
    T[np.tril_indices(m, k=-1)] = (
        np.random.standard_normal(m * (m - 1) // 2)
        + 1.0j * np.random.standard_normal(m * (m - 1) // 2)
    ) / 2 ** 0.5

    # Gamma variables on the diagonal
    for i in range(m):
        T[i, i] = gamma.rvs(n - i) ** 0.5

    # Return the square to get the Wishart matrix
    return np.dot(T, T.T.conj())

def draw_complex_wishart(C, n):
    """Draw a complex Wishart matrix.

    Parameters
    ----------
    C_exp : np.ndarray[:, :]
        Expected covaraince matrix.

    n : integer
        Number of measurements the covariance matrix is estimated from.

    Returns
    -------
    C_samp : np.ndarray
        Sample covariance matrix.
    """

    import scipy.linalg as la

    # Find Cholesky of C
    L = la.cholesky(C, lower=True)

    # Generate a standard Wishart
    A = standard_complex_wishart(C.shape[0], n)

    # Transform to get the Wishart variable
    return np.dot(L, np.dot(A, L.T.conj()))

@contextlib.contextmanager
def mpi_random_seed(seed, extra=0):
    """Use a specific random seed for the context, and return to the original state on exit.

    This is designed to work for MPI computations, incrementing the actual seed
    of each process by the MPI rank. Overall each process gets the numpy seed:
    `numpy_seed = seed + mpi_rank + 4096 * extra`.

    Parameters
    ----------
    seed : int
        Base seed to set. If seed is :obj:`None`, re-seed randomly.
    extra : int, optional
        An extra part of the seed, which should be changed for calculations
        using the same seed, but that want different random sequences.
    """

    import numpy as np
    from caput import mpiutil

    # Copy the old state for restoration later.
    old_state = np.random.get_state()

    # Just choose a random number per process as the seed if nothing was set.
    if seed is None:
        seed = np.random.randint(2 ** 30)

    # Construct the new process specific seed
    new_seed = seed + mpiutil.rank + 4096 * extra
    np.random.seed(new_seed)

    # Enter the context block, and reset the state on exit.
    try:
        yield
    finally:
        np.random.set_state(old_state)
<commit_msg>fix(GaussianNoise): fix issue when adding noise to stacked data<commit_after>"""Add the effects of instrumental noise into the simulation.

This is separated out into two tasks. The first, :class:`ReceiverTemperature`
adds in the effects of instrumental noise bias into the data. The second,
:class:`SampleNoise`, takes a timestream which is assumed to be the expected (or
average) value and returns an observed time stream.

Tasks
=====

.. autosummary::
    :toctree:

    ReceiverTemperature
    SampleNoise
"""
# === Start Python 2/3 compatibility
from __future__ import absolute_import, division, print_function, unicode_literals
from future.builtins import *  # noqa  pylint: disable=W0401, W0614
from future.builtins.disabled import *  # noqa  pylint: disable=W0401, W0614

# === End Python 2/3 compatibility

import contextlib

import numpy as np

from caput import config
from caput.time import STELLAR_S
from cora.util import nputil

from ..core import task, containers, io
from ..util import tools

class ReceiverTemperature(task.SingleTask):
    """Add a basic receiver temperature term into the data.

    This class adds in an uncorrelated, frequency and time independent receiver
    noise temperature to the data. As it is uncorrelated this will only affect
    the auto-correlations. Note this only adds in the offset to the visibility,
    to add the corresponding random fluctuations to subsequently use the
    :class:`SampleNoise` task.

    Attributes
    ----------
    recv_temp : float
        The receiver temperature in Kelvin.
    """

    recv_temp = config.Property(proptype=float, default=0.0)

    def process(self, data):

        # Iterate over the products to find the auto-correlations and add the noise into them
        for pi, prod in enumerate(data.prodstack):

            # Great an auto!
            if prod[0] == prod[1]:
                data.vis[:, pi] += self.recv_temp

        return data

class GaussianNoise(task.SingleTask):
    """Add Gaussian distributed noise to a visibility dataset.

    Note that this is an approximation to the actual noise distribution good only
    when T_recv >> T_sky and delta_time * delta_freq >> 1.

    Attributes
    ----------
    ndays : float
        Multiplies the number of samples in each measurement.
    seed : int
        Random seed for the noise generation.
    set_weights : bool
        Set the weights to the appropriate values.
    recv_temp : bool
        The temperature of the noise to add.
    """

    recv_temp = config.Property(proptype=float, default=50.0)
    ndays = config.Property(proptype=float, default=733.0)
    seed = config.Property(proptype=int, default=None)
    set_weights = config.Property(proptype=bool, default=True)

    def setup(self, manager=None):
        """Set the telescope instance if a manager object is given.

        This is used to simulate noise for visibilities that are stacked
        over redundant baselines.

        Parameters
        ----------
        manager : manager.ProductManager, optional
            The telescope/manager used to set the `redundancy`. If not set,
            `redundancy` is derived from the data.
        """
        if manager is not None:
            self.telescope = io.get_telescope(manager)
        else:
            self.telescope = None

    def process(self, data):
        """Generate a noisy dataset.

        Parameters
        ----------
        data : :class:`containers.SiderealStream` or :class:`containers.TimeStream`
            The expected (i.e. noiseless) visibility dataset. Note the modification
            is done in place.

        Returns
        -------
        data_noise : same as :param:`data`
            The sampled (i.e. noisy) visibility dataset.
        """

        data.redistribute("freq")

        visdata = data.vis[:]

        # Get the time interval
        if isinstance(data, containers.SiderealStream):
            dt = 240 * (data.ra[1] - data.ra[0]) * STELLAR_S
            ntime = len(data.ra)
        else:
            dt = data.time[1] - data.time[0]
            ntime = len(data.time)

        # TODO: this assumes uniform channels
        df = data.index_map["freq"]["width"][0] * 1e6
        nfreq = data.vis.local_shape[0]
        nprod = len(data.prodstack)
        ninput = len(data.index_map["input"])

        # Consider if this data is stacked over redundant baselines or not.
        if (self.telescope is not None) and (nprod == self.telescope.nbase):
            redundancy = self.telescope.redundancy
        elif nprod == ninput * (ninput + 1) / 2:
            redundancy = np.ones(nprod)
        else:
            raise ValueError("Unexpected number of products")

        # Calculate the number of samples, this is a 1D array for the prod axis.
        nsamp = int(self.ndays * dt * df) * redundancy
        std = self.recv_temp / np.sqrt(nsamp)

        with mpi_random_seed(self.seed):
            noise = std[np.newaxis, :, np.newaxis] * nputil.complex_std_normal(
                (nfreq, nprod, ntime)
            )

        # Iterate over the products to find the auto-correlations and add the noise into them
        for pi, prod in enumerate(data.prodstack):

            # Auto: multiply by sqrt(2) because auto has twice the variance
            if prod[0] == prod[1]:
                visdata[:, pi].real += np.sqrt(2) * noise[:, pi].real

            else:
                visdata[:, pi] += noise[:, pi]

        # Construct and set the correct weights in place
        if self.set_weights:
            for lfi, fi in visdata.enumerate(0):
                data.weight[fi] = 1.0 / std[:, np.newaxis] ** 2

        return data

class SampleNoise(task.SingleTask):
    """Add properly distributed noise to a visibility dataset.

    This task draws properly (complex Wishart) distributed samples from an input
    visibility dataset which is assumed to represent the expectation.

    See http://link.springer.com/article/10.1007%2Fs10440-010-9599-x for a
    discussion of the Bartlett decomposition for complex Wishart distributed
    quantities.

    Attributes
    ----------
    sample_frac : float
        Multiplies the number of samples in each measurement. For instance this
        could be a duty cycle if the correlator was not keeping up, or could be
        larger than one if multiple measurements were combined.
    seed : int
        Random seed for the noise generation.
    set_weights : bool
        Set the weights to the appropriate values.
    """

    sample_frac = config.Property(proptype=float, default=1.0)
    seed = config.Property(proptype=int, default=None)
    set_weights = config.Property(proptype=bool, default=True)

    def process(self, data_exp):
        """Generate a noisy dataset.

        Parameters
        ----------
        data_exp : :class:`containers.SiderealStream` or :class:`containers.TimeStream`
            The expected (i.e. noiseless) visibility dataset. Must be the full
            triangle. Make sure you have added an instrumental noise bias if you
            want instrumental noise.

        Returns
        -------
        data_samp : same as :param:`data_exp`
            The sampled (i.e. noisy) visibility dataset.
        """

        from caput.time import STELLAR_S
        from ..util import _fast_tools

        data_exp.redistribute("freq")

        nfeed = len(data_exp.index_map["input"])

        # Get a reference to the base MPIArray. Attempting to do this in the
        # loop fails if not all ranks enter the loop (as there is an implied MPI
        # Barrier)
        vis_data = data_exp.vis[:]

        # Get the time interval
        if isinstance(data_exp, containers.SiderealStream):
            dt = 240 * (data_exp.ra[1] - data_exp.ra[0]) * STELLAR_S
        else:
            dt = data_exp.time[1] - data_exp.time[0]

        with mpi_random_seed(self.seed):

            # Iterate over frequencies
            for lfi, fi in vis_data.enumerate(0):

                # Get the frequency interval
                df = data_exp.index_map["freq"]["width"][fi] * 1e6

                # Calculate the number of samples
                nsamp = int(self.sample_frac * dt * df)

                # Iterate over time
                for lti, ti in vis_data.enumerate(2):

                    # Unpack visibilites into full matrix
                    vis_utv = vis_data[lfi, :, lti].view(np.ndarray).copy()
                    vis_mat = np.zeros((nfeed, nfeed), dtype=vis_utv.dtype)
                    _fast_tools._unpack_product_array_fast(
                        vis_utv, vis_mat, np.arange(nfeed), nfeed
                    )

                    vis_samp = draw_complex_wishart(vis_mat, nsamp) / nsamp

                    vis_data[lfi, :, lti] = vis_samp[np.triu_indices(nfeed)]

                # Construct and set the correct weights in place
                if self.set_weights:
                    autos = tools.extract_diagonal(vis_data[lfi], axis=0).real
                    weight_fac = nsamp ** 0.5 / autos
                    tools.apply_gain(
                        data_exp.weight[fi][np.newaxis, ...],
                        weight_fac[np.newaxis, ...],
                        out=data_exp.weight[fi][np.newaxis, ...],
                    )

        return data_exp

def standard_complex_wishart(m, n):
    """Draw a standard Wishart matrix.

    Parameters
    ----------
    m : integer
        Number of variables (i.e. size of matrix).
    n : integer
        Number of measurements the covariance matrix is estimated from.

    Returns
    -------
    B : np.ndarray[m, m]
    """

    from scipy.stats import gamma

    # Fill in normal variables in the lower triangle
    T = np.zeros((m, m), dtype=np.complex128)
    T[np.tril_indices(m, k=-1)] = (
        np.random.standard_normal(m * (m - 1) // 2)
        + 1.0j * np.random.standard_normal(m * (m - 1) // 2)
    ) / 2 ** 0.5

    # Gamma variables on the diagonal
    for i in range(m):
        T[i, i] = gamma.rvs(n - i) ** 0.5

    # Return the square to get the Wishart matrix
    return np.dot(T, T.T.conj())

def draw_complex_wishart(C, n):
    """Draw a complex Wishart matrix.

    Parameters
    ----------
    C_exp : np.ndarray[:, :]
        Expected covaraince matrix.

    n : integer
        Number of measurements the covariance matrix is estimated from.

    Returns
    -------
    C_samp : np.ndarray
        Sample covariance matrix.
    """

    import scipy.linalg as la

    # Find Cholesky of C
    L = la.cholesky(C, lower=True)

    # Generate a standard Wishart
    A = standard_complex_wishart(C.shape[0], n)

    # Transform to get the Wishart variable
    return np.dot(L, np.dot(A, L.T.conj()))

@contextlib.contextmanager
def mpi_random_seed(seed, extra=0):
    """Use a specific random seed for the context, and return to the original state on exit.

    This is designed to work for MPI computations, incrementing the actual seed
    of each process by the MPI rank. Overall each process gets the numpy seed:
    `numpy_seed = seed + mpi_rank + 4096 * extra`.

    Parameters
    ----------
    seed : int
        Base seed to set. If seed is :obj:`None`, re-seed randomly.
    extra : int, optional
        An extra part of the seed, which should be changed for calculations
        using the same seed, but that want different random sequences.
    """

    import numpy as np
    from caput import mpiutil

    # Copy the old state for restoration later.
    old_state = np.random.get_state()

    # Just choose a random number per process as the seed if nothing was set.
    if seed is None:
        seed = np.random.randint(2 ** 30)

    # Construct the new process specific seed
    new_seed = seed + mpiutil.rank + 4096 * extra
    np.random.seed(new_seed)

    # Enter the context block, and reset the state on exit.
    try:
        yield
    finally:
        np.random.set_state(old_state)
<|endoftext|>


<commit_before>'''
Created on Feb 7, 2012

@author: Kristof
'''
from clawpack.pyclaw.solver import Solver
from peanoclaw.callbacks.initializationcallback import InitializationCallback
from peanoclaw.callbacks.solvercallback import SolverCallback
from peanoclaw.callbacks.refinementcriterioncallback import RefinementCriterionCallback
from peanoclaw.callbacks.boundaryconditioncallback import BoundaryConditionCallback
from peanoclaw.callbacks.interpolationcallback import InterpolationCallback
from peanoclaw.callbacks.restrictioncallback import RestrictionCallback
from peanoclaw.callbacks.fluxcorrectioncallback import FluxCorrectionCallback
from peanoclaw.peano import Peano
from peanoclaw.internalsettings import InternalSettings

class Solver(Solver):
    r"""
        This solver class wraps the AMR functionality of Peano. It holds a normal PyClaw-solver
        to advance separate subgrids in time. Therefore it provides the callbacks callback_solver(...)
        and callback_boundary_conditions(...) as an interface for Peano to use the PyClaw-solver.

        A Solver is typically instantiated as follows::

        >>> import pyclaw
        >>> solver = pyclaw.ClawSolver2D()
        >>> import peanoclaw
        >>> peanoclaw_solver = peanoclaw.Solver(solver, 1.0/18.0)
    """

    def __init__(
                 self,
                 solver,
                 initial_minimal_mesh_width,
                 q_initialization,
                 aux_initialization=None,
                 refinement_criterion=None,
                 interpolation=None,
                 restriction=None,
                 flux_correction=None,
                 internal_settings=InternalSettings()
                 ):
        r"""
        Initializes the Peano-solver. This keeps the Peano-spacetree internally and wraps the given PyClaw-solver.

        :Input:
         -  *solver* - (:class:`pyclaw.Solver`) The PyClaw-solver used internally.
         -  *initial_minimal_mesh_width* - The initial mesh width for the Peano mesh. I.e. Peano refines the mesh regularly
                                             until it is at least as fine as stated in this parameter.
        """
        #Initialize PeanoClaw solver
        self.solver = solver
        self.initial_minimal_mesh_width = initial_minimal_mesh_width
        self.q_initialization = q_initialization
        self.aux_initialization = aux_initialization
        self.refinement_criterion = refinement_criterion
        self.dt_initial = solver.dt_initial
        self.num_ghost = solver.num_ghost
        self.rp = solver.rp
        self.did_run_as_worker = False

        #Create callbacks
        self.initialization_callback = InitializationCallback(self, refinement_criterion, q_initialization, aux_initialization, initial_minimal_mesh_width)
        self.solver_callback = SolverCallback(self, refinement_criterion, initial_minimal_mesh_width, internal_settings.fixed_timestep_size)
        self.refinement_criterion_callback = RefinementCriterionCallback(self, refinement_criterion, initial_minimal_mesh_width)
        self.boundary_condition_callback = BoundaryConditionCallback(self)
        self.interpolation_callback = InterpolationCallback(interpolation, self)
        self.restriction_callback = RestrictionCallback(restriction, self)
        self.flux_correction_callback = FluxCorrectionCallback(flux_correction, self)

        self.internal_settings = internal_settings

        self._is_set_up = False

    def _del_(self):
      self.teardown()

    def setup(self, solution):
        r"""
        Initialize a Solver object. This method loads the library of Peano and prepares the initial mesh.

        See :class:`Solver` for full documentation
        """
        self.bc_lower = self.solver.bc_lower[:]
        self.bc_upper = self.solver.bc_upper[:]
        self.user_bc_lower = self.solver.user_bc_lower
        self.user_bc_upper = self.solver.user_bc_upper

        self.solution = solution

        self.peano = self.setup_peano(solution)

        self._is_set_up = True

    def run_tests(self, solution):
      self.peano = self.setup_peano(solution)
      self.peano.run_tests()

    def setup_peano(self, solution):
      return Peano(
         solution,
         self.initial_minimal_mesh_width,
         self.num_ghost,
         self.solver.dt_initial,
         self.initialization_callback,
         self.solver_callback,
         self.refinement_criterion_callback,
         self.boundary_condition_callback,
         self.interpolation_callback,
         self.restriction_callback,
         self.flux_correction_callback,
         self.internal_settings
         )

    def teardown(self):
        r"""
        See :class:`Solver` for full documentation
        """
        self.peano.teardown()
        #self.solver.teardown()

    def evolve_to_time(self, solution, tend=None):
        r"""
        Performs one global timestep until all patches in the mesh reach the given end time.

        See :class:`Solver` for full documentation
        """
        if(tend == None) :
            raise Exception("Not yet implemented.")

        self.solution = solution
        if self.peano.rank == 0:
          self.peano.evolve_to_time(tend)
        elif not self.did_run_as_worker:
          self.peano.runWorker()
          self.did_run_as_worker = True

    def solve_one_timestep(self, q, qbc):
        r"""
        """
        self.solver.step(self.solution)

<commit_msg>fixed segfault on exit of pyclaw solver simulation<commit_after>'''
Created on Feb 7, 2012

@author: Kristof
'''
from clawpack.pyclaw.solver import Solver
from peanoclaw.callbacks.initializationcallback import InitializationCallback
from peanoclaw.callbacks.solvercallback import SolverCallback
from peanoclaw.callbacks.refinementcriterioncallback import RefinementCriterionCallback
from peanoclaw.callbacks.boundaryconditioncallback import BoundaryConditionCallback
from peanoclaw.callbacks.interpolationcallback import InterpolationCallback
from peanoclaw.callbacks.restrictioncallback import RestrictionCallback
from peanoclaw.callbacks.fluxcorrectioncallback import FluxCorrectionCallback
from peanoclaw.peano import Peano
from peanoclaw.internalsettings import InternalSettings

class Solver(Solver):
    r"""
        This solver class wraps the AMR functionality of Peano. It holds a normal PyClaw-solver
        to advance separate subgrids in time. Therefore it provides the callbacks callback_solver(...)
        and callback_boundary_conditions(...) as an interface for Peano to use the PyClaw-solver.

        A Solver is typically instantiated as follows::

        >>> import pyclaw
        >>> solver = pyclaw.ClawSolver2D()
        >>> import peanoclaw
        >>> peanoclaw_solver = peanoclaw.Solver(solver, 1.0/18.0)
    """

    def __init__(
                 self,
                 solver,
                 initial_minimal_mesh_width,
                 q_initialization,
                 aux_initialization=None,
                 refinement_criterion=None,
                 interpolation=None,
                 restriction=None,
                 flux_correction=None,
                 internal_settings=InternalSettings()
                 ):
        r"""
        Initializes the Peano-solver. This keeps the Peano-spacetree internally and wraps the given PyClaw-solver.

        :Input:
         -  *solver* - (:class:`pyclaw.Solver`) The PyClaw-solver used internally.
         -  *initial_minimal_mesh_width* - The initial mesh width for the Peano mesh. I.e. Peano refines the mesh regularly
                                             until it is at least as fine as stated in this parameter.
        """
        #Initialize PeanoClaw solver
        self.solver = solver
        self.initial_minimal_mesh_width = initial_minimal_mesh_width
        self.q_initialization = q_initialization
        self.aux_initialization = aux_initialization
        self.refinement_criterion = refinement_criterion
        self.dt_initial = solver.dt_initial
        self.num_ghost = solver.num_ghost
        self.rp = solver.rp
        self.did_run_as_worker = False

        #Create callbacks
        self.initialization_callback = InitializationCallback(self, refinement_criterion, q_initialization, aux_initialization, initial_minimal_mesh_width)
        self.solver_callback = SolverCallback(self, refinement_criterion, initial_minimal_mesh_width, internal_settings.fixed_timestep_size)
        self.refinement_criterion_callback = RefinementCriterionCallback(self, refinement_criterion, initial_minimal_mesh_width)
        self.boundary_condition_callback = BoundaryConditionCallback(self)
        self.interpolation_callback = InterpolationCallback(interpolation, self)
        self.restriction_callback = RestrictionCallback(restriction, self)
        self.flux_correction_callback = FluxCorrectionCallback(flux_correction, self)

        self.internal_settings = internal_settings

        self._is_set_up = False

    def _del_(self):
      self.teardown()

    def setup(self, solution):
        r"""
        Initialize a Solver object. This method loads the library of Peano and prepares the initial mesh.

        See :class:`Solver` for full documentation
        """
        self.bc_lower = self.solver.bc_lower[:]
        self.bc_upper = self.solver.bc_upper[:]
        self.user_bc_lower = self.solver.user_bc_lower
        self.user_bc_upper = self.solver.user_bc_upper

        self.solution = solution

        self.peano = self.setup_peano(solution)

        self._is_set_up = True

    def run_tests(self, solution):
      self.peano = self.setup_peano(solution)
      self.peano.run_tests()

    def setup_peano(self, solution):
      return Peano(
         solution,
         self.initial_minimal_mesh_width,
         self.num_ghost,
         self.solver.dt_initial,
         self.initialization_callback,
         self.solver_callback,
         self.refinement_criterion_callback,
         self.boundary_condition_callback,
         self.interpolation_callback,
         self.restriction_callback,
         self.flux_correction_callback,
         self.internal_settings
         )

    def teardown(self):
        r"""
        See :class:`Solver` for full documentation
        """
        if self._is_set_up:
          self.peano.teardown()
          self._is_set_up = False

    def evolve_to_time(self, solution, tend=None):
        r"""
        Performs one global timestep until all patches in the mesh reach the given end time.

        See :class:`Solver` for full documentation
        """
        if(tend == None) :
            raise Exception("Not yet implemented.")

        self.solution = solution
        if self.peano.rank == 0:
          self.peano.evolve_to_time(tend)
        elif not self.did_run_as_worker:
          self.peano.runWorker()
          self.did_run_as_worker = True

    def solve_one_timestep(self, q, qbc):
        r"""
        """
        self.solver.step(self.solution)

<|endoftext|>


<commit_before>"""EayunOS Version plugin."""

import os

from otopi import plugin, util

from ovirt_engine_setup import constants as osetupcons
from ovirt_engine_setup.engine import constants as oenginecons
from ovirt_engine_setup.engine_common import constants as oengcommcons

@util.export
class Plugin(plugin.PluginBase):
    """EayunOS Version plugin."""

    def __init__(self, context):
        super(Plugin, self).__init__(context=context)

    @plugin.event(
        stage=plugin.Stages.STAGE_MISC,
        before=(
            oengcommcons.Stages.DB_SCHEMA,
        ),
    )
    def _customization(self):
        version = self.environment.get(
            oenginecons.ConfigEnv.EAYUNOS_VERSION
        )
        if version == 'Enterprise':
            self.enterprise_version_setup()
            system_operations()
            self.dialog.note(text="EayunOS version: Enterprise")

    def enterprise_version_setup(self):
        # update ovirt-engine files
        os.system("sed -i 's/4\.2 Basic/4\.2 Enterprise/' /usr/share/ovirt-engine/branding/ovirt.brand/messages.properties")
        os.system("sed -i 's/\\\u57FA\\\u7840\\\u7248/\\\u4f01\\\u4e1a\\\u7248/g' /usr/share/ovirt-engine/branding/ovirt.brand/messages_zh_CN.properties")
        os.system("sed -i 's/EayunOS_top_logo_basic\.png/EayunOS_top_logo_enterprise\.png/' /usr/share/ovirt-engine/branding/ovirt.brand/common.css")
        os.system("sed -i '/EayunOSVersion/d' /usr/share/ovirt-engine/dbscripts/upgrade/pre_upgrade/0000_config.sql")
        os.system("echo \"select fn_db_add_config_value('EayunOSVersion','Enterprise','general');\" >> /usr/share/ovirt-engine/dbscripts/upgrade/pre_upgrade/0000_config.sql")
        os.system("echo \"select fn_db_update_config_value('EayunOSVersion','Enterprise','general');\" >> /usr/share/ovirt-engine/dbscripts/upgrade/pre_upgrade/0000_config.sql")
        # make product uuid readable
        os.system("echo \"#! /bin/bash\" > /etc/init.d/systemuuid")
        os.system("echo \"# chkconfig: 2345 10 90\" >> /etc/init.d/systemuuid")
        os.system("echo \"chmod a+r /sys/class/dmi/id/product_uuid\" >> /etc/init.d/systemuuid")
        os.system("chmod a+x /etc/init.d/systemuuid")
        os.system("chkconfig systemuuid on")
        os.system("chmod a+r /sys/class/dmi/id/product_uuid")

    @plugin.event(
        stage=plugin.Stages.STAGE_MISC,
        before=(
            oengcommcons.Stages.DB_SCHEMA,
        ),
        condition=lambda self: (
            not self.environment[oenginecons.EngineDBEnv.NEW_DATABASE]
        ),
    )
    def _setup_installation_time(self):
        os.system("sed -i '/InstallationTime/d' /usr/share/ovirt-engine/dbscripts/upgrade/pre_upgrade/0000_config.sql")
        os.system("echo \"select fn_db_add_config_value('InstallationTime',to_char(current_timestamp,'yyyy-MM-dd HH24:mm:ss'),'general');\" >> /usr/share/ovirt-engine/dbscripts/upgrade/pre_upgrade/0000_config.sql")

<commit_msg>Fix python error of eayunos version setup<commit_after>"""EayunOS Version plugin."""

import os

from otopi import plugin, util

from ovirt_engine_setup import constants as osetupcons
from ovirt_engine_setup.engine import constants as oenginecons
from ovirt_engine_setup.engine_common import constants as oengcommcons

@util.export
class Plugin(plugin.PluginBase):
    """EayunOS Version plugin."""

    def __init__(self, context):
        super(Plugin, self).__init__(context=context)

    @plugin.event(
        stage=plugin.Stages.STAGE_MISC,
        before=(
            oengcommcons.Stages.DB_SCHEMA,
        ),
    )
    def _customization(self):
        version = self.environment.get(
            oenginecons.ConfigEnv.EAYUNOS_VERSION
        )
        if version == 'Enterprise':
            self.enterprise_version_setup()
            self.dialog.note(text="EayunOS version: Enterprise")

    def enterprise_version_setup(self):
        # update ovirt-engine files
        os.system("sed -i 's/4\.2 Basic/4\.2 Enterprise/' /usr/share/ovirt-engine/branding/ovirt.brand/messages.properties")
        os.system("sed -i 's/\\\u57FA\\\u7840\\\u7248/\\\u4f01\\\u4e1a\\\u7248/g' /usr/share/ovirt-engine/branding/ovirt.brand/messages_zh_CN.properties")
        os.system("sed -i 's/EayunOS_top_logo_basic\.png/EayunOS_top_logo_enterprise\.png/' /usr/share/ovirt-engine/branding/ovirt.brand/common.css")
        os.system("sed -i '/EayunOSVersion/d' /usr/share/ovirt-engine/dbscripts/upgrade/pre_upgrade/0000_config.sql")
        os.system("echo \"select fn_db_add_config_value('EayunOSVersion','Enterprise','general');\" >> /usr/share/ovirt-engine/dbscripts/upgrade/pre_upgrade/0000_config.sql")
        os.system("echo \"select fn_db_update_config_value('EayunOSVersion','Enterprise','general');\" >> /usr/share/ovirt-engine/dbscripts/upgrade/pre_upgrade/0000_config.sql")
        # make product uuid readable
        os.system("echo \"#! /bin/bash\" > /etc/init.d/systemuuid")
        os.system("echo \"# chkconfig: 2345 10 90\" >> /etc/init.d/systemuuid")
        os.system("echo \"chmod a+r /sys/class/dmi/id/product_uuid\" >> /etc/init.d/systemuuid")
        os.system("chmod a+x /etc/init.d/systemuuid")
        os.system("chkconfig systemuuid on")
        os.system("chmod a+r /sys/class/dmi/id/product_uuid")

    @plugin.event(
        stage=plugin.Stages.STAGE_MISC,
        before=(
            oengcommcons.Stages.DB_SCHEMA,
        ),
        condition=lambda self: (
            not self.environment[oenginecons.EngineDBEnv.NEW_DATABASE]
        ),
    )
    def _setup_installation_time(self):
        os.system("sed -i '/InstallationTime/d' /usr/share/ovirt-engine/dbscripts/upgrade/pre_upgrade/0000_config.sql")
        os.system("echo \"select fn_db_add_config_value('InstallationTime',to_char(current_timestamp,'yyyy-MM-dd HH24:mm:ss'),'general');\" >> /usr/share/ovirt-engine/dbscripts/upgrade/pre_upgrade/0000_config.sql")

<|endoftext|>


<commit_before>from django.conf import settings
from django.urls import include, path
from django.contrib import admin
from django.conf.urls import url

$PH_URL_IMPORTS

admin.autodiscover()

urlpatterns = [
    path('admin/', admin.site.urls),
    $PH_URLS
    $PH_CATCH_ALL_URL
]

# -----------------------------------------------------------------------------
# Django Debug Toolbar URLS
# -----------------------------------------------------------------------------
try:
    if settings.DEBUG:
        import debug_toolbar
        urlpatterns = [
            url(r'^__debug__/', include(debug_toolbar.urls)),
        ] + urlpatterns
except ImportError:
    pass

# -----------------------------------------------------------------------------
# Static file DEBUGGING
# -----------------------------------------------------------------------------
if settings.DEBUG:
    from django.conf.urls.static import static
    from django.contrib.staticfiles.urls import staticfiles_urlpatterns
    import os.path

    urlpatterns += staticfiles_urlpatterns()
    urlpatterns += static(settings.MEDIA_URL + 'images/',
                          document_root=os.path.join(settings.MEDIA_ROOT,
                                                     'images'))
<commit_msg>Removed deprecated URL import<commit_after>from django.conf import settings
from django.urls import include, path
from django.contrib import admin

$PH_URL_IMPORTS

admin.autodiscover()

urlpatterns = [
    path('admin/', admin.site.urls),
    $PH_URLS
    $PH_CATCH_ALL_URL
]

# -----------------------------------------------------------------------------
# Django Debug Toolbar URLS
# -----------------------------------------------------------------------------
try:
    if settings.DEBUG:
        import debug_toolbar
        urlpatterns = [
            url(r'^__debug__/', include(debug_toolbar.urls)),
        ] + urlpatterns
except ImportError:
    pass

# -----------------------------------------------------------------------------
# Static file DEBUGGING
# -----------------------------------------------------------------------------
if settings.DEBUG:
    from django.conf.urls.static import static
    from django.contrib.staticfiles.urls import staticfiles_urlpatterns
    import os.path

    urlpatterns += staticfiles_urlpatterns()
    urlpatterns += static(settings.MEDIA_URL + 'images/',
                          document_root=os.path.join(settings.MEDIA_ROOT,
                                                     'images'))
<|endoftext|>


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

import aloobj
import alofbx
import argparse
import collections
import json
import pprint
import struct
import sys

def load_format(file):
    return {int(k): v for k,v in json.load(file).items()}

def parse_chunked(format, buf):
    chunk_data = collections.defaultdict(list)

    while buf:
        chunk_id, size = unpack('<Ii', buf)

        sub_chunks = size < 0
        # Clear the sign bit (used to indicate if a chunk contains sub-chunks)
        size &= 0x7fffffff

        chunk_type = format.get(chunk_id)

        if chunk_type:
            if sub_chunks:
                chunk_data[chunk_type['name']].append(parse_chunked(format, buf[:size]))
            else:
                chunk_data[chunk_type['name']].append(parse_chunk(
                    chunk_type, buf[:size], chunk_data))
        del buf[:size]

    return chunk_data

def unpack(format, buf):
    """Both unpack and delete. Convert single-element tuples to their element"""
    result = struct.unpack_from(format, buf)
    if len(result) == 1:
        result = result[0]
    del buf[:struct.calcsize(format)]
    return result

def unpack_asciiz(buf):
    l = buf.find(b'\x00')
    if l < 0:
        # should not happen (famous last words), but if it does, interpret the whole
        # bytearray as a string and delete all of its contents (by setting the end
        # character to past the end
        l = len(buf)

    result = buf[:l].decode(encoding='ascii')
    del buf[:l+1]
    return result

def parse_chunk(format, buf, parent):
    result = {}
    content = format['content']

    for c in content:
        name = c.get('name')
        t = c['type']

        if c.get('head'):
            del buf[:2]

        if name is None:
            del buf[:struct.calcsize(t)]
            continue

        ct = c.get('count')
        if isinstance(ct, dict):
            # always take the first element of the given chunk type.
            ct = parent[ct['chunk_name']][0][ct['property']] * ct.get('scale', 1)

        if ct is None:
            if t == 'asciiz':
                result[name] = unpack_asciiz(buf)
            elif t == 'struct':
                result[name] = parse_chunk(c, buf, parent)
            else:
                result[name] = unpack(t, buf)
        elif ct == 'max':
            result[name] = []
            while buf:
                if t == 'asciiz':
                    result[name].append(unpack_asciiz(buf))
                elif t == 'struct':
                    result[name].append(parse_chunk(c, buf, parent))
                else:
                    result[name].append(unpack(t, buf))
        else:
            result[name] = []
            for _ in range(ct):
                if t == 'asciiz':
                    result[name].append(unpack_asciiz(buff))
                elif t == 'struct':
                    result[name].append(parse_chunk(c, buf, parent))
                else:
                    result[name].append(unpack(t, buf))

    return result

def main(args):
    with args.json_file as json_file, args.chunked_file as chunked_file,\
        args.output_file as output_file:

        format = load_format(json_file)
        buf = bytearray(chunked_file.read())

        parse_result = parse_chunked(format, buf)

        if args.output_format == 'dict':
            print(parse_result, file=args.output_file)
        elif args.output_format == 'json':
            json.dump(parse_result, output_file)
            print(file=args.output_file)
        elif args.output_format == 'obj':
            aloobj.dump(parse_result, output_file)
        elif args.output_format == 'fbx':
            alofbx.dump(parse_result, output_file)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Load chunked files based on a json descripton')
    parser.add_argument(
        'json_file', type=argparse.FileType('r'),
        help='The json file which describes the chunked format to be used')
    parser.add_argument(
        'chunked_file', type=argparse.FileType('rb'),
        help='The chunked file to be red using the specified format')
    parser.add_argument(
        '--output-format', '-f', type=str, choices=('dict', 'json', 'obj', 'fbx'),
        default='dict', help='The output format of the resulting data')
    parser.add_argument(
        '--output-file', '-o', type=argparse.FileType('w'), default=sys.stdout,
        help='where to store the output of the operation (default: stdout)')

    main(parser.parse_args())
<commit_msg>removed fbx stuff because it doesn't work on other python versions<commit_after>#!/usr/bin/env python3

import aloobj
import argparse
import collections
import json
import pprint
import struct
import sys

def load_format(file):
    return {int(k): v for k,v in json.load(file).items()}

def parse_chunked(format, buf):
    chunk_data = collections.defaultdict(list)

    while buf:
        chunk_id, size = unpack('<Ii', buf)

        sub_chunks = size < 0
        # Clear the sign bit (used to indicate if a chunk contains sub-chunks)
        size &= 0x7fffffff

        chunk_type = format.get(chunk_id)

        if chunk_type:
            if sub_chunks:
                chunk_data[chunk_type['name']].append(parse_chunked(format, buf[:size]))
            else:
                chunk_data[chunk_type['name']].append(parse_chunk(
                    chunk_type, buf[:size], chunk_data))
        del buf[:size]

    return chunk_data

def unpack(format, buf):
    """Both unpack and delete. Convert single-element tuples to their element"""
    result = struct.unpack_from(format, buf)
    if len(result) == 1:
        result = result[0]
    del buf[:struct.calcsize(format)]
    return result

def unpack_asciiz(buf):
    l = buf.find(b'\x00')
    if l < 0:
        # should not happen (famous last words), but if it does, interpret the whole
        # bytearray as a string and delete all of its contents (by setting the end
        # character to past the end
        l = len(buf)

    result = buf[:l].decode(encoding='ascii')
    del buf[:l+1]
    return result

def parse_chunk(format, buf, parent):
    result = {}
    content = format['content']

    for c in content:
        name = c.get('name')
        t = c['type']

        if c.get('head'):
            del buf[:2]

        if name is None:
            del buf[:struct.calcsize(t)]
            continue

        ct = c.get('count')
        if isinstance(ct, dict):
            # always take the first element of the given chunk type.
            ct = parent[ct['chunk_name']][0][ct['property']] * ct.get('scale', 1)

        if ct is None:
            if t == 'asciiz':
                result[name] = unpack_asciiz(buf)
            elif t == 'struct':
                result[name] = parse_chunk(c, buf, parent)
            else:
                result[name] = unpack(t, buf)
        elif ct == 'max':
            result[name] = []
            while buf:
                if t == 'asciiz':
                    result[name].append(unpack_asciiz(buf))
                elif t == 'struct':
                    result[name].append(parse_chunk(c, buf, parent))
                else:
                    result[name].append(unpack(t, buf))
        else:
            result[name] = []
            for _ in range(ct):
                if t == 'asciiz':
                    result[name].append(unpack_asciiz(buff))
                elif t == 'struct':
                    result[name].append(parse_chunk(c, buf, parent))
                else:
                    result[name].append(unpack(t, buf))

    return result

def main(args):
    with args.json_file as json_file, args.chunked_file as chunked_file,\
        args.output_file as output_file:

        format = load_format(json_file)
        buf = bytearray(chunked_file.read())

        parse_result = parse_chunked(format, buf)

        if args.output_format == 'dict':
            print(parse_result, file=args.output_file)
        elif args.output_format == 'json':
            json.dump(parse_result, output_file)
            print(file=args.output_file)
        elif args.output_format == 'obj':
            aloobj.dump(parse_result, output_file)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description='Load chunked files based on a json descripton')
    parser.add_argument(
        'json_file', type=argparse.FileType('r'),
        help='The json file which describes the chunked format to be used')
    parser.add_argument(
        'chunked_file', type=argparse.FileType('rb'),
        help='The chunked file to be red using the specified format')
    parser.add_argument(
        '--output-format', '-f', type=str, choices=('dict', 'json', 'obj'),
        default='dict', help='The output format of the resulting data')
    parser.add_argument(
        '--output-file', '-o', type=argparse.FileType('w'), default=sys.stdout,
        help='where to store the output of the operation (default: stdout)')

    main(parser.parse_args())
<|endoftext|>