<commit_before># -*- coding: utf-8 -*-

import os
import re
import textwrap

import click

from .pycompat import relpath
from .catalog import load_po

# ==================================
# settings

# To avoid using invalid resource name, append underscore to such names.
# As a limitation, append `_` doesn't care about collision to other resources.
# e.g. 'glossary' and 'glossary_' are pushed as a 'glossary_'.
IGNORED_RESOURCE_NAMES = (
    'glossary',  # transifex reject this name
)

TRANSIFEXRC_TEMPLATE = """\
[https://www.transifex.com]
hostname = https://www.transifex.com
password = %(transifex_password)s
username = %(transifex_username)s
token =
"""

TXCONFIG_TEMPLATE = """\
[main]
host = https://www.transifex.com
"""

# ==================================
# utility functions

def get_tx_root():
    import txclib.utils
    tx_root = txclib.utils.find_dot_tx()
    if tx_root is None:
        msg = "'.tx/config' not found. You need 'create-txconfig' first."
        raise click.BadParameter(msg)
    return tx_root

def normalize_resource_name(name):
    # replace path separator with '--'
    name = re.sub(r'[\\/]', '--', name)

    # replace unusable characters (not: -, _ ascii, digit) with '_'
    name = re.sub(r'[^\-\w]', '_', name)

    # append `_` for ignored resource names
    while name in IGNORED_RESOURCE_NAMES:
        name += '_'

    return name

# ==================================
# commands

def create_transifexrc(transifex_username, transifex_password):
    """
    Create `$HOME/.transifexrc`
    """
    target = os.path.normpath(os.path.expanduser('~/.transifexrc'))

    if os.path.exists(target):
        click.echo('{0} already exists, skipped.'.format(target))
        return

    if not transifex_username or not 'transifex_password':
        msg = textwrap.dedent("""\
        You need transifex username/password by command option or environment.
        command option: --transifex-username, --transifex-password
        """)
        raise click.BadParameter(msg, param_hint='transifex_username,transifex_password')

    with open(target, 'wt') as rc:
        rc.write(TRANSIFEXRC_TEMPLATE % locals())
    click.echo('Create: {0}'.format(target))

def create_txconfig():
    """
    Create `./.tx/config`
    """
    target = os.path.normpath('.tx/config')
    if os.path.exists(target):
        click.echo('{0} already exists, skipped.'.format(target))
        return

    if not os.path.exists('.tx'):
        os.mkdir('.tx')

    with open(target, 'wt') as f:
        f.write(TXCONFIG_TEMPLATE)

    click.echo('Create: {0}'.format(target))

def update_txconfig_resources(transifex_project_name, locale_dir, pot_dir):
    """
    Update resource sections of `./.tx/config`.
    """
    try:
        import txclib
        import txclib.utils
    except ImportError:
        msg = textwrap.dedent("""\
            Could not import 'txclib.utils'.
            You need install transifex_client external library.
            Please install below command if you want to this action:

                $ pip install sphinx-intl[transifex]
            """)
        raise click.BadParameter(msg)

    tx_root = get_tx_root()

    tx_version = getattr(txclib, '__version__', '0.0')
    if tx_version < '0.13':
        args_tmpl = (
            '--auto-local', '-r', '%(transifex_project_name)s.%(resource_name)s',
            '%(locale_dir)s/<lang>/LC_MESSAGES/%(resource_path)s.po',
            '--source-lang', 'en',
            '--source-file', '%(pot_dir)s/%(resource_path)s.pot',
            '--execute'
        )
    else:
        args_tmpl = (
            'mapping',
            '-r', '%(transifex_project_name)s.%(resource_name)s',
            '-t', 'PO',
            '-s', 'en',
            '-f', '%(pot_dir)s/%(resource_path)s.pot',
            '--execute',
            '%(locale_dir)s/<lang>/LC_MESSAGES/%(resource_path)s.po',
        )

    # convert transifex_project_name to internal name
    transifex_project_name = transifex_project_name.replace(' ', '-')
    transifex_project_name = re.sub(r'[^\-_\w]', '', transifex_project_name)

    for dirpath, dirnames, filenames in os.walk(pot_dir):
        for filename in filenames:
            pot_file = os.path.join(dirpath, filename)
            base, ext = os.path.splitext(pot_file)
            if ext != ".pot":
                continue
            resource_path = relpath(base, pot_dir)
            pot = load_po(pot_file)
            if len(pot):
                resource_name = normalize_resource_name(resource_path)
                l = locals()
                args = [arg % l for arg in args_tmpl]
                # print('set', args, tx_root)
                txclib.utils.exec_command('set', args, tx_root)
            else:
                click.echo('{0} is empty, skipped'.format(pot_file))
<commit_msg>fix for flake8 E741.<commit_after># -*- coding: utf-8 -*-

import os
import re
import textwrap

import click

from .pycompat import relpath
from .catalog import load_po

# ==================================
# settings

# To avoid using invalid resource name, append underscore to such names.
# As a limitation, append `_` doesn't care about collision to other resources.
# e.g. 'glossary' and 'glossary_' are pushed as a 'glossary_'.
IGNORED_RESOURCE_NAMES = (
    'glossary',  # transifex reject this name
)

TRANSIFEXRC_TEMPLATE = """\
[https://www.transifex.com]
hostname = https://www.transifex.com
password = %(transifex_password)s
username = %(transifex_username)s
token =
"""

TXCONFIG_TEMPLATE = """\
[main]
host = https://www.transifex.com
"""

# ==================================
# utility functions

def get_tx_root():
    import txclib.utils
    tx_root = txclib.utils.find_dot_tx()
    if tx_root is None:
        msg = "'.tx/config' not found. You need 'create-txconfig' first."
        raise click.BadParameter(msg)
    return tx_root

def normalize_resource_name(name):
    # replace path separator with '--'
    name = re.sub(r'[\\/]', '--', name)

    # replace unusable characters (not: -, _ ascii, digit) with '_'
    name = re.sub(r'[^\-\w]', '_', name)

    # append `_` for ignored resource names
    while name in IGNORED_RESOURCE_NAMES:
        name += '_'

    return name

# ==================================
# commands

def create_transifexrc(transifex_username, transifex_password):
    """
    Create `$HOME/.transifexrc`
    """
    target = os.path.normpath(os.path.expanduser('~/.transifexrc'))

    if os.path.exists(target):
        click.echo('{0} already exists, skipped.'.format(target))
        return

    if not transifex_username or not 'transifex_password':
        msg = textwrap.dedent("""\
        You need transifex username/password by command option or environment.
        command option: --transifex-username, --transifex-password
        """)
        raise click.BadParameter(msg, param_hint='transifex_username,transifex_password')

    with open(target, 'wt') as rc:
        rc.write(TRANSIFEXRC_TEMPLATE % locals())
    click.echo('Create: {0}'.format(target))

def create_txconfig():
    """
    Create `./.tx/config`
    """
    target = os.path.normpath('.tx/config')
    if os.path.exists(target):
        click.echo('{0} already exists, skipped.'.format(target))
        return

    if not os.path.exists('.tx'):
        os.mkdir('.tx')

    with open(target, 'wt') as f:
        f.write(TXCONFIG_TEMPLATE)

    click.echo('Create: {0}'.format(target))

def update_txconfig_resources(transifex_project_name, locale_dir, pot_dir):
    """
    Update resource sections of `./.tx/config`.
    """
    try:
        import txclib
        import txclib.utils
    except ImportError:
        msg = textwrap.dedent("""\
            Could not import 'txclib.utils'.
            You need install transifex_client external library.
            Please install below command if you want to this action:

                $ pip install sphinx-intl[transifex]
            """)
        raise click.BadParameter(msg)

    tx_root = get_tx_root()

    tx_version = getattr(txclib, '__version__', '0.0')
    if tx_version < '0.13':
        args_tmpl = (
            '--auto-local', '-r', '%(transifex_project_name)s.%(resource_name)s',
            '%(locale_dir)s/<lang>/LC_MESSAGES/%(resource_path)s.po',
            '--source-lang', 'en',
            '--source-file', '%(pot_dir)s/%(resource_path)s.pot',
            '--execute'
        )
    else:
        args_tmpl = (
            'mapping',
            '-r', '%(transifex_project_name)s.%(resource_name)s',
            '-t', 'PO',
            '-s', 'en',
            '-f', '%(pot_dir)s/%(resource_path)s.pot',
            '--execute',
            '%(locale_dir)s/<lang>/LC_MESSAGES/%(resource_path)s.po',
        )

    # convert transifex_project_name to internal name
    transifex_project_name = transifex_project_name.replace(' ', '-')
    transifex_project_name = re.sub(r'[^\-_\w]', '', transifex_project_name)

    for dirpath, dirnames, filenames in os.walk(pot_dir):
        for filename in filenames:
            pot_file = os.path.join(dirpath, filename)
            base, ext = os.path.splitext(pot_file)
            if ext != ".pot":
                continue
            resource_path = relpath(base, pot_dir)
            pot = load_po(pot_file)
            if len(pot):
                resource_name = normalize_resource_name(resource_path)
                lv = locals()
                args = [arg % lv for arg in args_tmpl]
                # print('set', args, tx_root)
                txclib.utils.exec_command('set', args, tx_root)
            else:
                click.echo('{0} is empty, skipped'.format(pot_file))
