<commit_before>"""
django-helpdesk - A Django powered ticket tracker for small enterprise.

(c) Copyright 2008 Jutda. All Rights Reserved. See LICENSE for details.

lib.py - Common functions (eg multipart e-mail)
"""

chart_colours = ('80C65A', '990066', 'FF9900', '3399CC', 'BBCCED', '3399CC', 'FFCC33')

try:
    from base64 import urlsafe_b64encode as b64encode
except ImportError:
    from base64 import encodestring as b64encode
try:
    from base64 import urlsafe_b64decode as b64decode
except ImportError:
    from base64 import decodestring as b64decode

def send_templated_mail(template_name, email_context, recipients, sender=None, bcc=None, fail_silently=False, files=None):
    """
    send_templated_mail() is a warpper around Django's e-mail routines that
    allows us to easily send multipart (text/plain & text/html) e-mails using
    templates that are stored in the database. This lets the admin provide
    both a text and a HTML template for each message.

    template_name is the slug of the template to use for this message (see
        models.EmailTemplate)

    email_context is a dictionary to be used when rendering the template

    recipients can be either a string, eg 'a@b.com', or a list of strings.

    sender should contain a string, eg 'My Site <me@z.com>'. If you leave it
        blank, it'll use settings.DEFAULT_FROM_EMAIL as a fallback.

    bcc is an optional list of addresses that will receive this message as a
        blind carbon copy.

    fail_silently is passed to Django's mail routine. Set to 'True' to ignore
        any errors at send time.

    files can be a list of file paths to be attached, or it can be left blank.
        eg ('/tmp/file1.txt', '/tmp/image.png')

    """
    from django.conf import settings
    from django.core.mail import EmailMultiAlternatives
    from django.template import loader, Context

    from helpdesk.models import EmailTemplate
    import os

    context = Context(email_context)
    locale = getattr(context['queue'], 'locale', '')
    if not locale:
        locale = 'en'

    t = None
    try:
        t = EmailTemplate.objects.get(template_name__iexact=template_name, locale=locale)
    except EmailTemplate.DoesNotExist:
        pass

    if not t:
        try:
            t = EmailTemplate.objects.get(template_name__iexact=template_name, locale__isnull=True)
        except EmailTemplate.DoesNotExist:
            return # just ignore if template doesn't exist

    if not sender:
        sender = settings.DEFAULT_FROM_EMAIL

    footer_file = os.path.join('helpdesk', locale, 'email_text_footer.txt')

    text_part = loader.get_template_from_string(
        "%s{%% include '%s' %%}" % (t.plain_text, footer_file)
        ).render(context)

    email_html_base_file = os.path.join('helpdesk', locale, 'email_html_base.html')

    ''' keep new lines in html emails '''
    from django.utils.safestring import mark_safe

    if context.has_key('comment'):
        html_txt = context['comment']
        html_txt = html_txt.replace('\r\n', '
')
        context['comment'] = mark_safe(html_txt)

    html_part = loader.get_template_from_string(
        "{%% extends '%s' %%}{%% block title %%}%s{%% endblock %%}{%% block content %%}%s{%% endblock %%}" % (email_html_base_file, t.heading, t.html)
        ).render(context)

    subject_part = loader.get_template_from_string(
        "{{ ticket.ticket }} {{ ticket.title|safe }} %s" % t.subject
        ).render(context)

    if type(recipients) == str:
        if recipients.find(','):
            recipients = recipients.split(',')
    elif type(recipients) != list:
        recipients = [recipients,]

    msg = EmailMultiAlternatives(   subject_part,
                                    text_part,
                                    sender,
                                    recipients,
                                    bcc=bcc)
    msg.attach_alternative(html_part, "text/html")

    if files:
        if type(files) != list:
            files = [files,]

        for file in files:
            msg.attach_file(file)

    return msg.send(fail_silently)

def normalise_data(data, to=100):
    """
    Used for normalising data prior to graphing with Google charting API. EG:

        [1, 4, 10] becomes [10, 40, 100]
        [36, 54, 240] becomes [15, 23, 100]
    """
    max_value = max(data)
    if max_value > to:
        new_data = []
        for d in data:
            new_data.append(int(d/float(max_value)*to))
        data = new_data
    return data

def line_chart(data):
    """
    'data' is a list of lists making a table.
    Row 1, columns 2-n are data headings (the time periods)
    Rows 2-n are data, with column 1 being the line labels
    """

    column_headings = data[0][1:]
    max = 0
    for row in data[1:]:
        for field in row[1:]:
            if field > max:
                max = field

    # Set width to '65px * number of months + 100 for headings.'.
    chart_url = 'http://chart.apis.google.com/chart?cht=lc&chs=%sx150&chd=t:' % (min(len(column_headings)*65+100, 1000))
    first_row = True
    row_headings = []
    for row in data[1:]:
        # Add data to URL, normalised to the maximum for all lines on this chart
        norm = normalise_data(row[1:], max)
        if not first_row:
            chart_url += '|'
        chart_url += ','.join([str(num) for num in norm])
        row_headings.append(row[0])
        first_row = False

    chart_url += '&chds='
    rows = len(data)-1
    first = True
    for row in range(rows):
        # Set maximum data ranges to '0:x' where 'x' is the maximum number in use.
        if not first:
            chart_url += ','
        chart_url += '0,%s' % max
        first = False
    chart_url += '&chdl=%s' % '|'.join(row_headings) # Display legend/labels
    chart_url += '&chco=%s' % ','.join(chart_colours) # Default colour set
    chart_url += '&chxt=x,y' # Turn on axis labels
    chart_url += '&chxl=0:|%s|1:|0|%s' % ('|'.join(column_headings), max) # Axis Label Text

    return chart_url

def bar_chart(data):
    """
    'data' is a list of lists making a table.
    Row 1, columns 2-n are data headings
    Rows 2-n are data, with column 1 being the line labels
    """

    column_headings = data[0][1:]
    max = 0
    for row in data[1:]:
        for field in row[1:]:
            if field > max:
                max = field

    # Set width to '220px * number of months'.
    chart_url = 'http://chart.apis.google.com/chart?cht=bvg&chs=%sx150&chd=t:' % (min(len(column_headings) * 220, 1000))
    first_row = True
    row_headings = []
    for row in data[1:]:
        # Add data to URL, normalised to the maximum for all lines on this chart
        norm = normalise_data(row[1:], max)
        if not first_row:
            chart_url += '|'
        chart_url += ','.join([str(num) for num in norm])
        row_headings.append(row[0])
        first_row = False

    chart_url += '&chds=0,%s' % max
    chart_url += '&chdl=%s' % '|'.join(row_headings) # Display legend/labels
    chart_url += '&chco=%s' % ','.join(chart_colours) # Default colour set
    chart_url += '&chxt=x,y' # Turn on axis labels
    chart_url += '&chxl=0:|%s|1:|0|%s' % ('|'.join(column_headings), max) # Axis Label Text

    return chart_url

def query_to_dict(results, descriptions):
    """
    Replacement method for cursor.dictfetchall() as that method no longer
    exists in psycopg2, and I'm guessing in other backends too.

    Converts the results of a raw SQL query into a list of dictionaries, suitable
    for use in templates etc.
    """

    output = []
    for data in results:
        row = {}
        i = 0
        for column in descriptions:
            row[column[0]] = data[i]
            i += 1

        output.append(row)
    return output

def apply_query(queryset, params):
    """
    Apply a dict-based set of filters & paramaters to a queryset.

    queryset is a Django queryset, eg MyModel.objects.all() or
             MyModel.objects.filter(user=request.user)

    params is a dictionary that contains the following:
        filtering: A dict of Django ORM filters, eg:
            {'user__id__in': [1, 3, 103], 'title__contains': 'foo'}
        other_filter: Another filter of some type, most likely a
            set of Q() objects.
        sorting: The name of the column to sort by
    """
    for key in params['filtering'].keys():
        filter = {key: params['filtering'][key]}
        queryset = queryset.filter(**filter)

    if params.get('other_filter', None):
        # eg a Q() set
        queryset = queryset.filter(params['other_filter'])

    if params.get('sorting', None):
        if params.get('sortreverse', None):
            params['sorting'] = "-%s" % params['sorting']
        queryset = queryset.order_by(params['sorting'])

    return queryset

def safe_template_context(ticket):
    """
    Return a dictionary that can be used as a template context to render
    comments and other details with ticket or queue paramaters. Note that
    we don't just provide the Ticket & Queue objects to the template as
    they could reveal confidential information. Just imagine these two options:
        * {{ ticket.queue.email_box_password }}
        * {{ ticket.assigned_to.password }}

    Ouch!

    The downside to this is that if we make changes to the model, we will also
    have to update this code. Perhaps we can find a better way in the future.
    """

    context = {
        'queue': {},
        'ticket': {},
        }
    queue = ticket.queue

    for field in (  'title', 'slug', 'email_address', 'from_address'):
        attr = getattr(queue, field, None)
        if callable(attr):
            context['queue'][field] = attr()
        else:
            context['queue'][field] = attr

    for field in (  'title', 'created', 'modified', 'submitter_email',
                    'status', 'get_status_display', 'on_hold', 'description',
                    'resolution', 'priority', 'get_priority_display',
                    'last_escalation', 'ticket', 'ticket_for_url',
                    'get_status', 'ticket_url', 'staff_url', '_get_assigned_to'
                 ):
        attr = getattr(ticket, field, None)
        if callable(attr):
            context['ticket'][field] = '%s' % attr()
        else:
            context['ticket'][field] = attr

    context['ticket']['queue'] = context['queue']
    context['ticket']['assigned_to'] = context['ticket']['_get_assigned_to']

    return context

def text_is_spam(text, request):
    # Based on a blog post by 'sciyoshi':
    # http://sciyoshi.com/blog/2008/aug/27/using-akismet-djangos-new-comments-framework/
    # This will return 'True' is the given text is deemed to be spam, or
    # False if it is not spam. If it cannot be checked for some reason, we
    # assume it isn't spam.
    from django.contrib.sites.models import Site
    from django.conf import settings
    try:
        from helpdesk.akismet import Akismet
    except:
        return False

    ak = Akismet(
        blog_url='http://%s/' % Site.objects.get(pk=settings.SITE_ID).domain,
        agent='django-helpdesk',
    )

    if hasattr(settings, 'TYPEPAD_ANTISPAM_API_KEY'):
        ak.setAPIKey(key = settings.TYPEPAD_ANTISPAM_API_KEY)
        ak.baseurl = 'api.antispam.typepad.com/1.1/'
    elif hasattr(settings, 'AKISMET_API_KEY'):
        ak.setAPIKey(key = settings.AKISMET_API_KEY)
    else:
        return False

    if ak.verify_key():
        ak_data = {
            'user_ip': request.META.get('REMOTE_ADDR', '127.0.0.1'),
            'user_agent': request.META.get('HTTP_USER_AGENT', ''),
            'referrer': request.META.get('HTTP_REFERER', ''),
            'comment_type': 'comment',
            'comment_author': '',
        }

        return ak.comment_check(text, data=ak_data)

    return False
<commit_msg>Ensure queue locale is included in safe_template_context as reported by vvarp. This should close GH-56 without further issues - please reopen if required..<commit_after>"""
django-helpdesk - A Django powered ticket tracker for small enterprise.

(c) Copyright 2008 Jutda. All Rights Reserved. See LICENSE for details.

lib.py - Common functions (eg multipart e-mail)
"""

chart_colours = ('80C65A', '990066', 'FF9900', '3399CC', 'BBCCED', '3399CC', 'FFCC33')

try:
    from base64 import urlsafe_b64encode as b64encode
except ImportError:
    from base64 import encodestring as b64encode
try:
    from base64 import urlsafe_b64decode as b64decode
except ImportError:
    from base64 import decodestring as b64decode

def send_templated_mail(template_name, email_context, recipients, sender=None, bcc=None, fail_silently=False, files=None):
    """
    send_templated_mail() is a warpper around Django's e-mail routines that
    allows us to easily send multipart (text/plain & text/html) e-mails using
    templates that are stored in the database. This lets the admin provide
    both a text and a HTML template for each message.

    template_name is the slug of the template to use for this message (see
        models.EmailTemplate)

    email_context is a dictionary to be used when rendering the template

    recipients can be either a string, eg 'a@b.com', or a list of strings.

    sender should contain a string, eg 'My Site <me@z.com>'. If you leave it
        blank, it'll use settings.DEFAULT_FROM_EMAIL as a fallback.

    bcc is an optional list of addresses that will receive this message as a
        blind carbon copy.

    fail_silently is passed to Django's mail routine. Set to 'True' to ignore
        any errors at send time.

    files can be a list of file paths to be attached, or it can be left blank.
        eg ('/tmp/file1.txt', '/tmp/image.png')

    """
    from django.conf import settings
    from django.core.mail import EmailMultiAlternatives
    from django.template import loader, Context

    from helpdesk.models import EmailTemplate
    import os

    context = Context(email_context)
    locale = getattr(context['queue'], 'locale', '')
    if not locale:
        locale = 'en'

    t = None
    try:
        t = EmailTemplate.objects.get(template_name__iexact=template_name, locale=locale)
    except EmailTemplate.DoesNotExist:
        pass

    if not t:
        try:
            t = EmailTemplate.objects.get(template_name__iexact=template_name, locale__isnull=True)
        except EmailTemplate.DoesNotExist:
            return # just ignore if template doesn't exist

    if not sender:
        sender = settings.DEFAULT_FROM_EMAIL

    footer_file = os.path.join('helpdesk', locale, 'email_text_footer.txt')

    text_part = loader.get_template_from_string(
        "%s{%% include '%s' %%}" % (t.plain_text, footer_file)
        ).render(context)

    email_html_base_file = os.path.join('helpdesk', locale, 'email_html_base.html')

    ''' keep new lines in html emails '''
    from django.utils.safestring import mark_safe

    if context.has_key('comment'):
        html_txt = context['comment']
        html_txt = html_txt.replace('\r\n', '
')
        context['comment'] = mark_safe(html_txt)

    html_part = loader.get_template_from_string(
        "{%% extends '%s' %%}{%% block title %%}%s{%% endblock %%}{%% block content %%}%s{%% endblock %%}" % (email_html_base_file, t.heading, t.html)
        ).render(context)

    subject_part = loader.get_template_from_string(
        "{{ ticket.ticket }} {{ ticket.title|safe }} %s" % t.subject
        ).render(context)

    if type(recipients) == str:
        if recipients.find(','):
            recipients = recipients.split(',')
    elif type(recipients) != list:
        recipients = [recipients,]

    msg = EmailMultiAlternatives(   subject_part,
                                    text_part,
                                    sender,
                                    recipients,
                                    bcc=bcc)
    msg.attach_alternative(html_part, "text/html")

    if files:
        if type(files) != list:
            files = [files,]

        for file in files:
            msg.attach_file(file)

    return msg.send(fail_silently)

def normalise_data(data, to=100):
    """
    Used for normalising data prior to graphing with Google charting API. EG:

        [1, 4, 10] becomes [10, 40, 100]
        [36, 54, 240] becomes [15, 23, 100]
    """
    max_value = max(data)
    if max_value > to:
        new_data = []
        for d in data:
            new_data.append(int(d/float(max_value)*to))
        data = new_data
    return data

def line_chart(data):
    """
    'data' is a list of lists making a table.
    Row 1, columns 2-n are data headings (the time periods)
    Rows 2-n are data, with column 1 being the line labels
    """

    column_headings = data[0][1:]
    max = 0
    for row in data[1:]:
        for field in row[1:]:
            if field > max:
                max = field

    # Set width to '65px * number of months + 100 for headings.'.
    chart_url = 'http://chart.apis.google.com/chart?cht=lc&chs=%sx150&chd=t:' % (min(len(column_headings)*65+100, 1000))
    first_row = True
    row_headings = []
    for row in data[1:]:
        # Add data to URL, normalised to the maximum for all lines on this chart
        norm = normalise_data(row[1:], max)
        if not first_row:
            chart_url += '|'
        chart_url += ','.join([str(num) for num in norm])
        row_headings.append(row[0])
        first_row = False

    chart_url += '&chds='
    rows = len(data)-1
    first = True
    for row in range(rows):
        # Set maximum data ranges to '0:x' where 'x' is the maximum number in use.
        if not first:
            chart_url += ','
        chart_url += '0,%s' % max
        first = False
    chart_url += '&chdl=%s' % '|'.join(row_headings) # Display legend/labels
    chart_url += '&chco=%s' % ','.join(chart_colours) # Default colour set
    chart_url += '&chxt=x,y' # Turn on axis labels
    chart_url += '&chxl=0:|%s|1:|0|%s' % ('|'.join(column_headings), max) # Axis Label Text

    return chart_url

def bar_chart(data):
    """
    'data' is a list of lists making a table.
    Row 1, columns 2-n are data headings
    Rows 2-n are data, with column 1 being the line labels
    """

    column_headings = data[0][1:]
    max = 0
    for row in data[1:]:
        for field in row[1:]:
            if field > max:
                max = field

    # Set width to '220px * number of months'.
    chart_url = 'http://chart.apis.google.com/chart?cht=bvg&chs=%sx150&chd=t:' % (min(len(column_headings) * 220, 1000))
    first_row = True
    row_headings = []
    for row in data[1:]:
        # Add data to URL, normalised to the maximum for all lines on this chart
        norm = normalise_data(row[1:], max)
        if not first_row:
            chart_url += '|'
        chart_url += ','.join([str(num) for num in norm])
        row_headings.append(row[0])
        first_row = False

    chart_url += '&chds=0,%s' % max
    chart_url += '&chdl=%s' % '|'.join(row_headings) # Display legend/labels
    chart_url += '&chco=%s' % ','.join(chart_colours) # Default colour set
    chart_url += '&chxt=x,y' # Turn on axis labels
    chart_url += '&chxl=0:|%s|1:|0|%s' % ('|'.join(column_headings), max) # Axis Label Text

    return chart_url

def query_to_dict(results, descriptions):
    """
    Replacement method for cursor.dictfetchall() as that method no longer
    exists in psycopg2, and I'm guessing in other backends too.

    Converts the results of a raw SQL query into a list of dictionaries, suitable
    for use in templates etc.
    """

    output = []
    for data in results:
        row = {}
        i = 0
        for column in descriptions:
            row[column[0]] = data[i]
            i += 1

        output.append(row)
    return output

def apply_query(queryset, params):
    """
    Apply a dict-based set of filters & paramaters to a queryset.

    queryset is a Django queryset, eg MyModel.objects.all() or
             MyModel.objects.filter(user=request.user)

    params is a dictionary that contains the following:
        filtering: A dict of Django ORM filters, eg:
            {'user__id__in': [1, 3, 103], 'title__contains': 'foo'}
        other_filter: Another filter of some type, most likely a
            set of Q() objects.
        sorting: The name of the column to sort by
    """
    for key in params['filtering'].keys():
        filter = {key: params['filtering'][key]}
        queryset = queryset.filter(**filter)

    if params.get('other_filter', None):
        # eg a Q() set
        queryset = queryset.filter(params['other_filter'])

    if params.get('sorting', None):
        if params.get('sortreverse', None):
            params['sorting'] = "-%s" % params['sorting']
        queryset = queryset.order_by(params['sorting'])

    return queryset

def safe_template_context(ticket):
    """
    Return a dictionary that can be used as a template context to render
    comments and other details with ticket or queue paramaters. Note that
    we don't just provide the Ticket & Queue objects to the template as
    they could reveal confidential information. Just imagine these two options:
        * {{ ticket.queue.email_box_password }}
        * {{ ticket.assigned_to.password }}

    Ouch!

    The downside to this is that if we make changes to the model, we will also
    have to update this code. Perhaps we can find a better way in the future.
    """

    context = {
        'queue': {},
        'ticket': {},
        }
    queue = ticket.queue

    for field in (  'title', 'slug', 'email_address', 'from_address', 'locale'):
        attr = getattr(queue, field, None)
        if callable(attr):
            context['queue'][field] = attr()
        else:
            context['queue'][field] = attr

    for field in (  'title', 'created', 'modified', 'submitter_email',
                    'status', 'get_status_display', 'on_hold', 'description',
                    'resolution', 'priority', 'get_priority_display',
                    'last_escalation', 'ticket', 'ticket_for_url',
                    'get_status', 'ticket_url', 'staff_url', '_get_assigned_to'
                 ):
        attr = getattr(ticket, field, None)
        if callable(attr):
            context['ticket'][field] = '%s' % attr()
        else:
            context['ticket'][field] = attr

    context['ticket']['queue'] = context['queue']
    context['ticket']['assigned_to'] = context['ticket']['_get_assigned_to']

    return context

def text_is_spam(text, request):
    # Based on a blog post by 'sciyoshi':
    # http://sciyoshi.com/blog/2008/aug/27/using-akismet-djangos-new-comments-framework/
    # This will return 'True' is the given text is deemed to be spam, or
    # False if it is not spam. If it cannot be checked for some reason, we
    # assume it isn't spam.
    from django.contrib.sites.models import Site
    from django.conf import settings
    try:
        from helpdesk.akismet import Akismet
    except:
        return False

    ak = Akismet(
        blog_url='http://%s/' % Site.objects.get(pk=settings.SITE_ID).domain,
        agent='django-helpdesk',
    )

    if hasattr(settings, 'TYPEPAD_ANTISPAM_API_KEY'):
        ak.setAPIKey(key = settings.TYPEPAD_ANTISPAM_API_KEY)
        ak.baseurl = 'api.antispam.typepad.com/1.1/'
    elif hasattr(settings, 'AKISMET_API_KEY'):
        ak.setAPIKey(key = settings.AKISMET_API_KEY)
    else:
        return False

    if ak.verify_key():
        ak_data = {
            'user_ip': request.META.get('REMOTE_ADDR', '127.0.0.1'),
            'user_agent': request.META.get('HTTP_USER_AGENT', ''),
            'referrer': request.META.get('HTTP_REFERER', ''),
            'comment_type': 'comment',
            'comment_author': '',
        }

        return ak.comment_check(text, data=ak_data)

    return False
<|endoftext|>


<commit_before># -*- coding: utf8 -*-
#
# Module RATIONAL
#
# Part of Nutils: open source numerical utilities for Python. Jointly developed
# by HvZ Computational Engineering, TU/e Multiscale Engineering Fluid Dynamics,
# and others. More info at http://nutils.org <info@nutils.org>. (c) 2014

"""
The rational module.
"""

from __future__ import division
from . import cache, log
import numpy, collections

masks = [ 0, 0 ] # primes per number encoded as bitmask

class Primes( object ):

  def __init__( self ):
    self.__primes = [] # all prime numbers found thus far
    self.__masks = collections.deque([ 0 ]) # primes per number in progress

  def __iter__( self ):
    for prime in self.__primes:
      yield prime
    while True:
      mask = self.__masks.popleft()
      if mask:
        masks.append( mask )
        iprime = 0
        while mask:
          if mask & 1:
            self.__mark( iprime )
          mask >>= 1
          iprime += 1
      else: # found prime
        iprime = len(self.__primes)
        prime = len(masks)
        self.__primes.append( prime )
        self.__mark( iprime )
        masks.append( 1 << iprime )
        yield prime

  def __mark( self, iprime ):
    i = self.__primes[iprime] - 1
    if i < len( self.__masks ):
      self.__masks[ i ] |= 1 << iprime
    else:
      self.__masks.extend( [0] * ( i - len(self.__masks) ) + [ 1 << iprime ] )

primes = Primes()

class Rational( object ):

  __array_priority__ = 1

  def __init__( self, init=() ):
    while len(init) and init[-1] == 0:
      init = init[:-1]
    if isinstance( init, numpy.ndarray ):
      assert init.dtype == int
    else:
      assert all( isinstance(n,int) for n in init )
    self.__n = tuple( init )

  def __hash__( self ):
    return hash( self.__n )

  def __eq__( self, other ):
    return self is other or isrational(other) and self.__n == other.__n

  def __op( self, other, op ):
    assert self.__n or other.__n, 'trivial cases should be handled by caller'
    return Rational( numpy.hstack([
      op( self.__n, other.__n[:len(self.__n)] ),
      op( 0, other.__n[len(self.__n):] ) ]) if len(self.__n) < len(other.__n)
                else numpy.hstack([
      op( self.__n[:len(other.__n)], other.__n ),
      op( self.__n[len(other.__n):], 0 ) ]) if len(self.__n) > len(other.__n)
                else
      op( self.__n, other.__n ) )

  def __mul__( self, other ):
    if not isexact( other ):
      return float(self) * other
    if isarray( other ):
      return other * self
    assert isrational( other )
    return other if not self.__n \
      else self if not other.__n \
      else self.__op( other, numpy.add )

  def __neg__( self ):
    raise TypeError, 'rational cannot be negated'

  def __add__( self, other ):
    if not isexact( other ):
      return float(self) + other
    if isarray( other ):
      return other + self
    raise NotImplementedError

  def __sub__( self, other ):
    if not isexact( other ):
      return float(self) - other
    if isarray( other ):
      return other - self
    raise NotImplementedError

  def __div__( self, other ):
    if not isexact( other ):
      return float(self) / other
    assert isrational( other )
    return self if not other.__n \
      else Rational( numpy.negative(other.__n) ) if not self.__n \
      else self.__op( other, numpy.subtract )

  def __rdiv__( self, other ):
    return Rational( numpy.negative(self.__n) ) * other

  __rmul__ = __mul__
  __radd__ = __add__
  __truediv__ = __div__
  __rtruediv__ = __rdiv__

  def __pow__( self, n ):
    assert isinstance( n, int )
    if n == 0 or not self.__n:
      return Rational()
    return Rational( numpy.multiply( self.__n, n ) )

  def __float__( self ):
    numer, denom = self.frac
    return numer / denom

  def __int__( self ):
    numer, denom = self.frac
    assert denom == 1
    return numer

  def __str__( self ):
    return '%d/%d' % self.frac

  def __repr__( self ):
    return 'Rational(%s)' % str(self)

  def factor( self ):
    return zip( primes, self.__n )

  @cache.property
  def frac( self ):
    numer = denom = 1
    for p, n in zip( primes, self.__n ):
      if n > 0:
        numer *= p**+n
      elif n < 0:
        denom *= p**-n
    return numer, denom

  @property
  def numer( self ):
    return Rational( self.__n and numpy.maximum( 0, self.__n ) )

  @property
  def denom( self ):
    return Rational( self.__n and numpy.minimum( 0, self.__n ) )

unit = Rational()
half = Rational((-1,))

class Array( object ):

  __array_priority__ = 1

  def __init__( self, array, factor=unit, isfactored=False ):
    assert isrational( factor )
    array = numpy.asarray( array )
    assert array.dtype == int
    if isfactored:
      self.__array = array
      self.__factor = factor
    else:
      numbers, common = gcd( array.ravel() )
      self.__array = numbers.reshape( array.shape )
      self.__factor = factor * common

  @property
  def ndim( self ):
    return self.__array.ndim

  @property
  def shape( self ):
    return self.__array.shape

  @property
  def T( self ):
    return Array( self.__array.T, self.__factor, True )

  def __len__( self ):
    return len(self.__array)

  @property
  def __cmpdata( self ):
    return self.__array.shape, tuple(self.__array.flat), self.__factor

  def __hash__( self ):
    return hash( self.__cmpdata )

  def __eq__( self, other ):
    return self is other or isarray(other) and self.__cmpdata == other.__cmpdata

  def __neg__( self ):
    return Array( -self.__array, self.__factor, True )

  def __add__( self, other ):
    if not isexact( other ):
      return self.__array * float(self.__factor) + other
    raise NotImplementedError

  def __sub__( self, other ):
    if not isexact( other ):
      return self.__array * float(self.__factor) - other
    raise NotImplementedError

  def __mul__( self, other ):
    if not isexact( other ):
      return self.__array * ( other * float(self.__factor) )
    if isrational( other ):
      return Array( self.__array, self.__factor*other, True )
    raise NotImplementedError

  def __div__( self, other ):
    if not isexact( other ):
      return self.__array * ( float(self.__factor) / other )
    if isrational( other ):
      return Array( self.__array, self.__factor/other, True )
    raise NotImplementedError

  def __rdiv__( self, other ):
    if not isexact( other ):
      return ( other / float(self.__factor) ) / self.__array
    raise NotImplementedError

  __rmul__ = __mul__
  __radd__ = __add__
  __truediv__ = __div__
  __rtruediv__ = __rdiv__

  def decompose( self ):
    return self.__array, self.__factor

  def __str__( self ):
    return str(self.__factor) + str(self.__array.tolist()).replace( ' ', '' )

def det( array ):
  assert isarray(array)
  if array.shape == (2,2):
    ((a,b),(c,d)), scale = array.decompose()
    det = a*d - b*c
  elif array.shape == (3,3):
    ((a,b,c),(d,e,f),(g,h,i)), scale = array.decompose()
    det = a*e*i + b*f*g + c*d*h - c*e*g - b*d*i - a*f*h
  else:
    raise NotImplementedError, 'shape=' + str(array.shape)
  return factor(det) * scale**len(array)

def inv( array ):
  assert isarray(array)
  if array.shape == (2,2):
    ((a,b),(c,d)), scale = array.decompose()
    inv = Array( ((d,-b),(-c,a)), scale / det(array), True )
  elif array.shape == (3,3):
    raise NotImplementedError
    ((a,b,c),(d,e,f),(g,h,i)), scale = array.decompose()
    inv = Array( ((e*i-f*h,c*h-b*i,b*f-c*e),(f*g-d*i,a*i-c*g,c*d-a*f),(d*h-e*g,b*g-a*h,a*e-b*d)), scale / det(array), False )
  else:
    raise NotImplementedError, 'shape=' + tuple(array.shape)
  return inv

def factor( n ):
  assert isinstance( n, int ) and n > 0
  factors = []
  if n >= len(masks):
    __log__ = log.iter( 'extending primes', primes )
    for prime in __log__:
      assert n >= prime
      count = 0
      while n % prime == 0:
        n //= prime
        count += 1
      factors.append( count )
      if n < len(masks):
        break
  factor = Rational( factors )
  while n > 1:
    mask = masks[n]
    nums = []
    while mask:
      nums.append( int(mask&1) )
      mask >>= 1
    r = Rational( nums )
    factor *= r
    n //= int(r)
  return factor

def isrational( num ):
  return isinstance( num, Rational )

def asrational( num ):
  if isinstance( num, int ):
    num = factor( num )
  assert isrational( num )
  return num

def isarray( arr ):
  return isinstance( arr, Array )

def asarray( arr ):
  if isarray( arr ):
    return arr
  arr = numpy.asarray( arr )
  assert arr.dtype == int
  return Array( arr, unit, False )

def isexact( obj ):
  return isrational(obj) or isarray(obj)

def asexact( obj ):
  return obj if isexact( obj ) \
    else asrational( obj ) if isinstance( obj, int ) \
    else asarray( obj )

def frac( numer, denom ):
  return asexact( numer ) / asrational( denom )

def gcd( numbers ):
  numbers = numpy.asarray( numbers )
  assert numbers.dtype == int
  if not numbers.size:
    return numbers, Rational()
  unique = numpy.unique( numpy.abs(numbers.flat) )
  if unique[0] == 0:
    unique = unique[1:]
  if not unique.size:
    return numbers, Rational()
  common = factor( unique[0] )
  unique = unique[1:]
  if unique.size:
    factors = []
    for prime, count in common.factor():
      n = 0
      while count > n and not numpy.any( unique % prime ):
        unique //= prime
        n += 1
      factors.append( n )
    common = Rational( factors )
  return numbers // int(common), common

def asfloat( obj ):
  if isinstance(obj,int) or isrational(obj):
    return float(obj)
  if isarray( obj ):
    array, factor = obj.decompose()
    return array * float(factor)
  if isinstance( obj, numpy.ndarray ) and obj.dtype == float:
    return obj # don't touch, might be of derived type
  return numpy.asarray( obj, dtype=float )

def dot( A, B ):
  if not isexact( A ) or not isexact( B ):
    return numpy.dot( asfloat(A), asfloat(B) )
  A, a = asarray( A ).decompose()
  B, b = asarray( B ).decompose()
  return Array( numpy.dot( A, B ), a * b, False )

# vim:shiftwidth=2:foldmethod=indent:foldnestmax=2
<commit_msg>fix numpy corner case<commit_after># -*- coding: utf8 -*-
#
# Module RATIONAL
#
# Part of Nutils: open source numerical utilities for Python. Jointly developed
# by HvZ Computational Engineering, TU/e Multiscale Engineering Fluid Dynamics,
# and others. More info at http://nutils.org <info@nutils.org>. (c) 2014

"""
The rational module.
"""

from __future__ import division
from . import cache, log
import numpy, collections

masks = [ 0, 0 ] # primes per number encoded as bitmask

class Primes( object ):

  def __init__( self ):
    self.__primes = [] # all prime numbers found thus far
    self.__masks = collections.deque([ 0 ]) # primes per number in progress

  def __iter__( self ):
    for prime in self.__primes:
      yield prime
    while True:
      mask = self.__masks.popleft()
      if mask:
        masks.append( mask )
        iprime = 0
        while mask:
          if mask & 1:
            self.__mark( iprime )
          mask >>= 1
          iprime += 1
      else: # found prime
        iprime = len(self.__primes)
        prime = len(masks)
        self.__primes.append( prime )
        self.__mark( iprime )
        masks.append( 1 << iprime )
        yield prime

  def __mark( self, iprime ):
    i = self.__primes[iprime] - 1
    if i < len( self.__masks ):
      self.__masks[ i ] |= 1 << iprime
    else:
      self.__masks.extend( [0] * ( i - len(self.__masks) ) + [ 1 << iprime ] )

primes = Primes()

class Rational( object ):

  __array_priority__ = 1

  def __init__( self, init=() ):
    while len(init) and init[-1] == 0:
      init = init[:-1]
    if isinstance( init, numpy.ndarray ):
      assert init.dtype == int
    else:
      assert all( isinstance(n,int) for n in init )
    self.__n = tuple( init )

  def __hash__( self ):
    return hash( self.__n )

  def __eq__( self, other ):
    return self is other or isrational(other) and self.__n == other.__n

  def __op( self, other, op ):
    assert self.__n or other.__n, 'trivial cases should be handled by caller'
    return Rational( numpy.hstack([
      op( self.__n, other.__n[:len(self.__n)] ),
      op( 0, other.__n[len(self.__n):] ) ]) if len(self.__n) < len(other.__n)
                else numpy.hstack([
      op( self.__n[:len(other.__n)], other.__n ),
      op( self.__n[len(other.__n):], 0 ) ]) if len(self.__n) > len(other.__n)
                else
      op( self.__n, other.__n ) )

  def __mul__( self, other ):
    if not isexact( other ):
      return float(self) * other
    if isarray( other ):
      return other * self
    assert isrational( other )
    return other if not self.__n \
      else self if not other.__n \
      else self.__op( other, numpy.add )

  def __neg__( self ):
    raise TypeError, 'rational cannot be negated'

  def __add__( self, other ):
    if not isexact( other ):
      return float(self) + other
    if isarray( other ):
      return other + self
    raise NotImplementedError

  def __sub__( self, other ):
    if not isexact( other ):
      return float(self) - other
    if isarray( other ):
      return other - self
    raise NotImplementedError

  def __div__( self, other ):
    if not isexact( other ):
      return float(self) / other
    assert isrational( other )
    return self if not other.__n \
      else Rational( numpy.negative(other.__n) ) if not self.__n \
      else self.__op( other, numpy.subtract )

  def __rdiv__( self, other ):
    return Rational( numpy.negative(self.__n) ) * other

  __rmul__ = __mul__
  __radd__ = __add__
  __truediv__ = __div__
  __rtruediv__ = __rdiv__

  def __pow__( self, n ):
    assert isinstance( n, int )
    if n == 0 or not self.__n:
      return Rational()
    return Rational( numpy.multiply( self.__n, n ) )

  def __float__( self ):
    numer, denom = self.frac
    return numer / denom

  def __int__( self ):
    numer, denom = self.frac
    assert denom == 1
    return numer

  def __str__( self ):
    return '%d/%d' % self.frac

  def __repr__( self ):
    return 'Rational(%s)' % str(self)

  def factor( self ):
    return zip( primes, self.__n )

  @cache.property
  def frac( self ):
    numer = denom = 1
    for p, n in zip( primes, self.__n ):
      if n > 0:
        numer *= p**+n
      elif n < 0:
        denom *= p**-n
    return numer, denom

  @property
  def numer( self ):
    return Rational( self.__n and numpy.maximum( 0, self.__n ) )

  @property
  def denom( self ):
    return Rational( self.__n and numpy.minimum( 0, self.__n ) )

unit = Rational()
half = Rational((-1,))

class Array( object ):

  __array_priority__ = 1

  def __init__( self, array, factor=unit, isfactored=False ):
    assert isrational( factor )
    array = numpy.asarray( array )
    assert array.dtype == int
    if isfactored:
      self.__array = array
      self.__factor = factor
    else:
      numbers, common = gcd( array.ravel() )
      self.__array = numbers.reshape( array.shape )
      self.__factor = factor * common

  @property
  def ndim( self ):
    return self.__array.ndim

  @property
  def shape( self ):
    return self.__array.shape

  @property
  def T( self ):
    return Array( self.__array.T, self.__factor, True )

  def __len__( self ):
    return len(self.__array)

  @property
  def __cmpdata( self ):
    return self.__array.shape, tuple(self.__array.flat), self.__factor

  def __hash__( self ):
    return hash( self.__cmpdata )

  def __eq__( self, other ):
    return self is other or isarray(other) and self.__cmpdata == other.__cmpdata

  def __neg__( self ):
    return Array( -self.__array, self.__factor, True )

  def __add__( self, other ):
    if not isexact( other ):
      return self.__array * float(self.__factor) + other
    raise NotImplementedError

  def __sub__( self, other ):
    if not isexact( other ):
      return self.__array * float(self.__factor) - other
    raise NotImplementedError

  def __mul__( self, other ):
    if not isexact( other ):
      return self.__array * ( other * float(self.__factor) )
    if isrational( other ):
      return Array( self.__array, self.__factor*other, True )
    raise NotImplementedError

  def __div__( self, other ):
    if not isexact( other ):
      return self.__array * ( float(self.__factor) / other )
    if isrational( other ):
      return Array( self.__array, self.__factor/other, True )
    raise NotImplementedError

  def __rdiv__( self, other ):
    if not isexact( other ):
      return ( other / float(self.__factor) ) / self.__array
    raise NotImplementedError

  __rmul__ = __mul__
  __radd__ = __add__
  __truediv__ = __div__
  __rtruediv__ = __rdiv__

  def decompose( self ):
    return self.__array, self.__factor

  def __str__( self ):
    return str(self.__factor) + str(self.__array.tolist()).replace( ' ', '' )

def det( array ):
  assert isarray(array)
  if array.shape == (2,2):
    ((a,b),(c,d)), scale = array.decompose()
    det = a*d - b*c
  elif array.shape == (3,3):
    ((a,b,c),(d,e,f),(g,h,i)), scale = array.decompose()
    det = a*e*i + b*f*g + c*d*h - c*e*g - b*d*i - a*f*h
  else:
    raise NotImplementedError, 'shape=' + str(array.shape)
  return factor(det) * scale**len(array)

def inv( array ):
  assert isarray(array)
  if array.shape == (2,2):
    ((a,b),(c,d)), scale = array.decompose()
    inv = Array( ((d,-b),(-c,a)), scale / det(array), True )
  elif array.shape == (3,3):
    raise NotImplementedError
    ((a,b,c),(d,e,f),(g,h,i)), scale = array.decompose()
    inv = Array( ((e*i-f*h,c*h-b*i,b*f-c*e),(f*g-d*i,a*i-c*g,c*d-a*f),(d*h-e*g,b*g-a*h,a*e-b*d)), scale / det(array), False )
  else:
    raise NotImplementedError, 'shape=' + tuple(array.shape)
  return inv

def factor( n ):
  assert isinstance( n, int ) and n > 0
  factors = []
  if n >= len(masks):
    __log__ = log.iter( 'extending primes', primes )
    for prime in __log__:
      assert n >= prime
      count = 0
      while n % prime == 0:
        n //= prime
        count += 1
      factors.append( count )
      if n < len(masks):
        break
  factor = Rational( factors )
  while n > 1:
    mask = masks[n]
    nums = []
    while mask:
      nums.append( int(mask&1) )
      mask >>= 1
    r = Rational( nums )
    factor *= r
    n //= int(r)
  return factor

def isrational( num ):
  return isinstance( num, Rational )

def asrational( num ):
  if isinstance( num, int ):
    num = factor( num )
  assert isrational( num )
  return num

def isarray( arr ):
  return isinstance( arr, Array )

def asarray( arr ):
  if isarray( arr ):
    return arr
  arr = numpy.asarray( arr )
  assert arr.dtype == int
  return Array( arr, unit, False )

def isexact( obj ):
  return isrational(obj) or isarray(obj)

def asexact( obj ):
  return obj if isexact( obj ) \
    else asrational( obj ) if isinstance( obj, int ) \
    else asarray( obj )

def frac( numer, denom ):
  return asexact( numer ) / asrational( denom )

def gcd( numbers ):
  numbers = numpy.asarray( numbers )
  assert numbers.dtype == int
  if not numbers.size:
    return numbers, Rational()
  unique = numpy.unique( numpy.abs(numbers.flat) )
  if unique[0] == 0:
    unique = unique[1:]
  if not unique.size:
    return numbers, Rational()
  common = factor( unique[0] )
  unique = unique[1:]
  if unique.size:
    factors = []
    for prime, count in common.factor():
      n = 0
      while count > n and not numpy.any( unique % prime ):
        unique //= prime
        n += 1
      factors.append( n )
    common = Rational( factors )
  return numbers // int(common), common

def asfloat( obj ):
  if isinstance(obj,int) or isrational(obj):
    return float(obj)
  if isarray( obj ):
    array, factor = obj.decompose()
    return array * float(factor)
  if isinstance( obj, numpy.ndarray ) and obj.dtype == float:
    return obj # don't touch, might be of derived type
  return numpy.asarray( obj, dtype=float )

def dot( A, B ):
  if not isexact( A ) or not isexact( B ):
    Aflt = asfloat(A)
    Bflt = asfloat(B)
    return numpy.dot( Aflt, Bflt ).view( Aflt.__class__ ) # .view necesssary for 1D Aflt (numpy exception)
  A, a = asarray( A ).decompose()
  B, b = asarray( B ).decompose()
  return Array( numpy.dot( A, B ), a * b, False )

# vim:shiftwidth=2:foldmethod=indent:foldnestmax=2
<|endoftext|>


<commit_before># Licensed under a 3-clause BSD style license - see LICENSE.rst
# -*- coding: utf-8 -*-
"""Test desimodel.io.
"""
import os
import uuid
import numpy as np
from astropy.table import Table
import unittest
from .. import io
#
# Try to import specter.
#
specter_available = True
specter_message = "The specter package was not detected."
try:
    import specter
except ImportError:
    specter_available = False
#
# Try to load the DESIMODEL environment variable
#
desimodel_available = True
desimodel_message = "The desimodel data set was not detected."
try:
    spam = os.environ['DESIMODEL']
except KeyError:
    desimodel_available = False
    specter_available = False
    specter_message = desimodel_message

class _TestIO(unittest.TestCase):
    """Test desimodel.io.
    """
    @classmethod
    def setUpClass(cls):
        global specter_available, desimodel_available
        cls.specter_available = specter_available
        cls.desimodel_available = desimodel_available
        cls.trimdir = 'test-'+uuid.uuid4().hex
        cls.testfile = 'test-'+uuid.uuid4().hex+'.fits'

    @classmethod
    def tearDownClass(cls):
        if os.path.exists(cls.trimdir):
            import shutil
            shutil.rmtree(cls.trimdir)

    def setUp(self):
        """Ensure that any desimodel.io caches are clear before running
        any test.
        """
        io._thru = dict()
        io._psf = dict()
        io._params = None
        io._gfa = None
        io._fiberpos = None
        io._tiles = dict()
        if os.path.exists(self.testfile):
            os.remove(self.testfile)

    def tearDown(self):
        pass

    @unittest.skipUnless(specter_available, specter_message)
    def test_load_throughput(self):
        """Test loading of throughput files.
        """
        for channel in ('b', 'r', 'z'):
            t = io.load_throughput(channel)

    @unittest.skipUnless(specter_available, specter_message)
    def test_load_psf(self):
        """Test loading of PSF files.
        """
        for channel in ('b', 'r', 'z'):
            t = io.load_psf(channel)

    @unittest.skipUnless(desimodel_available, desimodel_message)
    def test_load_desiparams(self):
        """Test loading of basic DESI parameters.
        """
        p = io.load_desiparams()
        self.assertTrue(isinstance(p, dict))

    @unittest.skipUnless(desimodel_available, desimodel_message)
    def test_load_gfa(self):
        """Test loading of gfa location data.
            """
        gfa = io.load_gfa()
        # length of the GFA table should be 40 since there are each corner of the 10 GFAs are included
        self.assertEqual(len(gfa), 40)
        for key in ('PETAL', 'X', 'Y', 'Z', 'Q', 'RADIUS_DEG', 'RADIUS_MM'):
            self.assertIn(key, gfa.dtype.names)

    @unittest.skipUnless(desimodel_available, desimodel_message)
    def test_load_fiberpos(self):
        """Test loading of fiber positioner data.
        """
        fiberpos = io.load_fiberpos()
        self.assertEqual(len(fiberpos), 5000)
        # Check workaround for astropy 1.0.x bug for lower -> upper col names.
        for key in ('FIBER', 'LOCATION', 'SPECTRO', 'X', 'Y', 'Z'):
            self.assertIn(key, fiberpos.dtype.names)
            x = fiberpos[key]

    @unittest.skipUnless(desimodel_available, desimodel_message)
    def test_load_platescale(self):
        """Test loading platescale.txt file.
        """
        p1 = io.load_platescale()
        p2 = io.load_platescale()
        self.assertTrue(p1 is p2)  #- caching worked

    @unittest.skipUnless(desimodel_available, desimodel_message)
    def test_load_targets(self):
        """Test loading of tile files.
        """
        data = io.load_target_info()
        #- Test a few keys, but not everything
        self.assertIn('ntarget_lrg', data.keys())
        self.assertIn('nobs_elg', data.keys())
        self.assertIn('success_qso', data.keys())

    @unittest.skip('Skip *until* the pixel weights file is in the DESIMODEL directory')
#    @unittest.skipUnless(desimodel_available, desimodel_message)
    def test_load_pix_file(self):
        """Test loading of the file of HEALPixel weights.
        """
        nside = 16
        pix = io.load_pixweight(nside)
        npix = len(pix)
        #ADM Test the length of the returned array is appropriate to the passed nside
        self.assertEqual(npix,12*nside*nside)

    @unittest.skipUnless(desimodel_available, desimodel_message)
    def test_load_tiles(self):
        """Test loading of tile files.
        """
        # starting clean
        self.assertEqual(io._tiles, {})
        t0 = io.load_tiles(cache=False)
        self.assertEqual(len(io._tiles), 0)
        # loading tiles fills the cache with one items
        t1 = io.load_tiles(onlydesi=False)
        self.assertEqual(len(io._tiles), 1)
        tile_cache_id1 = id(list(io._tiles.items())[0])
        # reloading, even with a filter, shoudln't change cache
        t2 = io.load_tiles(onlydesi=True)
        self.assertEqual(len(io._tiles), 1)
        tile_cache_id2 = id(list(io._tiles.items())[0])
        self.assertEqual(tile_cache_id1, tile_cache_id2)
        #- Temporarily support OBSCONDITIONS as u2 (old) or i4 (new)
        self.assertTrue(np.issubdtype(t1['OBSCONDITIONS'].dtype, 'i4') or \
                        np.issubdtype(t1['OBSCONDITIONS'].dtype, 'u2') )
        self.assertTrue(np.issubdtype(t2['OBSCONDITIONS'].dtype, 'i4') or \
                        np.issubdtype(t2['OBSCONDITIONS'].dtype, 'u2') )
        self.assertLess(len(t2), len(t1))
        # All tiles in DESI are also in full set.
        self.assertTrue(np.all(np.in1d(t2['TILEID'], t1['TILEID'])))
        # I think this is the exact same test as above, except using set theory.
        self.assertEqual(len(set(t2.TILEID) - set(t1.TILEID)), 0)
        t3 = io.load_tiles(onlydesi=False)
        tile_cache_id3 = id(list(io._tiles.items())[0])
        self.assertEqual(tile_cache_id1, tile_cache_id3)
        self.assertTrue(np.issubdtype(t3['OBSCONDITIONS'].dtype, 'i4') or \
                        np.issubdtype(t3['OBSCONDITIONS'].dtype, 'u2') )
        # Check for extra tiles.
        a = io.load_tiles(extra=False)
        self.assertEqual(np.sum(np.char.startswith(a['PROGRAM'], 'EXTRA')), 0)
        b = io.load_tiles(extra=True)
        self.assertGreater(np.sum(np.char.startswith(b['PROGRAM'], 'EXTRA')), 0)
        self.assertLess(len(a), len(b))

    @unittest.skipUnless(desimodel_available, desimodel_message)
    def test_load_tiles_alt(self):
        # starting clean
        self.assertEqual(io._tiles, {})
        t1 = Table(io.load_tiles())
        t1.write(self.testfile)
        #- no path; should fail since that file isn't in $DESIMODEL/data/footprint/
        with self.assertRaises(FileNotFoundError):
            t2 = io.load_tiles(tilesfile=self.testfile)

        #- with path, should work:
        t2 = Table(io.load_tiles(tilesfile='./'+self.testfile))
        self.assertTrue(np.all(t1 == t2))

        #- cache should have two items now
        self.assertEqual(len(io._tiles), 2)

    @unittest.skipUnless(desimodel_available, desimodel_message)
    def test_tiles_consistency(self):
        """Test consistency of tile files.

        - Validate different tile loading schemes.
        - Test numerical consistency between FITS & ECSV versions.
        - Assure that PROGRAM column has no trailing whitespace.
        """
        fitstiles = io.findfile('footprint/desi-tiles.fits')
        ecsvtiles = io.findfile('footprint/desi-tiles.ecsv')
        # tf=TilesFits  tt=TilesTable  te=TilesEcsv
        tf = io.load_tiles(onlydesi=False, extra=True)
        tt = Table.read(fitstiles)
        te = Table.read(ecsvtiles, format='ascii.ecsv')
        self.assertEqual(sorted(tf.dtype.names), sorted(tt.colnames))
        self.assertEqual(sorted(tf.dtype.names), sorted(te.colnames))

        for col in tt.colnames:
            self.assertTrue(np.all(tf[col]==tt[col]), 'fits[{col}] != table[{col}]'.format(col=col))
            if np.issubdtype(tf[col].dtype, float):
                self.assertTrue(np.allclose(tf[col], te[col], atol=1e-4, rtol=1e-4), 'fits[{col}] != ecsv[{col}]'.format(col=col))
            else:
                self.assertTrue(np.all(tf[col]==te[col]), 'fits[{col}] != ecsv[{col}]'.format(col=col))

        self.assertTrue(not np.any(np.char.endswith(tf['PROGRAM'], ' ')))
        self.assertTrue(not np.any(np.char.endswith(tt['PROGRAM'], ' ')))
        self.assertTrue(not np.any(np.char.endswith(te['PROGRAM'], ' ')))

    @unittest.skipUnless(desimodel_available, desimodel_message)
    def test_trim_data(self):
        '''Test trimming data files for lightweight tests'''
        psffile = io.findfile('specpsf/psf-b.fits')
        if os.path.getsize(psffile) > 20e6:
            from .. import trim
            indir = os.path.join(os.getenv('DESIMODEL'), 'data')
            trim.trim_data(indir, self.trimdir)
            self.assertTrue(os.path.isdir(self.trimdir))
            self.assertGreater(len(list(os.walk(self.trimdir))), 1)

def test_suite():
    """Allows testing of only this module with the command::
        python setup.py test -m <modulename>
    """
    return unittest.defaultTestLoader.loadTestsFromName(__name__)
<commit_msg>remove extraneous underscore<commit_after># Licensed under a 3-clause BSD style license - see LICENSE.rst
# -*- coding: utf-8 -*-
"""Test desimodel.io.
"""
import os
import uuid
import numpy as np
from astropy.table import Table
import unittest
from .. import io
#
# Try to import specter.
#
specter_available = True
specter_message = "The specter package was not detected."
try:
    import specter
except ImportError:
    specter_available = False
#
# Try to load the DESIMODEL environment variable
#
desimodel_available = True
desimodel_message = "The desimodel data set was not detected."
try:
    spam = os.environ['DESIMODEL']
except KeyError:
    desimodel_available = False
    specter_available = False
    specter_message = desimodel_message

class TestIO(unittest.TestCase):
    """Test desimodel.io.
    """
    @classmethod
    def setUpClass(cls):
        global specter_available, desimodel_available
        cls.specter_available = specter_available
        cls.desimodel_available = desimodel_available
        cls.trimdir = 'test-'+uuid.uuid4().hex
        cls.testfile = 'test-'+uuid.uuid4().hex+'.fits'

    @classmethod
    def tearDownClass(cls):
        if os.path.exists(cls.trimdir):
            import shutil
            shutil.rmtree(cls.trimdir)

    def setUp(self):
        """Ensure that any desimodel.io caches are clear before running
        any test.
        """
        io._thru = dict()
        io._psf = dict()
        io._params = None
        io._gfa = None
        io._fiberpos = None
        io._tiles = dict()
        if os.path.exists(self.testfile):
            os.remove(self.testfile)

    def tearDown(self):
        pass

    @unittest.skipUnless(specter_available, specter_message)
    def test_load_throughput(self):
        """Test loading of throughput files.
        """
        for channel in ('b', 'r', 'z'):
            t = io.load_throughput(channel)

    @unittest.skipUnless(specter_available, specter_message)
    def test_load_psf(self):
        """Test loading of PSF files.
        """
        for channel in ('b', 'r', 'z'):
            t = io.load_psf(channel)

    @unittest.skipUnless(desimodel_available, desimodel_message)
    def test_load_desiparams(self):
        """Test loading of basic DESI parameters.
        """
        p = io.load_desiparams()
        self.assertTrue(isinstance(p, dict))

    @unittest.skipUnless(desimodel_available, desimodel_message)
    def test_load_gfa(self):
        """Test loading of gfa location data.
            """
        gfa = io.load_gfa()
        # length of the GFA table should be 40 since there are each corner of the 10 GFAs are included
        self.assertEqual(len(gfa), 40)
        for key in ('PETAL', 'X', 'Y', 'Z', 'Q', 'RADIUS_DEG', 'RADIUS_MM'):
            self.assertIn(key, gfa.dtype.names)

    @unittest.skipUnless(desimodel_available, desimodel_message)
    def test_load_fiberpos(self):
        """Test loading of fiber positioner data.
        """
        fiberpos = io.load_fiberpos()
        self.assertEqual(len(fiberpos), 5000)
        # Check workaround for astropy 1.0.x bug for lower -> upper col names.
        for key in ('FIBER', 'LOCATION', 'SPECTRO', 'X', 'Y', 'Z'):
            self.assertIn(key, fiberpos.dtype.names)
            x = fiberpos[key]

    @unittest.skipUnless(desimodel_available, desimodel_message)
    def test_load_platescale(self):
        """Test loading platescale.txt file.
        """
        p1 = io.load_platescale()
        p2 = io.load_platescale()
        self.assertTrue(p1 is p2)  #- caching worked

    @unittest.skipUnless(desimodel_available, desimodel_message)
    def test_load_targets(self):
        """Test loading of tile files.
        """
        data = io.load_target_info()
        #- Test a few keys, but not everything
        self.assertIn('ntarget_lrg', data.keys())
        self.assertIn('nobs_elg', data.keys())
        self.assertIn('success_qso', data.keys())

    @unittest.skip('Skip *until* the pixel weights file is in the DESIMODEL directory')
#    @unittest.skipUnless(desimodel_available, desimodel_message)
    def test_load_pix_file(self):
        """Test loading of the file of HEALPixel weights.
        """
        nside = 16
        pix = io.load_pixweight(nside)
        npix = len(pix)
        #ADM Test the length of the returned array is appropriate to the passed nside
        self.assertEqual(npix,12*nside*nside)

    @unittest.skipUnless(desimodel_available, desimodel_message)
    def test_load_tiles(self):
        """Test loading of tile files.
        """
        # starting clean
        self.assertEqual(io._tiles, {})
        t0 = io.load_tiles(cache=False)
        self.assertEqual(len(io._tiles), 0)
        # loading tiles fills the cache with one items
        t1 = io.load_tiles(onlydesi=False)
        self.assertEqual(len(io._tiles), 1)
        tile_cache_id1 = id(list(io._tiles.items())[0])
        # reloading, even with a filter, shoudln't change cache
        t2 = io.load_tiles(onlydesi=True)
        self.assertEqual(len(io._tiles), 1)
        tile_cache_id2 = id(list(io._tiles.items())[0])
        self.assertEqual(tile_cache_id1, tile_cache_id2)
        #- Temporarily support OBSCONDITIONS as u2 (old) or i4 (new)
        self.assertTrue(np.issubdtype(t1['OBSCONDITIONS'].dtype, 'i4') or \
                        np.issubdtype(t1['OBSCONDITIONS'].dtype, 'u2') )
        self.assertTrue(np.issubdtype(t2['OBSCONDITIONS'].dtype, 'i4') or \
                        np.issubdtype(t2['OBSCONDITIONS'].dtype, 'u2') )
        self.assertLess(len(t2), len(t1))
        # All tiles in DESI are also in full set.
        self.assertTrue(np.all(np.in1d(t2['TILEID'], t1['TILEID'])))
        # I think this is the exact same test as above, except using set theory.
        self.assertEqual(len(set(t2.TILEID) - set(t1.TILEID)), 0)
        t3 = io.load_tiles(onlydesi=False)
        tile_cache_id3 = id(list(io._tiles.items())[0])
        self.assertEqual(tile_cache_id1, tile_cache_id3)
        self.assertTrue(np.issubdtype(t3['OBSCONDITIONS'].dtype, 'i4') or \
                        np.issubdtype(t3['OBSCONDITIONS'].dtype, 'u2') )
        # Check for extra tiles.
        a = io.load_tiles(extra=False)
        self.assertEqual(np.sum(np.char.startswith(a['PROGRAM'], 'EXTRA')), 0)
        b = io.load_tiles(extra=True)
        self.assertGreater(np.sum(np.char.startswith(b['PROGRAM'], 'EXTRA')), 0)
        self.assertLess(len(a), len(b))

    @unittest.skipUnless(desimodel_available, desimodel_message)
    def test_load_tiles_alt(self):
        # starting clean
        self.assertEqual(io._tiles, {})
        t1 = Table(io.load_tiles())
        t1.write(self.testfile)
        #- no path; should fail since that file isn't in $DESIMODEL/data/footprint/
        with self.assertRaises(FileNotFoundError):
            t2 = io.load_tiles(tilesfile=self.testfile)

        #- with path, should work:
        t2 = Table(io.load_tiles(tilesfile='./'+self.testfile))
        self.assertTrue(np.all(t1 == t2))

        #- cache should have two items now
        self.assertEqual(len(io._tiles), 2)

    @unittest.skipUnless(desimodel_available, desimodel_message)
    def test_tiles_consistency(self):
        """Test consistency of tile files.

        - Validate different tile loading schemes.
        - Test numerical consistency between FITS & ECSV versions.
        - Assure that PROGRAM column has no trailing whitespace.
        """
        fitstiles = io.findfile('footprint/desi-tiles.fits')
        ecsvtiles = io.findfile('footprint/desi-tiles.ecsv')
        # tf=TilesFits  tt=TilesTable  te=TilesEcsv
        tf = io.load_tiles(onlydesi=False, extra=True)
        tt = Table.read(fitstiles)
        te = Table.read(ecsvtiles, format='ascii.ecsv')
        self.assertEqual(sorted(tf.dtype.names), sorted(tt.colnames))
        self.assertEqual(sorted(tf.dtype.names), sorted(te.colnames))

        for col in tt.colnames:
            self.assertTrue(np.all(tf[col]==tt[col]), 'fits[{col}] != table[{col}]'.format(col=col))
            if np.issubdtype(tf[col].dtype, float):
                self.assertTrue(np.allclose(tf[col], te[col], atol=1e-4, rtol=1e-4), 'fits[{col}] != ecsv[{col}]'.format(col=col))
            else:
                self.assertTrue(np.all(tf[col]==te[col]), 'fits[{col}] != ecsv[{col}]'.format(col=col))

        self.assertTrue(not np.any(np.char.endswith(tf['PROGRAM'], ' ')))
        self.assertTrue(not np.any(np.char.endswith(tt['PROGRAM'], ' ')))
        self.assertTrue(not np.any(np.char.endswith(te['PROGRAM'], ' ')))

    @unittest.skipUnless(desimodel_available, desimodel_message)
    def test_trim_data(self):
        '''Test trimming data files for lightweight tests'''
        psffile = io.findfile('specpsf/psf-b.fits')
        if os.path.getsize(psffile) > 20e6:
            from .. import trim
            indir = os.path.join(os.getenv('DESIMODEL'), 'data')
            trim.trim_data(indir, self.trimdir)
            self.assertTrue(os.path.isdir(self.trimdir))
            self.assertGreater(len(list(os.walk(self.trimdir))), 1)

def test_suite():
    """Allows testing of only this module with the command::
        python setup.py test -m <modulename>
    """
    return unittest.defaultTestLoader.loadTestsFromName(__name__)
<|endoftext|>


<commit_before>
import cPickle
from glob import glob
import gzip
import hashlib
from hashlib import md5
import json
import os
import re
import shutil

import numpy as np

from ..interfaces.traits_extension import isdefined
from .misc import is_container

from .. import logging, config
fmlogger = logging.getLogger("filemanip")

def update_minchd_json(filename, data_in, var, attr):
    data=load_json(filename)
    if data.get(var,'None') == 'None':
    	data[var]={attr:data_in}
    else:
    	if data.get(var).get(attr,'None') == 'None':
	    data[var][attr]=data_in
    	else:
	    data[var][attr]=[data[var][attr],data_in]

    fp = file(filename, 'w')
    json.dump(data, fp, sort_keys=True, indent=4)
    fp.close()
<commit_msg>Remove update_minchd_json.py<commit_after><|endoftext|>


<commit_before># Copyright (c) 2007 Ferran Pegueroles <ferran@pegueroles.com>
# Copyright (c) 2009 Albert Cervera i Areny <albert@nan-tic.com>
# Copyright (C) 2011 Agile Business Group sagl (<http://www.agilebg.com>)
# Copyright (C) 2011 Domsense srl (<http://www.domsense.com>)
# Copyright (C) 2013-2014 Camptocamp (<http://www.camptocamp.com>)
# License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl).

{
    "name": "Report to printer",
    "version": "14.0.1.1.0",
    "category": "Generic Modules/Base",
    "author": "Agile Business Group & Domsense, Pegueroles SCP, NaN,"
    " LasLabs, Camptocamp, Odoo Community Association (OCA),"
    " Open for Small Business Ltd",
    "website": "https://github.com/OCA/report-print-send",
    "license": "AGPL-3",
    "depends": ["web"],
    "data": [
        "data/printing_data.xml",
        "security/security.xml",
        "views/assets.xml",
        "views/printing_printer.xml",
        "views/printing_server.xml",
        "views/printing_job.xml",
        "views/printing_report.xml",
        "views/res_users.xml",
        "views/ir_actions_report.xml",
        "wizards/printing_printer_update_wizard_view.xml",
    ],
    "installable": True,
    "application": False,
    "external_dependencies": {"python": ["pycups"]},
}
<commit_msg>base_report_to_printer 14.0.1.1.1<commit_after># Copyright (c) 2007 Ferran Pegueroles <ferran@pegueroles.com>
# Copyright (c) 2009 Albert Cervera i Areny <albert@nan-tic.com>
# Copyright (C) 2011 Agile Business Group sagl (<http://www.agilebg.com>)
# Copyright (C) 2011 Domsense srl (<http://www.domsense.com>)
# Copyright (C) 2013-2014 Camptocamp (<http://www.camptocamp.com>)
# License AGPL-3.0 or later (http://www.gnu.org/licenses/agpl).

{
    "name": "Report to printer",
    "version": "14.0.1.1.1",
    "category": "Generic Modules/Base",
    "author": "Agile Business Group & Domsense, Pegueroles SCP, NaN,"
    " LasLabs, Camptocamp, Odoo Community Association (OCA),"
    " Open for Small Business Ltd",
    "website": "https://github.com/OCA/report-print-send",
    "license": "AGPL-3",
    "depends": ["web"],
    "data": [
        "data/printing_data.xml",
        "security/security.xml",
        "views/assets.xml",
        "views/printing_printer.xml",
        "views/printing_server.xml",
        "views/printing_job.xml",
        "views/printing_report.xml",
        "views/res_users.xml",
        "views/ir_actions_report.xml",
        "wizards/printing_printer_update_wizard_view.xml",
    ],
    "installable": True,
    "application": False,
    "external_dependencies": {"python": ["pycups"]},
}
<|endoftext|>


<commit_before>#!/usr/bin/env python
# Copyright (C) 2010-2013 Claudio Guarnieri.
# Copyright (C) 2014-2016 Cuckoo Foundation.
# This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
# See the file 'docs/LICENSE' for copying permission.

import os
import sys
import shutil
import urllib2
import argparse
import tempfile
from tarfile import TarFile

sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)), ".."))

