<commit_before># Copyright (c) 2007-2011 gocept gmbh & co. kg
# See also LICENSE.txt

from zeit.cms.i18n import MessageFactory as _
import PIL.Image
import transaction
import urlparse
import z3c.conditionalviews
import zeit.cms.browser.interfaces
import zeit.cms.browser.listing
import zeit.cms.browser.view
import zeit.cms.content.property
import zeit.cms.repository.interfaces
import zeit.cms.settings.interfaces
import zeit.connector.interfaces
import zeit.content.image.imagereference
import zeit.content.image.interfaces
import zope.cachedescriptors.property
import zope.component
import zope.file.download
import zope.publisher.interfaces

def get_img_tag(image, request, view=None):
    """Render -tag."""
    url = zope.component.getMultiAdapter(
        (image, request), name='absolute_url')
    width, height = image.getImageSize()
    if view:
        view = '/' + view
    else:
        view = ''
    return (
        '' % (
        url, view, height, width))

class Image(zope.file.download.Display):

    def __call__(self):
        self.request.response.setHeader('Content-Type', self.context.mimeType)
        return self.stream_image()

    @z3c.conditionalviews.ConditionalView
    def stream_image(self):
        return super(Image, self).__call__()

class ImageView(zeit.cms.browser.view.Base):

    title = _('View image')

    @zope.cachedescriptors.property.Lazy
    def metadata(self):
        return zeit.content.image.interfaces.IImageMetadata(self.context)

    def tag(self):
        return get_img_tag(self.context, self.request, view='@@raw')

    @property
    def width(self):
        return self.context.getImageSize()[0]

    @property
    def height(self):
        return self.context.getImageSize()[1]

    @property
    def copyrights(self):
        result = []
        for copyright, url, nofollow in self.metadata.copyrights:
            result.append(dict(
                copyright=copyright,
                url=url,
                nofollow=nofollow))
        return result

class Scaled(object):

    filter = PIL.Image.ANTIALIAS

    def __call__(self):
        return self.scaled()

    def tag(self):
        return get_img_tag(self.scaled.context, self.request)

    @zope.cachedescriptors.property.Lazy
    def scaled(self):
        try:
            image = zeit.content.image.interfaces.ITransform(self.context)
        except TypeError:
            image = self.context
        else:
            image = image.thumbnail(self.width, self.height, self.filter)
            image.__name__ = self.__name__

            def cleanup(commited, image):
                # Releasing the last reference triggers the weakref cleanup of
                # ZODB.blob.Blob, since this local_data Blob never was part of
                # a ZODB connection, which will delete the temporary file.
                image.local_data = None
            transaction.get().addAfterCommitHook(cleanup, [image])
        image_view = zope.component.getMultiAdapter(
            (image, self.request), name='raw')
        return image_view

class Preview(Scaled):

    width = 500
    height = 500

class MetadataPreview(Scaled):

    width = 500
    height = 90

class Thumbnail(Scaled):

    width = height = 100

class ImageListRepresentation(
    zeit.cms.browser.listing.BaseListRepresentation):
    """Adapter for listing article content resources"""

    zope.interface.implements(zeit.cms.browser.interfaces.IListRepresentation)
    zope.component.adapts(zeit.content.image.interfaces.IImage,
                          zope.publisher.interfaces.IPublicationRequest)

    author = ressort = page = u''

    @property
    def title(self):
        title = self.image_metadata.title
        if not title:
            title = self.context.__name__
        return title

    @property
    def volume(self):
        return self.image_metadata.volume

    @property
    def year(self):
        return self.image_metadata.year

    @property
    def searchableText(self):
        # XXX
        return ''

    @zope.cachedescriptors.property.Lazy
    def image_metadata(self):
        return zeit.content.image.interfaces.IImageMetadata(self.context)

@zope.component.adapter(
    zeit.cms.repository.interfaces.IFolder,
    zeit.content.image.interfaces.IImageSource)
@zope.interface.implementer(
    zeit.cms.browser.interfaces.IDefaultBrowsingLocation)
def imagefolder_browse_location(context, source):
    """The image browse location is deduced from the current folder, i.e.

        for /online/2007/32 it is /bilder/2007/32

    """
    unique_id = context.uniqueId
    repository = zope.component.getUtility(
        zeit.cms.repository.interfaces.IRepository)
    base = image_folder = None
    try:
        obj_in_repository = repository.getContent(unique_id)
    except KeyError:
        pass
    else:
        # Try to get a base folder
        while base is None:
            properties = zeit.connector.interfaces.IWebDAVProperties(
                obj_in_repository, None)
            if properties is None:
                break
            base = properties.get(('base-folder',
                                   'http://namespaces.zeit.de/CMS/Image'))
            obj_in_repository = obj_in_repository.__parent__

    if base is not None:
        try:
            base_obj = repository.getContent(base)
        except KeyError:
            pass
        else:
            # Get from the base folder to the year/volume folder
            settings = zeit.cms.settings.interfaces.IGlobalSettings(context)
            try:
                image_folder = base_obj[
                    '%04d' % settings.default_year][
                    '%02d' % settings.default_volume]
            except KeyError:
                pass

    if image_folder is None:
        all_content_source = zope.component.getUtility(
            zeit.cms.content.interfaces.ICMSContentSource, name='all-types')
        image_folder = zope.component.queryMultiAdapter(
            (context, all_content_source),
            zeit.cms.browser.interfaces.IDefaultBrowsingLocation)

    return image_folder

@zope.component.adapter(
    zeit.content.image.imagereference.ImagesAdapter,
    zeit.content.image.interfaces.IImageSource)
@zope.interface.implementer(
    zeit.cms.browser.interfaces.IDefaultBrowsingLocation)
def imageadapter_browse_location(context, source):
    return zope.component.queryMultiAdapter(
        (context.__parent__, source),
        zeit.cms.browser.interfaces.IDefaultBrowsingLocation)

class MetadataPreviewHTML(object):

    @zope.cachedescriptors.property.Lazy
    def metadata(self):
        return zeit.content.image.interfaces.IImageMetadata(self.context)
<commit_msg>Simplify spelling<commit_after># Copyright (c) 2007-2011 gocept gmbh & co. kg
# See also LICENSE.txt

from zeit.cms.i18n import MessageFactory as _
from zope.cachedescriptors.property import Lazy as cachedproperty
import PIL.Image
import transaction
import z3c.conditionalviews
import zeit.cms.browser.interfaces
import zeit.cms.browser.listing
import zeit.cms.browser.view
import zeit.cms.repository.interfaces
import zeit.cms.settings.interfaces
import zeit.connector.interfaces
import zeit.content.image.imagereference
import zeit.content.image.interfaces
import zope.component
import zope.file.download
import zope.publisher.interfaces

def get_img_tag(image, request, view=None):
    """Render -tag."""
    url = zope.component.getMultiAdapter(
        (image, request), name='absolute_url')
    width, height = image.getImageSize()
    if view:
        view = '/' + view
    else:
        view = ''
    return (
        '' % (
        url, view, height, width))

class Image(zope.file.download.Display):

    def __call__(self):
        self.request.response.setHeader('Content-Type', self.context.mimeType)
        return self.stream_image()

    @z3c.conditionalviews.ConditionalView
    def stream_image(self):
        return super(Image, self).__call__()

class ImageView(zeit.cms.browser.view.Base):

    title = _('View image')

    @cachedproperty
    def metadata(self):
        return zeit.content.image.interfaces.IImageMetadata(self.context)

    def tag(self):
        return get_img_tag(self.context, self.request, view='@@raw')

    @property
    def width(self):
        return self.context.getImageSize()[0]

    @property
    def height(self):
        return self.context.getImageSize()[1]

    @property
    def copyrights(self):
        result = []
        for copyright, url, nofollow in self.metadata.copyrights:
            result.append(dict(
                copyright=copyright,
                url=url,
                nofollow=nofollow))
        return result

class Scaled(object):

    filter = PIL.Image.ANTIALIAS

    def __call__(self):
        return self.scaled()

    def tag(self):
        return get_img_tag(self.scaled.context, self.request)

    @cachedproperty
    def scaled(self):
        try:
            image = zeit.content.image.interfaces.ITransform(self.context)
        except TypeError:
            image = self.context
        else:
            image = image.thumbnail(self.width, self.height, self.filter)
            image.__name__ = self.__name__

            def cleanup(commited, image):
                # Releasing the last reference triggers the weakref cleanup of
                # ZODB.blob.Blob, since this local_data Blob never was part of
                # a ZODB connection, which will delete the temporary file.
                image.local_data = None
            transaction.get().addAfterCommitHook(cleanup, [image])
        image_view = zope.component.getMultiAdapter(
            (image, self.request), name='raw')
        return image_view

class Preview(Scaled):

    width = 500
    height = 500

class MetadataPreview(Scaled):

    width = 500
    height = 90

class Thumbnail(Scaled):

    width = height = 100

class ImageListRepresentation(
    zeit.cms.browser.listing.BaseListRepresentation):
    """Adapter for listing article content resources"""

    zope.interface.implements(zeit.cms.browser.interfaces.IListRepresentation)
    zope.component.adapts(zeit.content.image.interfaces.IImage,
                          zope.publisher.interfaces.IPublicationRequest)

    author = ressort = page = u''

    @property
    def title(self):
        title = self.image_metadata.title
        if not title:
            title = self.context.__name__
        return title

    @property
    def volume(self):
        return self.image_metadata.volume

    @property
    def year(self):
        return self.image_metadata.year

    @property
    def searchableText(self):
        # XXX
        return ''

    @cachedproperty
    def image_metadata(self):
        return zeit.content.image.interfaces.IImageMetadata(self.context)