<|endoftext|>


<commit_before>from docutils.parsers.rst import directives
from sphinx.domains.python import PyModulelevel, PyClassmember
from sphinx.ext.autodoc import FunctionDocumenter, MethodDocumenter
from asyncio import iscoroutinefunction

__version__ = '0.1.1'

def merge_dicts(*dcts):
    ret = {}
    for d in dcts:
        for k, v in d.items():
            ret[k] = v
    return ret

class PyCoroutineMixin(object):
    option_spec = {'coroutine': directives.flag,
                   'async-with': directives.flag,
                   'async-for': directives.flag}

    def get_signature_prefix(self, sig):
        ret = ''
        if 'staticmethod' in self.options:
            ret += 'staticmethod '
        if 'classmethod' in self.options:
            ret += 'classmethod '
        if 'coroutine' in self.options:
            coroutine = True
        else:
            coroutine = ('async-with' not in self.options and
                         'async-for' not in self.options)
        if coroutine:
            ret += 'coroutine '
        if 'async-with' in self.options:
            ret += 'async-with '
        if 'async-for' in self.options:
            ret += 'async-for '
        return ret

class PyCoroutineFunction(PyCoroutineMixin, PyModulelevel):
    option_spec = merge_dicts(PyCoroutineMixin.option_spec,
                              PyModulelevel.option_spec)

    def run(self):
        self.name = 'py:function'
        return super(PyCoroutineFunction, self).run()

class PyCoroutineMethod(PyCoroutineMixin, PyClassmember):
    option_spec = merge_dicts(PyCoroutineMixin.option_spec,
                              PyClassmember.option_spec,
                              {'staticmethod': directives.flag,
                               'classmethod': directives.flag})

    def run(self):
        self.name = 'py:method'
        return super(PyCoroutineMethod, self).run()

class CoFunctionDocumenter(FunctionDocumenter):
    """
    Specialized Documenter subclass for functions and coroutines.
    """
    objtype = "cofunction"
    directivetype = "cofunction"
    priority = 2

    @classmethod
    def can_document_member(cls, member, membername, isattr, parent):
        """Called to see if a member can be documented by this documenter."""
        if not super().can_document_member(member, membername, isattr, parent):
            return False
        return iscoroutinefunction(member)

class CoMethodDocumenter(MethodDocumenter):
    """
    Specialized Documenter subclass for methods and coroutines.
    """
    objtype = "comethod"
    priority = 3  # Higher than CoFunctionDocumenter

    @classmethod
    def can_document_member(cls, member, membername, isattr, parent):
        """Called to see if a member can be documented by this documenter."""
        if not super().can_document_member(member, membername, isattr, parent):
            return False
        return iscoroutinefunction(member)

    def import_object(self):
        ret = super().import_object()
        # Was overridden by method documenter, return to default
        self.directivetype = "comethod"
        return ret

def setup(app):
    app.add_directive_to_domain('py', 'coroutinefunction', PyCoroutineFunction)
    app.add_directive_to_domain('py', 'coroutinemethod', PyCoroutineMethod)
    app.add_directive_to_domain('py', 'corofunction', PyCoroutineFunction)
    app.add_directive_to_domain('py', 'coromethod', PyCoroutineMethod)
    app.add_directive_to_domain('py', 'cofunction', PyCoroutineFunction)
    app.add_directive_to_domain('py', 'comethod', PyCoroutineMethod)

    app.add_autodocumenter(CoFunctionDocumenter)
    app.add_autodocumenter(CoMethodDocumenter)
    return {'version': '1.0', 'parallel_read_safe': True}
<commit_msg>Make python < 3.4 compatioble<commit_after>from docutils.parsers.rst import directives
from sphinx.domains.python import PyModulelevel, PyClassmember
from sphinx.ext.autodoc import FunctionDocumenter, MethodDocumenter
try:
    from asyncio import iscoroutinefunction
except ImportError:
    def iscoroutinefunction(func):
        """Return True if func is a decorated coroutine function."""
        return getattr(func, '_is_coroutine', False)

__version__ = '0.1.1'

def merge_dicts(*dcts):
    ret = {}
    for d in dcts:
        for k, v in d.items():
            ret[k] = v
    return ret

class PyCoroutineMixin(object):
    option_spec = {'coroutine': directives.flag,
                   'async-with': directives.flag,
                   'async-for': directives.flag}

    def get_signature_prefix(self, sig):
        ret = ''
        if 'staticmethod' in self.options:
            ret += 'staticmethod '
        if 'classmethod' in self.options:
            ret += 'classmethod '
        if 'coroutine' in self.options:
            coroutine = True
        else:
            coroutine = ('async-with' not in self.options and
                         'async-for' not in self.options)
        if coroutine:
            ret += 'coroutine '
        if 'async-with' in self.options:
            ret += 'async-with '
        if 'async-for' in self.options:
            ret += 'async-for '
        return ret

class PyCoroutineFunction(PyCoroutineMixin, PyModulelevel):
    option_spec = merge_dicts(PyCoroutineMixin.option_spec,
                              PyModulelevel.option_spec)

    def run(self):
        self.name = 'py:function'
        return super(PyCoroutineFunction, self).run()

class PyCoroutineMethod(PyCoroutineMixin, PyClassmember):
    option_spec = merge_dicts(PyCoroutineMixin.option_spec,
                              PyClassmember.option_spec,
                              {'staticmethod': directives.flag,
                               'classmethod': directives.flag})

    def run(self):
        self.name = 'py:method'
        return super(PyCoroutineMethod, self).run()

class CoFunctionDocumenter(FunctionDocumenter):
    """
    Specialized Documenter subclass for functions and coroutines.
    """
    objtype = "cofunction"
    directivetype = "cofunction"
    priority = 2

    @classmethod
    def can_document_member(cls, member, membername, isattr, parent):
        """Called to see if a member can be documented by this documenter."""
        if not super().can_document_member(member, membername, isattr, parent):
            return False
        return iscoroutinefunction(member)

class CoMethodDocumenter(MethodDocumenter):
    """
    Specialized Documenter subclass for methods and coroutines.
    """
    objtype = "comethod"
    priority = 3  # Higher than CoFunctionDocumenter

    @classmethod
    def can_document_member(cls, member, membername, isattr, parent):
        """Called to see if a member can be documented by this documenter."""
        if not super().can_document_member(member, membername, isattr, parent):
            return False
        return iscoroutinefunction(member)

    def import_object(self):
        ret = super().import_object()
        # Was overridden by method documenter, return to default
        self.directivetype = "comethod"
        return ret

def setup(app):
    app.add_directive_to_domain('py', 'coroutinefunction', PyCoroutineFunction)
    app.add_directive_to_domain('py', 'coroutinemethod', PyCoroutineMethod)
    app.add_directive_to_domain('py', 'corofunction', PyCoroutineFunction)
    app.add_directive_to_domain('py', 'coromethod', PyCoroutineMethod)
    app.add_directive_to_domain('py', 'cofunction', PyCoroutineFunction)
    app.add_directive_to_domain('py', 'comethod', PyCoroutineMethod)

    app.add_autodocumenter(CoFunctionDocumenter)
    app.add_autodocumenter(CoMethodDocumenter)
    return {'version': '1.0', 'parallel_read_safe': True}
<|endoftext|>


<commit_before>'''
update dataset DB, and then close dataset and start Activator if needed

'''

import re
import sys
import time
import urllib
import commands
import threading
from DDM import ddm
import Notifier
import RetryMaker
from Activator import Activator
from pandalogger.PandaLogger import PandaLogger
from taskbuffer.JobSpec import JobSpec
from taskbuffer.FileSpec import FileSpec
from taskbuffer.DatasetSpec import DatasetSpec
from brokerage.SiteMapper import SiteMapper
from config import panda_config
import brokerage.broker_util

# logger
_logger = PandaLogger().getLogger('Closer')

def initLogger(pLogger):
    # redirect logging to parent as it doesn't work in nested threads
    global _logger
    _logger = pLogger
    Notifier.initLogger(_logger)
    RetryMaker.initLogger(_logger)

