<commit_before>"""
This module provides a plugin architecture for security methods.

Each security plugin can be distributed as separate package by providing
modules inside the namespace-package "puresnmp_plugins.security". Note that
in order to be a valid namespace-package, such a package *must not* have a
``__init__.py`` file!

Example folder-structure for a privacy plugin:

.. code-block:: text

    my-security-plugin/
     +- setup.py (or pyproject.toml)
     +- puresnmp/
         +- security/
             +- mymodule.py
             +- myothermodule.py

Note that there is no ``__init__.py`` file!

In order for modules to be detected as plugin, they must follow the following
rules:

* Have a no-arg function ``create`` returning a instance of the class (or
  subclass) of :py:class:`puresnmp.security.SecurityModel`.
* Contain a int-variable ``IDENTIFIER``. This variable should map to the
  security model identifiers as defined in the SNMPv3 standard. See
  https://www.iana.org/assignments/snmp-number-spaces/snmp-number-spaces.xhtml
  and :rfc:`3411`
"""
from types import ModuleType
from typing import Any, Awaitable, Callable, Dict, Generic, TypeVar

from puresnmp.credentials import Credentials
from puresnmp.exc import UnknownSecurityModel
from puresnmp.plugins.pluginbase import Loader

try:
    from typing import Protocol
except ImportError:
    from typing_extensions import Protocol  # type: ignore

#: The type of an *unsecured* message inside of "puresnmp"
TPureSNMPType = TypeVar("TPureSNMPType", bound=Any)

#: The type of a *secured* message outside of "puresnmp"
TX690Type = TypeVar("TX690Type", bound=Any)

class TSecurityPlugin(Protocol):
    """
    Protocol for security plugins
    """

    # pylint: disable=too-few-public-methods

    def create(
        self,
    ) -> "SecurityModel[TPureSNMPType, TX690Type]":  # pragma: no cover
        """
        Create a new instance of a security model
        """
        ...

class SecurityModel(Generic[TPureSNMPType, TX690Type]):
    """
    Each Security Model defines the applied protecion on SNMP PDUs
    """

    #: The "Local Configuration Datastore" (LCD). This contains contextual
    #: information which may be needed for some security models (as defined by
    #: the SNMPv3 architecture.)
    local_config: Dict[bytes, Dict[str, Any]]

    def __init__(self) -> None:
        self.local_config = {}

    def generate_request_message(
        self,
        message: TPureSNMPType,
        security_engine_id: bytes,
        credentials: Credentials,
    ) -> TX690Type:  # pragma: no cover
        """
        Take a plain unprocessed message and applies security to the message
        as defined by the concrete security model.

        It returns the processed message including security modifications.
        Further processing is provided by the "message-processing-model" to
        prepare it for sending out to the network.

        The kind of processing applied to the message depends on the
        credential type.

        :param message: The original message
        :param security_engine_id: The engine-id of the receiving SNMP engine
        :param credentials: A credential object giving information on what
            kind of operations are needed on the message.
        """
        raise NotImplementedError(f"Not yet implemented in {self.__class__}")

    def process_incoming_message(
        self,
        message: TX690Type,
        credentials: Credentials,
    ) -> TPureSNMPType:  # pragma: no cover
        """
        Takes a message which included potential security modifications (like
        encryption) and "undoes" these modifications in order to make the
        message usable again.

        Returns an unprocessed message.

        The kind of processing applied to the message depends on the
        credential type.

        :param message: The original message
        :param credentials: A credential object giving information on what
            kind of operations are needed on the message.
        """
        raise NotImplementedError(f"Not yet implemented in {self.__class__}")

    def set_engine_timing(
        self,
        engine_id: bytes,
        engine_boots: int,
        engine_time: int,
    ) -> None:  # pragma: no cover
        """
        Update the security model with timing information of the remote-engine.

        The timing information is required by SNMPv3 to authenticat/encrypt
        messages.
        """
        raise NotImplementedError(f"Not yet implemented in {self.__class__}")

    async def send_discovery_message(
        self,
        transport_handler: Callable[[bytes], Awaitable[bytes]],
    ) -> Any:  # pragma: no cover
        """
        Send a discovery message to the remote engine.

        :param handler: A callable that is responsible to send data to the
            remote engine. This callable should already be aware of how to
            connect. It takes the bytes to be sent to the remote and should
            return the response bytes without any processing.
        """
        raise NotImplementedError(f"Not yet implemented in {self.__class__}")

def is_valid_sec_plugin(mod: ModuleType) -> bool:
    return all(
        [
            hasattr(mod, "create"),
            hasattr(mod, "IDENTIFIER"),
        ]
    )

def create(identifier: int) -> SecurityModel[TPureSNMPType, TX690Type]:
    """
    Return an instance of the given security module by identifier.

    This looks up the module by "IDENTIFIER" as specified in the given plugin.

    If no plugin with the given identifier is found, a *KeyError* is raised
    """
    namespace = "puresnmp_plugins.security"
    loader = Loader(namespace, is_valid_sec_plugin)
    result = loader.create(identifier)
    if not result:
        raise UnknownSecurityModel(
            namespace,
            identifier,
            sorted(loader.discovered_plugins.keys()),
        )
    return result.create()  # type: ignore
<commit_msg>Change docstring markup style<commit_after>"""
This module provides a plugin architecture for security methods.

Each security plugin can be distributed as separate package by providing
modules inside the namespace-package "puresnmp_plugins.security". Note that
in order to be a valid namespace-package, such a package *must not* have a
``__init__.py`` file!

Example folder-structure for a privacy plugin:

.. code-block:: text

    my-security-plugin/
     +- setup.py (or pyproject.toml)
     +- puresnmp/
         +- security/
             +- mymodule.py
             +- myothermodule.py

Note that there is no ``__init__.py`` file!

In order for modules to be detected as plugin, they must follow the following
rules:

* Have a no-arg function ``create`` returning a instance of the class (or
  subclass) of :py:class:`puresnmp.security.SecurityModel`.
* Contain a int-variable ``IDENTIFIER``. This variable should map to the
  security model identifiers as defined in the SNMPv3 standard. See
  https://www.iana.org/assignments/snmp-number-spaces/snmp-number-spaces.xhtml
  and :rfc:`3411`
"""
from types import ModuleType
from typing import Any, Awaitable, Callable, Dict, Generic, TypeVar

from puresnmp.credentials import Credentials
from puresnmp.exc import UnknownSecurityModel
from puresnmp.plugins.pluginbase import Loader

try:
    from typing import Protocol
except ImportError:
    from typing_extensions import Protocol  # type: ignore

#: The type of an *unsecured* message inside of "puresnmp"
TPureSNMPType = TypeVar("TPureSNMPType", bound=Any)

#: The type of a *secured* message outside of "puresnmp"
TX690Type = TypeVar("TX690Type", bound=Any)

class TSecurityPlugin(Protocol):
    """
    Protocol for security plugins
    """

    # pylint: disable=too-few-public-methods

    def create(
        self,
    ) -> "SecurityModel[TPureSNMPType, TX690Type]":  # pragma: no cover
        """
        Create a new instance of a security model
        """
        ...

class SecurityModel(Generic[TPureSNMPType, TX690Type]):
    """
    Each Security Model defines the applied protecion on SNMP PDUs
    """

    local_config: Dict[bytes, Dict[str, Any]]
    """
    The "Local Configuration Datastore" (LCD). This contains contextual
    information which may be needed for some security models (as defined by
    the SNMPv3 architecture.)
    """

    def __init__(self) -> None:
        self.local_config = {}

    def generate_request_message(
        self,
        message: TPureSNMPType,
        security_engine_id: bytes,
        credentials: Credentials,
    ) -> TX690Type:  # pragma: no cover
        """
        Take a plain unprocessed message and applies security to the message
        as defined by the concrete security model.

        It returns the processed message including security modifications.
        Further processing is provided by the "message-processing-model" to
        prepare it for sending out to the network.

        The kind of processing applied to the message depends on the
        credential type.

        :param message: The original message
        :param security_engine_id: The engine-id of the receiving SNMP engine
        :param credentials: A credential object giving information on what
            kind of operations are needed on the message.
        """
        raise NotImplementedError(f"Not yet implemented in {self.__class__}")

    def process_incoming_message(
        self,
        message: TX690Type,
        credentials: Credentials,
    ) -> TPureSNMPType:  # pragma: no cover
        """
        Takes a message which included potential security modifications (like
        encryption) and "undoes" these modifications in order to make the
        message usable again.

        Returns an unprocessed message.

        The kind of processing applied to the message depends on the
        credential type.

        :param message: The original message
        :param credentials: A credential object giving information on what
            kind of operations are needed on the message.
        """
        raise NotImplementedError(f"Not yet implemented in {self.__class__}")

    def set_engine_timing(
        self,
        engine_id: bytes,
        engine_boots: int,
        engine_time: int,
    ) -> None:  # pragma: no cover
        """
        Update the security model with timing information of the remote-engine.

        The timing information is required by SNMPv3 to authenticat/encrypt
        messages.
        """
        raise NotImplementedError(f"Not yet implemented in {self.__class__}")

    async def send_discovery_message(
        self,
        transport_handler: Callable[[bytes], Awaitable[bytes]],
    ) -> Any:  # pragma: no cover
        """
        Send a discovery message to the remote engine.

        :param handler: A callable that is responsible to send data to the
            remote engine. This callable should already be aware of how to
            connect. It takes the bytes to be sent to the remote and should
            return the response bytes without any processing.
        """
        raise NotImplementedError(f"Not yet implemented in {self.__class__}")