import lib.cuckoo.common.colors as colors
from lib.cuckoo.common.constants import CUCKOO_ROOT

URL = "https://github.com/cuckoosandbox/community/archive/{0}.tar.gz"

def download_archive():
    print("Downloading modules from {0}".format(URL))

    try:
        data = urllib2.urlopen(URL).read()
    except Exception as e:
        print("ERROR: Unable to download archive: %s" % e)
        sys.exit(-1)

    return data

def extract_archive(data):
    fd, filepath = tempfile.mkstemp()
    os.write(fd, data)
    os.close(fd)

    archive = TarFile.open(filepath, mode="r:gz")
    temp_dir = tempfile.mkdtemp()
    archive.extractall(temp_dir)
    archive.close()
    os.unlink(filepath)
    return temp_dir

def installdir(src, dst, force, rewrite, origin=[]):
    for file_name in os.listdir(src):
        if file_name == ".gitignore":
            continue

        destination = os.path.join(dst, file_name)

        if not rewrite:
            if os.path.exists(destination):
                print("File \"{0}\" already exists, "
                      "{1}".format(file_name, colors.yellow("skipped")))
                continue

        install = False

        if not force:
            while 1:
                choice = raw_input("Do you want to install file "
                                   "\"{0}\"? [yes/no] ".format(file_name))
                if choice.lower() == "yes":
                    install = True
                    break
                elif choice.lower() == "no":
                    break
                else:
                    continue
        else:
            install = True

        if install:
            srcpath = os.path.join(src, file_name)
            if os.path.islink(srcpath):
                if os.path.lexists(destination):
                    try:
                        shutil.rmtree(destination)
                    except OSError:
                        os.unlink(destination)
                os.symlink(os.readlink(srcpath), destination)
                print "Symbolic link \"%s/%s\" -> \"%s\" %s" % (
                    "/".join(origin), file_name, os.readlink(srcpath),
                    colors.green("installed"))

            elif os.path.isdir(srcpath):
                installdir(srcpath, destination, force, rewrite,
                           origin + [file_name])
            else:
                if not os.path.isdir(os.path.dirname(destination)):
                    os.makedirs(os.path.dirname(destination))

                shutil.copy(srcpath, destination)
                print "File \"%s/%s\" %s" % (
                    "/".join(origin), file_name, colors.green("installed"))