class Closer (threading.Thread):
    # constructor
    def __init__(self,taskBuffer,destinationDBlocks,job,pandaDDM=False,datasetMap={}):
        threading.Thread.__init__(self)
        self.taskBuffer = taskBuffer
        self.destinationDBlocks = destinationDBlocks
        self.job = job
        self.pandaID = job.PandaID
        self.pandaDDM = pandaDDM
        self.siteMapper = None
        self.datasetMap = datasetMap

    # main
    def run(self):
        try:
            _logger.debug('%s Start %s' % (self.pandaID,self.job.jobStatus))
            flagComplete    = True
            ddmJobs         = []
            topUserDsList   = []
            usingMerger     = False
            disableNotifier = False
            firstIndvDS     = True
            for destinationDBlock in self.destinationDBlocks:
                dsList = []
                _logger.debug('%s start %s' % (self.pandaID,destinationDBlock))
                # ignore tid datasets
                if re.search('_tid[\d_]+$',destinationDBlock):
                    _logger.debug('%s skip %s' % (self.pandaID,destinationDBlock))
                    continue
                # query dataset
                if self.datasetMap.has_key(destinationDBlock):
                    dataset = self.datasetMap[destinationDBlock]
                else:
                    dataset = self.taskBuffer.queryDatasetWithMap({'name':destinationDBlock})
                if dataset == None:
                    _logger.error('%s Not found : %s' % (self.pandaID,destinationDBlock))
                    flagComplete = False
                    continue
                # skip tobedeleted
                if dataset.status in ['cleanup']:
                    _logger.debug('%s skip cleanup %s' % (self.pandaID,destinationDBlock))
                    continue
                dsList.append(dataset)
                # sort
                dsList.sort()
                # count number of completed files
                notFinish = self.taskBuffer.countFilesWithMap({'destinationDBlock':destinationDBlock,
                                                               'status':'unknown'})
                if notFinish < 0:
                    _logger.error('%s Invalid DB return : %s' % (self.pandaID,notFinish))
                    flagComplete = False
                    continue
                # check if completed
                _logger.debug('%s notFinish:%s' % (self.pandaID,notFinish))
                if self.job.destinationSE == 'local' and self.job.prodSourceLabel in ['user','panda']:
                    # close non-DQ2 destinationDBlock immediately
                    finalStatus = 'closed'
                elif self.job.prodSourceLabel in ['user'] and "--mergeOutput" in self.job.jobParameters \
                         and self.job.processingType != 'usermerge' and self.job.commandToPilot != 'tobekilled':
                    # merge output files
                    if firstIndvDS:
                        # set 'tobemerged' to only the first dataset to avoid triggering many Mergers for --individualOutDS
                        finalStatus = 'tobemerged'
                        firstIndvDS = False
                    else:
                        finalStatus = 'tobeclosed'
                    # set merging to top dataset
                    usingMerger = True
                    # disable Notifier
                    disableNotifier = True
                else:
                    # set status to 'tobeclosed' to trigger DQ2 closing
                    finalStatus = 'tobeclosed'
                if notFinish==0:
                    _logger.debug('%s set %s to dataset : %s' % (self.pandaID,finalStatus,destinationDBlock))
                    # set status
                    dataset.status = finalStatus
                    # update dataset in DB
                    retT = self.taskBuffer.updateDatasets(dsList,withLock=True,withCriteria="status<>:crStatus AND status<>:lockStatus ",
                                                          criteriaMap={':crStatus':finalStatus,':lockStatus':'locked'})
                    if len(retT) > 0 and retT[0]==1:
                        # close user datasets
                        if self.job.prodSourceLabel in ['user'] and self.job.destinationDBlock.endswith('/') \
                               and (dataset.name.startswith('user') or dataset.name.startswith('group')):
                            # get top-level user dataset
                            topUserDsName = re.sub('_sub\d+$','',dataset.name)
                            # update if it is the first attempt
                            if topUserDsName != dataset.name and not topUserDsName in topUserDsList:
                                topUserDs = self.taskBuffer.queryDatasetWithMap({'name':topUserDsName})
                                if topUserDs != None:
                                    # check status
                                    if topUserDs.status in ['completed','cleanup','tobeclosed',
                                                            'tobemerged','merging']:
                                        _logger.debug('%s skip %s due to status=%s' % (self.pandaID,topUserDsName,topUserDs.status))
                                    else:
                                        # set status
                                        if self.job.processingType in ['gangarobot-rctest']:
                                            # not trigger freezing for HC datasets so that files can be appended
                                            topUserDs.status = 'completed'
                                        elif not usingMerger:
                                            topUserDs.status = finalStatus
                                        else:
                                            topUserDs.status = 'merging'
                                        # append to avoid repetition
                                        topUserDsList.append(topUserDsName)
                                        # update DB
                                        retTopT = self.taskBuffer.updateDatasets([topUserDs],withLock=True,withCriteria="status<>:crStatus",
                                                                                criteriaMap={':crStatus':topUserDs.status})
                                        if len(retTopT) > 0 and retTopT[0]==1:
                                            _logger.debug('%s set %s to top dataset : %s' % (self.pandaID,topUserDs.status,topUserDsName))
                                        else:
                                            _logger.debug('%s failed to update top dataset : %s' % (self.pandaID,topUserDsName))
                            # get parent dataset for merge job
                            if self.job.processingType == 'usermerge':
                                tmpMatch = re.search('--parentDS ([^ \'\"]+)',self.job.jobParameters)
                                if tmpMatch == None:
                                    _logger.error('%s failed to extract parentDS' % self.pandaID)
                                else:
                                    unmergedDsName = tmpMatch.group(1)
                                    # update if it is the first attempt
                                    if not unmergedDsName in topUserDsList:
                                        unmergedDs = self.taskBuffer.queryDatasetWithMap({'name':unmergedDsName})
                                        if unmergedDs == None:
                                            _logger.error('%s failed to get parentDS=%s from DB' % (self.pandaID,unmergedDsName))
                                        else:
                                            # check status
                                            if unmergedDs.status in ['completed','cleanup','tobeclosed']:
                                                _logger.debug('%s skip %s due to status=%s' % (self.pandaID,unmergedDsName,unmergedDs.status))
                                            else:
                                                # set status
                                                unmergedDs.status = finalStatus
                                                # append to avoid repetition
                                                topUserDsList.append(unmergedDsName)
                                                # update DB
                                                retTopT = self.taskBuffer.updateDatasets([unmergedDs],withLock=True,withCriteria="status<>:crStatus",
                                                                                criteriaMap={':crStatus':unmergedDs.status})
                                                if len(retTopT) > 0 and retTopT[0]==1:
                                                    _logger.debug('%s set %s to parent dataset : %s' % (self.pandaID,unmergedDs.status,unmergedDsName))
                                                else:
                                                    _logger.debug('%s failed to update parent dataset : %s' % (self.pandaID,unmergedDsName))
                        if self.pandaDDM and self.job.prodSourceLabel=='managed':
                            # instantiate SiteMapper
                            if self.siteMapper == None:
                                self.siteMapper = SiteMapper(self.taskBuffer)
                            # get file list for PandaDDM
                            retList = self.taskBuffer.queryFilesWithMap({'destinationDBlock':destinationDBlock})
                            lfnsStr = ''
                            guidStr = ''
                            for tmpFile in retList:
                                if tmpFile.type in ['log','output']:
                                    lfnsStr += '%s,' % tmpFile.lfn
                                    guidStr += '%s,' % tmpFile.GUID
                            if lfnsStr != '':
                                guidStr = guidStr[:-1]
                                lfnsStr = lfnsStr[:-1]
                                # create a DDM job
                                ddmjob = JobSpec()
                                ddmjob.jobDefinitionID   = int(time.time()) % 10000
                                ddmjob.jobName           = "%s" % commands.getoutput('uuidgen')
                                ddmjob.transformation    = 'http://pandaserver.cern.ch:25080/trf/mover/run_dq2_cr'
                                ddmjob.destinationDBlock = 'testpanda.%s' % ddmjob.jobName
                                ddmjob.computingSite     = "BNL_ATLAS_DDM"
                                ddmjob.destinationSE     = ddmjob.computingSite
                                ddmjob.currentPriority   = 200000
                                ddmjob.prodSourceLabel   = 'ddm'
                                ddmjob.transferType      = 'sub'
                                # append log file
                                fileOL = FileSpec()
                                fileOL.lfn = "%s.job.log.tgz" % ddmjob.jobName
                                fileOL.destinationDBlock = ddmjob.destinationDBlock
                                fileOL.destinationSE     = ddmjob.destinationSE
                                fileOL.dataset           = ddmjob.destinationDBlock
                                fileOL.type = 'log'
                                ddmjob.addFile(fileOL)
                                # make arguments
                                dstDQ2ID = 'BNLPANDA'
                                srcDQ2ID = self.siteMapper.getSite(self.job.computingSite).ddm
                                callBackURL = 'https://%s:%s/server/panda/datasetCompleted?vuid=%s&site=%s' % \
                                              (panda_config.pserverhost,panda_config.pserverport,
                                               dataset.vuid,dstDQ2ID)
                                _logger.debug(callBackURL)
                                # set src/dest
                                ddmjob.sourceSite      = srcDQ2ID
                                ddmjob.destinationSite = dstDQ2ID
                                # if src==dst, send callback without ddm job
                                if dstDQ2ID == srcDQ2ID:
                                    comout = commands.getoutput('curl -k %s' % callBackURL)
                                    _logger.debug(comout)
                                else:
                                    # run dq2_cr
                                    callBackURL = urllib.quote(callBackURL)
                                    # get destination dir
                                    destDir = brokerage.broker_util._getDefaultStorage(self.siteMapper.getSite(self.job.computingSite).dq2url)
                                    argStr = "-s %s -r %s --guids %s --lfns %s --callBack %s -d %s/%s %s" % \
                                             (srcDQ2ID,dstDQ2ID,guidStr,lfnsStr,callBackURL,destDir,
                                              destinationDBlock,destinationDBlock)
                                    # set job parameters
                                    ddmjob.jobParameters = argStr
                                    _logger.debug('%s pdq2_cr %s' % (self.pandaID,ddmjob.jobParameters))
                                    ddmJobs.append(ddmjob)
                        # start Activator
                        if re.search('_sub\d+$',dataset.name) == None:
                            if self.job.prodSourceLabel=='panda' and self.job.processingType in ['merge','unmerge']:
                                # don't trigger Activator for merge jobs
                                pass
                            else:
                                if self.job.jobStatus == 'finished':
                                    aThr = Activator(self.taskBuffer,dataset)
                                    aThr.start()
                                    aThr.join()
                    else:
                        # unset flag since another thread already updated
                        flagComplete = False
                else:
                    # update dataset in DB
                    self.taskBuffer.updateDatasets(dsList,withLock=True,withCriteria="status<>:crStatus AND status<>:lockStatus ",
                                                   criteriaMap={':crStatus':finalStatus,':lockStatus':'locked'})
                    # unset flag
                    flagComplete = False
                # end
                _logger.debug('%s end %s' % (self.pandaID,destinationDBlock))
            # start DDM jobs
            if ddmJobs != []:
                self.taskBuffer.storeJobs(ddmJobs,self.job.prodUserID,joinThr=True)
            # change pending jobs to failed
            if flagComplete and self.job.prodSourceLabel=='user':
                #_logger.debug('%s call RetryMaker for %s %s' % (self.pandaID,self.job.prodUserName,self.job.jobDefinitionID))
                #retryMaker = RetryMaker.RetryMaker(self.taskBuffer,self.job)
                #retryMaker.run()
                _logger.debug('%s finalize %s %s' % (self.pandaID,self.job.prodUserName,self.job.jobDefinitionID))
                self.taskBuffer.finalizePendingJobs(self.job.prodUserName,self.job.jobDefinitionID)
            # start notifier
            _logger.debug('%s source:%s complete:%s' % (self.pandaID,self.job.prodSourceLabel,flagComplete))
            if (self.job.jobStatus != 'transferring') and ((flagComplete and self.job.prodSourceLabel=='user') or \
               (self.job.jobStatus=='failed' and self.job.prodSourceLabel=='panda')):
                # don't send email for merge jobs
                if (not disableNotifier) and not self.job.processingType in ['merge','unmerge']:
                    useNotifier = True
                    summaryInfo = {}
                    # check all jobDefIDs in jobsetID
                    if not self.job.jobsetID in [0,None,'NULL']:
                        useNotifier,summaryInfo = self.taskBuffer.checkDatasetStatusForNotifier(self.job.jobsetID,self.job.jobDefinitionID,
                                                                                self.job.prodUserName)
                        _logger.debug('%s useNotifier:%s' % (self.pandaID,useNotifier))
                    if useNotifier:
                        _logger.debug('%s start Notifier' % self.pandaID)
                        nThr = Notifier.Notifier(self.taskBuffer,self.job,self.destinationDBlocks,summaryInfo)
                        nThr.run()
                        _logger.debug('%s end Notifier' % self.pandaID)
            _logger.debug('%s End' % self.pandaID)
        except:
            errType,errValue = sys.exc_info()[:2]
            _logger.error("%s %s" % (errType,errValue))