@zope.component.adapter(
    zeit.cms.repository.interfaces.IFolder,
    zeit.content.image.interfaces.IImageSource)
@zope.interface.implementer(
    zeit.cms.browser.interfaces.IDefaultBrowsingLocation)
def imagefolder_browse_location(context, source):
    """The image browse location is deduced from the current folder, i.e.

        for /online/2007/32 it is /bilder/2007/32

    """
    unique_id = context.uniqueId
    repository = zope.component.getUtility(
        zeit.cms.repository.interfaces.IRepository)
    base = image_folder = None
    try:
        obj_in_repository = repository.getContent(unique_id)
    except KeyError:
        pass
    else:
        # Try to get a base folder
        while base is None:
            properties = zeit.connector.interfaces.IWebDAVProperties(
                obj_in_repository, None)
            if properties is None:
                break
            base = properties.get(('base-folder',
                                   'http://namespaces.zeit.de/CMS/Image'))
            obj_in_repository = obj_in_repository.__parent__

    if base is not None:
        try:
            base_obj = repository.getContent(base)
        except KeyError:
            pass
        else:
            # Get from the base folder to the year/volume folder
            settings = zeit.cms.settings.interfaces.IGlobalSettings(context)
            try:
                image_folder = base_obj[
                    '%04d' % settings.default_year][
                    '%02d' % settings.default_volume]
            except KeyError:
                pass

    if image_folder is None:
        all_content_source = zope.component.getUtility(
            zeit.cms.content.interfaces.ICMSContentSource, name='all-types')
        image_folder = zope.component.queryMultiAdapter(
            (context, all_content_source),
            zeit.cms.browser.interfaces.IDefaultBrowsingLocation)

    return image_folder

@zope.component.adapter(
    zeit.content.image.imagereference.ImagesAdapter,
    zeit.content.image.interfaces.IImageSource)
@zope.interface.implementer(
    zeit.cms.browser.interfaces.IDefaultBrowsingLocation)
def imageadapter_browse_location(context, source):
    return zope.component.queryMultiAdapter(
        (context.__parent__, source),
        zeit.cms.browser.interfaces.IDefaultBrowsingLocation)

class MetadataPreviewHTML(object):

    @cachedproperty
    def metadata(self):
        return zeit.content.image.interfaces.IImageMetadata(self.context)
<|endoftext|>


<commit_before>from datetime import datetime, timedelta
import json
import logging
import os
from os.path import splitext, basename, exists, dirname
import sys
import time

import datadog

logging.basicConfig(level=logging.INFO, stream=sys.stdout)
logging.getLogger("requests").setLevel(logging.WARNING)
logging.getLogger("dd").setLevel(logging.WARNING)
LOGGER = logging.getLogger(__name__)

"""
Originally written by 'Jharrod LaFon'
#https://github.com/jlafon/ansible-profile/blob/master/callback_plugins/profile_tasks.py

"""

ANSIBLE_TIMER_LOG = os.environ.get('ANSIBLE_TIMER_LOG')

class Timestamp(object):
    def __init__(self):
        self.start = datetime.utcnow()
        self.end = None

    def stop(self):
        self.end = datetime.utcnow()

    @property
    def duration(self):
        return self.end - self.start

class Formatter(object):
    pass

class DatadogFormatter(Formatter):
    def __init__(self):
        super(DatadogFormatter, self).__init__()

        self.datadog_api_key = os.getenv('DATADOG_API_KEY')
        self.datadog_api_initialized = False

        if self.datadog_api_key:
            datadog.initialize(
                api_key=self.datadog_api_key,
                app_key=None
            )
            self.datadog_api_initialized = True

    def clean_tag_value(self, value):
        return value.replace(" | ", ".").replace(" ", "-").lower()

    def log_play(self, playbook_name, playbook_timestamp, results):
        if not self.datadog_api_initialized:
            return

        datadog_tasks_metrics = []
        for name, timestamp in results.items():
            datadog_tasks_metrics.append({
                'metric': 'edx.ansible.task_duration',
                'date_happened': timestamp.start,
                'points': timestamp.duration.total_seconds(),
                'tags': [
                    'task:{0}'.format(self.clean_tag_value(name)),
                    'playbook:{0}'.format(self.clean_tag_value(playbook_name))
                ]
            })
        try:
            datadog.api.Metric.send(datadog_tasks_metrics)
            datadog.api.Metric.send(
                metric="edx.ansible.playbook_duration",
                date_happened=time.time(),
                points=playbook_timestamp.duration.total_seconds(),
                tags=["playbook:{0}".format(self.clean_tag_value(playbook_name))]
            )
        except Exception:
            LOGGER.exception("Failed to log timing data to datadog")

class JsonFormatter(Formatter):
    def log_play(self, playbook_name, playbook_timestamp, results):
        # N.B. This is intended to provide a consistent interface and message format
        # across all of Open edX tooling, so it deliberately eschews standard
        # python logging infrastructure.
        if ANSIBLE_TIMER_LOG is None:
            return

        messages = []
        for name, timestamp in results.items():
            messages.append({
                'task': name,
                'playbook': playbook_name,
                'started_at': timestamp.start.isoformat(),
                'ended_at': timestamp.end.isoformat(),
                'duration': timestamp.duration.total_seconds(),
            })

        messages.append({
            'playbook': playbook_name,
            'started_at': playbook_timestamp.start.isoformat(),
            'ended_at': playbook_timestamp.end.isoformat(),
            'duration': playbook_timestamp.duration.total_seconds(),
        })

        log_path = playbook_timestamp.start.strftime(ANSIBLE_TIMER_LOG)

        try:
            if not exists(dirname(log_path)):
                os.makedirs(dirname(log_path))

            with open(log_path, 'a') as outfile:
                for log_message in messages:
                    json.dump(
                        log_message,
                        outfile,
                        separators=(',', ':'),
                        sort_keys=True,
                    )
                    outfile.write('\n')
        except Exception:
            LOGGER.exception("Unable to write json timing log messages")

class LoggingFormatter(Formatter):
    def log_play(self, playbook_name, playbook_timestamp, results):

        # Sort the tasks by their running time
        sorted_results = sorted(
            results.items(),
            key=lambda (task, timestamp): timestamp.duration,
            reverse=True
        )

        for name, timestamp in sorted_results[:10]:
            LOGGER.info(
                "{0:-<80}{1:->8}".format(
                    '{0} '.format(name),
                    ' {0:.02f}s'.format(timestamp.duration.total_seconds()),
                )
            )

        LOGGER.info("\nPlaybook {0} finished: {1}, {2} total tasks.  {3} elapsed. \n".format(
            playbook_name,
            time.asctime(),
            len(results),
            timedelta(seconds=(int(playbook_timestamp.duration.total_seconds())))
        ))

class CallbackModule(object):

    """
    Ansible plugin get the time of each task and total time
    to run the complete playbook
    """
    def __init__(self):
        self.stats = {}
        self.current_task = None
        self.playbook_name = None
        self.playbook_timestamp = None

        self.formatters = [
            DatadogFormatter(),
            LoggingFormatter(),
            JsonFormatter(),
        ]

    def playbook_on_play_start(self, pattern):
        self.playbook_name, _ = splitext(
            basename(self.play.playbook.filename)
        )
        self.playbook_timestamp = Timestamp()

    def playbook_on_task_start(self, name, is_conditional):
        """
        Logs the start of each task
        """

        if self.current_task is not None:
            # Record the running time of the last executed task
            self.stats[self.current_task].stop()

        # Record the start time of the current task
        self.current_task = name
        self.stats[self.current_task] = Timestamp()

    def playbook_on_stats(self, stats):
        """
        Prints the timing of each task and total time to
        run the complete playbook
        """
        # Record the timing of the very last task, we use it here, because we
        # don't have stop task function by default
        if self.current_task is not None:
            self.stats[self.current_task].stop()

        self.playbook_timestamp.stop()

        for formatter in self.formatters:
            formatter.log_play(
                self.playbook_name,
                self.playbook_timestamp,
                self.stats,
            )
<commit_msg>Remove extraneous timedelta conversion<commit_after>from datetime import datetime, timedelta
import json
import logging
import os
from os.path import splitext, basename, exists, dirname
import sys
import time

import datadog

logging.basicConfig(level=logging.INFO, stream=sys.stdout)
logging.getLogger("requests").setLevel(logging.WARNING)
logging.getLogger("dd").setLevel(logging.WARNING)
LOGGER = logging.getLogger(__name__)

"""
Originally written by 'Jharrod LaFon'
#https://github.com/jlafon/ansible-profile/blob/master/callback_plugins/profile_tasks.py

"""

ANSIBLE_TIMER_LOG = os.environ.get('ANSIBLE_TIMER_LOG')

class Timestamp(object):
    def __init__(self):
        self.start = datetime.utcnow()
        self.end = None

    def stop(self):
        self.end = datetime.utcnow()

    @property
    def duration(self):
        return self.end - self.start

class Formatter(object):
    pass