def is_valid_sec_plugin(mod: ModuleType) -> bool:
    return all(
        [
            hasattr(mod, "create"),
            hasattr(mod, "IDENTIFIER"),
        ]
    )

def create(identifier: int) -> SecurityModel[TPureSNMPType, TX690Type]:
    """
    Return an instance of the given security module by identifier.

    This looks up the module by "IDENTIFIER" as specified in the given plugin.

    If no plugin with the given identifier is found, a *KeyError* is raised
    """
    namespace = "puresnmp_plugins.security"
    loader = Loader(namespace, is_valid_sec_plugin)
    result = loader.create(identifier)
    if not result:
        raise UnknownSecurityModel(
            namespace,
            identifier,
            sorted(loader.discovered_plugins.keys()),
        )
    return result.create()  # type: ignore
<|endoftext|>


<commit_before>#!/usr/bin/env python3
# -*-coding:UTF-8 -*

import os
import re
import sys
import time
import redis
import datetime

sys.path.append(os.path.join(os.environ['AIL_BIN'], 'lib/'))
import ConfigLoader

new_version = 'v2.5'

if __name__ == '__main__':

    start_deb = time.time()

    config_loader = ConfigLoader.ConfigLoader()
    r_serv = config_loader.get_redis_conn("ARDB_DB")
    config_loader = None

    r_serv.zadd('ail:all_role', 3, 'user')
    r_serv.zadd('ail:all_role', 4, 'user_no_api')
    r_serv.zadd('ail:all_role', 5, 'read_only')

    for user in r_serv.hkeys(user:all):
        r_serv.sadd('user_role:user', user)
        r_serv.sadd('user_role:user_no_api', user)
        r_serv.sadd('user_role:read_only', user)

    #Set current ail version
    r_serv.set('ail:version', new_version)

    #Set current ail version
    r_serv.hset('ail:update_date', new_version, datetime.datetime.now().strftime("%Y%m%d"))
<commit_msg>fix: [update v2.5]<commit_after>#!/usr/bin/env python3
# -*-coding:UTF-8 -*

import os
import re
import sys
import time
import redis
import datetime

sys.path.append(os.path.join(os.environ['AIL_BIN'], 'lib/'))
import ConfigLoader

new_version = 'v2.5'

if __name__ == '__main__':

    start_deb = time.time()

    config_loader = ConfigLoader.ConfigLoader()
    r_serv = config_loader.get_redis_conn("ARDB_DB")
    config_loader = None

    r_serv.zadd('ail:all_role', 3, 'user')
    r_serv.zadd('ail:all_role', 4, 'user_no_api')
    r_serv.zadd('ail:all_role', 5, 'read_only')

    for user in r_serv.hkeys('user:all'):
        r_serv.sadd('user_role:user', user)
        r_serv.sadd('user_role:user_no_api', user)
        r_serv.sadd('user_role:read_only', user)

    #Set current ail version
    r_serv.set('ail:version', new_version)

    #Set current ail version
    r_serv.hset('ail:update_date', new_version, datetime.datetime.now().strftime("%Y%m%d"))
<|endoftext|>


<commit_before>from __future__ import unicode_literals

from django.conf import settings
from django.core.exceptions import ValidationError
from django.db import models
from django.utils.functional import cached_property
from django.utils.translation import ugettext_lazy as _

from .apns import (apns_fetch_inactive_ids, apns_send_bulk_message,
					apns_send_message)
from .fields import HexIntegerField

SERVICE_GCM = 0
SERVICE_APNS = 1
SERVICE_INACTIVE = 2

SERVICES = (
	(SERVICE_GCM, "GCM"),
	(SERVICE_APNS, "APNS"),
	(SERVICE_INACTIVE, "Inactive"),
)

class DeviceManager(models.Manager):
	def get_queryset(self):
		return DeviceQuerySet(self.model)

	def invalidate(self, registration_ids):
		""" Called when some registration ids are deemed invalid. """
		self.filter(registration_id__in=registration_ids).update(service=SERVICE_INACTIVE)

class DeviceQuerySet(models.query.QuerySet):
	def send_message(self, message, **kwargs):
		if self:
			gcmDevices = []
			apnsDevices = []

			for device in self:
				if device.service == SERVICE_APNS:
					apnsDevices.append(device)
				elif device.service == SERVICE_GCM:
					gcmDevices.append(device)

			if apnsDevices:
				apns_send_bulk_message(
					devices=apnsDevices,
					alert=message,
					certificate=apnsDevices[0].APNS_CERTIFICATE,
					**kwargs
				)

			if gcmDevices:
				data = kwargs.pop("extra", {})
				if message is not None:
					data["message"] = message

				from .gcm import gcm_send_bulk_message
				return gcm_send_bulk_message(
					devices=gcmDevices,
					data=data,
					api_key=apnsDevices[0].APNS_CERTIFICATE,
					**kwargs
				)
			return None

class BareDevice(models.Model):
	service = models.IntegerField(
		choices=SERVICES,
		verbose_name=_("Notification service"),
		default=SERVICE_INACTIVE
	)
	registration_id = models.TextField(
		blank=True,
		verbose_name=_("Registration ID")
	)

	objects = DeviceManager()
	APNS_CERTIFICATE = settings.PUSH_NOTIFICATIONS_SETTINGS.get("APNS_CERTIFICATE", None) if hasattr(settings, 'PUSH_NOTIFICATIONS_SETTINGS') else None
	GCM_API_KEY = settings.PUSH_NOTIFICATIONS_SETTINGS.get("GCM_API_KEY", None) if hasattr(settings, 'PUSH_NOTIFICATIONS_SETTINGS') else None

	class Meta:
		abstract = True

	@cached_property
	def active(self):
		return self.service != SERVICE_INACTIVE

	def save(self, *args, **kwargs):
		if (self.service != SERVICE_INACTIVE
			and not self.registration_id):
			self.invalidate()
		else:
			if (self.service == SERVICE_APNS
				and len(self.registration_id) > 64):
				raise ValidationError("APNS registration_id's max length is 64.")
			super(BareDevice, self).save(*args, **kwargs)

	def send_message(self, message, **kwargs):
		if self.active:
			if self.service == SERVICE_APNS:
				return apns_send_message(
					device=self,
					alert=message,
					certificate=self.APNS_CERTIFICATE,
					**kwargs
				)
			else:
				data = kwargs.pop("extra", {})
				if message is not None:
					data["message"] = message
				from .gcm import gcm_send_message
				return gcm_send_message(
					device=self,
					data=data,
					api_key=self.GCM_API_KEY,
					**kwargs
				)

	def invalidate(self):
		""" Called when the registration_id is deemed invalid. """
		self.service = SERVICE_INACTIVE
		self.save()

	def __unicode__(self):
		return "{service}: {registration_id}".format(
			service=self.get_service_display(),
			registration_id=self.registration_id
		)

class SimpleDevice(BareDevice):
	pass

class Device(models.Model):
	name = models.CharField(max_length=255, verbose_name=_("Name"), blank=True, null=True)
	active = models.BooleanField(verbose_name=_("Is active"), default=True,
		help_text=_("Inactive devices will not be sent notifications"))
	user = models.ForeignKey(settings.AUTH_USER_MODEL, blank=True, null=True)
	date_created = models.DateTimeField(verbose_name=_("Creation date"), auto_now_add=True, null=True)

	class Meta:
		abstract = True

	def __unicode__(self):
		return self.name or \
			str(self.device_id or "") or \
			"%s for %s" % (self.__class__.__name__, self.user or "unknown user")

class GCMDeviceManager(models.Manager):
	def get_queryset(self):
		return GCMDeviceQuerySet(self.model)

class GCMDeviceQuerySet(models.query.QuerySet):
	def send_message(self, message, **kwargs):
		if self:
			from .gcm import gcm_send_bulk_message

			data = kwargs.pop("extra", {})
			if message is not None:
				data["message"] = message

			return gcm_send_bulk_message(devices=self, data=data, **kwargs)

class GCMDevice(Device):
	# device_id cannot be a reliable primary key as fragmentation between different devices
	# can make it turn out to be null and such:
	# http://android-developers.blogspot.co.uk/2011/03/identifying-app-installations.html
	device_id = HexIntegerField(verbose_name=_("Device ID"), blank=True, null=True, db_index=True,
		help_text=_("ANDROID_ID / TelephonyManager.getDeviceId() (always as hex)"))
	registration_id = models.TextField(verbose_name=_("Registration ID"))

	objects = GCMDeviceManager()

	class Meta:
		verbose_name = _("GCM device")

	def send_message(self, message, **kwargs):
		from .gcm import gcm_send_message
		data = kwargs.pop("extra", {})
		if message is not None:
			data["message"] = message
		return gcm_send_message(device=self, data=data, **kwargs)