def install(enabled, force, rewrite, archive):
    if archive:
        if not os.path.isfile(archive):
            print("ERROR: Provided archive not found!")
            sys.exit(-1)

        data = open(archive, "rb").read()
    else:
        data = download_archive()

    source = extract_archive(data)

    folders = {
        "signatures": os.path.join("modules", "signatures"),
        "processing": os.path.join("modules", "processing"),
        "reporting": os.path.join("modules", "reporting"),
        "machinery": os.path.join("modules", "machinery"),
        "analyzer": os.path.join("analyzer"),
        "monitor": os.path.join("data", "monitor"),
        "agent": os.path.join("agent"),
    }

    for category in enabled:
        folder = folders[category]

        print("\nInstalling {0}".format(colors.cyan(category.upper())))

        origin = os.path.join(source, folder)
        if not os.path.isdir(origin):
            print "  No candidates available, continuing."
            continue

        installdir(origin, os.path.join(CUCKOO_ROOT, folder), force, rewrite)

    shutil.rmtree(source)

def main():
    global URL

    parser = argparse.ArgumentParser()
    parser.add_argument("-a", "--all", help="Download everything", action="store_true", required=False)
    parser.add_argument("-s", "--signatures", help="Download Cuckoo signatures", action="store_true", required=False)
    parser.add_argument("-p", "--processing", help="Download processing modules", action="store_true", required=False)
    parser.add_argument("-m", "--machinery", help="Download machine managers", action="store_true", required=False)
    parser.add_argument("-n", "--analyzer", help="Download analyzer modules", action="store_true", required=False)
    parser.add_argument("-M", "--monitor", help="Download monitoring binaries", action="store_true", required=False)
    parser.add_argument("-g", "--agent", help="Download agent modules", action="store_true", required=False)
    parser.add_argument("-r", "--reporting", help="Download reporting modules", action="store_true", required=False)
    parser.add_argument("-f", "--force", help="Install files without confirmation", action="store_true", required=False)
    parser.add_argument("-w", "--rewrite", help="Rewrite existing files", action="store_true", required=False)
    parser.add_argument("-b", "--branch", help="Specify a different branch", action="store", default="master", required=False)
    parser.add_argument("archive", help="Install a stored archive", nargs="?")
    args = parser.parse_args()

    enabled = []
    force = False
    rewrite = False

    if args.all:
        enabled.append("processing")
        enabled.append("signatures")
        enabled.append("reporting")
        enabled.append("machinery")
        enabled.append("analyzer")
        enabled.append("monitor")
        enabled.append("agent")
    else:
        if args.signatures:
            enabled.append("signatures")
        if args.processing:
            enabled.append("processing")
        if args.reporting:
            enabled.append("reporting")
        if args.machinery:
            enabled.append("machinery")
        if args.analyzer:
            enabled.append("analyzer")
        if args.agent:
            enabled.append("agent")

    if not enabled:
        print(colors.red("You need to enable some category!\n"))
        parser.print_help()
        return

    if args.force:
        force = True
    if args.rewrite:
        rewrite = True

    URL = URL.format(args.branch)

    install(enabled, force, rewrite, args.archive)

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        pass
<commit_msg>Revert "fix age-old community tarfile unpack issue"<commit_after>#!/usr/bin/env python
# Copyright (C) 2010-2013 Claudio Guarnieri.
# Copyright (C) 2014-2016 Cuckoo Foundation.
# This file is part of Cuckoo Sandbox - http://www.cuckoosandbox.org
# See the file 'docs/LICENSE' for copying permission.