class DatadogFormatter(Formatter):
    def __init__(self):
        super(DatadogFormatter, self).__init__()

        self.datadog_api_key = os.getenv('DATADOG_API_KEY')
        self.datadog_api_initialized = False

        if self.datadog_api_key:
            datadog.initialize(
                api_key=self.datadog_api_key,
                app_key=None
            )
            self.datadog_api_initialized = True

    def clean_tag_value(self, value):
        return value.replace(" | ", ".").replace(" ", "-").lower()

    def log_play(self, playbook_name, playbook_timestamp, results):
        if not self.datadog_api_initialized:
            return

        datadog_tasks_metrics = []
        for name, timestamp in results.items():
            datadog_tasks_metrics.append({
                'metric': 'edx.ansible.task_duration',
                'date_happened': timestamp.start,
                'points': timestamp.duration.total_seconds(),
                'tags': [
                    'task:{0}'.format(self.clean_tag_value(name)),
                    'playbook:{0}'.format(self.clean_tag_value(playbook_name))
                ]
            })
        try:
            datadog.api.Metric.send(datadog_tasks_metrics)
            datadog.api.Metric.send(
                metric="edx.ansible.playbook_duration",
                date_happened=time.time(),
                points=playbook_timestamp.duration.total_seconds(),
                tags=["playbook:{0}".format(self.clean_tag_value(playbook_name))]
            )
        except Exception:
            LOGGER.exception("Failed to log timing data to datadog")

class JsonFormatter(Formatter):
    def log_play(self, playbook_name, playbook_timestamp, results):
        # N.B. This is intended to provide a consistent interface and message format
        # across all of Open edX tooling, so it deliberately eschews standard
        # python logging infrastructure.
        if ANSIBLE_TIMER_LOG is None:
            return

        messages = []
        for name, timestamp in results.items():
            messages.append({
                'task': name,
                'playbook': playbook_name,
                'started_at': timestamp.start.isoformat(),
                'ended_at': timestamp.end.isoformat(),
                'duration': timestamp.duration.total_seconds(),
            })

        messages.append({
            'playbook': playbook_name,
            'started_at': playbook_timestamp.start.isoformat(),
            'ended_at': playbook_timestamp.end.isoformat(),
            'duration': playbook_timestamp.duration.total_seconds(),
        })

        log_path = playbook_timestamp.start.strftime(ANSIBLE_TIMER_LOG)

        try:
            if not exists(dirname(log_path)):
                os.makedirs(dirname(log_path))

            with open(log_path, 'a') as outfile:
                for log_message in messages:
                    json.dump(
                        log_message,
                        outfile,
                        separators=(',', ':'),
                        sort_keys=True,
                    )
                    outfile.write('\n')
        except Exception:
            LOGGER.exception("Unable to write json timing log messages")

class LoggingFormatter(Formatter):
    def log_play(self, playbook_name, playbook_timestamp, results):

        # Sort the tasks by their running time
        sorted_results = sorted(
            results.items(),
            key=lambda (task, timestamp): timestamp.duration,
            reverse=True
        )

        for name, timestamp in sorted_results[:10]:
            LOGGER.info(
                "{0:-<80}{1:->8}".format(
                    '{0} '.format(name),
                    ' {0:.02f}s'.format(timestamp.duration.total_seconds()),
                )
            )

        LOGGER.info("\nPlaybook {0} finished: {1}, {2} total tasks.  {3} elapsed. \n".format(
            playbook_name,
            time.asctime(),
            len(results),
            playbook_timestamp,
        ))

class CallbackModule(object):

    """
    Ansible plugin get the time of each task and total time
    to run the complete playbook
    """
    def __init__(self):
        self.stats = {}
        self.current_task = None
        self.playbook_name = None
        self.playbook_timestamp = None

        self.formatters = [
            DatadogFormatter(),
            LoggingFormatter(),
            JsonFormatter(),
        ]

    def playbook_on_play_start(self, pattern):
        self.playbook_name, _ = splitext(
            basename(self.play.playbook.filename)
        )
        self.playbook_timestamp = Timestamp()

    def playbook_on_task_start(self, name, is_conditional):
        """
        Logs the start of each task
        """

        if self.current_task is not None:
            # Record the running time of the last executed task
            self.stats[self.current_task].stop()

        # Record the start time of the current task
        self.current_task = name
        self.stats[self.current_task] = Timestamp()

    def playbook_on_stats(self, stats):
        """
        Prints the timing of each task and total time to
        run the complete playbook
        """
        # Record the timing of the very last task, we use it here, because we
        # don't have stop task function by default
        if self.current_task is not None:
            self.stats[self.current_task].stop()

        self.playbook_timestamp.stop()

        for formatter in self.formatters:
            formatter.log_play(
                self.playbook_name,
                self.playbook_timestamp,
                self.stats,
            )
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import os
import time
import unittest

import intelmq.lib.test as test
from intelmq.bots.outputs.elasticsearch.output import ElasticsearchOutputBot

if os.environ.get('INTELMQ_TEST_DATABASES'):
    import elasticsearch

INPUT1 = {"__type": "Event",
          "classification.type": "botnet drone",
          "source.asn": 64496,
          "source.ip": "192.0.2.1",
          "feed.name": "Example Feed",
          "extra": '{"foo.bar": "test"}'
          }
OUTPUT1 = {'classification_type': 'botnet drone',
           'extra_foo_bar': 'test',
           'feed_name': 'Example Feed',
           'source_asn': 64496,
           'source_ip': '192.0.2.1',
           }
ES_SEARCH = {"query": {
    "constant_score": {
        "filter": {
            "term": {
                "source_asn": 64496
            }
        }
    }
}
}

@test.skip_database()
class TestElasticsearchOutputBot(test.BotTestCase, unittest.TestCase):

    @classmethod
    def set_bot(cls):
        cls.bot_reference = ElasticsearchOutputBot
        cls.default_input_message = INPUT1
        cls.sysconfig = {"flatten_fields": "extra"}
        if os.environ.get('INTELMQ_TEST_DATABASES'):
            cls.con = elasticsearch.Elasticsearch()

    def test_event(self):
        self.run_bot()
        time.sleep(1)  # ES needs some time between inserting and searching
        result = self.con.search(index='intelmq', body=ES_SEARCH)['hits']['hits'][0]
        self.con.delete(index='intelmq', doc_type='events', id=result['_id'])
        self.assertDictEqual(OUTPUT1, result['_source'])

TIMESTAMP_1 = "1869-12-02T00:00:00+00:00"
TIMESTAMP_2 = "2020-02-02T01:23:45+00:00"

@test.skip_database()
class TestElasticsearchRotatingIndices(test.BotTestCase, unittest.TestCase):

    @classmethod
    def set_bot(cls):
        cls.bot_reference = ElasticsearchOutputBot
        cls.default_input_message = INPUT1
        cls.sysconfig = {"flatten_fields": "extra",
                         "elastic_index": "intelmq",
                         "elastic_doctype": "events",
                         "rotate_index": "true"}
        if os.environ.get('INTELMQ_TEST_DATABASES'):
            cls.con = elasticsearch.Elasticsearch()

    def test_index_detected_from_time_source(self):
        """
        Tests whether an input event with a time.source field is indexed according to its time.source date.
        """

        # Use the sample input, but set the source timestamp
        self.input_message = INPUT1.update({"time_source": TIMESTAMP_1})
        expected_index_name = "{}-1869-12-02".format(self.sysconfig.get('elastic_index'))

        self.assertFalse(self.con.indices.exists(expected_index_name))

        self.run_bot()
        time.sleep(1)  # Let ES store the event. Can also force this with ES API

        self.assertTrue(self.con.indices.exists(expected_index_name))

        result = self.con.search(index=expected_index_name, body=ES_SEARCH)['hits']['hits'][0]
        result_index_name = result["_index"]

        # Clean up test event and check that the index name was set correctly
        self.con.delete(index=result_index_name, doc_type=self.sysconfig.get('doc_type'), id=result['_id'])
        self.assertEqual(result_index_name, expected_index_name)

    def test_index_detected_from_time_observation(self):
        """
        Tests whether an input event with a time.observation field (and no time.source field) is indexed according to
        its time.observation date.
        """

        # Use the sample input, but set the observation timestamp
        self.input_message = INPUT1.update({"time_observation": TIMESTAMP_2})
        expected_index_name = "{}-2020-02-02".format(self.sysconfig.get('elastic_index'))

        self.assertFalse(self.con.indices.exists(expected_index_name))

        self.run_bot()
        time.sleep(1)  # Let ES store the event. Can also force this with ES API

        self.assertTrue(self.con.indices.exists(expected_index_name))

        result = self.con.search(index=expected_index_name, body=ES_SEARCH)['hits']['hits'][0]
        result_index_name = result["_index"]

        # Clean up test event and check that the index name was set correctly
        self.con.delete(index=result_index_name, doc_type=self.sysconfig.get('doc_type'), id=result['_id'])
        self.assertEqual(result_index_name, expected_index_name)

    # def test_index_detected_from_current_date(self):
    #     """
    #     Tests whether an input event without a time.source or time.observation field is indexed using the current date.
    #     """
    #
    #     self.input_message = INPUT1
    #     expected_index_name = "{}-2020-02-02".format(self.sysconfig.get('elastic_index'))
    #
    #     self.assertFalse(self.con.indices.exists(expected_index_name))
    #
    #     self.run_bot()
    #     time.sleep(1)  # Let ES store the event. Can also force this with ES API
    #
    #     self.assertTrue(self.con.indices.exists(expected_index_name))
    #
    #     result = self.con.search(index=expected_index_name, body=ES_SEARCH)['hits']['hits'][0]
    #     result_index_name = result["_index"]
    #
    #     # Clean up test event and check that the index name was set correctly
    #     self.con.delete(index=result_index_name, doc_type=self.sysconfig.get('doc_type'), id=result['_id'])
    #     self.assertEqual(result_index_name, expected_index_name)