class APNSDeviceManager(models.Manager):
	def get_queryset(self):
		return APNSDeviceQuerySet(self.model)

class APNSDeviceQuerySet(models.query.QuerySet):
	def send_message(self, message, **kwargs):
		if self:
			return apns_send_bulk_message(devices=self, alert=message, **kwargs)

class APNSDevice(Device):
	device_id = models.UUIDField(verbose_name=_("Device ID"), blank=True, null=True, db_index=True,
		help_text="UDID / UIDevice.identifierForVendor()")
	registration_id = models.CharField(verbose_name=_("Registration ID"), max_length=64, unique=True)

	objects = APNSDeviceManager()

	class Meta:
		verbose_name = _("APNS device")

	def send_message(self, message, **kwargs):
		return apns_send_message(device=self, alert=message, **kwargs)

# This is an APNS-only function right now, but maybe GCM will implement it
# in the future.  But the definition of 'expired' may not be the same. Whatevs
def get_expired_tokens():
	return apns_fetch_inactive_ids()
<commit_msg>Fixed DeviceQuerySet gcm using APNS instead of GCM key.<commit_after>from __future__ import unicode_literals

from django.conf import settings
from django.core.exceptions import ValidationError
from django.db import models
from django.utils.functional import cached_property
from django.utils.translation import ugettext_lazy as _

from .apns import (apns_fetch_inactive_ids, apns_send_bulk_message,
					apns_send_message)
from .fields import HexIntegerField

SERVICE_GCM = 0
SERVICE_APNS = 1
SERVICE_INACTIVE = 2

SERVICES = (
	(SERVICE_GCM, "GCM"),
	(SERVICE_APNS, "APNS"),
	(SERVICE_INACTIVE, "Inactive"),
)

class DeviceManager(models.Manager):
	def get_queryset(self):
		return DeviceQuerySet(self.model)

	def invalidate(self, registration_ids):
		""" Called when some registration ids are deemed invalid. """
		self.filter(registration_id__in=registration_ids).update(service=SERVICE_INACTIVE)

class DeviceQuerySet(models.query.QuerySet):
	def send_message(self, message, **kwargs):
		if self:
			gcmDevices = []
			apnsDevices = []

			for device in self:
				if device.service == SERVICE_APNS:
					apnsDevices.append(device)
				elif device.service == SERVICE_GCM:
					gcmDevices.append(device)

			if apnsDevices:
				apns_send_bulk_message(
					devices=apnsDevices,
					alert=message,
					certificate=apnsDevices[0].APNS_CERTIFICATE,
					**kwargs
				)

			if gcmDevices:
				data = kwargs.pop("extra", {})
				if message is not None:
					data["message"] = message

				from .gcm import gcm_send_bulk_message
				return gcm_send_bulk_message(
					devices=gcmDevices,
					data=data,
					api_key=gcmDevices[0].GCM_API_KEY,
					**kwargs
				)
			return None

class BareDevice(models.Model):
	service = models.IntegerField(
		choices=SERVICES,
		verbose_name=_("Notification service"),
		default=SERVICE_INACTIVE
	)
	registration_id = models.TextField(
		blank=True,
		verbose_name=_("Registration ID")
	)

	objects = DeviceManager()
	APNS_CERTIFICATE = settings.PUSH_NOTIFICATIONS_SETTINGS.get("APNS_CERTIFICATE", None) if hasattr(settings, 'PUSH_NOTIFICATIONS_SETTINGS') else None
	GCM_API_KEY = settings.PUSH_NOTIFICATIONS_SETTINGS.get("GCM_API_KEY", None) if hasattr(settings, 'PUSH_NOTIFICATIONS_SETTINGS') else None

	class Meta:
		abstract = True

	@cached_property
	def active(self):
		return self.service != SERVICE_INACTIVE

	def save(self, *args, **kwargs):
		if (self.service != SERVICE_INACTIVE
			and not self.registration_id):
			self.invalidate()
		else:
			if (self.service == SERVICE_APNS
				and len(self.registration_id) > 64):
				raise ValidationError("APNS registration_id's max length is 64.")
			super(BareDevice, self).save(*args, **kwargs)

	def send_message(self, message, **kwargs):
		if self.active:
			if self.service == SERVICE_APNS:
				return apns_send_message(
					device=self,
					alert=message,
					certificate=self.APNS_CERTIFICATE,
					**kwargs
				)
			else:
				data = kwargs.pop("extra", {})
				if message is not None:
					data["message"] = message
				from .gcm import gcm_send_message
				return gcm_send_message(
					device=self,
					data=data,
					api_key=self.GCM_API_KEY,
					**kwargs
				)

	def invalidate(self):
		""" Called when the registration_id is deemed invalid. """
		self.service = SERVICE_INACTIVE
		self.save()

	def __unicode__(self):
		return "{service}: {registration_id}".format(
			service=self.get_service_display(),
			registration_id=self.registration_id
		)

class SimpleDevice(BareDevice):
	pass

class Device(models.Model):
	name = models.CharField(max_length=255, verbose_name=_("Name"), blank=True, null=True)
	active = models.BooleanField(verbose_name=_("Is active"), default=True,
		help_text=_("Inactive devices will not be sent notifications"))
	user = models.ForeignKey(settings.AUTH_USER_MODEL, blank=True, null=True)
	date_created = models.DateTimeField(verbose_name=_("Creation date"), auto_now_add=True, null=True)

	class Meta:
		abstract = True

	def __unicode__(self):
		return self.name or \
			str(self.device_id or "") or \
			"%s for %s" % (self.__class__.__name__, self.user or "unknown user")

class GCMDeviceManager(models.Manager):
	def get_queryset(self):
		return GCMDeviceQuerySet(self.model)

class GCMDeviceQuerySet(models.query.QuerySet):
	def send_message(self, message, **kwargs):
		if self:
			from .gcm import gcm_send_bulk_message

			data = kwargs.pop("extra", {})
			if message is not None:
				data["message"] = message

			return gcm_send_bulk_message(devices=self, data=data, **kwargs)

class GCMDevice(Device):
	# device_id cannot be a reliable primary key as fragmentation between different devices
	# can make it turn out to be null and such:
	# http://android-developers.blogspot.co.uk/2011/03/identifying-app-installations.html
	device_id = HexIntegerField(verbose_name=_("Device ID"), blank=True, null=True, db_index=True,
		help_text=_("ANDROID_ID / TelephonyManager.getDeviceId() (always as hex)"))
	registration_id = models.TextField(verbose_name=_("Registration ID"))

	objects = GCMDeviceManager()

	class Meta:
		verbose_name = _("GCM device")

	def send_message(self, message, **kwargs):
		from .gcm import gcm_send_message
		data = kwargs.pop("extra", {})
		if message is not None:
			data["message"] = message
		return gcm_send_message(device=self, data=data, **kwargs)

class APNSDeviceManager(models.Manager):
	def get_queryset(self):
		return APNSDeviceQuerySet(self.model)

class APNSDeviceQuerySet(models.query.QuerySet):
	def send_message(self, message, **kwargs):
		if self:
			return apns_send_bulk_message(devices=self, alert=message, **kwargs)

class APNSDevice(Device):
	device_id = models.UUIDField(verbose_name=_("Device ID"), blank=True, null=True, db_index=True,
		help_text="UDID / UIDevice.identifierForVendor()")
	registration_id = models.CharField(verbose_name=_("Registration ID"), max_length=64, unique=True)

	objects = APNSDeviceManager()

	class Meta:
		verbose_name = _("APNS device")

	def send_message(self, message, **kwargs):
		return apns_send_message(device=self, alert=message, **kwargs)

# This is an APNS-only function right now, but maybe GCM will implement it
# in the future.  But the definition of 'expired' may not be the same. Whatevs
def get_expired_tokens():
	return apns_fetch_inactive_ids()
<|endoftext|>


<commit_before>#!/usr/bin/env python
from django.core.management import execute_manager
try:
    import settings # Assumed to be in the same directory.
except ImportError:
    import sys
    sys.stderr.write("Error: Can't find the file 'settings.py' in the directory containing %r. It appears you've customized things.\nYou'll have to run django-admin.py, passing it your settings module.\n(If the file settings.py does indeed exist, it's causing an ImportError somehow.)\n" % __file__)
    sys.exit(1)

if __name__ == "__main__":
    execute_manager(settings)
<commit_msg>Removed custom modifications from manage.py<commit_after>#!/usr/bin/env python
import os
import sys

if __name__ == "__main__":
    os.environ.setdefault("DJANGO_SETTINGS_MODULE", "settings")

    from django.core.management import execute_from_command_line

    execute_from_command_line(sys.argv)