import os
import sys
import shutil
import urllib2
import argparse
import tempfile
from tarfile import TarFile

sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)), ".."))

import lib.cuckoo.common.colors as colors
from lib.cuckoo.common.constants import CUCKOO_ROOT

URL = "https://github.com/cuckoosandbox/community/archive/{0}.tar.gz"

def download_archive():
    print("Downloading modules from {0}".format(URL))

    try:
        data = urllib2.urlopen(URL).read()
    except Exception as e:
        print("ERROR: Unable to download archive: %s" % e)
        sys.exit(-1)

    return data

def extract_archive(data):
    fd, filepath = tempfile.mkstemp()
    os.write(fd, data)
    os.close(fd)

    archive = TarFile.open(filepath, mode="r:gz")
    temp_dir = tempfile.mkdtemp()
    archive.extractall(temp_dir)
    archive.close()
    os.unlink(filepath)
    final_dir = os.path.join(temp_dir, os.listdir(temp_dir)[0])

    return temp_dir, final_dir

def installdir(src, dst, force, rewrite, origin=[]):
    for file_name in os.listdir(src):
        if file_name == ".gitignore":
            continue

        destination = os.path.join(dst, file_name)

        if not rewrite:
            if os.path.exists(destination):
                print("File \"{0}\" already exists, "
                      "{1}".format(file_name, colors.yellow("skipped")))
                continue

        install = False

        if not force:
            while 1:
                choice = raw_input("Do you want to install file "
                                   "\"{0}\"? [yes/no] ".format(file_name))
                if choice.lower() == "yes":
                    install = True
                    break
                elif choice.lower() == "no":
                    break
                else:
                    continue
        else:
            install = True

        if install:
            srcpath = os.path.join(src, file_name)
            if os.path.islink(srcpath):
                if os.path.lexists(destination):
                    try:
                        shutil.rmtree(destination)
                    except OSError:
                        os.unlink(destination)
                os.symlink(os.readlink(srcpath), destination)
                print "Symbolic link \"%s/%s\" -> \"%s\" %s" % (
                    "/".join(origin), file_name, os.readlink(srcpath),
                    colors.green("installed"))

            elif os.path.isdir(srcpath):
                installdir(srcpath, destination, force, rewrite,
                           origin + [file_name])
            else:
                if not os.path.isdir(os.path.dirname(destination)):
                    os.makedirs(os.path.dirname(destination))

                shutil.copy(srcpath, destination)
                print "File \"%s/%s\" %s" % (
                    "/".join(origin), file_name, colors.green("installed"))