if __name__ == '__main__':  # pragma: no cover
    unittest.main()
<commit_msg>Fix field names in rotating index tests<commit_after># -*- coding: utf-8 -*-
import os
import time
import unittest

import intelmq.lib.test as test
from intelmq.bots.outputs.elasticsearch.output import ElasticsearchOutputBot

if os.environ.get('INTELMQ_TEST_DATABASES'):
    import elasticsearch

INPUT1 = {"__type": "Event",
          "classification.type": "botnet drone",
          "source.asn": 64496,
          "source.ip": "192.0.2.1",
          "feed.name": "Example Feed",
          "extra": '{"foo.bar": "test"}'
          }
OUTPUT1 = {'classification_type': 'botnet drone',
           'extra_foo_bar': 'test',
           'feed_name': 'Example Feed',
           'source_asn': 64496,
           'source_ip': '192.0.2.1',
           }
ES_SEARCH = {"query": {
    "constant_score": {
        "filter": {
            "term": {
                "source_asn": 64496
            }
        }
    }
}
}

@test.skip_database()
class TestElasticsearchOutputBot(test.BotTestCase, unittest.TestCase):

    @classmethod
    def set_bot(cls):
        cls.bot_reference = ElasticsearchOutputBot
        cls.default_input_message = INPUT1
        cls.sysconfig = {"flatten_fields": "extra"}
        if os.environ.get('INTELMQ_TEST_DATABASES'):
            cls.con = elasticsearch.Elasticsearch()

    def test_event(self):
        self.run_bot()
        time.sleep(1)  # ES needs some time between inserting and searching
        result = self.con.search(index='intelmq', body=ES_SEARCH)['hits']['hits'][0]
        self.con.delete(index='intelmq', doc_type='events', id=result['_id'])
        self.assertDictEqual(OUTPUT1, result['_source'])

TIMESTAMP_1 = "1869-12-02T00:00:00+00:00"
TIMESTAMP_2 = "2020-02-02T01:23:45+00:00"

@test.skip_database()
class TestElasticsearchRotatingIndices(test.BotTestCase, unittest.TestCase):

    @classmethod
    def set_bot(cls):
        cls.bot_reference = ElasticsearchOutputBot
        cls.default_input_message = INPUT1
        cls.sysconfig = {"flatten_fields": "extra",
                         "elastic_index": "intelmq",
                         "elastic_doctype": "events",
                         "rotate_index": "true"}
        if os.environ.get('INTELMQ_TEST_DATABASES'):
            cls.con = elasticsearch.Elasticsearch()

    def test_index_detected_from_time_source(self):
        """
        Tests whether an input event with a time.source field is indexed according to its time.source date.
        """

        # Use the sample input, but set the source timestamp
        # sample_input = INPUT1.copy().update({"time.source": TIMESTAMP_1})
        self.input_message = INPUT1.update({"time.source": TIMESTAMP_1})

        expected_index_name = "{}-1869-12-02".format(self.sysconfig.get('elastic_index'))

        self.assertFalse(self.con.indices.exists(expected_index_name))

        self.run_bot()
        time.sleep(1)  # Let ES store the event. Can also force this with ES API

        self.assertTrue(self.con.indices.exists(expected_index_name))

        result = self.con.search(index=expected_index_name, body=ES_SEARCH)['hits']['hits'][0]
        result_index_name = result["_index"]

        # Clean up test event and check that the index name was set correctly
        self.con.delete(index=result_index_name, doc_type=self.sysconfig.get('doc_type'), id=result['_id'])
        self.assertEqual(result_index_name, expected_index_name)

    def test_index_detected_from_time_observation(self):
        """
        Tests whether an input event with a time.observation field (and no time.source field) is indexed according to
        its time.observation date.
        """

        # Use the sample input, but set the observation timestamp
        # sample_input = INPUT1.copy().update({"time.source": None, "time.observation": TIMESTAMP_2})
        self.input_message = INPUT1.update({"time.source": None, "time.observation": TIMESTAMP_2})

        expected_index_name = "{}-2020-02-02".format(self.sysconfig.get('elastic_index'))

        self.assertFalse(self.con.indices.exists(expected_index_name))

        self.run_bot()
        time.sleep(1)  # Let ES store the event. Can also force this with ES API

        self.assertTrue(self.con.indices.exists(expected_index_name))

        result = self.con.search(index=expected_index_name, body=ES_SEARCH)['hits']['hits'][0]
        result_index_name = result["_index"]

        # Clean up test event and check that the index name was set correctly
        self.con.delete(index=result_index_name, doc_type=self.sysconfig.get('doc_type'), id=result['_id'])
        self.assertEqual(result_index_name, expected_index_name)

    # def test_index_detected_from_current_date(self):
    #     """
    #     Tests whether an input event without a time.source or time.observation field is indexed using the current date.
    #     """
    #
    #     self.input_message = INPUT1
    #     expected_index_name = "{}-2020-02-02".format(self.sysconfig.get('elastic_index'))
    #
    #     self.assertFalse(self.con.indices.exists(expected_index_name))
    #
    #     self.run_bot()
    #     time.sleep(1)  # Let ES store the event. Can also force this with ES API
    #
    #     self.assertTrue(self.con.indices.exists(expected_index_name))
    #
    #     result = self.con.search(index=expected_index_name, body=ES_SEARCH)['hits']['hits'][0]
    #     result_index_name = result["_index"]
    #
    #     # Clean up test event and check that the index name was set correctly
    #     self.con.delete(index=result_index_name, doc_type=self.sysconfig.get('doc_type'), id=result['_id'])
    #     self.assertEqual(result_index_name, expected_index_name)

if __name__ == '__main__':  # pragma: no cover
    unittest.main()
<|endoftext|>


<commit_before>import os
from os.path import splitext, basename, exists, dirname
from datetime import datetime, timedelta
import json
import time
import logging
import datadog
import sys

logging.basicConfig(level=logging.INFO, stream=sys.stdout)
logging.getLogger("requests").setLevel(logging.WARNING)
logging.getLogger("dd").setLevel(logging.WARNING)
logger = logging.getLogger(__name__)

"""
Originally written by 'Jharrod LaFon'
#https://github.com/jlafon/ansible-profile/blob/master/callback_plugins/profile_tasks.py

"""

ANSIBLE_TIMER_LOG = os.environ.get('ANSIBLE_TIMER_LOG')

class Timestamp(object):
    def __init__(self):
        self.start = datetime.utcnow()
        self.end = None

    def stop(self):
        self.end = datetime.utcnow()

    @property
    def duration(self):
        return self.end - self.start

class Formatter(object):
    def __init__(self, callback_module):
        self.callback_module = callback_module

class DatadogFormatter(Formatter):
    def __init__(self, callback_module):
        super(DatadogFormatter, self).__init__(callback_module)

        self.datadog_api_key = os.getenv('DATADOG_API_KEY')
        self.datadog_api_initialized = False

        if self.datadog_api_key:
            datadog.initialize(
                api_key=self.datadog_api_key,
                app_key=None
            )
            self.datadog_api_initialized = True

    def clean_tag_value(self, value):
        return value.replace(" | ", ".").replace(" ", "-").lower()

    def log_tasks(self, playbook_name, results):
        if self.datadog_api_initialized:
            datadog_tasks_metrics = []
            for name, timestamp in results:
                datadog_tasks_metrics.append({
                    'metric': 'edx.ansible.task_duration',
                    'date_happened': timestamp.start,
                    'points': timestamp.duration.total_seconds(),
                    'tags': [
                        'task:{0}'.format(self.clean_tag_value(name)),
                        'playbook:{0}'.format(self.clean_tag_value(playbook_name))
                    ]
                })
            try:
                datadog.api.Metric.send(datadog_tasks_metrics)
            except Exception as ex:
                logger.error(ex.message)

    def log_play(self, playbook_name, duration, task_count):
        if self.datadog_api_initialized:
            try:
                datadog.api.Metric.send(
                    metric="edx.ansible.playbook_duration",
                    date_happened=time.time(),
                    points=duration,
                    tags=["playbook:{0}".format(self.clean_tag_value(playbook_name))]
                )
            except Exception as ex:
                logger.error(ex.message)

class JsonFormatter(Formatter):
    def log_tasks(self, playbook_name, results):
        if ANSIBLE_TIMER_LOG is not None:
            log_path = self.callback_module.playbook_timestamp.start.strftime(ANSIBLE_TIMER_LOG)

            if not exists(dirname(log_path)):
                os.makedirs(dirname(log_path))

            with open(log_path, 'a') as outfile:
                for name, timestamp in results:
                    log_message = {
                        'task': name,
                        'playbook': playbook_name,
                        'started_at': timestamp.start.isoformat(),
                        'ended_at': timestamp.end.isoformat(),
                        'duration': timestamp.duration.total_seconds(),
                    }

                    json.dump(
                        log_message,
                        outfile,
                        separators=(',', ':'),
                        sort_keys=True,
                    )
                    outfile.write('\n')

    def log_play(self, playbook_name, duration, task_count):
        # N.B. This is intended to provide a consistent interface and message format
        # across all of Open edX tooling, so it deliberately eschews standard
        # python logging infrastructure.
        if ANSIBLE_TIMER_LOG is not None:
            log_path = self.callback_module.playbook_timestamp.start.strftime(ANSIBLE_TIMER_LOG)

            if not exists(dirname(log_path)):
                os.makedirs(dirname(log_path))

            with open(log_path, 'a') as outfile:
                log_message = {
                    'playbook': playbook_name,
                    'started_at': self.callback_module.playbook_timestamp.start.isoformat(),
                    'ended_at': self.callback_module.playbook_timestamp.end.isoformat(),
                    'duration': duration,
                }

                json.dump(
                    log_message,
                    outfile,
                    separators=(',', ':'),
                    sort_keys=True,
                )
                outfile.write('\n')