<|endoftext|>


<commit_before>"""Report the state of arcyd as it updates repositories."""
# =============================================================================
# CONTENTS
# -----------------------------------------------------------------------------
# abdt_arcydreporter
#
# Public Classes:
#   SharedFileDictOutput
#    .write
#   SharedDictOutput
#    .write
#   ArcydReporter
#    .start_sleep
#    .update_sleep
#    .finish_sleep
#    .start_repo
#    .fail_repo
#    .finish_repo
#    .close
#
# Public Assignments:
#   ARCYD_STATUS
#   ARCYD_CURRENT_REPO
#   ARCYD_REPOS
#   ARCYD_STATISTICS
#   ARCYD_LIST_ATTRIB
#   ARCYD_STATUS_STARTING
#   ARCYD_STATUS_UPDATING
#   ARCYD_STATUS_SLEEPING
#   ARCYD_STATUS_STOPPED
#   ARCYD_STATUS_IDLE
#   ARCYD_LIST_STATUS
#   REPO_ATTRIB_NAME
#   REPO_ATTRIB_HUMAN_NAME
#   REPO_ATTRIB_STATUS
#   REPO_LIST_ATTRIB
#   REPO_STATUS_UPDATING
#   REPO_STATUS_FAILED
#   REPO_STATUS_OK
#   REPO_STATUSES
#   ARCYD_STAT_CURRENT_CYCLE_TIME
#   ARCYD_STAT_LAST_CYCLE_TIME
#   ARCYD_LIST_STATISTICS
#
# -----------------------------------------------------------------------------
# (this contents block is generated, edits will be lost)
# =============================================================================

from __future__ import absolute_import

import datetime
import json

import phlsys_fs

ARCYD_STATUS = 'status'
ARCYD_CURRENT_REPO = 'current-repo'
ARCYD_REPOS = 'repos'
ARCYD_STATISTICS = 'statistics'

ARCYD_LIST_ATTRIB = [
    ARCYD_CURRENT_REPO,
    ARCYD_STATUS,
    ARCYD_REPOS,
    ARCYD_STATISTICS,
]

ARCYD_STATUS_STARTING = 'starting'
ARCYD_STATUS_UPDATING = 'updating'
ARCYD_STATUS_SLEEPING = 'sleeping'
ARCYD_STATUS_STOPPED = 'stopped'
ARCYD_STATUS_IDLE = 'idle'

ARCYD_LIST_STATUS = [
    ARCYD_STATUS_UPDATING,
    ARCYD_STATUS_SLEEPING,
    ARCYD_STATUS_STARTING,
    ARCYD_STATUS_STOPPED,
    ARCYD_STATUS_IDLE,
]

REPO_ATTRIB_NAME = 'name'
REPO_ATTRIB_HUMAN_NAME = 'human-name'
REPO_ATTRIB_STATUS = 'status'

REPO_LIST_ATTRIB = [
    REPO_ATTRIB_NAME,
    REPO_ATTRIB_HUMAN_NAME,
    REPO_ATTRIB_STATUS,
]

REPO_STATUS_UPDATING = 'updating'
REPO_STATUS_FAILED = 'failed'
REPO_STATUS_OK = 'ok'

REPO_STATUSES = [
    REPO_STATUS_UPDATING,
    REPO_STATUS_FAILED,
    REPO_STATUS_OK,
]

ARCYD_STAT_CURRENT_CYCLE_TIME = 'current-cycle-time'
ARCYD_STAT_LAST_CYCLE_TIME = 'last-cycle-time'

ARCYD_LIST_STATISTICS = [
    ARCYD_STAT_CURRENT_CYCLE_TIME,
    ARCYD_STAT_LAST_CYCLE_TIME,
]

class SharedFileDictOutput(object):

    def __init__(self, filename):
        super(SharedFileDictOutput, self).__init__()
        self._filename = filename

    def write(self, d):
        assert isinstance(d, dict)
        with phlsys_fs.write_file_lock_context(self._filename) as f:
            f.write(json.dumps(d))

class SharedDictOutput(object):

    def __init__(self, shared_d):
        super(SharedDictOutput, self).__init__()
        self._shared_d = shared_d
        assert isinstance(self._shared_d, dict)

    def write(self, d):
        assert isinstance(d, dict)
        # copy contents to other dict
        self._shared_d.clear()
        self._shared_d.update(d)

class _CycleTimer(object):

    def __init__(self):
        self._current_start = None
        self._last_duration = None

    def start_cycle(self):
        self._current_start = datetime.datetime.utcnow()

    def stop_cycle(self):
        self._last_duration = self.current_duration()
        self._current_start = None

    def current_duration(self):
        if self._current_start is None:
            return None
        duration = datetime.datetime.utcnow() - self._current_start
        return duration.total_seconds()

    @property
    def last_duration(self):
        if self._last_duration is not None:
            return self._last_duration.total_seconds()
        else:
            return None

class ArcydReporter(object):

    def __init__(self, output):
        """Initialise a new reporter to report to the specified outputs.

        :output: output to write status to

        """
        super(ArcydReporter, self).__init__()
        self._output = output

        assert self._output

        self._repos = []
        self._repo = None

        self._cycle_timer = _CycleTimer()

        self._write_status(ARCYD_STATUS_STARTING)

    def start_sleep(self, duration):
        _ = duration  # NOQA
        self._cycle_timer.stop_cycle()
        self._write_status(ARCYD_STATUS_SLEEPING)

    def update_sleep(self, duration):
        _ = duration  # NOQA
        self._write_status(ARCYD_STATUS_SLEEPING)

    def finish_sleep(self):
        self._write_status(ARCYD_STATUS_IDLE)
        self._cycle_timer.start_cycle()

    def start_repo(self, name, human_name):
        self._repo = {
            REPO_ATTRIB_NAME: name,
            REPO_ATTRIB_HUMAN_NAME: human_name,
            REPO_ATTRIB_STATUS: REPO_STATUS_UPDATING,
        }
        self._write_status(ARCYD_STATUS_UPDATING)

    def fail_repo(self):
        self._repo[REPO_ATTRIB_STATUS] = REPO_STATUS_FAILED
        self._repos.append(self._repo)
        self._repo = None
        self._write_status(ARCYD_STATUS_UPDATING)

    def finish_repo(self):
        self._repo[REPO_ATTRIB_STATUS] = REPO_STATUS_OK
        self._repos.append(self._repo)
        self._repo = None
        self._write_status(ARCYD_STATUS_IDLE)

    def _write_status(self, status):
        timer = self._cycle_timer
        statistics = {
            ARCYD_STAT_CURRENT_CYCLE_TIME: timer.current_duration(),
            ARCYD_STAT_LAST_CYCLE_TIME: timer.last_duration,
        }
        d = {
            ARCYD_STATUS: status,
            ARCYD_CURRENT_REPO: self._repo,
            ARCYD_REPOS: self._repos,
            ARCYD_STATISTICS: statistics,
        }
        assert set(d.keys()) == set(ARCYD_LIST_ATTRIB)
        self._output.write(d)

    def close(self):
        self._write_status(ARCYD_STATUS_STOPPED)

#------------------------------------------------------------------------------
# Copyright (C) 2012 Bloomberg L.P.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
#------------------------------- END-OF-FILE ----------------------------------
<commit_msg>Arcyd: fix crash on last cycle duration calc<commit_after>"""Report the state of arcyd as it updates repositories."""
# =============================================================================
# CONTENTS
# -----------------------------------------------------------------------------
# abdt_arcydreporter
#
# Public Classes:
#   SharedFileDictOutput
#    .write
#   SharedDictOutput
#    .write
#   ArcydReporter
#    .start_sleep
#    .update_sleep
#    .finish_sleep
#    .start_repo
#    .fail_repo
#    .finish_repo
#    .close
#
# Public Assignments:
#   ARCYD_STATUS
#   ARCYD_CURRENT_REPO
#   ARCYD_REPOS
#   ARCYD_STATISTICS
#   ARCYD_LIST_ATTRIB
#   ARCYD_STATUS_STARTING
#   ARCYD_STATUS_UPDATING
#   ARCYD_STATUS_SLEEPING
#   ARCYD_STATUS_STOPPED
#   ARCYD_STATUS_IDLE
#   ARCYD_LIST_STATUS
#   REPO_ATTRIB_NAME
#   REPO_ATTRIB_HUMAN_NAME
#   REPO_ATTRIB_STATUS
#   REPO_LIST_ATTRIB
#   REPO_STATUS_UPDATING
#   REPO_STATUS_FAILED
#   REPO_STATUS_OK
#   REPO_STATUSES
#   ARCYD_STAT_CURRENT_CYCLE_TIME
#   ARCYD_STAT_LAST_CYCLE_TIME
#   ARCYD_LIST_STATISTICS
#
# -----------------------------------------------------------------------------
# (this contents block is generated, edits will be lost)
# =============================================================================

from __future__ import absolute_import

import datetime
import json

import phlsys_fs