<commit_msg>not freeze hammercloud datasets<commit_after>'''
update dataset DB, and then close dataset and start Activator if needed

'''

import re
import sys
import time
import urllib
import commands
import threading
from DDM import ddm
import Notifier
import RetryMaker
from Activator import Activator
from pandalogger.PandaLogger import PandaLogger
from taskbuffer.JobSpec import JobSpec
from taskbuffer.FileSpec import FileSpec
from taskbuffer.DatasetSpec import DatasetSpec
from brokerage.SiteMapper import SiteMapper
from config import panda_config
import brokerage.broker_util

# logger
_logger = PandaLogger().getLogger('Closer')

def initLogger(pLogger):
    # redirect logging to parent as it doesn't work in nested threads
    global _logger
    _logger = pLogger
    Notifier.initLogger(_logger)
    RetryMaker.initLogger(_logger)

class Closer (threading.Thread):
    # constructor
    def __init__(self,taskBuffer,destinationDBlocks,job,pandaDDM=False,datasetMap={}):
        threading.Thread.__init__(self)
        self.taskBuffer = taskBuffer
        self.destinationDBlocks = destinationDBlocks
        self.job = job
        self.pandaID = job.PandaID
        self.pandaDDM = pandaDDM
        self.siteMapper = None
        self.datasetMap = datasetMap

    # main
    def run(self):
        try:
            _logger.debug('%s Start %s' % (self.pandaID,self.job.jobStatus))
            flagComplete    = True
            ddmJobs         = []
            topUserDsList   = []
            usingMerger     = False
            disableNotifier = False
            firstIndvDS     = True
            for destinationDBlock in self.destinationDBlocks:
                dsList = []
                _logger.debug('%s start %s' % (self.pandaID,destinationDBlock))
                # ignore tid datasets
                if re.search('_tid[\d_]+$',destinationDBlock):
                    _logger.debug('%s skip %s' % (self.pandaID,destinationDBlock))
                    continue
                # query dataset
                if self.datasetMap.has_key(destinationDBlock):
                    dataset = self.datasetMap[destinationDBlock]
                else:
                    dataset = self.taskBuffer.queryDatasetWithMap({'name':destinationDBlock})
                if dataset == None:
                    _logger.error('%s Not found : %s' % (self.pandaID,destinationDBlock))
                    flagComplete = False
                    continue
                # skip tobedeleted
                if dataset.status in ['cleanup']:
                    _logger.debug('%s skip cleanup %s' % (self.pandaID,destinationDBlock))
                    continue
                dsList.append(dataset)
                # sort
                dsList.sort()
                # count number of completed files
                notFinish = self.taskBuffer.countFilesWithMap({'destinationDBlock':destinationDBlock,
                                                               'status':'unknown'})
                if notFinish < 0:
                    _logger.error('%s Invalid DB return : %s' % (self.pandaID,notFinish))
                    flagComplete = False
                    continue
                # check if completed
                _logger.debug('%s notFinish:%s' % (self.pandaID,notFinish))
                if self.job.destinationSE == 'local' and self.job.prodSourceLabel in ['user','panda']:
                    # close non-DQ2 destinationDBlock immediately
                    finalStatus = 'closed'
                elif self.job.prodSourceLabel in ['user'] and "--mergeOutput" in self.job.jobParameters \
                         and self.job.processingType != 'usermerge' and self.job.commandToPilot != 'tobekilled':
                    # merge output files
                    if firstIndvDS:
                        # set 'tobemerged' to only the first dataset to avoid triggering many Mergers for --individualOutDS
                        finalStatus = 'tobemerged'
                        firstIndvDS = False
                    else:
                        finalStatus = 'tobeclosed'
                    # set merging to top dataset
                    usingMerger = True
                    # disable Notifier
                    disableNotifier = True
                else:
                    # set status to 'tobeclosed' to trigger DQ2 closing
                    finalStatus = 'tobeclosed'
                if notFinish==0:
                    _logger.debug('%s set %s to dataset : %s' % (self.pandaID,finalStatus,destinationDBlock))
                    # set status
                    dataset.status = finalStatus
                    # update dataset in DB
                    retT = self.taskBuffer.updateDatasets(dsList,withLock=True,withCriteria="status<>:crStatus AND status<>:lockStatus ",
                                                          criteriaMap={':crStatus':finalStatus,':lockStatus':'locked'})
                    if len(retT) > 0 and retT[0]==1:
                        # close user datasets
                        if self.job.prodSourceLabel in ['user'] and self.job.destinationDBlock.endswith('/') \
                               and (dataset.name.startswith('user') or dataset.name.startswith('group')):
                            # get top-level user dataset
                            topUserDsName = re.sub('_sub\d+$','',dataset.name)
                            # update if it is the first attempt
                            if topUserDsName != dataset.name and not topUserDsName in topUserDsList:
                                topUserDs = self.taskBuffer.queryDatasetWithMap({'name':topUserDsName})
                                if topUserDs != None:
                                    # check status
                                    if topUserDs.status in ['completed','cleanup','tobeclosed',
                                                            'tobemerged','merging']:
                                        _logger.debug('%s skip %s due to status=%s' % (self.pandaID,topUserDsName,topUserDs.status))
                                    else:
                                        # set status
                                        if self.job.processingType in ['gangarobot-rctest','hammercloud']:
                                            # not trigger freezing for HC datasets so that files can be appended
                                            topUserDs.status = 'completed'
                                        elif not usingMerger:
                                            topUserDs.status = finalStatus
                                        else:
                                            topUserDs.status = 'merging'
                                        # append to avoid repetition
                                        topUserDsList.append(topUserDsName)
                                        # update DB
                                        retTopT = self.taskBuffer.updateDatasets([topUserDs],withLock=True,withCriteria="status<>:crStatus",
                                                                                criteriaMap={':crStatus':topUserDs.status})
                                        if len(retTopT) > 0 and retTopT[0]==1:
                                            _logger.debug('%s set %s to top dataset : %s' % (self.pandaID,topUserDs.status,topUserDsName))
                                        else:
                                            _logger.debug('%s failed to update top dataset : %s' % (self.pandaID,topUserDsName))
                            # get parent dataset for merge job
                            if self.job.processingType == 'usermerge':
                                tmpMatch = re.search('--parentDS ([^ \'\"]+)',self.job.jobParameters)
                                if tmpMatch == None:
                                    _logger.error('%s failed to extract parentDS' % self.pandaID)
                                else:
                                    unmergedDsName = tmpMatch.group(1)
                                    # update if it is the first attempt
                                    if not unmergedDsName in topUserDsList:
                                        unmergedDs = self.taskBuffer.queryDatasetWithMap({'name':unmergedDsName})
                                        if unmergedDs == None:
                                            _logger.error('%s failed to get parentDS=%s from DB' % (self.pandaID,unmergedDsName))
                                        else:
                                            # check status
                                            if unmergedDs.status in ['completed','cleanup','tobeclosed']:
                                                _logger.debug('%s skip %s due to status=%s' % (self.pandaID,unmergedDsName,unmergedDs.status))
                                            else:
                                                # set status
                                                unmergedDs.status = finalStatus
                                                # append to avoid repetition
                                                topUserDsList.append(unmergedDsName)
                                                # update DB
                                                retTopT = self.taskBuffer.updateDatasets([unmergedDs],withLock=True,withCriteria="status<>:crStatus",
                                                                                criteriaMap={':crStatus':unmergedDs.status})
                                                if len(retTopT) > 0 and retTopT[0]==1:
                                                    _logger.debug('%s set %s to parent dataset : %s' % (self.pandaID,unmergedDs.status,unmergedDsName))
                                                else:
                                                    _logger.debug('%s failed to update parent dataset : %s' % (self.pandaID,unmergedDsName))
                        if self.pandaDDM and self.job.prodSourceLabel=='managed':
                            # instantiate SiteMapper
                            if self.siteMapper == None:
                                self.siteMapper = SiteMapper(self.taskBuffer)
                            # get file list for PandaDDM
                            retList = self.taskBuffer.queryFilesWithMap({'destinationDBlock':destinationDBlock})
                            lfnsStr = ''
                            guidStr = ''
                            for tmpFile in retList:
                                if tmpFile.type in ['log','output']:
                                    lfnsStr += '%s,' % tmpFile.lfn
                                    guidStr += '%s,' % tmpFile.GUID
                            if lfnsStr != '':
                                guidStr = guidStr[:-1]
                                lfnsStr = lfnsStr[:-1]
                                # create a DDM job
                                ddmjob = JobSpec()
                                ddmjob.jobDefinitionID   = int(time.time()) % 10000
                                ddmjob.jobName           = "%s" % commands.getoutput('uuidgen')
                                ddmjob.transformation    = 'http://pandaserver.cern.ch:25080/trf/mover/run_dq2_cr'
                                ddmjob.destinationDBlock = 'testpanda.%s' % ddmjob.jobName
                                ddmjob.computingSite     = "BNL_ATLAS_DDM"
                                ddmjob.destinationSE     = ddmjob.computingSite
                                ddmjob.currentPriority   = 200000
                                ddmjob.prodSourceLabel   = 'ddm'
                                ddmjob.transferType      = 'sub'
                                # append log file
                                fileOL = FileSpec()
                                fileOL.lfn = "%s.job.log.tgz" % ddmjob.jobName
                                fileOL.destinationDBlock = ddmjob.destinationDBlock
                                fileOL.destinationSE     = ddmjob.destinationSE
                                fileOL.dataset           = ddmjob.destinationDBlock
                                fileOL.type = 'log'
                                ddmjob.addFile(fileOL)
                                # make arguments
                                dstDQ2ID = 'BNLPANDA'
                                srcDQ2ID = self.siteMapper.getSite(self.job.computingSite).ddm
                                callBackURL = 'https://%s:%s/server/panda/datasetCompleted?vuid=%s&site=%s' % \
                                              (panda_config.pserverhost,panda_config.pserverport,
                                               dataset.vuid,dstDQ2ID)
                                _logger.debug(callBackURL)
                                # set src/dest
                                ddmjob.sourceSite      = srcDQ2ID
                                ddmjob.destinationSite = dstDQ2ID
                                # if src==dst, send callback without ddm job
                                if dstDQ2ID == srcDQ2ID:
                                    comout = commands.getoutput('curl -k %s' % callBackURL)
                                    _logger.debug(comout)
                                else:
                                    # run dq2_cr
                                    callBackURL = urllib.quote(callBackURL)
                                    # get destination dir
                                    destDir = brokerage.broker_util._getDefaultStorage(self.siteMapper.getSite(self.job.computingSite).dq2url)
                                    argStr = "-s %s -r %s --guids %s --lfns %s --callBack %s -d %s/%s %s" % \
                                             (srcDQ2ID,dstDQ2ID,guidStr,lfnsStr,callBackURL,destDir,
                                              destinationDBlock,destinationDBlock)
                                    # set job parameters
                                    ddmjob.jobParameters = argStr
                                    _logger.debug('%s pdq2_cr %s' % (self.pandaID,ddmjob.jobParameters))
                                    ddmJobs.append(ddmjob)
                        # start Activator
                        if re.search('_sub\d+$',dataset.name) == None:
                            if self.job.prodSourceLabel=='panda' and self.job.processingType in ['merge','unmerge']:
                                # don't trigger Activator for merge jobs
                                pass
                            else:
                                if self.job.jobStatus == 'finished':
                                    aThr = Activator(self.taskBuffer,dataset)
                                    aThr.start()
                                    aThr.join()
                    else:
                        # unset flag since another thread already updated
                        flagComplete = False
                else:
                    # update dataset in DB
                    self.taskBuffer.updateDatasets(dsList,withLock=True,withCriteria="status<>:crStatus AND status<>:lockStatus ",
                                                   criteriaMap={':crStatus':finalStatus,':lockStatus':'locked'})
                    # unset flag
                    flagComplete = False
                # end
                _logger.debug('%s end %s' % (self.pandaID,destinationDBlock))
            # start DDM jobs
            if ddmJobs != []:
                self.taskBuffer.storeJobs(ddmJobs,self.job.prodUserID,joinThr=True)
            # change pending jobs to failed
            if flagComplete and self.job.prodSourceLabel=='user':
                #_logger.debug('%s call RetryMaker for %s %s' % (self.pandaID,self.job.prodUserName,self.job.jobDefinitionID))
                #retryMaker = RetryMaker.RetryMaker(self.taskBuffer,self.job)
                #retryMaker.run()
                _logger.debug('%s finalize %s %s' % (self.pandaID,self.job.prodUserName,self.job.jobDefinitionID))
                self.taskBuffer.finalizePendingJobs(self.job.prodUserName,self.job.jobDefinitionID)
            # start notifier
            _logger.debug('%s source:%s complete:%s' % (self.pandaID,self.job.prodSourceLabel,flagComplete))
            if (self.job.jobStatus != 'transferring') and ((flagComplete and self.job.prodSourceLabel=='user') or \
               (self.job.jobStatus=='failed' and self.job.prodSourceLabel=='panda')):
                # don't send email for merge jobs
                if (not disableNotifier) and not self.job.processingType in ['merge','unmerge']:
                    useNotifier = True
                    summaryInfo = {}
                    # check all jobDefIDs in jobsetID
                    if not self.job.jobsetID in [0,None,'NULL']:
                        useNotifier,summaryInfo = self.taskBuffer.checkDatasetStatusForNotifier(self.job.jobsetID,self.job.jobDefinitionID,
                                                                                self.job.prodUserName)
                        _logger.debug('%s useNotifier:%s' % (self.pandaID,useNotifier))
                    if useNotifier:
                        _logger.debug('%s start Notifier' % self.pandaID)
                        nThr = Notifier.Notifier(self.taskBuffer,self.job,self.destinationDBlocks,summaryInfo)
                        nThr.run()
                        _logger.debug('%s end Notifier' % self.pandaID)
            _logger.debug('%s End' % self.pandaID)
        except:
            errType,errValue = sys.exc_info()[:2]
            _logger.error("%s %s" % (errType,errValue))

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Author: Óscar Nájera
# License: 3-clause BSD
"""
Generate the rst files for the examples by iterating over the python
example files.

Files that generate images should start with 'plot'

"""
from __future__ import division, print_function, absolute_import
from time import time
import ast
import os
import re
import shutil
import traceback
import sys
import subprocess
import warnings
from . import glr_path_static
from .backreferences import write_backreferences, _thumbnail_div