class LoggingFormatter(Formatter):
    def log_tasks(self, playbook_name, results):
        for name, timestamp in results[:10]:
            logger.info(
                "{0:-<80}{1:->8}".format(
                    '{0} '.format(name),
                    ' {0:.02f}s'.format(timestamp.duration.total_seconds()),
                )
            )

    def log_play(self, playbook_name, duration, task_count):
        logger.info("\nPlaybook {0} finished: {1}, {2} total tasks.  {3} elapsed. \n".format(
            playbook_name,
            time.asctime(),
            task_count,
            timedelta(seconds=(int(duration)))
        ))

class CallbackModule(object):

    """
    Ansible plugin get the time of each task and total time
    to run the complete playbook
    """
    def __init__(self):
        self.stats = {}
        self.current_task = None
        self.playbook_name = None
        self.playbook_timestamp = None

        self.formatters = [
            DatadogFormatter,
            LoggingFormatter,
            JsonFormatter,
        ]

    def playbook_on_play_start(self, pattern):
        self.playbook_name, _ = splitext(
            basename(self.play.playbook.filename)
        )
        self.playbook_timestamp = Timestamp()

    def playbook_on_task_start(self, name, is_conditional):
        """
        Logs the start of each task
        """

        if self.current_task is not None:
            # Record the running time of the last executed task
            self.stats[self.current_task].stop()

        # Record the start time of the current task
        self.current_task = name
        self.stats[self.current_task] = Timestamp()

    def playbook_on_stats(self, stats):
        """
        Prints the timing of each task and total time to
        run the complete playbook
        """
        # Record the timing of the very last task, we use it here, because we
        # don't have stop task function by default
        if self.current_task is not None:
            self.stats[self.current_task].stop()

        self.playbook_timestamp.stop()

        # Sort the tasks by their running time
        results = sorted(
            self.stats.items(),
            key=lambda (task, timestamp): timestamp.duration,
            reverse=True
        )

        total_seconds = self.playbook_timestamp.duration.total_seconds()

        for fmt_class in self.formatters:
            formatter = fmt_class(self)
            formatter.log_tasks(self.playbook_name, results)
            formatter.log_play(self.playbook_name, total_seconds, len(self.stats))
<commit_msg># This is a combination of 2 commits. # The first commit's message is:<commit_after>import os
from os.path import splitext, basename, exists, dirname
from datetime import datetime, timedelta
import json
import time
import logging
import datadog
import sys

logging.basicConfig(level=logging.INFO, stream=sys.stdout)
logging.getLogger("requests").setLevel(logging.WARNING)
logging.getLogger("dd").setLevel(logging.WARNING)
logger = logging.getLogger(__name__)

"""
Originally written by 'Jharrod LaFon'
#https://github.com/jlafon/ansible-profile/blob/master/callback_plugins/profile_tasks.py

"""

ANSIBLE_TIMER_LOG = os.environ.get('ANSIBLE_TIMER_LOG')

class Timestamp(object):
    def __init__(self):
        self.start = datetime.utcnow()
        self.end = None

    def stop(self):
        self.end = datetime.utcnow()

    @property
    def duration(self):
        return self.end - self.start

class Formatter(object):
    def __init__(self, callback_module, playbook_timestamp):
        self.callback_module = callback_module
        self.playbook_timestamp = playbook_timestamp

class DatadogFormatter(Formatter):
    def __init__(self, callback_module, playbook_timestamp):
        super(DatadogFormatter, self).__init__(callback_module, playbook_timestamp)

        self.datadog_api_key = os.getenv('DATADOG_API_KEY')
        self.datadog_api_initialized = False

        if self.datadog_api_key:
            datadog.initialize(
                api_key=self.datadog_api_key,
                app_key=None
            )
            self.datadog_api_initialized = True

    def clean_tag_value(self, value):
        return value.replace(" | ", ".").replace(" ", "-").lower()

    def log_tasks(self, playbook_name, results):
        if self.datadog_api_initialized:
            datadog_tasks_metrics = []
            for name, timestamp in results:
                datadog_tasks_metrics.append({
                    'metric': 'edx.ansible.task_duration',
                    'date_happened': timestamp.start,
                    'points': timestamp.duration.total_seconds(),
                    'tags': [
                        'task:{0}'.format(self.clean_tag_value(name)),
                        'playbook:{0}'.format(self.clean_tag_value(playbook_name))
                    ]
                })
            try:
                datadog.api.Metric.send(datadog_tasks_metrics)
            except Exception as ex:
                logger.error(ex.message)

    def log_play(self, playbook_name, duration, task_count):
        if self.datadog_api_initialized:
            try:
                datadog.api.Metric.send(
                    metric="edx.ansible.playbook_duration",
                    date_happened=time.time(),
                    points=duration,
                    tags=["playbook:{0}".format(self.clean_tag_value(playbook_name))]
                )
            except Exception as ex:
                logger.error(ex.message)

class JsonFormatter(Formatter):
    def log_tasks(self, playbook_name, results):
        if ANSIBLE_TIMER_LOG is not None:
            log_path = self.playbook_timestamp.start.strftime(ANSIBLE_TIMER_LOG)

            if not exists(dirname(log_path)):
                os.makedirs(dirname(log_path))

            with open(log_path, 'a') as outfile:
                for name, timestamp in results:
                    log_message = {
                        'task': name,
                        'playbook': playbook_name,
                        'started_at': timestamp.start.isoformat(),
                        'ended_at': timestamp.end.isoformat(),
                        'duration': timestamp.duration.total_seconds(),
                    }

                    json.dump(
                        log_message,
                        outfile,
                        separators=(',', ':'),
                        sort_keys=True,
                    )
                    outfile.write('\n')

    def log_play(self, playbook_name, duration, task_count):
        # N.B. This is intended to provide a consistent interface and message format
        # across all of Open edX tooling, so it deliberately eschews standard
        # python logging infrastructure.
        if ANSIBLE_TIMER_LOG is not None:
            log_path = self.playbook_timestamp.start.strftime(ANSIBLE_TIMER_LOG)

            if not exists(dirname(log_path)):
                os.makedirs(dirname(log_path))

            with open(log_path, 'a') as outfile:
                log_message = {
                    'playbook': playbook_name,
                    'started_at': self.playbook_timestamp.start.isoformat(),
                    'ended_at': self.playbook_timestamp.end.isoformat(),
                    'duration': duration,
                }

                json.dump(
                    log_message,
                    outfile,
                    separators=(',', ':'),
                    sort_keys=True,
                )
                outfile.write('\n')

class LoggingFormatter(Formatter):
    def log_tasks(self, playbook_name, results):
        for name, timestamp in results[:10]:
            logger.info(
                "{0:-<80}{1:->8}".format(
                    '{0} '.format(name),
                    ' {0:.02f}s'.format(timestamp.duration.total_seconds()),
                )
            )

    def log_play(self, playbook_name, duration, task_count):
        logger.info("\nPlaybook {0} finished: {1}, {2} total tasks.  {3} elapsed. \n".format(
            playbook_name,
            time.asctime(),
            task_count,
            timedelta(seconds=(int(duration)))
        ))

class CallbackModule(object):

    """
    Ansible plugin get the time of each task and total time
    to run the complete playbook
    """
    def __init__(self):
        self.stats = {}
        self.current_task = None
        self.playbook_name = None
        self.playbook_timestamp = None

        self.formatters = [
            DatadogFormatter,
            LoggingFormatter,
            JsonFormatter,
        ]

    def playbook_on_play_start(self, pattern):
        self.playbook_name, _ = splitext(
            basename(self.play.playbook.filename)
        )
        self.playbook_timestamp = Timestamp()

    def playbook_on_task_start(self, name, is_conditional):
        """
        Logs the start of each task
        """

        if self.current_task is not None:
            # Record the running time of the last executed task
            self.stats[self.current_task].stop()

        # Record the start time of the current task
        self.current_task = name
        self.stats[self.current_task] = Timestamp()

    def playbook_on_stats(self, stats):
        """
        Prints the timing of each task and total time to
        run the complete playbook
        """
        # Record the timing of the very last task, we use it here, because we
        # don't have stop task function by default
        if self.current_task is not None:
            self.stats[self.current_task].stop()

        self.playbook_timestamp.stop()

        # Sort the tasks by their running time
        results = sorted(
            self.stats.items(),
            key=lambda (task, timestamp): timestamp.duration,
            reverse=True
        )

        total_seconds = self.playbook_timestamp.duration.total_seconds()

        for fmt_class in self.formatters:
            formatter = fmt_class(self, playbook_timestamp)
            formatter.log_tasks(self.playbook_name, results)
            formatter.log_play(self.playbook_name, total_seconds, len(self.stats))
<|endoftext|>


<commit_before># -*- python -*-
from twisted.application import service, internet
from nevow import appserver, inevow
from ldaptor import config
from ldaptor.apps.webui import main, i18n