ARCYD_STATUS = 'status'
ARCYD_CURRENT_REPO = 'current-repo'
ARCYD_REPOS = 'repos'
ARCYD_STATISTICS = 'statistics'

ARCYD_LIST_ATTRIB = [
    ARCYD_CURRENT_REPO,
    ARCYD_STATUS,
    ARCYD_REPOS,
    ARCYD_STATISTICS,
]

ARCYD_STATUS_STARTING = 'starting'
ARCYD_STATUS_UPDATING = 'updating'
ARCYD_STATUS_SLEEPING = 'sleeping'
ARCYD_STATUS_STOPPED = 'stopped'
ARCYD_STATUS_IDLE = 'idle'

ARCYD_LIST_STATUS = [
    ARCYD_STATUS_UPDATING,
    ARCYD_STATUS_SLEEPING,
    ARCYD_STATUS_STARTING,
    ARCYD_STATUS_STOPPED,
    ARCYD_STATUS_IDLE,
]

REPO_ATTRIB_NAME = 'name'
REPO_ATTRIB_HUMAN_NAME = 'human-name'
REPO_ATTRIB_STATUS = 'status'

REPO_LIST_ATTRIB = [
    REPO_ATTRIB_NAME,
    REPO_ATTRIB_HUMAN_NAME,
    REPO_ATTRIB_STATUS,
]

REPO_STATUS_UPDATING = 'updating'
REPO_STATUS_FAILED = 'failed'
REPO_STATUS_OK = 'ok'

REPO_STATUSES = [
    REPO_STATUS_UPDATING,
    REPO_STATUS_FAILED,
    REPO_STATUS_OK,
]

ARCYD_STAT_CURRENT_CYCLE_TIME = 'current-cycle-time'
ARCYD_STAT_LAST_CYCLE_TIME = 'last-cycle-time'

ARCYD_LIST_STATISTICS = [
    ARCYD_STAT_CURRENT_CYCLE_TIME,
    ARCYD_STAT_LAST_CYCLE_TIME,
]

class SharedFileDictOutput(object):

    def __init__(self, filename):
        super(SharedFileDictOutput, self).__init__()
        self._filename = filename

    def write(self, d):
        assert isinstance(d, dict)
        with phlsys_fs.write_file_lock_context(self._filename) as f:
            f.write(json.dumps(d))

class SharedDictOutput(object):

    def __init__(self, shared_d):
        super(SharedDictOutput, self).__init__()
        self._shared_d = shared_d
        assert isinstance(self._shared_d, dict)

    def write(self, d):
        assert isinstance(d, dict)
        # copy contents to other dict
        self._shared_d.clear()
        self._shared_d.update(d)

class _CycleTimer(object):

    def __init__(self):
        self._current_start = None
        self._last_duration = None

    def start_cycle(self):
        self._current_start = datetime.datetime.utcnow()

    def stop_cycle(self):
        self._last_duration = self.current_duration()
        self._current_start = None

    def current_duration(self):
        if self._current_start is None:
            return None
        duration = datetime.datetime.utcnow() - self._current_start
        return duration.total_seconds()

    @property
    def last_duration(self):
        return self._last_duration

class ArcydReporter(object):

    def __init__(self, output):
        """Initialise a new reporter to report to the specified outputs.

        :output: output to write status to

        """
        super(ArcydReporter, self).__init__()
        self._output = output

        assert self._output

        self._repos = []
        self._repo = None

        self._cycle_timer = _CycleTimer()

        self._write_status(ARCYD_STATUS_STARTING)

    def start_sleep(self, duration):
        _ = duration  # NOQA
        self._cycle_timer.stop_cycle()
        self._write_status(ARCYD_STATUS_SLEEPING)

    def update_sleep(self, duration):
        _ = duration  # NOQA
        self._write_status(ARCYD_STATUS_SLEEPING)

    def finish_sleep(self):
        self._write_status(ARCYD_STATUS_IDLE)
        self._cycle_timer.start_cycle()

    def start_repo(self, name, human_name):
        self._repo = {
            REPO_ATTRIB_NAME: name,
            REPO_ATTRIB_HUMAN_NAME: human_name,
            REPO_ATTRIB_STATUS: REPO_STATUS_UPDATING,
        }
        self._write_status(ARCYD_STATUS_UPDATING)

    def fail_repo(self):
        self._repo[REPO_ATTRIB_STATUS] = REPO_STATUS_FAILED
        self._repos.append(self._repo)
        self._repo = None
        self._write_status(ARCYD_STATUS_UPDATING)

    def finish_repo(self):
        self._repo[REPO_ATTRIB_STATUS] = REPO_STATUS_OK
        self._repos.append(self._repo)
        self._repo = None
        self._write_status(ARCYD_STATUS_IDLE)

    def _write_status(self, status):
        timer = self._cycle_timer
        statistics = {
            ARCYD_STAT_CURRENT_CYCLE_TIME: timer.current_duration(),
            ARCYD_STAT_LAST_CYCLE_TIME: timer.last_duration,
        }
        d = {
            ARCYD_STATUS: status,
            ARCYD_CURRENT_REPO: self._repo,
            ARCYD_REPOS: self._repos,
            ARCYD_STATISTICS: statistics,
        }
        assert set(d.keys()) == set(ARCYD_LIST_ATTRIB)
        self._output.write(d)

    def close(self):
        self._write_status(ARCYD_STATUS_STOPPED)

#------------------------------------------------------------------------------
# Copyright (C) 2012 Bloomberg L.P.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
#------------------------------- END-OF-FILE ----------------------------------
<|endoftext|>


<commit_before># vim: set et ts=4 sw=4 fdm=marker
"""
MIT License

Copyright (c) 2016 Jesse Hogan

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""
from pdb import set_trace; B=set_trace
from random import randint
class entities(object):
    def __init__(self, initial=None):
        self.clear()
        if initial != None:
            self.append(initial)

    def clear(self):
        self._ls=[]

    def __call__(self, ix):
        """
        Allow collections to be called providing similar functionality to the
        way they can be indexed. The difference is that None is returned if ix
        is not a valid index. Only positve indexes are considered valid.
        """
        try:
            if ix < 0:
                raise IndexError
            return self[ix]
        except IndexError:
            return None

    def __iter__(self):
        for t in self._list:
            yield t

    def getrandom(self, returnIndex=False):
        if self.isempty: return None
        ix = randint(0, self.ubound)
        if returnIndex:
            return self[ix], ix
        else:
            return self[ix]

    def where(self, qry):
        if type(qry) == type:
            cls = self.__class__
            return cls([x for x in self if type(x) == qry])

        fn = qry
        es = entities()
        for e in self:
            if fn(e): es += e
        return es

    def sort(self, key):
        self._ls.sort(key=key)

    def sorted(self, key):
        es = type(self)()
        ls = sorted(self._ls, key=key)
        for e in ls:
            es += e
        return es

    def tail(self, number):
        es = type(self)()
        for e in self[-number:]:
            es += e
        return es

    def remove(self, e):
        if isinstance(e, entities):
            rms = e
        elif callable(e):
            rms = self.where(e)
        else:
            rms = [e]

        for i in range(self.ubound, -1, -1):
            for rm in rms:
                if rm is self[i]:
                    del self._list[i]
                    break

    def reversed(self):
        r = type(self)()
        for e in reversed(self._ls):
            r += e
        return r

    def reverse(self):
        self._ls.reverse()

    @property
    def ubound(self):
        if self.isempty: return None
        return self.count - 1

    def insert(self, ix, e):
        self.insertBefore(ix, e)

    def insertBefore(self, ix, e):
        self._ls.insert(ix, e)

    def insertAfter(self, ix, e):
        self._ls.insert(ix + 1, e)

    def move(self, srcix, dstix):
        # TODO: This is untested
        if srcix == dstix:
            raise Exception('Source and destination are the same: {}'.format((srcix, dstix)))

        e = self.pop(srcix)
        self.insert(dstix, e)

    def shift(self):
        return self._ls.pop(0)

    def unshift(self, t):
        # TODO: Return entities object to indicate what was unshifted
        return self._ls.insert(0, t)

    def pop(self, ix=None):
        if ix == None: return self._ls.pop()
        return self._ls.pop(ix)

    def push(self, e):
        self += e

    def __lshift__(self, a):
        self.unshift(a)

    def append(self, obj, uniq=False, r=None):
        if not r: r = []
        if isinstance(obj, entity):
            t = obj
        elif isinstance(obj, entities) or hasattr(obj, '__iter__'):
            for t in obj:
                if uniq:
                    for t1 in self:
                        if t is t1: break
                    else:
                        self.append(t, r=r)
                        continue
                    break
                else:
                    self.append(t, r=r)
            return r
        else:
            raise ValueError('Unsupported object appended')

        if uniq:
            for t1 in self:
                if t is t1: return r

        r.append(t)
        for t in r:
            self._list.append(t)

        return r

    def __iadd__(self, t):
        self.append(t)
        return self

    def __iand__(self, t):
        self.append(t, uniq=True)
        return self

    def add(self, ts):
        self += ts
        return self

    def __add__(self, t):
        return self.add(t)

    @property
    def _list(self):
        if not hasattr(self, '_ls'):
            self._ls = []
        return self._ls

    @property
    def count(self):
        return len(self._list)

    def __len__(self):
        return self.count

    @property
    def isempty(self):
        return self.count == 0

    @property
    def hasone(self):
        return self.count == 1

    @property
    def ispopulated(self):
        return not self.isempty

    def __str__(self):
        if not self.isempty:
            r=''
            for i, t in enumerate(self):
                if i > 0: r += "\n"
                r += str(t)
            return r
        return ''

    def __setitem__(self, key, item):
        self._ls[key]=item

    def __getitem__(self, key):
        if type(key) == int or type(key) == slice:
            return self._list[key]

        for e in self._list:
            if hasattr(e, 'id'):
                if e.id == key:   return e
            elif hasattr(e, 'name'):
                if e.name == key: return e

    def getindex(self, e):
        """ Return the firsnt index of e in the collection.

        This is similar to list.index except here we use the `is` operator for
        comarison instead of the `==` operator."""

        for ix, e1 in enumerate(self):
            if e is e1: return ix
        raise ValueError("'{}' is not in the collection " + e)

    def _getbyindex(self, ix):
        # TODO This can be replace with a call to enttities.__call__
        try: return self[ix]
        except IndexError: return None

    @property
    def first(self): return self._getbyindex(0)

    @property
    def second(self): return self._getbyindex(1)

    @second.setter
    def second(self, v): self[1] = v

    @property
    def third(self): return self._getbyindex(2)

    @property
    def fourth(self): return self._getbyindex(3)

    @property
    def fifth(self): return self._getbyindex(4)

    @property
    def sixth(self): return self._getbyindex(5)

    @property
    def last(self): return self._getbyindex(-1)

    @last.setter
    def last(self, v): self[-1] = v

    @property
    def brokenrules(self):
        r = brokenrules()
        for ent in self:
            r += ent.brokenrules
        return r

    @property
    def isvalid(self):
        return self.brokenrules.isempty

class entity():
    def __init__(self):
        pass

    def add(self, t):
        th = entities()
        th += self
        th += t
        return th

    def isin(self, es):
        """Test if self is in entities object `es`. This is like the the `in`
        operator (__contains__()) except it tests for object identity (with `is`)
        instead of object equality (`==` or __eq__()).
        """
        for e in es:
            if self is e: return True
        return False

    def __add__(self, t):
        return self.add(t)

    @property
    def brokenrules(self):
        return brokenrules()

    @property
    def isvalid(self):
        return self.brokenrules.isempty

class brokenrules(entities):
    def append(self, o, r=None):
        if isinstance(o, str):
            o = brokenrule(o)
        super().append(o, r)

class brokenrule(entity):
    def __init__(self, msg):
        self.message = msg

    def __str__(self):
        return self.message

class event(entities):
    def __call__(self, src, e):
        for callable in self:
            callable(src, e)

    def append(self, fn):
        if not callable(fn):
            raise ValueError('Event must be callable')
        self._list.append(fn)

class eventargs(entity):
    pass

<commit_msg>Reimplement e.isin() using any()<commit_after># vim: set et ts=4 sw=4 fdm=marker
"""
MIT License