# Try Python 2 first, otherwise load from Python 3
try:
    from StringIO import StringIO
except ImportError:
    from io import StringIO

try:
    basestring
except NameError:
    basestring = str

import token
import tokenize

try:
    # make sure that the Agg backend is set before importing any
    # matplotlib
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
except ImportError:
    # this script can be imported by nosetest to find tests to run: we should
    # not impose the matplotlib requirement in that case.
    pass

###############################################################################

class Tee(object):
    """A tee object to redirect streams to multiple outputs"""
    def __init__(self, file1, file2):
        self.file1 = file1
        self.file2 = file2

    def write(self, data):
        self.file1.write(data)
        self.file2.write(data)

    def flush(self):
        self.file1.flush()
        self.file2.flush()

###############################################################################
CODE_DOWNLOAD = """**Total running time of the script:**
({0:.0f} minutes {1:.3f} seconds)\n\n
\n**Download Python source code:** :download:`{2} <{2}>`\n"""

# The following strings are used when we have several pictures: we use
# an html div tag that our CSS uses to turn the lists into horizontal
# lists.
HLIST_HEADER = """
.. rst-class:: sphx-glr-horizontal

"""

HLIST_IMAGE_TEMPLATE = """
    *

      .. image:: /%s
            :scale: 47
"""