def install(enabled, force, rewrite, archive):
    if archive:
        if not os.path.isfile(archive):
            print("ERROR: Provided archive not found!")
            sys.exit(-1)

        data = open(archive, "rb").read()
    else:
        data = download_archive()

    temp, source = extract_archive(data)

    folders = {
        "signatures": os.path.join("modules", "signatures"),
        "processing": os.path.join("modules", "processing"),
        "reporting": os.path.join("modules", "reporting"),
        "machinery": os.path.join("modules", "machinery"),
        "analyzer": os.path.join("analyzer"),
        "monitor": os.path.join("data", "monitor"),
        "agent": os.path.join("agent"),
    }

    for category in enabled:
        folder = folders[category]

        print("\nInstalling {0}".format(colors.cyan(category.upper())))

        origin = os.path.join(source, folder)
        if not os.path.isdir(origin):
            print "  No candidates available, continuing."
            continue

        installdir(origin, os.path.join(CUCKOO_ROOT, folder), force, rewrite)

    shutil.rmtree(temp)

def main():
    global URL

    parser = argparse.ArgumentParser()
    parser.add_argument("-a", "--all", help="Download everything", action="store_true", required=False)
    parser.add_argument("-s", "--signatures", help="Download Cuckoo signatures", action="store_true", required=False)
    parser.add_argument("-p", "--processing", help="Download processing modules", action="store_true", required=False)
    parser.add_argument("-m", "--machinery", help="Download machine managers", action="store_true", required=False)
    parser.add_argument("-n", "--analyzer", help="Download analyzer modules", action="store_true", required=False)
    parser.add_argument("-M", "--monitor", help="Download monitoring binaries", action="store_true", required=False)
    parser.add_argument("-g", "--agent", help="Download agent modules", action="store_true", required=False)
    parser.add_argument("-r", "--reporting", help="Download reporting modules", action="store_true", required=False)
    parser.add_argument("-f", "--force", help="Install files without confirmation", action="store_true", required=False)
    parser.add_argument("-w", "--rewrite", help="Rewrite existing files", action="store_true", required=False)
    parser.add_argument("-b", "--branch", help="Specify a different branch", action="store", default="master", required=False)
    parser.add_argument("archive", help="Install a stored archive", nargs="?")
    args = parser.parse_args()

    enabled = []
    force = False
    rewrite = False

    if args.all:
        enabled.append("processing")
        enabled.append("signatures")
        enabled.append("reporting")
        enabled.append("machinery")
        enabled.append("analyzer")
        enabled.append("monitor")
        enabled.append("agent")
    else:
        if args.signatures:
            enabled.append("signatures")
        if args.processing:
            enabled.append("processing")
        if args.reporting:
            enabled.append("reporting")
        if args.machinery:
            enabled.append("machinery")
        if args.analyzer:
            enabled.append("analyzer")
        if args.agent:
            enabled.append("agent")

    if not enabled:
        print(colors.red("You need to enable some category!\n"))
        parser.print_help()
        return

    if args.force:
        force = True
    if args.rewrite:
        rewrite = True

    URL = URL.format(args.branch)

    install(enabled, force, rewrite, args.archive)