Copyright (c) 2016 Jesse Hogan

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
"""
from pdb import set_trace; B=set_trace
from random import randint
class entities(object):
    def __init__(self, initial=None):
        self.clear()
        if initial != None:
            self.append(initial)

    def clear(self):
        self._ls=[]

    def __call__(self, ix):
        """
        Allow collections to be called providing similar functionality to the
        way they can be indexed. The difference is that None is returned if ix
        is not a valid index. Only positve indexes are considered valid.
        """
        try:
            if ix < 0:
                raise IndexError
            return self[ix]
        except IndexError:
            return None

    def __iter__(self):
        for t in self._list:
            yield t

    def getrandom(self, returnIndex=False):
        if self.isempty: return None
        ix = randint(0, self.ubound)
        if returnIndex:
            return self[ix], ix
        else:
            return self[ix]

    def where(self, qry):
        if type(qry) == type:
            cls = self.__class__
            return cls([x for x in self if type(x) == qry])

        fn = qry
        es = entities()
        for e in self:
            if fn(e): es += e
        return es

    def sort(self, key):
        self._ls.sort(key=key)

    def sorted(self, key):
        es = type(self)()
        ls = sorted(self._ls, key=key)
        for e in ls:
            es += e
        return es

    def tail(self, number):
        es = type(self)()
        for e in self[-number:]:
            es += e
        return es

    def remove(self, e):
        if isinstance(e, entities):
            rms = e
        elif callable(e):
            rms = self.where(e)
        else:
            rms = [e]

        for i in range(self.ubound, -1, -1):
            for rm in rms:
                if rm is self[i]:
                    del self._list[i]
                    break

    def reversed(self):
        r = type(self)()
        for e in reversed(self._ls):
            r += e
        return r

    def reverse(self):
        self._ls.reverse()

    @property
    def ubound(self):
        if self.isempty: return None
        return self.count - 1

    def insert(self, ix, e):
        self.insertBefore(ix, e)

    def insertBefore(self, ix, e):
        self._ls.insert(ix, e)

    def insertAfter(self, ix, e):
        self._ls.insert(ix + 1, e)

    def move(self, srcix, dstix):
        # TODO: This is untested
        if srcix == dstix:
            raise Exception('Source and destination are the same: {}'.format((srcix, dstix)))

        e = self.pop(srcix)
        self.insert(dstix, e)

    def shift(self):
        return self._ls.pop(0)

    def unshift(self, t):
        # TODO: Return entities object to indicate what was unshifted
        return self._ls.insert(0, t)

    def pop(self, ix=None):
        if ix == None: return self._ls.pop()
        return self._ls.pop(ix)

    def push(self, e):
        self += e

    def __lshift__(self, a):
        self.unshift(a)

    def append(self, obj, uniq=False, r=None):
        if not r: r = []
        if isinstance(obj, entity):
            t = obj
        elif isinstance(obj, entities) or hasattr(obj, '__iter__'):
            for t in obj:
                if uniq:
                    for t1 in self:
                        if t is t1: break
                    else:
                        self.append(t, r=r)
                        continue
                    break
                else:
                    self.append(t, r=r)
            return r
        else:
            raise ValueError('Unsupported object appended')

        if uniq:
            for t1 in self:
                if t is t1: return r

        r.append(t)
        for t in r:
            self._list.append(t)

        return r

    def __iadd__(self, t):
        self.append(t)
        return self

    def __iand__(self, t):
        self.append(t, uniq=True)
        return self

    def add(self, ts):
        self += ts
        return self

    def __add__(self, t):
        return self.add(t)

    @property
    def _list(self):
        if not hasattr(self, '_ls'):
            self._ls = []
        return self._ls

    @property
    def count(self):
        return len(self._list)

    def __len__(self):
        return self.count

    @property
    def isempty(self):
        return self.count == 0

    @property
    def hasone(self):
        return self.count == 1

    @property
    def ispopulated(self):
        return not self.isempty

    def __str__(self):
        if not self.isempty:
            r=''
            for i, t in enumerate(self):
                if i > 0: r += "\n"
                r += str(t)
            return r
        return ''

    def __setitem__(self, key, item):
        self._ls[key]=item

    def __getitem__(self, key):
        if type(key) == int or type(key) == slice:
            return self._list[key]

        for e in self._list:
            if hasattr(e, 'id'):
                if e.id == key:   return e
            elif hasattr(e, 'name'):
                if e.name == key: return e

    def getindex(self, e):
        """ Return the firsnt index of e in the collection.

        This is similar to list.index except here we use the `is` operator for
        comarison instead of the `==` operator."""

        for ix, e1 in enumerate(self):
            if e is e1: return ix
        raise ValueError("'{}' is not in the collection " + e)

    def _getbyindex(self, ix):
        # TODO This can be replace with a call to enttities.__call__
        try: return self[ix]
        except IndexError: return None

    @property
    def first(self): return self._getbyindex(0)

    @property
    def second(self): return self._getbyindex(1)

    @second.setter
    def second(self, v): self[1] = v

    @property
    def third(self): return self._getbyindex(2)

    @property
    def fourth(self): return self._getbyindex(3)

    @property
    def fifth(self): return self._getbyindex(4)

    @property
    def sixth(self): return self._getbyindex(5)

    @property
    def last(self): return self._getbyindex(-1)

    @last.setter
    def last(self, v): self[-1] = v

    @property
    def brokenrules(self):
        r = brokenrules()
        for ent in self:
            r += ent.brokenrules
        return r

    @property
    def isvalid(self):
        return self.brokenrules.isempty

class entity():
    def __init__(self):
        pass

    def add(self, t):
        th = entities()
        th += self
        th += t
        return th

    def isin(self, es):
        """Test if self is in entities object `es`. This is like the the `in`
        operator (__contains__()) except it tests for object identity (with `is`)
        instead of object equality (`==` or __eq__()).
        """
        return any(self is e for e in es)

    def __add__(self, t):
        return self.add(t)

    @property
    def brokenrules(self):
        return brokenrules()

    @property
    def isvalid(self):
        return self.brokenrules.isempty

class brokenrules(entities):
    def append(self, o, r=None):
        if isinstance(o, str):
            o = brokenrule(o)
        super().append(o, r)

class brokenrule(entity):
    def __init__(self, msg):
        self.message = msg

    def __str__(self):
        return self.message

class event(entities):
    def __call__(self, src, e):
        for callable in self:
            callable(src, e)

    def append(self, fn):
        if not callable(fn):
            raise ValueError('Event must be callable')
        self._list.append(fn)

class eventargs(entity):
    pass

<|endoftext|>


<commit_before>"""
Implementations of the Harmonic Balance Method.