application = service.Application("ldaptor-webui")
myService = service.IServiceCollection(application)

cp = config.loadConfig(configFiles=[])
cp.add_section('webui')
cp.set('webui', 'search-field 1 Name',
       '(|(cn=%(input)s)(uid=%(input)s))')
cp.add_section('authentication')
cp.set('authentication', 'identity-base',
       'dc=example,dc=com')

cfg = config.LDAPConfig(serviceLocationOverrides={
    'dc=example,dc=com': ('localhost', 10389),
    })
resource = main.getResource(cfg)

site = appserver.NevowSite(resource)

myServer = internet.TCPServer(38980, site)
myServer.setServiceParent(myService)

############################################################################
import sys, trace
class Coverage(service.Service):
    def startService(self):

        # begin monkey patch ---------------------------
        def find_executable_linenos(filename):
            """Return dict where keys are line numbers in the line number table."""
            #assert filename.endswith('.py') # YOU BASTARDS
            try:
                prog = open(filename).read()
                prog = '\n'.join(prog.splitlines()) + '\n'
            except IOError, err:
                sys.stderr.write("Not printing coverage data for %r: %s\n" % (filename, err))
                sys.stderr.flush()
                return {}
            code = compile(prog, filename, "exec")
            strs = trace.find_strings(filename)
            return trace.find_lines(code, strs)

        trace.find_executable_linenos = find_executable_linenos
        # end monkey patch ------------------------------

        service.Service.startService(self)
        self.tracer = trace.Trace(count=1, trace=0)
        sys.settrace(self.tracer.globaltrace)

    def stopService(self):
        sys.settrace(None)
        results = self.tracer.results()
        results.write_results(show_missing=1,
                              summary=False,
                              coverdir='coverage')
        service.Service.stopService(self)

svc = Coverage()
svc.setServiceParent(application)
############################################################################
<commit_msg>Make test webui tac set up a base DN. Otherwise, adding posixAccounts will fail because the search for uidNumbers requires a base.<commit_after># -*- python -*-
from twisted.application import service, internet
from nevow import appserver, inevow
from ldaptor import config
from ldaptor.apps.webui import main, i18n

application = service.Application("ldaptor-webui")
myService = service.IServiceCollection(application)

cp = config.loadConfig(configFiles=[])
cp.add_section('webui')
cp.set('webui', 'search-field 1 Name',
       '(|(cn=%(input)s)(uid=%(input)s))')
cp.add_section('ldap')
cp.set('ldap', 'base',
       'dc=example,dc=com')

cfg = config.LDAPConfig(serviceLocationOverrides={
    'dc=example,dc=com': ('localhost', 10389),
    })
resource = main.getResource(cfg)

site = appserver.NevowSite(resource)

myServer = internet.TCPServer(38980, site)
myServer.setServiceParent(myService)

############################################################################
import sys, trace
class Coverage(service.Service):
    def startService(self):

        # begin monkey patch ---------------------------
        def find_executable_linenos(filename):
            """Return dict where keys are line numbers in the line number table."""
            #assert filename.endswith('.py') # YOU BASTARDS
            try:
                prog = open(filename).read()
                prog = '\n'.join(prog.splitlines()) + '\n'
            except IOError, err:
                sys.stderr.write("Not printing coverage data for %r: %s\n" % (filename, err))
                sys.stderr.flush()
                return {}
            code = compile(prog, filename, "exec")
            strs = trace.find_strings(filename)
            return trace.find_lines(code, strs)

        trace.find_executable_linenos = find_executable_linenos
        # end monkey patch ------------------------------

        service.Service.startService(self)
        self.tracer = trace.Trace(count=1, trace=0)
        sys.settrace(self.tracer.globaltrace)

    def stopService(self):
        sys.settrace(None)
        results = self.tracer.results()
        results.write_results(show_missing=1,
                              summary=False,
                              coverdir='coverage')
        service.Service.stopService(self)

svc = Coverage()
svc.setServiceParent(application)
############################################################################
<|endoftext|>


<commit_before>"""Utility functions for the j2cl_* build rules / macros"""

def _get_java_root_index(pkg_name):
  """Returns the index of the java_root within a build package"""
  # Find the java folder in the beginning, middle or end of a path.
  if pkg_name.endswith("java"):
    java_index = pkg_name.rfind("java")
  else:
    java_index = pkg_name.rfind("java/")

  # Find the javatests folder in the beginning, middle or end of a path.
  if pkg_name.endswith("javatests"):
    javatests_index = pkg_name.rfind("javatests")
  else:
    javatests_index = pkg_name.rfind("javatests/")

  if java_index == -1 and javatests_index == -1:
    fail("can not find java root: " + pkg_name)

  if java_index > javatests_index:
    index = java_index + len("java/")
  else:
    index = javatests_index + len("javatests/")
  return index

def get_java_root(pkg_name):
  """Extract the path to java root from the build package"""
  return pkg_name[:_get_java_root_index(pkg_name)]

def get_java_path(pkg_name):
  """Extract the java path from the build package"""
  return pkg_name[len(get_java_root(pkg_name)):]

def get_java_package(pkg_name):
  """Extract the java package from the build package"""
  return get_java_path(pkg_name).replace("/", ".")

def generate_zip(name, srcs, pkg, testonly = None):
  """Generates a zip target with given srcs.

  See j2cl_library for details of pkg handling
  """
  # Exit early to avoid parse errors when running under bazel
  if not hasattr(native, "genzip"):
    return

  native.genzip(name=name, deps=[name + "_pkg_library"], testonly = testonly)

  if pkg == "RELATIVE":
    flatten = 0
    package_dir = None
    strip_prefix = None
  elif pkg == "ABSOLUTE":
    flatten = 0
    package_dir = None
    strip_prefix = ""
  elif pkg == "CONVENTION":
    flatten = 1
    package_dir = get_java_path(PACKAGE_NAME)
    strip_prefix = None
  else:
    flatten = 1
    package_dir = pkg
    strip_prefix = None

  native.pkg_library(
      name=name + "_pkg_library",
      srcs=srcs,
      flatten=flatten,
      package_dir=package_dir,
      strip_prefix=strip_prefix,
      testonly = testonly,
  )

load(
    "//javascript/tools/jscompiler/builddefs:flags.bzl",
    "ADVANCED_OPTIMIZATIONS_FLAGS",
    "USE_TYPES_FOR_OPTIMIZATIONS_FLAGS",
    "JS_TEST_FLAGS",
)

J2CL_UNOPTIMIZED_DEFS = [
    "--j2cl_pass=true",
    "--language_in=ECMASCRIPT6_STRICT",
    "--language_out=ECMASCRIPT5",
]

# TODO(goktug): Switch to RECOMMENDED_FLAGS and opt-out from checks as needed.
J2CL_OPTIMIZED_DEFS = (J2CL_UNOPTIMIZED_DEFS +
                       ADVANCED_OPTIMIZATIONS_FLAGS +
                       USE_TYPES_FOR_OPTIMIZATIONS_FLAGS +
                       ["--extra_smart_name_removal=true"])

J2CL_TEST_DEFS = JS_TEST_FLAGS + [
    # Manage closure deps will strip our outputs in some tests
    "--manage_closure_dependencies=false",
    "--j2cl_pass=true",
    "--extra_smart_name_removal=true",
]
<commit_msg>Remove explicitly set language mode, the default is now ECMASCRIPT6_STRICT (aka EcmaScript 2015) and the language versions will be advancing much more quickly in coming years.<commit_after>"""Utility functions for the j2cl_* build rules / macros"""

def _get_java_root_index(pkg_name):
  """Returns the index of the java_root within a build package"""
  # Find the java folder in the beginning, middle or end of a path.
  if pkg_name.endswith("java"):
    java_index = pkg_name.rfind("java")
  else:
    java_index = pkg_name.rfind("java/")

  # Find the javatests folder in the beginning, middle or end of a path.
  if pkg_name.endswith("javatests"):
    javatests_index = pkg_name.rfind("javatests")
  else:
    javatests_index = pkg_name.rfind("javatests/")

  if java_index == -1 and javatests_index == -1:
    fail("can not find java root: " + pkg_name)

  if java_index > javatests_index:
    index = java_index + len("java/")
  else:
    index = javatests_index + len("javatests/")
  return index

def get_java_root(pkg_name):
  """Extract the path to java root from the build package"""
  return pkg_name[:_get_java_root_index(pkg_name)]

def get_java_path(pkg_name):
  """Extract the java path from the build package"""
  return pkg_name[len(get_java_root(pkg_name)):]

def get_java_package(pkg_name):
  """Extract the java package from the build package"""
  return get_java_path(pkg_name).replace("/", ".")

def generate_zip(name, srcs, pkg, testonly = None):
  """Generates a zip target with given srcs.

  See j2cl_library for details of pkg handling
  """
  # Exit early to avoid parse errors when running under bazel
  if not hasattr(native, "genzip"):
    return

  native.genzip(name=name, deps=[name + "_pkg_library"], testonly = testonly)

  if pkg == "RELATIVE":
    flatten = 0
    package_dir = None
    strip_prefix = None
  elif pkg == "ABSOLUTE":
    flatten = 0
    package_dir = None
    strip_prefix = ""
  elif pkg == "CONVENTION":
    flatten = 1
    package_dir = get_java_path(PACKAGE_NAME)
    strip_prefix = None
  else:
    flatten = 1
    package_dir = pkg
    strip_prefix = None

  native.pkg_library(
      name=name + "_pkg_library",
      srcs=srcs,
      flatten=flatten,
      package_dir=package_dir,
      strip_prefix=strip_prefix,
      testonly = testonly,
  )