if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        pass
<|endoftext|>


<commit_before>import asyncio
import discord
from discord.ext import commands
from sys import argv

class Lockdown:
    """
    Channel lockdown commands.
    """
    def __init__(self, bot):
        self.bot = bot
        print('Addon "{}" loaded'.format(self.__class__.__name__))

    @commands.has_permissions(manage_messages=True)
    @commands.command(pass_context=True, name="lock")
    async def lock(self, ctx):
        """Lock message sending in the channel. Staff only."""
        try:
            if ctx.message.channel in self.bot.community_channels:
                roles = (self.bot.community_role, self.bot.hacker_role)
            else:
                roles = (self.bot.everyone_role,)
            overwrites = ctx.message.channel.overwrites_for(roles[0])
            if overwrites.send_messages is False:
                await self.bot.say("🔒 Channel is already locked down. Use `.unlock [community]` to unlock.")
                return
            overwrites.send_messages = False
            overwrites.add_reactions = False
            await asyncio.gather(*(self.bot.edit_channel_permissions(ctx.message.channel, role, overwrites) for role in roles))
            await self.bot.say("🔒 Channel locked down. Only staff members may speak. Do not bring the topic to other channels or risk disciplinary actions.")
            msg = "🔒 **Lockdown**: {0} by {1} | {2}#{3}".format(ctx.message.channel.mention, ctx.message.author.mention, ctx.message.author.name, ctx.message.author.discriminator)
            await self.bot.send_message(self.bot.modlogs_channel, msg)

        except discord.errors.Forbidden:
            await self.bot.say("💢 I don't have permission to do this.")

    @commands.has_permissions(manage_messages=True)
    @commands.command(pass_context=True, name="softlock")
    async def softlock(self, ctx):
        """Lock message sending in the channel, without the "disciplinary action" note. Staff only."""
        try:
            if ctx.message.channel in self.bot.community_channels:
                roles = (self.bot.community_role, self.bot.hacker_role)
            else:
                roles = (self.bot.everyone_role,)
            overwrites = ctx.message.channel.overwrites_for(roles[0])
            if overwrites.send_messages is False:
                await self.bot.say("🔒 Channel is already locked down. Use `.unlock [community]` to unlock.")
                return
            overwrites.send_messages = False
            overwrites.add_reactions = False
            await asyncio.gather(*(self.bot.edit_channel_permissions(ctx.message.channel, role, overwrites) for role in roles))
            await self.bot.say("🔒 Channel locked.")
            msg = "🔒 **Soft-lock**: {0} by {1} | {2}#{3}".format(ctx.message.channel.mention, ctx.message.author.mention, ctx.message.author.name, ctx.message.author.discriminator)
            await self.bot.send_message(self.bot.modlogs_channel, msg)
        except discord.errors.Forbidden:
            await self.bot.say("💢 I don't have permission to do this.")

    @commands.has_permissions(manage_messages=True)
    @commands.command(pass_context=True, name="unlock")
    async def unlock(self, ctx):
        """Unlock message sending in the channel. Staff only."""
        try:
            if ctx.message.channel in self.bot.community_channels:
                roles = (self.bot.community_role, self.bot.hacker_role)
            else:
                roles = (self.bot.everyone_role,)
            overwrites = ctx.message.channel.overwrites_for(roles[0])
            if overwrites.send_messages is True:
                await self.bot.say("🔓 Channel is already unlocked.")
                return
            overwrites.send_messages = True
            overwrites.add_reactions = True
            await asyncio.gather(*(self.bot.edit_channel_permissions(ctx.message.channel, role, overwrites) for role in roles))
            await self.bot.say("🔓 Channel unlocked.")
            msg = "🔓 **Unlock**: {0} by {1} | {2}#{3}".format(ctx.message.channel.mention, ctx.message.author.mention, ctx.message.author.name, ctx.message.author.discriminator)
            await self.bot.send_message(self.bot.modlogs_channel, msg)

        except discord.errors.Forbidden:
            await self.bot.say("💢 I don't have permission to do this.")