Mousai provides tools for harmonic balance solutions.

Solvers can be coupled with external codes through connectivity modules.

Documentation is available in the docstrings and online at
https://github.com/josephcslater/mousai.
Joseph C. Slater

.. code-block:: python

    >>> import mousai as ms

Copyright (c) Joseph C. Slater

Distributed under the terms of the BSD-3-Clause License.
"""

__title__ = 'mousai'
# 1) Version may have other than numerical digits after decimal point.
# 2) 1.11 is actually a higher release than 1.2 (confusing).
# 3) Let's just increment with single digits.

# Moved to Semantic Versioning. 0.3 was last non-semantic release.
__version__ = '0.3.0'
__author__ = 'Joseph C. Slater'
__license__ = 'BSD-3-Clause'
__copyright__ = 'Copyright 2017 Joseph C. Slater'
# __name__ = 'mousai'
# __package__ = 'mousai'

import sys
import matplotlib as mpl

from .har_bal import *

# Accomodate using Travis-ci and potential matlotlib results.
if 'pytest' in sys.argv[0]:
    print('Setting backend to agg to run tests')
    mpl.use('agg')

# This would have value for using
# from mousai import *
# That's bad practice and I'd rather not support it
# __all__ = ['har_bal', '__version__']
<commit_msg>semantic versioning incorporated.<commit_after>"""
Implementations of the Harmonic Balance Method.

Mousai provides tools for harmonic balance solutions.

Solvers can be coupled with external codes through connectivity modules.

Documentation is available in the docstrings and online at
https://github.com/josephcslater/mousai.
Joseph C. Slater

.. code-block:: python

    >>> import mousai as ms

Copyright (c) Joseph C. Slater