load(
    "//javascript/tools/jscompiler/builddefs:flags.bzl",
    "ADVANCED_OPTIMIZATIONS_FLAGS",
    "USE_TYPES_FOR_OPTIMIZATIONS_FLAGS",
    "JS_TEST_FLAGS",
)

J2CL_UNOPTIMIZED_DEFS = [
    "--j2cl_pass=true",
    "--language_out=ECMASCRIPT5",
]

# TODO(goktug): Switch to RECOMMENDED_FLAGS and opt-out from checks as needed.
J2CL_OPTIMIZED_DEFS = (J2CL_UNOPTIMIZED_DEFS +
                       ADVANCED_OPTIMIZATIONS_FLAGS +
                       USE_TYPES_FOR_OPTIMIZATIONS_FLAGS +
                       ["--extra_smart_name_removal=true"])

J2CL_TEST_DEFS = JS_TEST_FLAGS + [
    # Manage closure deps will strip our outputs in some tests
    "--manage_closure_dependencies=false",
    "--j2cl_pass=true",
    "--extra_smart_name_removal=true",
]
<|endoftext|>


<commit_before>#
# This file is part of Plinth.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#

"""
Plinth module for running diagnostics
"""

import collections
from django.http import Http404
from django.template.response import TemplateResponse
from django.views.decorators.http import require_POST
from django.utils.translation import ugettext_lazy as _
import logging
import threading

from plinth import module_loader
from plinth.modules import diagnostics

logger = logging.Logger(__name__)

current_results = {}

_running_task = None

def index(request):
    """Serve the index page"""
    if request.method == 'POST' and not _running_task:
        _start_task()

    return TemplateResponse(request, 'diagnostics.html',
                            {'title': diagnostics.title,
                             'description': diagnostics.description,
                             'is_running': _running_task is not None,
                             'results': current_results})

@require_POST
def module(request, module_name):
    """Return diagnostics for a particular module."""
    try:
        module = module_loader.loaded_modules[module_name]
    except KeyError:
        raise Http404('Module does not exist or not loaded')

    results = []
    if hasattr(module, 'diagnose'):
        results = module.diagnose()

    return TemplateResponse(request, 'diagnostics_module.html',
                            {'title': _('Diagnostic Test'),
                             'module_name': module_name,
                             'results': results})

def _start_task():
    """Start the run task in a separate thread."""
    global _running_task
    if _running_task:
        raise Exception('Task already running')

    _running_task = threading.Thread(target=_run_on_all_modules_wrapper)
    _running_task.start()

def _run_on_all_modules_wrapper():
    """Wrapper over actual task to catch exceptions."""
    try:
        run_on_all_modules()
    except Exception as exception:
        logger.exception('Error running diagnostics - %s', exception)
        current_results['error'] = str(exception)

    global _running_task
    _running_task = None

def run_on_all_modules():
    """Run diagnostics on all modules and store the result."""
    global current_results
    current_results = {'modules': [],
                       'results': collections.OrderedDict(),
                       'progress_percentage': 0}

    modules = []
    for module_name, module in module_loader.loaded_modules.items():
        if not hasattr(module, 'diagnose'):
            continue

        modules.append((module_name, module))
        current_results['results'][module_name] = None

    current_results['modules'] = modules
    for current_index, (module_name, module) in enumerate(modules):
        current_results['results'][module_name] = module.diagnose()
        current_results['progress_percentage'] = \
            int((current_index + 1) * 100 / len(modules))
<commit_msg>diagnostics: Don't run on modules not yet setup<commit_after>#
# This file is part of Plinth.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#

"""
Plinth module for running diagnostics
"""

import collections
from django.http import Http404
from django.template.response import TemplateResponse
from django.views.decorators.http import require_POST
from django.utils.translation import ugettext_lazy as _
import logging
import threading

from plinth import module_loader
from plinth.modules import diagnostics

logger = logging.Logger(__name__)

current_results = {}

_running_task = None

def index(request):
    """Serve the index page"""
    if request.method == 'POST' and not _running_task:
        _start_task()

    return TemplateResponse(request, 'diagnostics.html',
                            {'title': diagnostics.title,
                             'description': diagnostics.description,
                             'is_running': _running_task is not None,
                             'results': current_results})

@require_POST
def module(request, module_name):
    """Return diagnostics for a particular module."""
    try:
        module = module_loader.loaded_modules[module_name]
    except KeyError:
        raise Http404('Module does not exist or not loaded')

    results = []
    if hasattr(module, 'diagnose'):
        results = module.diagnose()

    return TemplateResponse(request, 'diagnostics_module.html',
                            {'title': _('Diagnostic Test'),
                             'module_name': module_name,
                             'results': results})

def _start_task():
    """Start the run task in a separate thread."""
    global _running_task
    if _running_task:
        raise Exception('Task already running')

    _running_task = threading.Thread(target=_run_on_all_modules_wrapper)
    _running_task.start()

def _run_on_all_modules_wrapper():
    """Wrapper over actual task to catch exceptions."""
    try:
        run_on_all_modules()
    except Exception as exception:
        logger.exception('Error running diagnostics - %s', exception)
        current_results['error'] = str(exception)

    global _running_task
    _running_task = None

def run_on_all_modules():
    """Run diagnostics on all modules and store the result."""
    global current_results
    current_results = {'modules': [],
                       'results': collections.OrderedDict(),
                       'progress_percentage': 0}

    modules = []
    for module_name, module in module_loader.loaded_modules.items():
        if not hasattr(module, 'diagnose'):
            continue

        # Don't run setup on modules have not been setup yet.
        # However, run on modules that need an upgrade.
        if module.setup_helper.get_state() == 'needs-setup':
            continue

        modules.append((module_name, module))
        current_results['results'][module_name] = None

    current_results['modules'] = modules
    for current_index, (module_name, module) in enumerate(modules):
        current_results['results'][module_name] = module.diagnose()
        current_results['progress_percentage'] = \
            int((current_index + 1) * 100 / len(modules))
<|endoftext|>


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

"""Helper functions for creating partitioned variables.

This is a convenient abstraction to partition a large variable across
multiple smaller variables that can be assigned to different devices.

The full variable can be reconstructed by concatenating the smaller variables.
Using partitioned variables instead of a single variable is mostly a
performance choice.  It however also has an impact on:

1. Random initialization, as the random number generator is called once per
   slice
2. Updates, as they happen in parallel across slices

A key design goal is to allow a different graph to repartition a variable
with the same name but different slicings, including possibly no partitions.

TODO(touts): If an initializer provides a seed, the seed must be changed
deterministicaly for each slice, maybe by adding one to it, otherwise each
slice will use the same values.  Maybe this can be done by passing the
slice offsets to the initializer functions.

Typical usage:

```python
# Create a list of partitioned variables with:
vs = create_partitioned_variables(
    <shape>, <slicing>, <initializer>, name=<optional-name>)

# Pass the list as inputs to embedding_lookup for sharded, parallel lookup:
y = embedding_lookup(vs, ids, partition_strategy="div")

# Or fetch the variables in parallel to speed up large matmuls:
z = matmul(x, concat(slice_dim, vs))
```
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from six.moves import xrange  # pylint: disable=redefined-builtin

from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables

__all__ = ["create_partitioned_variables"]

def _compute_slice_dim_and_shape(full_shape, slicing):
  """Computes which dimension is being sliced and the typical slice shape."""

  slice_shape = [0] * len(full_shape)
  slice_dim = None
  for dim, num_slices in enumerate(slicing):
    dim_size = full_shape[dim]
    if num_slices <= 0 or dim_size < num_slices:
      raise ValueError("Cannot create %d slices for size %d. shape: %s, "
                       "slicing: %s" %
                       (num_slices, full_shape[dim], full_shape, slicing))
    if num_slices == 1:
      # Not slicing in this dimension.
      slice_shape[dim] = dim_size
    elif slice_dim is not None:
      # We only support slicing along one of the dimensions.
      raise ValueError("Can only slice a variable along one dimension: "
                       "shape: %s, slicing: %s" % (full_shape, slicing))
    else:
      # Note: We will add any extras onto the last slice, later.
      slice_dim = dim
      slice_shape[dim] = dim_size // num_slices

  # Degenerate case: If "slicing" was all ones, pretend we are slicing along
  # the first dimension.
  if slice_dim is None:
    slice_dim = 0
  return slice_dim, slice_shape

def create_partitioned_variables(
    shape, slicing, initializer, dtype=dtypes.float32,
    trainable=True, collections=None, name=None):
  """Create a list of partitioned variables according to the given `slicing`.

  Currently only one dimension of the full variable can be sliced, and the
  full variable can be reconstructed by the concatenation of the returned
  list along that dimension.

  Args:
    shape: List of integers.  The shape of the full variable.
    slicing: List of integers.  How to partition the variable.
      Must be of the same length as `shape`.  Each value
      indicate how many slices to create in the corresponding
      dimension.  Presently only one of the values can be more than 1;
      that is, the variable can only be sliced along one dimension.

      For convenience, The requested number of partitions does not have to
      divide the corresponding dimension evenly.  If it does not, the
      shapes of the partitions are incremented by 1 starting from partition
      0 until all slack is absorbed.  The adjustment rules may change in the
      future, but as you can save/restore these variables with different
      slicing specifications this should not be a problem.
    initializer: A `Tensor` of shape `shape` or a variable initializer
      function.  If a function, it will be called once for each slice,
      passing the shape and data type of the slice as parameters.  The
      function must return a tensor with the same shape as the slice.
    dtype: Type of the variables. Ignored if `initializer` is a `Tensor`.
    trainable: If True also add all the variables to the graph collection
      `GraphKeys.TRAINABLE_VARIABLES`.
    collections: List of graph collections keys to add the variables to.
      Defaults to `[GraphKeys.VARIABLES]`.
    name: Optional name for the full variable.  Defaults to
      `"PartitionedVariable"` and gets uniquified automatically.

  Returns:
    A list of Variables corresponding to the slicing.

  Raises:
    ValueError: If any of the arguments is malformed.
  """
  if len(shape) != len(slicing):
    raise ValueError("The 'shape' and 'slicing' of a partitioned Variable "
                     "must have the length: shape: %s, slicing: %s" %
                     (shape, slicing))
  if len(shape) < 1:
    raise ValueError("A partitioned Variable must have rank at least 1: "
                     "shape: %s" % shape)
  full_shape = tensor_shape.as_shape(shape)
  full_shape.assert_is_fully_defined()
  full_shape = full_shape.as_list()

  slice_dim, slice_shape = _compute_slice_dim_and_shape(full_shape, slicing)

  vs = []
  num_slices = slicing[slice_dim]
  num_slices_with_excess = full_shape[slice_dim] % num_slices

  with variable_scope.variable_op_scope([], name,
                                        "PartitionedVariable") as scope:
    full_name = scope.name
    slice_offset = [0] * len(full_shape)
    for i in xrange(num_slices):
      var_shape = slice_shape[:]
      var_offset = slice_offset[:]
      if i < num_slices_with_excess:
        var_shape[slice_dim] += 1
      slice_offset[slice_dim] += var_shape[slice_dim]

      if callable(initializer):
        init_val = initializer(var_shape, dtype=dtype)
        init_val = ops.convert_to_tensor(init_val, dtype=dtype)
      elif isinstance(initializer, ops.Tensor):
        init_val = array_ops.slice(initializer, var_offset, var_shape)
        # Use the dtype of the given tensor.
        dtype = init_val.dtype.base_dtype
      else:
        init_val = ops.convert_to_tensor(initializer, dtype=dtype)
        init_val = array_ops.slice(init_val, var_offset, var_shape)

      var = variable_scope.get_variable(name="part_%d" % i,
                                        dtype=dtype,
                                        initializer=init_val,
                                        trainable=trainable,
                                        collections=collections)

      # pylint: disable=protected-access
      var._set_save_slice_info(variables.Variable.SaveSliceInfo(
          full_name, full_shape, var_offset, var_shape))
      # pylint: enable=protected-access
      vs.append(var)
    assert slice_offset[slice_dim] == full_shape[slice_dim]
  return vs
<commit_msg>Allow reuse of partitioned_variables. Change: 119392087<commit_after># Copyright 2015 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================

"""Helper functions for creating partitioned variables.