def setup(bot):
    bot.add_cog(Lockdown(bot))
<commit_msg>Switch @everyone with participant for locks<commit_after>import asyncio
import discord
from discord.ext import commands
from sys import argv

class Lockdown:
    """
    Channel lockdown commands.
    """
    def __init__(self, bot):
        self.bot = bot
        print('Addon "{}" loaded'.format(self.__class__.__name__))

    @commands.has_permissions(manage_messages=True)
    @commands.command(pass_context=True, name="lock")
    async def lock(self, ctx):
        """Lock message sending in the channel. Staff only."""
        try:
            if ctx.message.channel in self.bot.community_channels:
                roles = (self.bot.community_role, self.bot.hacker_role)
            else:
                roles = (self.bot.unprobated_role,)
            overwrites = ctx.message.channel.overwrites_for(roles[0])
            if overwrites.send_messages is False:
                await self.bot.say("🔒 Channel is already locked down. Use `.unlock [community]` to unlock.")
                return
            overwrites.send_messages = False
            overwrites.add_reactions = False
            await asyncio.gather(*(self.bot.edit_channel_permissions(ctx.message.channel, role, overwrites) for role in roles))
            await self.bot.say("🔒 Channel locked down. Only staff members may speak. Do not bring the topic to other channels or risk disciplinary actions.")
            msg = "🔒 **Lockdown**: {0} by {1} | {2}#{3}".format(ctx.message.channel.mention, ctx.message.author.mention, ctx.message.author.name, ctx.message.author.discriminator)
            await self.bot.send_message(self.bot.modlogs_channel, msg)

        except discord.errors.Forbidden:
            await self.bot.say("💢 I don't have permission to do this.")

    @commands.has_permissions(manage_messages=True)
    @commands.command(pass_context=True, name="softlock")
    async def softlock(self, ctx):
        """Lock message sending in the channel, without the "disciplinary action" note. Staff only."""
        try:
            if ctx.message.channel in self.bot.community_channels:
                roles = (self.bot.community_role, self.bot.hacker_role)
            else:
                roles = (self.bot.unprobated_role,)
            overwrites = ctx.message.channel.overwrites_for(roles[0])
            if overwrites.send_messages is False:
                await self.bot.say("🔒 Channel is already locked down. Use `.unlock [community]` to unlock.")
                return
            overwrites.send_messages = False
            overwrites.add_reactions = False
            await asyncio.gather(*(self.bot.edit_channel_permissions(ctx.message.channel, role, overwrites) for role in roles))
            await self.bot.say("🔒 Channel locked.")
            msg = "🔒 **Soft-lock**: {0} by {1} | {2}#{3}".format(ctx.message.channel.mention, ctx.message.author.mention, ctx.message.author.name, ctx.message.author.discriminator)
            await self.bot.send_message(self.bot.modlogs_channel, msg)
        except discord.errors.Forbidden:
            await self.bot.say("💢 I don't have permission to do this.")

    @commands.has_permissions(manage_messages=True)
    @commands.command(pass_context=True, name="unlock")
    async def unlock(self, ctx):
        """Unlock message sending in the channel. Staff only."""
        try:
            if ctx.message.channel in self.bot.community_channels:
                roles = (self.bot.community_role, self.bot.hacker_role)
            else:
                roles = (self.bot.unprobated_role,)
            overwrites = ctx.message.channel.overwrites_for(roles[0])
            if overwrites.send_messages is True:
                await self.bot.say("🔓 Channel is already unlocked.")
                return
            overwrites.send_messages = True
            overwrites.add_reactions = True
            await asyncio.gather(*(self.bot.edit_channel_permissions(ctx.message.channel, role, overwrites) for role in roles))
            await self.bot.say("🔓 Channel unlocked.")
            msg = "🔓 **Unlock**: {0} by {1} | {2}#{3}".format(ctx.message.channel.mention, ctx.message.author.mention, ctx.message.author.name, ctx.message.author.discriminator)
            await self.bot.send_message(self.bot.modlogs_channel, msg)

        except discord.errors.Forbidden:
            await self.bot.say("💢 I don't have permission to do this.")

def setup(bot):
    bot.add_cog(Lockdown(bot))
<|endoftext|>