Distributed under the terms of the BSD-3-Clause License.
"""

__title__ = 'mousai'
# 1) Version may have other than numerical digits after decimal point.
# 2) 1.11 is actually a higher release than 1.2 (confusing).
# 3) Let's just increment with single digits.

# Moved to Semantic Versioning. 0.3 was last non-semantic release.
__version__ = '0.3.0'
__author__ = 'Joseph C. Slater'
__license__ = 'BSD-3-Clause'
__copyright__ = 'Copyright 2017 Joseph C. Slater'
# __name__ = 'mousai'
# __package__ = 'mousai'

import sys
import matplotlib as mpl

from .har_bal import *

# Accomodate using Travis-ci and potential matlotlib results.
if 'pytest' in sys.argv[0]:
    print('Setting backend to agg to run tests')
    mpl.use('agg')

# This would have value for using
# from mousai import *
# That's bad practice and I'd rather not support it
# __all__ = ['har_bal', '__version__']
<|endoftext|>


<commit_before>
"""
Fit redshifts and classifications on DESI bricks
"""

import sys
import os
import numpy as np
import multiprocessing

from desispec import io
from desispec.interpolation import resample_flux
from desispec.log import get_logger
from desispec.zfind.redmonster import RedMonsterZfind
from desispec.zfind import ZfindBase
from desispec.io.qa import load_qa_brick, write_qa_brick

import argparse

def parse(options=None):
    default_nproc = max(1, multiprocessing.cpu_count() // 2)

    parser = argparse.ArgumentParser(description="Fit redshifts and classifications on bricks.")
    parser.add_argument("-b", "--brick", type=str, required=False,
        help="input brickname")
    parser.add_argument("-n", "--nspec", type=int, required=False,
        help="number of spectra to fit [default: all]")
    parser.add_argument("-o", "--outfile", type=str, required=False,
        help="output file name")
    parser.add_argument("--objtype", type=str, required=False,
        help="only use templates for these objtypes (comma separated elg,lrg,qso,star)")
    parser.add_argument("--zspec", action="store_true",
        help="also include spectra in output file")
    parser.add_argument('--qafile', type=str,
        help='path of QA file.')
    parser.add_argument('--qafig', type=str,
        help='path of QA figure file')
    parser.add_argument("--nproc", type=int, default=default_nproc,
        help="number of parallel processes for multiprocessing")
    parser.add_argument("brickfiles", nargs="*")

    args = None
    if options is None:
        args = parser.parse_args()
    else:
        args = parser.parse_args(options)
    return args

#- function for multiprocessing
def _func(arg) :
    return RedMonsterZfind(**arg)

def main(args) :

    log = get_logger()

    if args.objtype is not None:
        args.objtype = args.objtype.split(',')

    #- Read brick files for each channel
    log.info("Reading bricks")
    brick = dict()
    if args.brick is not None:
        if len(args.brickfiles) != 0:
            raise RuntimeError('Give -b/--brick or input brickfiles but not both')
        for channel in ('b', 'r', 'z'):
            filename = io.findfile('brick', band=channel, brickname=args.brick)
            brick[channel] = io.Brick(filename)
    else:
        if len(args.brickfiles) != 3:
            raise RuntimeError('Exactly 3 brickfiles should be given, or use --brick')
        for filename in args.brickfiles:
            bx = io.Brick(filename)
            if bx.channel not in brick:
                brick[bx.channel] = bx
            else:
                log.error('Channel {} in multiple input files'.format(bx.channel))
                sys.exit(2)

    filters=brick.keys()
    for fil in filters:
        log.info("Filter found: "+fil)

    #- Assume all channels have the same number of targets
    #- TODO: generalize this to allow missing channels
    if args.nspec is None:
        args.nspec = brick['b'].get_num_targets()
        log.info("Fitting {} targets".format(args.nspec))
    else:
        log.info("Fitting {} of {} targets".format(args.nspec, brick['b'].get_num_targets()))

    nspec = args.nspec

    #- Coadd individual exposures and combine channels
    #- Full coadd code is a bit slow, so try something quick and dirty for
    #- now to get something going for redshifting
    log.info("Combining individual channels and exposures")
    wave=[]
    for fil in filters:
        wave=np.concatenate([wave,brick[fil].get_wavelength_grid()])
    np.ndarray.sort(wave)
    nwave = len(wave)

    #- flux and ivar arrays to fill for all targets
    flux = np.zeros((nspec, nwave))
    ivar = np.zeros((nspec, nwave))
    targetids = brick['b'].get_target_ids()[0:nspec]

    func_args = []

    for i, targetid in enumerate(targetids):
        #- wave, flux, and ivar for this target; concatenate
        xwave = list()
        xflux = list()
        xivar = list()
        for channel in filters:
            exp_flux, exp_ivar, resolution, info = brick[channel].get_target(targetid)
            weights = np.sum(exp_ivar, axis=0)
            ii, = np.where(weights > 0)
            xwave.extend(brick[channel].get_wavelength_grid()[ii])
            #- Average multiple exposures on the same wavelength grid for each channel
            xflux.extend(np.average(exp_flux[:,ii], weights=exp_ivar[:,ii], axis=0))
            xivar.extend(weights[ii])

        xwave = np.array(xwave)
        xivar = np.array(xivar)
        xflux = np.array(xflux)

        ii = np.argsort(xwave)
        flux[i], ivar[i] = resample_flux(wave, xwave[ii], xflux[ii], xivar[ii])

    #- distribute the spectra in nspec groups
    if args.nproc > nspec:
        args.nproc = nspec

    ii = np.linspace(0, nspec, args.nproc+1).astype(int)
    for i in range(args.nproc):
        lo, hi = ii[i], ii[i+1]
        log.debug('CPU {} spectra {}:{}'.format(i, lo, hi))
        arguments={"wave":wave, "flux":flux[lo:hi],"ivar":ivar[lo:hi],"objtype":args.objtype}
        func_args.append( arguments )

    #- Do the redshift fit

    if args.nproc==1 : # No parallelization, simple loop over arguments

        zf = list()
        for arg in func_args:
            zff = RedMonsterZfind(**arg)
            zf.append(zff)

    else: # Multiprocessing
        log.info("starting multiprocessing with {} cpus for {} spectra in {} groups".format(args.nproc, nspec, len(func_args)))
        pool = multiprocessing.Pool(args.nproc)
        zf = pool.map(_func, func_args)
        pool.close()
        pool.join()

    # reformat results
    dtype = list()

    dtype = [
        ('Z',         zf[0].z.dtype),
        ('ZERR',      zf[0].zerr.dtype),
        ('ZWARN',     zf[0].zwarn.dtype),
        ('TYPE',      zf[0].type.dtype),
        ('SUBTYPE',   zf[0].subtype.dtype),
    ]

    formatted_data = np.empty(nspec, dtype=dtype)

    i = 0
    for result in zf:
        n = result.nspec
        formatted_data['Z'][i:i+n]       = result.z
        formatted_data['ZERR'][i:i+n]    = result.zerr
        formatted_data['ZWARN'][i:i+n]   = result.zwarn
        formatted_data['TYPE'][i:i+n]    = result.type
        formatted_data['SUBTYPE'][i:i+n] = result.subtype
        i += n

    assert i == nspec

    # Create a ZfindBase object with formatted results
    zfi = ZfindBase(None, None, None, results=formatted_data)
    zfi.nspec = nspec

    # QA
    if (args.qafile is not None) or (args.qafig is not None):
        log.info("performing skysub QA")
        # Load
        qabrick = load_qa_brick(args.qafile)
        # Run
        qabrick.run_qa('ZBEST', (zfi,brick))
        # Write
        if args.qafile is not None:
            write_qa_brick(args.qafile, qabrick)
            log.info("successfully wrote {:s}".format(args.qafile))
        # Figure(s)
        if args.qafig is not None:
            raise IOError("Not yet implemented")
            qa_plots.brick_zbest(args.qafig, zfi, qabrick)

    #- Write some output
    if args.outfile is None:
        args.outfile = io.findfile('zbest', brickname=args.brick)

    log.info("Writing "+args.outfile)
    io.write_zbest(args.outfile, args.brick, targetids, zfi, zspec=args.zspec)

<commit_msg>Fix zfind logic.  When specifying the brick files manually, allow any combination of bands.<commit_after>
"""
Fit redshifts and classifications on DESI bricks
"""

import sys
import os
import numpy as np
import multiprocessing

from desispec import io
from desispec.interpolation import resample_flux
from desispec.log import get_logger
from desispec.zfind.redmonster import RedMonsterZfind
from desispec.zfind import ZfindBase
from desispec.io.qa import load_qa_brick, write_qa_brick

import argparse

def parse(options=None):
    default_nproc = max(1, multiprocessing.cpu_count() // 2)

    parser = argparse.ArgumentParser(description="Fit redshifts and classifications on bricks.")
    parser.add_argument("-b", "--brick", type=str, required=False,
        help="input brickname")
    parser.add_argument("-n", "--nspec", type=int, required=False,
        help="number of spectra to fit [default: all]")
    parser.add_argument("-o", "--outfile", type=str, required=False,
        help="output file name")
    parser.add_argument("--objtype", type=str, required=False,
        help="only use templates for these objtypes (comma separated elg,lrg,qso,star)")
    parser.add_argument("--zspec", action="store_true",
        help="also include spectra in output file")
    parser.add_argument('--qafile', type=str,
        help='path of QA file.')
    parser.add_argument('--qafig', type=str,
        help='path of QA figure file')
    parser.add_argument("--nproc", type=int, default=default_nproc,
        help="number of parallel processes for multiprocessing")
    parser.add_argument("brickfiles", nargs="*")

    args = None
    if options is None:
        args = parser.parse_args()
    else:
        args = parser.parse_args(options)
    return args

#- function for multiprocessing
def _func(arg) :
    return RedMonsterZfind(**arg)

def main(args) :

    log = get_logger()

    if args.objtype is not None:
        args.objtype = args.objtype.split(',')

    #- Read brick files for each channel
    log.info("Reading bricks")
    brick = dict()
    if args.brick is not None:
        if len(args.brickfiles) != 0:
            raise RuntimeError('Give -b/--brick or input brickfiles but not both')
        for channel in ('b', 'r', 'z'):
            filename = io.findfile('brick', band=channel, brickname=args.brick)
            brick[channel] = io.Brick(filename)
    else:
        for filename in args.brickfiles:
            bx = io.Brick(filename)
            if bx.channel not in brick:
                brick[bx.channel] = bx
            else:
                log.error('Channel {} in multiple input files'.format(bx.channel))
                sys.exit(2)

    filters=brick.keys()
    for fil in filters:
        log.info("Filter found: "+fil)

    #- Assume all channels have the same number of targets
    #- TODO: generalize this to allow missing channels
    if args.nspec is None:
        args.nspec = brick['b'].get_num_targets()
        log.info("Fitting {} targets".format(args.nspec))
    else:
        log.info("Fitting {} of {} targets".format(args.nspec, brick['b'].get_num_targets()))

    nspec = args.nspec

    #- Coadd individual exposures and combine channels
    #- Full coadd code is a bit slow, so try something quick and dirty for
    #- now to get something going for redshifting
    log.info("Combining individual channels and exposures")
    wave=[]
    for fil in filters:
        wave=np.concatenate([wave,brick[fil].get_wavelength_grid()])
    np.ndarray.sort(wave)
    nwave = len(wave)

    #- flux and ivar arrays to fill for all targets
    flux = np.zeros((nspec, nwave))
    ivar = np.zeros((nspec, nwave))
    targetids = brick['b'].get_target_ids()[0:nspec]

    func_args = []

    for i, targetid in enumerate(targetids):
        #- wave, flux, and ivar for this target; concatenate
        xwave = list()
        xflux = list()
        xivar = list()
        for channel in filters:
            exp_flux, exp_ivar, resolution, info = brick[channel].get_target(targetid)
            weights = np.sum(exp_ivar, axis=0)
            ii, = np.where(weights > 0)
            xwave.extend(brick[channel].get_wavelength_grid()[ii])
            #- Average multiple exposures on the same wavelength grid for each channel
            xflux.extend(np.average(exp_flux[:,ii], weights=exp_ivar[:,ii], axis=0))
            xivar.extend(weights[ii])

        xwave = np.array(xwave)
        xivar = np.array(xivar)
        xflux = np.array(xflux)

        ii = np.argsort(xwave)
        flux[i], ivar[i] = resample_flux(wave, xwave[ii], xflux[ii], xivar[ii])

    #- distribute the spectra in nspec groups
    if args.nproc > nspec:
        args.nproc = nspec

    ii = np.linspace(0, nspec, args.nproc+1).astype(int)
    for i in range(args.nproc):
        lo, hi = ii[i], ii[i+1]
        log.debug('CPU {} spectra {}:{}'.format(i, lo, hi))
        arguments={"wave":wave, "flux":flux[lo:hi],"ivar":ivar[lo:hi],"objtype":args.objtype}
        func_args.append( arguments )

    #- Do the redshift fit

    if args.nproc==1 : # No parallelization, simple loop over arguments

        zf = list()
        for arg in func_args:
            zff = RedMonsterZfind(**arg)
            zf.append(zff)

    else: # Multiprocessing
        log.info("starting multiprocessing with {} cpus for {} spectra in {} groups".format(args.nproc, nspec, len(func_args)))
        pool = multiprocessing.Pool(args.nproc)
        zf = pool.map(_func, func_args)
        pool.close()
        pool.join()

    # reformat results
    dtype = list()

    dtype = [
        ('Z',         zf[0].z.dtype),
        ('ZERR',      zf[0].zerr.dtype),
        ('ZWARN',     zf[0].zwarn.dtype),
        ('TYPE',      zf[0].type.dtype),
        ('SUBTYPE',   zf[0].subtype.dtype),
    ]

    formatted_data = np.empty(nspec, dtype=dtype)

    i = 0
    for result in zf:
        n = result.nspec
        formatted_data['Z'][i:i+n]       = result.z
        formatted_data['ZERR'][i:i+n]    = result.zerr
        formatted_data['ZWARN'][i:i+n]   = result.zwarn
        formatted_data['TYPE'][i:i+n]    = result.type
        formatted_data['SUBTYPE'][i:i+n] = result.subtype
        i += n

    assert i == nspec

    # Create a ZfindBase object with formatted results
    zfi = ZfindBase(None, None, None, results=formatted_data)
    zfi.nspec = nspec

    # QA
    if (args.qafile is not None) or (args.qafig is not None):
        log.info("performing skysub QA")
        # Load
        qabrick = load_qa_brick(args.qafile)
        # Run
        qabrick.run_qa('ZBEST', (zfi,brick))
        # Write
        if args.qafile is not None:
            write_qa_brick(args.qafile, qabrick)
            log.info("successfully wrote {:s}".format(args.qafile))
        # Figure(s)
        if args.qafig is not None:
            raise IOError("Not yet implemented")
            qa_plots.brick_zbest(args.qafig, zfi, qabrick)

    #- Write some output
    if args.outfile is None:
        args.outfile = io.findfile('zbest', brickname=args.brick)

    log.info("Writing "+args.outfile)
    io.write_zbest(args.outfile, args.brick, targetids, zfi, zspec=args.zspec)

<|endoftext|>