This is a convenient abstraction to partition a large variable across
multiple smaller variables that can be assigned to different devices.

The full variable can be reconstructed by concatenating the smaller variables.
Using partitioned variables instead of a single variable is mostly a
performance choice.  It however also has an impact on:

1. Random initialization, as the random number generator is called once per
   slice
2. Updates, as they happen in parallel across slices

A key design goal is to allow a different graph to repartition a variable
with the same name but different slicings, including possibly no partitions.

TODO(touts): If an initializer provides a seed, the seed must be changed
deterministicaly for each slice, maybe by adding one to it, otherwise each
slice will use the same values.  Maybe this can be done by passing the
slice offsets to the initializer functions.

Typical usage:

```python
# Create a list of partitioned variables with:
vs = create_partitioned_variables(
    <shape>, <slicing>, <initializer>, name=<optional-name>)

# Pass the list as inputs to embedding_lookup for sharded, parallel lookup:
y = embedding_lookup(vs, ids, partition_strategy="div")

# Or fetch the variables in parallel to speed up large matmuls:
z = matmul(x, concat(slice_dim, vs))
```
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from six.moves import xrange  # pylint: disable=redefined-builtin

from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import tensor_shape
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import variable_scope
from tensorflow.python.ops import variables

__all__ = ["create_partitioned_variables"]

def _compute_slice_dim_and_shape(full_shape, slicing):
  """Computes which dimension is being sliced and the typical slice shape."""

  slice_shape = [0] * len(full_shape)
  slice_dim = None
  for dim, num_slices in enumerate(slicing):
    dim_size = full_shape[dim]
    if num_slices <= 0 or dim_size < num_slices:
      raise ValueError("Cannot create %d slices for size %d. shape: %s, "
                       "slicing: %s" %
                       (num_slices, full_shape[dim], full_shape, slicing))
    if num_slices == 1:
      # Not slicing in this dimension.
      slice_shape[dim] = dim_size
    elif slice_dim is not None:
      # We only support slicing along one of the dimensions.
      raise ValueError("Can only slice a variable along one dimension: "
                       "shape: %s, slicing: %s" % (full_shape, slicing))
    else:
      # Note: We will add any extras onto the last slice, later.
      slice_dim = dim
      slice_shape[dim] = dim_size // num_slices

  # Degenerate case: If "slicing" was all ones, pretend we are slicing along
  # the first dimension.
  if slice_dim is None:
    slice_dim = 0
  return slice_dim, slice_shape

def create_partitioned_variables(
    shape, slicing, initializer, dtype=dtypes.float32,
    trainable=True, collections=None, name=None, reuse=None):
  """Create a list of partitioned variables according to the given `slicing`.

  Currently only one dimension of the full variable can be sliced, and the
  full variable can be reconstructed by the concatenation of the returned
  list along that dimension.

  Args:
    shape: List of integers.  The shape of the full variable.
    slicing: List of integers.  How to partition the variable.
      Must be of the same length as `shape`.  Each value
      indicate how many slices to create in the corresponding
      dimension.  Presently only one of the values can be more than 1;
      that is, the variable can only be sliced along one dimension.

      For convenience, The requested number of partitions does not have to
      divide the corresponding dimension evenly.  If it does not, the
      shapes of the partitions are incremented by 1 starting from partition
      0 until all slack is absorbed.  The adjustment rules may change in the
      future, but as you can save/restore these variables with different
      slicing specifications this should not be a problem.
    initializer: A `Tensor` of shape `shape` or a variable initializer
      function.  If a function, it will be called once for each slice,
      passing the shape and data type of the slice as parameters.  The
      function must return a tensor with the same shape as the slice.
    dtype: Type of the variables. Ignored if `initializer` is a `Tensor`.
    trainable: If True also add all the variables to the graph collection
      `GraphKeys.TRAINABLE_VARIABLES`.
    collections: List of graph collections keys to add the variables to.
      Defaults to `[GraphKeys.VARIABLES]`.
    name: Optional name for the full variable.  Defaults to
      `"PartitionedVariable"` and gets uniquified automatically.
    reuse: Boolean or `None`; if `True` and name is set, it would reuse
      previously created variables. if `False` it will create new variables.
      if `None`, it would inherit the parent scope reuse.

  Returns:
    A list of Variables corresponding to the slicing.

  Raises:
    ValueError: If any of the arguments is malformed.
  """
  if len(shape) != len(slicing):
    raise ValueError("The 'shape' and 'slicing' of a partitioned Variable "
                     "must have the length: shape: %s, slicing: %s" %
                     (shape, slicing))
  if len(shape) < 1:
    raise ValueError("A partitioned Variable must have rank at least 1: "
                     "shape: %s" % shape)
  full_shape = tensor_shape.as_shape(shape)
  full_shape.assert_is_fully_defined()
  full_shape = full_shape.as_list()

  slice_dim, slice_shape = _compute_slice_dim_and_shape(full_shape, slicing)

  vs = []
  num_slices = slicing[slice_dim]
  num_slices_with_excess = full_shape[slice_dim] % num_slices

  with variable_scope.variable_op_scope([], name,
                                        "PartitionedVariable",
                                        reuse=reuse) as scope:
    full_name = scope.name
    slice_offset = [0] * len(full_shape)
    for i in xrange(num_slices):
      var_shape = slice_shape[:]
      var_offset = slice_offset[:]
      if i < num_slices_with_excess:
        var_shape[slice_dim] += 1
      slice_offset[slice_dim] += var_shape[slice_dim]

      if callable(initializer):
        init_val = initializer(var_shape, dtype=dtype)
        init_val = ops.convert_to_tensor(init_val, dtype=dtype)
      elif isinstance(initializer, ops.Tensor):
        init_val = array_ops.slice(initializer, var_offset, var_shape)
        # Use the dtype of the given tensor.
        dtype = init_val.dtype.base_dtype
      else:
        init_val = ops.convert_to_tensor(initializer, dtype=dtype)
        init_val = array_ops.slice(init_val, var_offset, var_shape)

      var = variable_scope.get_variable(name="part_%d" % i,
                                        dtype=dtype,
                                        initializer=init_val,
                                        trainable=trainable,
                                        collections=collections)

      # pylint: disable=protected-access
      var._set_save_slice_info(variables.Variable.SaveSliceInfo(
          full_name, full_shape, var_offset, var_shape))
      # pylint: enable=protected-access
      vs.append(var)
    assert slice_offset[slice_dim] == full_shape[slice_dim]
  return vs
<|endoftext|>