SINGLE_IMAGE = """
.. image:: /%s
    :align: center
"""

CODE_OUTPUT = """**Script output**:\n
.. rst-class:: sphx-glr-script-out

  ::

    {0}\n"""

def split_code_and_text_blocks(source_file):
    f=open(source_file)

    blocks = []
    block = ''

    header_string = 0
    continue_text = False
    for line in f:
        # python docstring
        if line.startswith('"""') and not header_string:
            header_string += 1
            block = line
            continue
        elif line.startswith('"""') and header_string==1:
            header_string += 1
            block += line
            blocks.append(('text', block))
            block= ''
            continue

        # comment blocks
        if line.startswith('#') and not continue_text:
            if 20*'#' in line:
                continue_text = True
                if len(block) > 1:
                    blocks.append(('code', block.strip()))
                    block = '"""'
                continue
        if line.startswith('#') and continue_text:
            block += line[1:].lstrip()
            continue
        if not line.startswith('#') and continue_text:

            block += '"""'
            blocks.append(('text', block))
            block = ''
            continue_text = False

        # code blocks
        block += line

    # close  last block
    if len(block) > 1:
        if continue_text:
            blocks.append(('text', block))
        else:
            blocks.append(('code', block.strip()))

    return blocks

def codestr2rst(codestr):
    """Return reStructuredText code block from code string"""
    code_directive = "\n.. code-block:: python\n\n"
    indented_block = '    ' + codestr.replace('\n', '\n    ')
    return code_directive + indented_block

def extract_intro(filename):
    """ Extract the first paragraph of module-level docstring. max:95 char"""

    first_text = split_code_and_text_blocks(filename)[0][-1]

    paragraphs = first_text.split('\n\n')
    if len(first_text) > 1:
        first_par = re.sub('\n', ' ', paragraphs[1])
        first_par = ((first_par[:95] + '...')
                     if len(first_par) > 95 else first_par)
    else:
        raise ValueError("Missing first paragraph."
                         "Please check the layout of your"
                         " example file:\n {}\n and make sure"
                         " it's correct.".format(filename))

    return first_par

def _plots_are_current(src_file, image_file):
    """Test existence of image file and later touch time to source script"""

    first_image_file = image_file.format(1)
    needs_replot = (not os.path.exists(first_image_file) or
              os.stat(first_image_file).st_mtime <= os.stat(src_file).st_mtime)
    return not needs_replot

def save_figures(image_path, fig_count):
    """Save all open matlplotlib figures of the example code-block

    Parameters
    ----------
    image_path : str
        Path where plots are saved (format string which accepts figure number)
    fig_count : int
        Previous figure number count. Figure number add from this number
    """
    figure_list = []
    # In order to save every figure we have two solutions :
    # * iterate from 1 to infinity and call plt.fignum_exists(n)
    #   (this requires the figures to be numbered
    #    incrementally: 1, 2, 3 and not 1, 2, 5)
    # * iterate over [fig_mngr.num for fig_mngr in
    #   matplotlib._pylab_helpers.Gcf.get_all_fig_managers()]

    fig_managers = matplotlib._pylab_helpers.Gcf.get_all_fig_managers()
    for fig_mngr in fig_managers:
        # Set the fig_num figure as the current figure as we can't
        # save a figure that's not the current figure.
        fig = plt.figure(fig_mngr.num)
        kwargs = {}
        to_rgba = matplotlib.colors.colorConverter.to_rgba
        for attr in ['facecolor', 'edgecolor']:
            fig_attr = getattr(fig, 'get_' + attr)()
            default_attr = matplotlib.rcParams['figure.' + attr]
            if to_rgba(fig_attr) != to_rgba(default_attr):
                kwargs[attr] = fig_attr

        current_fig = image_path.format(fig_count + fig_mngr.num)
        fig.savefig(current_fig, **kwargs)
        figure_list.append(current_fig)
    return figure_list

def scale_image(in_fname, out_fname, max_width, max_height):
    """Scales an image with the same aspect ratio centered in an
       image with a given max_width and max_height
       if in_fname == out_fname the image can only be scaled down
    """
    # local import to avoid testing dependency on PIL:
    try:
        from PIL import Image
    except ImportError:
        import Image
    img = Image.open(in_fname)
    width_in, height_in = img.size
    scale_w = max_width / float(width_in)
    scale_h = max_height / float(height_in)

    if height_in * scale_w <= max_height:
        scale = scale_w
    else:
        scale = scale_h

    if scale >= 1.0 and in_fname == out_fname:
        return

    width_sc = int(round(scale * width_in))
    height_sc = int(round(scale * height_in))

    # resize the image
    img.thumbnail((width_sc, height_sc), Image.ANTIALIAS)

    # insert centered
    thumb = Image.new('RGB', (max_width, max_height), (255, 255, 255))
    pos_insert = ((max_width - width_sc) // 2, (max_height - height_sc) // 2)
    thumb.paste(img, pos_insert)

    thumb.save(out_fname)
    # Use optipng to perform lossless compression on the resized image if
    # software is installed
    if os.environ.get('SKLEARN_DOC_OPTIPNG', False):
        try:
            subprocess.call(["optipng", "-quiet", "-o", "9", out_fname])
        except Exception:
            warnings.warn('Install optipng to reduce the size of the \
                          generated images')

def save_thumbnail(image_path, base_image_name):
    """Save the thumbnail image"""
    first_image_file = image_path.format(1)
    thumb_dir = os.path.join(os.path.dirname(first_image_file), 'thumb')
    if not os.path.exists(thumb_dir):
        os.makedirs(thumb_dir)

    thumb_file = os.path.join(thumb_dir, 'sphx_glr_%s_thumb.png' % base_image_name)

    if os.path.exists(first_image_file):
        scale_image(first_image_file, thumb_file, 400, 280)
    elif not os.path.exists(thumb_file):
        # create something to replace the thumbnail
        scale_image(os.path.join(glr_path_static(), 'no_image.png'),
                    thumb_file, 200, 140)

def generate_dir_rst(src_dir, target_dir, gallery_conf, seen_backrefs):
    """Generate the gallery reStructuredText for an example directory"""
    if not os.path.exists(os.path.join(src_dir, 'README.txt')):
        print(80 * '_')
        print('Example directory %s does not have a README.txt file' %
              src_dir)
        print('Skipping this directory')
        print(80 * '_')
        return ""  # because string is an expected return type

    fhindex = open(os.path.join(src_dir, 'README.txt')).read()
    if not os.path.exists(target_dir):
        os.makedirs(target_dir)
    sorted_listdir = [fname for fname in sorted(os.listdir(src_dir))
                                if fname.endswith('py')]
    for fname in sorted_listdir:
        generate_file_rst(fname, target_dir, src_dir)
        new_fname = os.path.join(src_dir, fname)
        intro = extract_intro(new_fname)
        write_backreferences(seen_backrefs, gallery_conf,
                             target_dir, fname, intro)

        fhindex += _thumbnail_div(target_dir, fname, intro)
        fhindex += """

.. toctree::
   :hidden:

   /%s/%s\n""" % (target_dir, fname[:-3])

# clear at the end of the section
    fhindex += """.. raw:: html\n

\n\n"""

    return fhindex

def execute_script(code_block, example_globals, image_path, fig_count, src_file):
    """Executes the code block of the example file"""
    time_elapsed = 0
    stdout = ''

    # We need to execute the code
    print('plotting code blocks in %s' % src_file)

    plt.close('all')
    cwd = os.getcwd()
    # Redirect output to stdout and
    orig_stdout = sys.stdout

    try:
        # First CD in the original example dir, so that any file
        # created by the example get created in this directory
        os.chdir(os.path.dirname(src_file))
        my_buffer = StringIO()
        my_stdout = Tee(sys.stdout, my_buffer)
        sys.stdout = my_stdout

        t_start = time()
        exec(code_block, example_globals)
        time_elapsed = time() - t_start

        sys.stdout = orig_stdout

        my_stdout = my_buffer.getvalue().strip().expandtabs()
        if my_stdout:
            stdout = CODE_OUTPUT.format(my_stdout.replace('\n', '\n    '))
        os.chdir(cwd)
        figure_list = save_figures(image_path, fig_count)

        # Depending on whether we have one or more figures, we're using a
        # horizontal list or a single rst call to 'image'.
        if len(figure_list) == 1:
            figure_name = figure_list[0]
            image_list = SINGLE_IMAGE % figure_name.lstrip('/')
        else:
            image_list = HLIST_HEADER
            for figure_name in figure_list:
                image_list += HLIST_IMAGE_TEMPLATE % figure_name.lstrip('/')

    except Exception:
        image_list = '%s is not compiling:' % src_file
        print(80 * '_')
        print(image_list)
        traceback.print_exc()
        print(80 * '_')
    finally:
        os.chdir(cwd)
        sys.stdout = orig_stdout

    print(" - time elapsed : %.2g sec" % time_elapsed)
    code_output = "\n{0}\n\n{1}".format(image_list, stdout)

    return code_output, time_elapsed, fig_count + len(figure_list)

def generate_file_rst(fname, target_dir, src_dir):
    """ Generate the rst file for a given example."""

    src_file = os.path.join(src_dir, fname)
    example_file = os.path.join(target_dir, fname)
    shutil.copyfile(src_file, example_file)

    image_dir = os.path.join(target_dir, 'images')
    if not os.path.exists(image_dir):
        os.makedirs(image_dir)

    base_image_name = os.path.splitext(fname)[0]
    image_fname = 'sphx_glr_' + base_image_name + '_{0:03}.png'
    image_path = os.path.join(image_dir, image_fname)

    if _plots_are_current(src_file, image_path):
        return

    time_elapsed = 0
    script_blocks = split_code_and_text_blocks(example_file)

    ref_fname = example_file.replace(os.path.sep, '_')
    example_rst = """\n\n.. _example_{0}:\n\n""".format(ref_fname)

    if not fname.startswith('plot'):
        convert_func = dict(code=codestr2rst, text=ast.literal_eval)
        for blabel, bcontent in script_blocks:
            example_rst += convert_func[blabel](bcontent)+'\n'
    else:
        example_globals = {}
        fig_count = 0
        for blabel, bcontent in script_blocks:
            if blabel == 'code':
                code_output, rtime, fig_count = execute_script(bcontent,
                                                               example_globals,
                                                               image_path,
                                                               fig_count,
                                                               src_file)

                time_elapsed += rtime
                # Single example style
                if len(script_blocks)==2:
                    example_rst += code_output
                    example_rst += codestr2rst(bcontent)+'\n'
                # Notebook style
                else:
                    example_rst += codestr2rst(bcontent)+'\n'
                    example_rst += code_output
            else:
                example_rst += ast.literal_eval(bcontent)+'\n'

    save_thumbnail(image_path, base_image_name)

    time_m, time_s = divmod(time_elapsed, 60)
    with open(os.path.join(target_dir, base_image_name + '.rst'), 'w') as f:
        example_rst += CODE_DOWNLOAD.format(time_m, time_s, fname)
        f.write(example_rst)
<commit_msg>file closing and code style<commit_after># -*- coding: utf-8 -*-
# Author: Óscar Nájera
# License: 3-clause BSD
"""
Generate the rst files for the examples by iterating over the python
example files.

Files that generate images should start with 'plot'

"""
from __future__ import division, print_function, absolute_import
from time import time
import ast
import os
import re
import shutil
import traceback
import sys
import subprocess
import warnings
from . import glr_path_static
from .backreferences import write_backreferences, _thumbnail_div

# Try Python 2 first, otherwise load from Python 3
try:
    from StringIO import StringIO
except ImportError:
    from io import StringIO

try:
    basestring
except NameError:
    basestring = str

import token
import tokenize

try:
    # make sure that the Agg backend is set before importing any
    # matplotlib
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
except ImportError:
    # this script can be imported by nosetest to find tests to run: we should
    # not impose the matplotlib requirement in that case.
    pass

###############################################################################

class Tee(object):
    """A tee object to redirect streams to multiple outputs"""
    def __init__(self, file1, file2):
        self.file1 = file1
        self.file2 = file2

    def write(self, data):
        self.file1.write(data)
        self.file2.write(data)

    def flush(self):
        self.file1.flush()
        self.file2.flush()

###############################################################################
CODE_DOWNLOAD = """**Total running time of the script:**
({0:.0f} minutes {1:.3f} seconds)\n\n
\n**Download Python source code:** :download:`{2} <{2}>`\n"""

# The following strings are used when we have several pictures: we use
# an html div tag that our CSS uses to turn the lists into horizontal
# lists.
HLIST_HEADER = """
.. rst-class:: sphx-glr-horizontal

"""

HLIST_IMAGE_TEMPLATE = """
    *

      .. image:: /%s
            :scale: 47
"""

SINGLE_IMAGE = """
.. image:: /%s
    :align: center
"""

CODE_OUTPUT = """**Script output**:\n
.. rst-class:: sphx-glr-script-out

  ::

    {0}\n"""

def split_code_and_text_blocks(source_file):
    f = open(source_file)

    blocks = []
    block = ''

    has_header_string = 0
    continue_text = False
    for line in f:
        # python docstring
        if line.startswith('"""') and not has_header_string:
            has_header_string += 1
            block = line
            continue
        # because we only allow for one single docstring
        elif line.startswith('"""') and has_header_string == 1:
            has_header_string += 1
            block += line
            blocks.append(('text', block))
            block = ''
            continue

        # comment blocks
        if line.startswith('#') and not continue_text:
            if 20*'#' in line:
                continue_text = True
                if len(block) > 1:
                    blocks.append(('code', block.strip()))
                    block = '"""'
                continue
        if line.startswith('#') and continue_text:
            block += line[1:].lstrip()
            continue
        if not line.startswith('#') and continue_text:

            block += '"""'
            blocks.append(('text', block))
            block = ''
            continue_text = False

        # code blocks
        block += line

    # close  last block
    if len(block) > 1:
        if continue_text:
            blocks.append(('text', block))
        else:
            blocks.append(('code', block.strip()))

    f.close()
    return blocks

def codestr2rst(codestr):
    """Return reStructuredText code block from code string"""
    code_directive = "\n.. code-block:: python\n\n"
    indented_block = '    ' + codestr.replace('\n', '\n    ')
    return code_directive + indented_block

def extract_intro(filename):
    """ Extract the first paragraph of module-level docstring. max:95 char"""

    first_text = split_code_and_text_blocks(filename)[0][-1]

    paragraphs = first_text.split('\n\n')
    if len(first_text) > 1:
        first_par = re.sub('\n', ' ', paragraphs[1])
        first_par = ((first_par[:95] + '...')
                     if len(first_par) > 95 else first_par)
    else:
        raise ValueError("Missing first paragraph."
                         "Please check the layout of your"
                         " example file:\n {}\n and make sure"
                         " it's correct.".format(filename))

    return first_par

def _plots_are_current(src_file, image_file):
    """Test existence of image file and later touch time to source script"""

    first_image_file = image_file.format(1)
    needs_replot = (not os.path.exists(first_image_file) or
              os.stat(first_image_file).st_mtime <= os.stat(src_file).st_mtime)
    return not needs_replot

def save_figures(image_path, fig_count):
    """Save all open matlplotlib figures of the example code-block

    Parameters
    ----------
    image_path : str
        Path where plots are saved (format string which accepts figure number)
    fig_count : int
        Previous figure number count. Figure number add from this number
    """
    figure_list = []
    # In order to save every figure we have two solutions :
    # * iterate from 1 to infinity and call plt.fignum_exists(n)
    #   (this requires the figures to be numbered
    #    incrementally: 1, 2, 3 and not 1, 2, 5)
    # * iterate over [fig_mngr.num for fig_mngr in
    #   matplotlib._pylab_helpers.Gcf.get_all_fig_managers()]

    fig_managers = matplotlib._pylab_helpers.Gcf.get_all_fig_managers()
    for fig_mngr in fig_managers:
        # Set the fig_num figure as the current figure as we can't
        # save a figure that's not the current figure.
        fig = plt.figure(fig_mngr.num)
        kwargs = {}
        to_rgba = matplotlib.colors.colorConverter.to_rgba
        for attr in ['facecolor', 'edgecolor']:
            fig_attr = getattr(fig, 'get_' + attr)()
            default_attr = matplotlib.rcParams['figure.' + attr]
            if to_rgba(fig_attr) != to_rgba(default_attr):
                kwargs[attr] = fig_attr

        current_fig = image_path.format(fig_count + fig_mngr.num)
        fig.savefig(current_fig, **kwargs)
        figure_list.append(current_fig)
    return figure_list

def scale_image(in_fname, out_fname, max_width, max_height):
    """Scales an image with the same aspect ratio centered in an
       image with a given max_width and max_height
       if in_fname == out_fname the image can only be scaled down
    """
    # local import to avoid testing dependency on PIL:
    try:
        from PIL import Image
    except ImportError:
        import Image
    img = Image.open(in_fname)
    width_in, height_in = img.size
    scale_w = max_width / float(width_in)
    scale_h = max_height / float(height_in)

    if height_in * scale_w <= max_height:
        scale = scale_w
    else:
        scale = scale_h

    if scale >= 1.0 and in_fname == out_fname:
        return

    width_sc = int(round(scale * width_in))
    height_sc = int(round(scale * height_in))

    # resize the image
    img.thumbnail((width_sc, height_sc), Image.ANTIALIAS)

    # insert centered
    thumb = Image.new('RGB', (max_width, max_height), (255, 255, 255))
    pos_insert = ((max_width - width_sc) // 2, (max_height - height_sc) // 2)
    thumb.paste(img, pos_insert)

    thumb.save(out_fname)
    # Use optipng to perform lossless compression on the resized image if
    # software is installed
    if os.environ.get('SKLEARN_DOC_OPTIPNG', False):
        try:
            subprocess.call(["optipng", "-quiet", "-o", "9", out_fname])
        except Exception:
            warnings.warn('Install optipng to reduce the size of the \
                          generated images')

def save_thumbnail(image_path, base_image_name):
    """Save the thumbnail image"""
    first_image_file = image_path.format(1)
    thumb_dir = os.path.join(os.path.dirname(first_image_file), 'thumb')
    if not os.path.exists(thumb_dir):
        os.makedirs(thumb_dir)

    thumb_file = os.path.join(thumb_dir, 'sphx_glr_%s_thumb.png' % base_image_name)

    if os.path.exists(first_image_file):
        scale_image(first_image_file, thumb_file, 400, 280)
    elif not os.path.exists(thumb_file):
        # create something to replace the thumbnail
        scale_image(os.path.join(glr_path_static(), 'no_image.png'),
                    thumb_file, 200, 140)

def generate_dir_rst(src_dir, target_dir, gallery_conf, seen_backrefs):
    """Generate the gallery reStructuredText for an example directory"""
    if not os.path.exists(os.path.join(src_dir, 'README.txt')):
        print(80 * '_')
        print('Example directory %s does not have a README.txt file' %
              src_dir)
        print('Skipping this directory')
        print(80 * '_')
        return ""  # because string is an expected return type

    fhindex = open(os.path.join(src_dir, 'README.txt')).read()
    if not os.path.exists(target_dir):
        os.makedirs(target_dir)
    sorted_listdir = [fname for fname in sorted(os.listdir(src_dir))
                                if fname.endswith('py')]
    for fname in sorted_listdir:
        generate_file_rst(fname, target_dir, src_dir)
        new_fname = os.path.join(src_dir, fname)
        intro = extract_intro(new_fname)
        write_backreferences(seen_backrefs, gallery_conf,
                             target_dir, fname, intro)

        fhindex += _thumbnail_div(target_dir, fname, intro)
        fhindex += """

.. toctree::
   :hidden:

   /%s/%s\n""" % (target_dir, fname[:-3])

# clear at the end of the section
    fhindex += """.. raw:: html\n

\n\n"""

    return fhindex

def execute_script(code_block, example_globals, image_path, fig_count, src_file):
    """Executes the code block of the example file"""
    time_elapsed = 0
    stdout = ''

    # We need to execute the code
    print('plotting code blocks in %s' % src_file)

    plt.close('all')
    cwd = os.getcwd()
    # Redirect output to stdout and
    orig_stdout = sys.stdout

    try:
        # First CD in the original example dir, so that any file
        # created by the example get created in this directory
        os.chdir(os.path.dirname(src_file))
        my_buffer = StringIO()
        my_stdout = Tee(sys.stdout, my_buffer)
        sys.stdout = my_stdout

        t_start = time()
        exec(code_block, example_globals)
        time_elapsed = time() - t_start

        sys.stdout = orig_stdout

        my_stdout = my_buffer.getvalue().strip().expandtabs()
        if my_stdout:
            stdout = CODE_OUTPUT.format(my_stdout.replace('\n', '\n    '))
        os.chdir(cwd)
        figure_list = save_figures(image_path, fig_count)

        # Depending on whether we have one or more figures, we're using a
        # horizontal list or a single rst call to 'image'.
        if len(figure_list) == 1:
            figure_name = figure_list[0]
            image_list = SINGLE_IMAGE % figure_name.lstrip('/')
        else:
            image_list = HLIST_HEADER
            for figure_name in figure_list:
                image_list += HLIST_IMAGE_TEMPLATE % figure_name.lstrip('/')

    except Exception:
        image_list = '%s is not compiling:' % src_file
        print(80 * '_')
        print(image_list)
        traceback.print_exc()
        print(80 * '_')
    finally:
        os.chdir(cwd)
        sys.stdout = orig_stdout

    print(" - time elapsed : %.2g sec" % time_elapsed)
    code_output = "\n{0}\n\n{1}".format(image_list, stdout)

    return code_output, time_elapsed, fig_count + len(figure_list)

def generate_file_rst(fname, target_dir, src_dir):
    """ Generate the rst file for a given example."""

    src_file = os.path.join(src_dir, fname)
    example_file = os.path.join(target_dir, fname)
    shutil.copyfile(src_file, example_file)

    image_dir = os.path.join(target_dir, 'images')
    if not os.path.exists(image_dir):
        os.makedirs(image_dir)

    base_image_name = os.path.splitext(fname)[0]
    image_fname = 'sphx_glr_' + base_image_name + '_{0:03}.png'
    image_path = os.path.join(image_dir, image_fname)

    if _plots_are_current(src_file, image_path):
        return

    time_elapsed = 0
    script_blocks = split_code_and_text_blocks(example_file)

    ref_fname = example_file.replace(os.path.sep, '_')
    example_rst = """\n\n.. _example_{0}:\n\n""".format(ref_fname)

    if not fname.startswith('plot'):
        convert_func = dict(code=codestr2rst, text=ast.literal_eval)
        for blabel, bcontent in script_blocks:
            example_rst += convert_func[blabel](bcontent)+'\n'
    else:
        example_globals = {}
        fig_count = 0
        for blabel, bcontent in script_blocks:
            if blabel == 'code':
                code_output, rtime, fig_count = execute_script(bcontent,
                                                               example_globals,
                                                               image_path,
                                                               fig_count,
                                                               src_file)

                time_elapsed += rtime
                # Single example style
                if len(script_blocks)==2:
                    example_rst += code_output
                    example_rst += codestr2rst(bcontent)+'\n'
                # Notebook style
                else:
                    example_rst += codestr2rst(bcontent)+'\n'
                    example_rst += code_output
            else:
                example_rst += ast.literal_eval(bcontent)+'\n'

    save_thumbnail(image_path, base_image_name)

    time_m, time_s = divmod(time_elapsed, 60)
    with open(os.path.join(target_dir, base_image_name + '.rst'), 'w') as f:
        example_rst += CODE_DOWNLOAD.format(time_m, time_s, fname)
        f.write(example_rst)
<|endoftext|>