<commit_before>"""
The task runner.

Run it with `python -m iceprod.core.i3exec`.

optional arguments:
  -h, --help            show this help message and exit
  -f CFGFILE, --cfgfile CFGFILE
                        Specify config file
  -v VALIDATE, --validate VALIDATE
                        Specify if the config file should be validated
  -u URL, --url URL     URL of the iceprod server
  -p PASSKEY, --passkey PASSKEY
                        passkey for communcation with iceprod server
  --pilot_id PILOTID    ID of the pilot (if this is a pilot)
  -d, --debug           Enable debug actions and logging
  --offline             Enable offline mode (don't talk with server)
"""

import os
import sys
import logging
import logging.config
import time
import signal
from functools import partial
import tempfile
import shutil

from iceprod.core import to_file, constants
import iceprod.core.dataclasses
import iceprod.core.serialization
import iceprod.core.exe
import iceprod.core.exe_json

import iceprod.core.logger
logging.basicConfig()
logger = None

def handler(signum, frame):
    """Signal handler. Exit on SIGQUIT or SIGINT."""
    logger.warn('Signal handler called with signal %s' % signum)
    logger.warn('Exiting...')
    os._exit(0)

def load_config(cfgfile):
    """Load a config from file, serialized string, dictionary, etc"""
    config = None
    if isinstance(cfgfile,str):
        try:
            if os.path.exists(cfgfile):
                config = iceprod.core.serialization.serialize_json.load(cfgfile)
            else:
                config = iceprod.core.serialization.serialize_json.loads(cfgfile)
            if not config:
                raise Exception('Config not found')
        except Exception as e:
            logger.critical('Error loading configuration: %s' % str(e))
            raise
    elif isinstance(cfgfile,iceprod.core.dataclasses.Job):
        config = cfgfile
    elif isinstance(cfgfile,dict):
        config = iceprod.core.serialization.dict_to_dataclasses(cfgfile)
    else:
        logger.warn('cfgfile: %r',cfgfile)
        raise Exception('cfgfile is not a str or a Job')
    return config

def main(cfgfile=None, logfile=None, url=None, debug=False,
         passkey='', pilot_id=None, offline=False):
    """Main task runner for iceprod"""
    global logger
    # set up stdout and stderr
    stdout = partial(to_file,sys.stdout,constants['stdout'])
    stderr = partial(to_file,sys.stderr,constants['stderr'])
    with stdout(), stderr():
        # set up logger
        if debug:
            logl = 'INFO'
        else:
            logl = 'WARNING'
        if logfile:
            logf = os.path.abspath(logfile)
        else:
            logf = os.path.abspath(os.path.expandvars(constants['stdlog']))
        iceprod.core.logger.setlogger('i3exec',
                                      None,
                                      loglevel=logl,
                                      logfile=logf,
                                      logsize=67108864,
                                      lognum=1)
        iceprod.core.logger.removestdout()
        logger = logging.getLogger('i3exec')
        logger.warn('starting...%s ' % logger.name)

        signal.signal(signal.SIGQUIT, handler)
        signal.signal(signal.SIGINT, handler)

        if cfgfile is None:
            logger.critical('There is no cfgfile')
            raise Exception('missing cfgfile')
        config = load_config(cfgfile)
        logger.info('config: %r',config)

        if offline is True:
            # run in offline mode
            runner(config,url,debug,offline)
            return

        # setup jsonRPC
        kwargs = {}
        if 'username' in config['options']:
            kwargs['username'] = config['options']['username']
        if 'password' in config['options']:
            kwargs['password'] = config['options']['password']
        if 'ssl' in config['options'] and config['options']['ssl']:
            kwargs.update(config['options']['ssl'])
        iceprod.core.exe_json.setupjsonRPC(url+'/jsonrpc',passkey,**kwargs)

        if 'tasks' in config and config['tasks']:
            logger.info('default configuration - a single task')
            runner(config,url,debug)
        else:
            logger.info('pilot mode - get many tasks from server')
            if 'gridspec' not in config['options']:
                logger.critical('gridspec missing')
                raise Exception('gridspec missing')
            errors = 0
            while errors < 5:
                try:
                    kwargs = {}
                    task_config = iceprod.core.exe_json.downloadtask(config['options']['gridspec'])
                    task_id = task_config['options']['task_id']
                    iceprod.core.exe_json.update_pilot(pilot_id, tasks=task_id)
                except Exception:
                    errors += 1
                    logger.error('cannot download task. current error count is %d',
                                 errors, exc_info=True)
                    continue
                logger.info('task config: %r', task_config)

                if task_config is None:
                    break # assuming server wants client to exit
                else:
                    # add grid-specific config
                    for k in config['options']:
                        if k not in task_config['options']:
                            task_config['options'][k] = config['options'][k]
                    # clear stdout,stderr,log
                    for stream in (sys.stdout,sys.stderr):
                        stream.flush()
                        os.ftruncate(stream.fileno(), 0)
                    iceprod.core.logger.rotate()
                    # run task
                    main_dir = os.getcwd()
                    try:
                        tmpdir = tempfile.mkdtemp(dir=main_dir)
                        os.chdir(tmpdir)
                        for f in os.listdir(main_dir):
                            os.symlink(os.path.join(main_dir, f),
                                       os.path.join(tmpdir, f))
                        runner(task_config,url,debug)
                    except Exception:
                        errors += 1
                        logger.error('task encountered an error. current error count is %d',
                                     errors, exc_info=True)
                    finally:
                        os.chdir(main_dir)
                        shutil.rmtree(tmpdir)
                        iceprod.core.exe_json.update_pilot(pilot_id, tasks='')
            if errors >= 5:
                logger.critical('too many errors when running tasks')
        logger.warn('finished running normally; exiting...')

def runner(config,url,debug=False,offline=False):
    """Run a config"""
    # set logging verbosity
    if 'debug' not in config['options']:
        config['options']['debug'] = debug
    if ('debug' in config['options'] and config['options']['debug'] and
        'loglevel' not in config['options']):
        config['options']['loglevel'] = 'INFO'
    if ('loglevel' in config['options'] and
        config['options']['loglevel'].upper() in iceprod.core.logger.setlevel):
        try:
            logging.getLogger().setLevel(iceprod.core.logger.setlevel[config['options']['loglevel'].upper()])
        except Exception:
            logger.warn('failed to set a new log level', exc_info=True)

    # make sure some basic options are set
    if 'job' not in config['options']:
        config['options']['job'] = 0
    if 'jobs_submitted' not in config['options']:
        config['options']['jobs_submitted'] = 1
    if 'resource_url' not in config['options']:
        config['options']['resource_url'] = str(url)+'/download'
    if 'offline' not in config['options']:
        config['options']['offline'] = offline
    if 'data_url' not in config['options']:
        config['options']['data_url'] = 'gsiftp://gridftp.icecube.wisc.edu/'
    if 'svn_repository' not in config['options']:
        config['options']['svn_repository'] = 'http://code.icecube.wisc.edu/svn/'
    if 'site_temp' not in config['options']:
        config['options']['site_temp'] = 'gsiftp://gridftp.icecube.wisc.edu/data/sim/sim-new/tmp/dagtemp2'
    if 'dataset_temp' not in config['options']:
        config['options']['dataset_temp'] = os.path.join(config['options']['site_temp'],'$(dataset)')
    if 'job_temp' not in config['options']:
        config['options']['job_temp'] = os.path.join(config['options']['dataset_temp'],'$(job)')
    if 'task_temp' not in config['options']:
        config['options']['task_temp'] = 'file:'+os.path.join(os.getcwd(),'task_temp')
    if 'tray_temp' not in config['options']:
        config['options']['tray_temp'] = 'file:'+os.path.join(os.getcwd(),'tray_temp')
    if 'local_temp' not in config['options']:
        config['options']['local_temp'] = os.path.join(os.getcwd(),'local_temp')
    if 'stillrunninginterval' not in config['options']:
        config['options']['stillrunninginterval'] = 60

    if not config['steering']:
        # make sure steering exists in the config
        config['steering'] = iceprod.core.dataclasses.Steering()

    # make exe Config
    cfg = iceprod.core.exe.Config(config=config)

    # set up global env, based on config['options'] and config.steering
    env_opts = cfg.parseObject(config['options'], {})
    stats = {}
    try:
        with iceprod.core.exe.setupenv(cfg, config['steering'], {'options':env_opts}) as env:
            logger.warn("config options: %r",config['options'])

            if not offline:
                # tell the server that we are processing this task
                try:
                    iceprod.core.exe_json.processing(cfg)
                except Exception as e:
                    logging.error(e)

            # keep track of the start time
            start_time = time.time()

            # find tasks to run
            try:
                if 'task' in config['options']:
                    logger.warn('task specified: %r',config['options']['task'])
                    # run only this task name or number
                    name = config['options']['task']
                    if isinstance(name, iceprod.core.dataclasses.String) and name.isdigit():
                        name = int(name)
                    if isinstance(name, iceprod.core.dataclasses.String):
                        # find task by name
                        for task in config['tasks']:
                            if task['name'] == name:
                                stats = iceprod.core.exe.runtask(cfg, env, task)
                                break
                        else:
                            logger.critical('cannot find task named %r', name)
                            raise Exception('cannot find specified task')
                    elif isinstance(name, int):
                        # find task by index
                        if (name >= 0 and
                            name < len(config['tasks'])):
                            stats = iceprod.core.exe.runtask(cfg, env, config['tasks'][name])
                        else:
                            logger.critical('cannot find task index %d', name)
                            raise Exception('cannot find specified task')

                    else:
                        logger.critical('task specified in options is %r, but no task found',
                                        name)
                        raise Exception('cannot find specified task')
                elif offline:
                    # run all tasks in order
                    for task in config['tasks']:
                        iceprod.core.exe.runtask(cfg, env, task)
                else:
                    raise Exception('task to run not specified')
            except Exception as e:
                logger.error('task failed, exiting without running completion steps.',
                             exc_info=True)
                # set task status on server
                if not offline:
                    try:
                        iceprod.core.exe_json.taskerror(cfg, start_time=start_time)
                    except Exception as e:
                        logger.error(e)
                raise

    finally:
        # upload log files to server
        try:
            if 'upload' in config['options']:
                if isinstance(config['options']['upload'],
                              iceprod.core.dataclasses.String):
                    upload = config['options']['upload'].lower().split('|')
                elif isinstance(config['options']['upload'],(tuple,list)):
                    upload = [x.lower() for x in config['options']['upload']]
                else:
                    raise Exception('upload config is not a valid type')
                for up in upload:
                    if up.startswith('logging'):
                        # upload err,log,out files
                        iceprod.core.exe_json.uploadLogging(cfg)
                        break
                    elif up.startswith('log'):
                        # upload log files
                        iceprod.core.exe_json.uploadLog(cfg)
                    elif up.startswith('err'):
                        # upload err files
                        iceprod.core.exe_json.uploadErr(cfg)
                    elif up.startswith('out'):
                        # upload out files
                        iceprod.core.exe_json.uploadOut(cfg)
        except Exception as e:
            logger.error('failed when uploading logging info',exc_info=True)

    iceprod.core.exe_json.finishtask(cfg, stats)
    logger.warn('finished without error')

if __name__ == '__main__':
    # get arguments
    import argparse
    parser = argparse.ArgumentParser(description='IceProd Core')
    parser.add_argument('-f','--cfgfile', type=str,
                        help='Specify config file')
    parser.add_argument('-u','--url', type=str,
                        help='URL of the iceprod server')
    parser.add_argument('-p','--passkey', type=str,
                        help='passkey for communcation with iceprod server')
    parser.add_argument('--pilot_id', type=str, default=None,
                        help='ID of the pilot (if this is a pilot)')
    parser.add_argument('-d','--debug', action='store_true', default=False,
                        help='Enable debug actions and logging')
    parser.add_argument('--offline', action='store_true', default=False,
                        help='Enable offline mode (don\'t talk with server)')
    parser.add_argument('--logfile', type=str, default=None,
                        help='Specify the logfile to use')

    args = vars(parser.parse_args())
    print args

    # check cfgfile
    if args['cfgfile'] is not None and not os.path.isfile(args['cfgfile']):
        if os.path.isfile(os.path.join(os.getcwd(),args['cfgfile'])):
            args['cfgfile'] = os.path.join(os.getcwd(),args['cfgfile'])
        else:
            args['cfgfile'] = None

    # start iceprod
    main(**args)
<commit_msg>only update pilot once it's confirmed we have a task to run<commit_after>"""
The task runner.

Run it with `python -m iceprod.core.i3exec`.

optional arguments:
  -h, --help            show this help message and exit
  -f CFGFILE, --cfgfile CFGFILE
                        Specify config file
  -v VALIDATE, --validate VALIDATE
                        Specify if the config file should be validated
  -u URL, --url URL     URL of the iceprod server
  -p PASSKEY, --passkey PASSKEY
                        passkey for communcation with iceprod server
  --pilot_id PILOTID    ID of the pilot (if this is a pilot)
  -d, --debug           Enable debug actions and logging
  --offline             Enable offline mode (don't talk with server)
"""

import os
import sys
import logging
import logging.config
import time
import signal
from functools import partial
import tempfile
import shutil

from iceprod.core import to_file, constants
import iceprod.core.dataclasses
import iceprod.core.serialization
import iceprod.core.exe
import iceprod.core.exe_json

import iceprod.core.logger
logging.basicConfig()
logger = None

def handler(signum, frame):
    """Signal handler. Exit on SIGQUIT or SIGINT."""
    logger.warn('Signal handler called with signal %s' % signum)
    logger.warn('Exiting...')
    os._exit(0)

def load_config(cfgfile):
    """Load a config from file, serialized string, dictionary, etc"""
    config = None
    if isinstance(cfgfile,str):
        try:
            if os.path.exists(cfgfile):
                config = iceprod.core.serialization.serialize_json.load(cfgfile)
            else:
                config = iceprod.core.serialization.serialize_json.loads(cfgfile)
            if not config:
                raise Exception('Config not found')
        except Exception as e:
            logger.critical('Error loading configuration: %s' % str(e))
            raise
    elif isinstance(cfgfile,iceprod.core.dataclasses.Job):
        config = cfgfile
    elif isinstance(cfgfile,dict):
        config = iceprod.core.serialization.dict_to_dataclasses(cfgfile)
    else:
        logger.warn('cfgfile: %r',cfgfile)
        raise Exception('cfgfile is not a str or a Job')
    return config

def main(cfgfile=None, logfile=None, url=None, debug=False,
         passkey='', pilot_id=None, offline=False):
    """Main task runner for iceprod"""
    global logger
    # set up stdout and stderr
    stdout = partial(to_file,sys.stdout,constants['stdout'])
    stderr = partial(to_file,sys.stderr,constants['stderr'])
    with stdout(), stderr():
        # set up logger
        if debug:
            logl = 'INFO'
        else:
            logl = 'WARNING'
        if logfile:
            logf = os.path.abspath(logfile)
        else:
            logf = os.path.abspath(os.path.expandvars(constants['stdlog']))
        iceprod.core.logger.setlogger('i3exec',
                                      None,
                                      loglevel=logl,
                                      logfile=logf,
                                      logsize=67108864,
                                      lognum=1)
        iceprod.core.logger.removestdout()
        logger = logging.getLogger('i3exec')
        logger.warn('starting...%s ' % logger.name)

        signal.signal(signal.SIGQUIT, handler)
        signal.signal(signal.SIGINT, handler)

        if cfgfile is None:
            logger.critical('There is no cfgfile')
            raise Exception('missing cfgfile')
        config = load_config(cfgfile)
        logger.info('config: %r',config)

        if offline is True:
            # run in offline mode
            runner(config,url,debug,offline)
            return

        # setup jsonRPC
        kwargs = {}
        if 'username' in config['options']:
            kwargs['username'] = config['options']['username']
        if 'password' in config['options']:
            kwargs['password'] = config['options']['password']
        if 'ssl' in config['options'] and config['options']['ssl']:
            kwargs.update(config['options']['ssl'])
        iceprod.core.exe_json.setupjsonRPC(url+'/jsonrpc',passkey,**kwargs)

        if 'tasks' in config and config['tasks']:
            logger.info('default configuration - a single task')
            runner(config,url,debug)
        else:
            logger.info('pilot mode - get many tasks from server')
            if 'gridspec' not in config['options']:
                logger.critical('gridspec missing')
                raise Exception('gridspec missing')
            errors = 0
            while errors < 5:
                try:
                    kwargs = {}
                    task_config = iceprod.core.exe_json.downloadtask(config['options']['gridspec'])
                except Exception:
                    errors += 1
                    logger.error('cannot download task. current error count is %d',
                                 errors, exc_info=True)
                    continue
                logger.info('task config: %r', task_config)

                if task_config is None:
                    break # assuming server wants client to exit
                else:
                    task_id = task_config['options']['task_id']
                    iceprod.core.exe_json.update_pilot(pilot_id, tasks=task_id)

                    # add grid-specific config
                    for k in config['options']:
                        if k not in task_config['options']:
                            task_config['options'][k] = config['options'][k]
                    # clear stdout,stderr,log
                    for stream in (sys.stdout,sys.stderr):
                        stream.flush()
                        os.ftruncate(stream.fileno(), 0)
                    iceprod.core.logger.rotate()
                    # run task
                    main_dir = os.getcwd()
                    try:
                        tmpdir = tempfile.mkdtemp(dir=main_dir)
                        os.chdir(tmpdir)
                        for f in os.listdir(main_dir):
                            os.symlink(os.path.join(main_dir, f),
                                       os.path.join(tmpdir, f))
                        runner(task_config,url,debug)
                    except Exception:
                        errors += 1
                        logger.error('task encountered an error. current error count is %d',
                                     errors, exc_info=True)
                    finally:
                        os.chdir(main_dir)
                        shutil.rmtree(tmpdir)
                        iceprod.core.exe_json.update_pilot(pilot_id, tasks='')
            if errors >= 5:
                logger.critical('too many errors when running tasks')
        logger.warn('finished running normally; exiting...')

def runner(config,url,debug=False,offline=False):
    """Run a config"""
    # set logging verbosity
    if 'debug' not in config['options']:
        config['options']['debug'] = debug
    if ('debug' in config['options'] and config['options']['debug'] and
        'loglevel' not in config['options']):
        config['options']['loglevel'] = 'INFO'
    if ('loglevel' in config['options'] and
        config['options']['loglevel'].upper() in iceprod.core.logger.setlevel):
        try:
            logging.getLogger().setLevel(iceprod.core.logger.setlevel[config['options']['loglevel'].upper()])
        except Exception:
            logger.warn('failed to set a new log level', exc_info=True)

    # make sure some basic options are set
    if 'job' not in config['options']:
        config['options']['job'] = 0
    if 'jobs_submitted' not in config['options']:
        config['options']['jobs_submitted'] = 1
    if 'resource_url' not in config['options']:
        config['options']['resource_url'] = str(url)+'/download'
    if 'offline' not in config['options']:
        config['options']['offline'] = offline
    if 'data_url' not in config['options']:
        config['options']['data_url'] = 'gsiftp://gridftp.icecube.wisc.edu/'
    if 'svn_repository' not in config['options']:
        config['options']['svn_repository'] = 'http://code.icecube.wisc.edu/svn/'
    if 'site_temp' not in config['options']:
        config['options']['site_temp'] = 'gsiftp://gridftp.icecube.wisc.edu/data/sim/sim-new/tmp/dagtemp2'
    if 'dataset_temp' not in config['options']:
        config['options']['dataset_temp'] = os.path.join(config['options']['site_temp'],'$(dataset)')
    if 'job_temp' not in config['options']:
        config['options']['job_temp'] = os.path.join(config['options']['dataset_temp'],'$(job)')
    if 'task_temp' not in config['options']:
        config['options']['task_temp'] = 'file:'+os.path.join(os.getcwd(),'task_temp')
    if 'tray_temp' not in config['options']:
        config['options']['tray_temp'] = 'file:'+os.path.join(os.getcwd(),'tray_temp')
    if 'local_temp' not in config['options']:
        config['options']['local_temp'] = os.path.join(os.getcwd(),'local_temp')
    if 'stillrunninginterval' not in config['options']:
        config['options']['stillrunninginterval'] = 60

    if not config['steering']:
        # make sure steering exists in the config
        config['steering'] = iceprod.core.dataclasses.Steering()

    # make exe Config
    cfg = iceprod.core.exe.Config(config=config)

    # set up global env, based on config['options'] and config.steering
    env_opts = cfg.parseObject(config['options'], {})
    stats = {}
    try:
        with iceprod.core.exe.setupenv(cfg, config['steering'], {'options':env_opts}) as env:
            logger.warn("config options: %r",config['options'])

            if not offline:
                # tell the server that we are processing this task
                try:
                    iceprod.core.exe_json.processing(cfg)
                except Exception as e:
                    logging.error(e)

            # keep track of the start time
            start_time = time.time()

            # find tasks to run
            try:
                if 'task' in config['options']:
                    logger.warn('task specified: %r',config['options']['task'])
                    # run only this task name or number
                    name = config['options']['task']
                    if isinstance(name, iceprod.core.dataclasses.String) and name.isdigit():
                        name = int(name)
                    if isinstance(name, iceprod.core.dataclasses.String):
                        # find task by name
                        for task in config['tasks']:
                            if task['name'] == name:
                                stats = iceprod.core.exe.runtask(cfg, env, task)
                                break
                        else:
                            logger.critical('cannot find task named %r', name)
                            raise Exception('cannot find specified task')
                    elif isinstance(name, int):
                        # find task by index
                        if (name >= 0 and
                            name < len(config['tasks'])):
                            stats = iceprod.core.exe.runtask(cfg, env, config['tasks'][name])
                        else:
                            logger.critical('cannot find task index %d', name)
                            raise Exception('cannot find specified task')

                    else:
                        logger.critical('task specified in options is %r, but no task found',
                                        name)
                        raise Exception('cannot find specified task')
                elif offline:
                    # run all tasks in order
                    for task in config['tasks']:
                        iceprod.core.exe.runtask(cfg, env, task)
                else:
                    raise Exception('task to run not specified')
            except Exception as e:
                logger.error('task failed, exiting without running completion steps.',
                             exc_info=True)
                # set task status on server
                if not offline:
                    try:
                        iceprod.core.exe_json.taskerror(cfg, start_time=start_time)
                    except Exception as e:
                        logger.error(e)
                raise

    finally:
        # upload log files to server
        try:
            if 'upload' in config['options']:
                if isinstance(config['options']['upload'],
                              iceprod.core.dataclasses.String):
                    upload = config['options']['upload'].lower().split('|')
                elif isinstance(config['options']['upload'],(tuple,list)):
                    upload = [x.lower() for x in config['options']['upload']]
                else:
                    raise Exception('upload config is not a valid type')
                for up in upload:
                    if up.startswith('logging'):
                        # upload err,log,out files
                        iceprod.core.exe_json.uploadLogging(cfg)
                        break
                    elif up.startswith('log'):
                        # upload log files
                        iceprod.core.exe_json.uploadLog(cfg)
                    elif up.startswith('err'):
                        # upload err files
                        iceprod.core.exe_json.uploadErr(cfg)
                    elif up.startswith('out'):
                        # upload out files
                        iceprod.core.exe_json.uploadOut(cfg)
        except Exception as e:
            logger.error('failed when uploading logging info',exc_info=True)

    iceprod.core.exe_json.finishtask(cfg, stats)
    logger.warn('finished without error')

if __name__ == '__main__':
    # get arguments
    import argparse
    parser = argparse.ArgumentParser(description='IceProd Core')
    parser.add_argument('-f','--cfgfile', type=str,
                        help='Specify config file')
    parser.add_argument('-u','--url', type=str,
                        help='URL of the iceprod server')
    parser.add_argument('-p','--passkey', type=str,
                        help='passkey for communcation with iceprod server')
    parser.add_argument('--pilot_id', type=str, default=None,
                        help='ID of the pilot (if this is a pilot)')
    parser.add_argument('-d','--debug', action='store_true', default=False,
                        help='Enable debug actions and logging')
    parser.add_argument('--offline', action='store_true', default=False,
                        help='Enable offline mode (don\'t talk with server)')
    parser.add_argument('--logfile', type=str, default=None,
                        help='Specify the logfile to use')

    args = vars(parser.parse_args())
    print args

    # check cfgfile
    if args['cfgfile'] is not None and not os.path.isfile(args['cfgfile']):
        if os.path.isfile(os.path.join(os.getcwd(),args['cfgfile'])):
            args['cfgfile'] = os.path.join(os.getcwd(),args['cfgfile'])
        else:
            args['cfgfile'] = None

    # start iceprod
    main(**args)
<|endoftext|>


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

from threading import current_thread

from PIL import Image
from six import BytesIO
from six.moves import queue
from six.moves.urllib.parse import urlparse

from icrawler.utils import ThreadPool

class Downloader(ThreadPool):
    """Base class for downloader.

    A thread pool of downloader threads, in charge of downloading files and
    saving them in the corresponding paths.

    Attributes:
        task_queue (CachedQueue): A queue storing image downloading tasks,
            connecting :class:`Parser` and :class:`Downloader`.
        signal (Signal): A Signal object shared by all components.
        session (Session): A session object.
        logger: A logging.Logger object used for logging.
        workers (list): A list of downloader threads.
        thread_num (int): The number of downloader threads.
        lock (Lock): A threading.Lock object.
        storage (BaseStorage): storage backend.
    """

    def __init__(self, thread_num, signal, session, storage):
        """Init Parser with some shared variables."""
        super(Downloader, self).__init__(
            thread_num, out_queue=None, name='downloader')
        self.signal = signal
        self.session = session
        self.storage = storage
        self.file_idx_offset = 0
        self.clear_status()

    def clear_status(self):
        """Reset fetched_num to 0."""
        self.fetched_num = 0

    def set_file_idx_offset(self, file_idx_offset=0):
        """Set offset of file index.

        Args:
            file_idx_offset: It can be either an integer or 'auto'. If set
                to an integer, the filename will start from
                ``file_idx_offset`` + 1. If set to ``'auto'``, the filename
                will start from existing max file index plus 1.
        """
        if isinstance(file_idx_offset, int):
            self.file_idx_offset = file_idx_offset
        elif file_idx_offset == 'auto':
            self.file_idx_offset = self.storage.max_file_idx()
        else:
            raise ValueError('"file_idx_offset" must be an integer or `auto`')

    def get_filename(self, task, default_ext):
        """Set the path where the image will be saved.

        The default strategy is to use an increasing 6-digit number as
        the filename. You can override this method if you want to set custom
        naming rules. The file extension is kept if it can be obtained from
        the url, otherwise ``default_ext`` is used as extension.

        Args:
            task (dict): The task dict got from ``task_queue``.

        Output:
            Filename with extension.
        """
        url_path = urlparse(task['file_url'])[2]
        extension = url_path.split('.')[-1] if '.' in url_path else default_ext
        file_idx = self.fetched_num + self.file_idx_offset
        return '{:06d}.{}'.format(file_idx, extension)

    def reach_max_num(self):
        """Check if downloaded images reached max num.

        Returns:
            bool: if downloaded images reached max num.
        """
        if self.signal.get('reach_max_num'):
            return True
        if self.max_num > 0 and self.fetched_num >= self.max_num:
            return True
        else:
            return False

    def keep_file(self, task, response, **kwargs):
        return True

    def download(self, task, default_ext, timeout=5, max_retry=3, **kwargs):
        """Download the image and save it to the corresponding path.

        Args:
            task (dict): The task dict got from ``task_queue``.
            timeout (int): Timeout of making requests for downloading images.
            max_retry (int): the max retry times if the request fails.
            **kwargs: reserved arguments for overriding.
        """
        file_url = task['file_url']
        task['success'] = False
        task['filename'] = None
        retry = max_retry
        while retry > 0 and not self.signal.get('reach_max_num'):
            try:
                response = self.session.get(file_url, timeout=timeout)
            except Exception as e:
                self.logger.error('Exception caught when downloading file %s, '
                                  'error: %s, remaining retry times: %d',
                                  file_url, e, retry - 1)
            else:
                if self.reach_max_num():
                    self.signal.set(reach_max_num=True)
                    break
                elif response.status_code != 200:
                    self.logger.error('Response status code %d, file %s',
                                      response.status_code, file_url)
                    break
                elif not self.keep_file(task, response, **kwargs):
                    break
                with self.lock:
                    self.fetched_num += 1
                    filename = self.get_filename(task, default_ext)
                self.logger.info('image #%s\t%s', self.fetched_num, file_url)
                self.storage.write(filename, response.content)
                task['success'] = True
                task['filename'] = filename
                break
            finally:
                retry -= 1

    def process_meta(self, task):
        """Process some meta data of the images.

        This method should be overridden by users if wanting to do more things
        other than just downloading the image, such as saving annotations.

        Args:
            task (dict): The task dict got from task_queue. This method will
                make use of fields other than ``file_url`` in the dict.
        """
        pass

    def start(self, file_idx_offset=0, *args, **kwargs):
        self.clear_status()
        self.set_file_idx_offset(file_idx_offset)
        self.init_workers(*args, **kwargs)
        for worker in self.workers:
            worker.start()
            self.logger.debug('thread %s started', worker.name)

    def worker_exec(self,
                    max_num,
                    default_ext='',
                    queue_timeout=5,
                    req_timeout=5,
                    **kwargs):
        """Target method of workers.

        Get task from ``task_queue`` and then download files and process meta
        data. A downloader thread will exit in either of the following cases:

        1. All parser threads have exited and the task_queue is empty.
        2. Downloaded image number has reached required number(max_num).

        Args:
            queue_timeout (int): Timeout of getting tasks from ``task_queue``.
            req_timeout (int): Timeout of making requests for downloading pages.
            **kwargs: Arguments passed to the :func:`download` method.
        """
        self.max_num = max_num
        while True:
            if self.signal.get('reach_max_num'):
                self.logger.info('downloaded images reach max num, thread %s'
                                 ' is ready to exit', current_thread().name)
                break
            try:
                task = self.in_queue.get(timeout=queue_timeout)
            except queue.Empty:
                if self.signal.get('parser_exited'):
                    self.logger.info('no more download task for thread %s',
                                     current_thread().name)
                    break
                else:
                    self.logger.info('%s is waiting for new download tasks',
                                     current_thread().name)
            except:
                self.logger.error('exception in thread %s',
                                  current_thread().name)
            else:
                self.download(task, default_ext, req_timeout, **kwargs)
                self.process_meta(task)
                self.in_queue.task_done()
        self.logger.info('thread {} exit'.format(current_thread().name))

    def __exit__(self):
        self.logger.info('all downloader threads exited')

class ImageDownloader(Downloader):
    """Downloader specified for images.
    """

    def _size_lt(self, sz1, sz2):
        return max(sz1) <= max(sz2) and min(sz1) <= min(sz2)

    def _size_gt(self, sz1, sz2):
        return max(sz1) >= max(sz2) and min(sz1) >= min(sz2)

    def keep_file(self, task, response, min_size=None, max_size=None):
        """Decide whether to keep the image

        Compare image size with ``min_size`` and ``max_size`` to decide.

        Args:
            response (Response): response of requests.
            min_size (tuple or None): minimum size of required images.
            max_size (tuple or None): maximum size of required images.
        Returns:
            bool: whether to keep the image.
        """
        try:
            img = Image.open(BytesIO(response.content))
        except (IOError, OSError):
            return False
        task['img_size'] = img.size
        if min_size and not self._size_gt(img.size, min_size):
            return False
        if max_size and not self._size_lt(img.size, max_size):
            return False
        return True

    def get_filename(self, task, default_ext):
        url_path = urlparse(task['file_url'])[2]
        if '.' in url_path:
            extension = url_path.split('.')[-1]
            if extension.lower() not in [
                    'jpg', 'jpeg', 'png', 'bmp', 'tiff', 'gif', 'ppm', 'pgm'
            ]:
                extension = default_ext
        else:
            extension = default_ext
        file_idx = self.fetched_num + self.file_idx_offset
        return '{:06d}.{}'.format(file_idx, extension)

    def worker_exec(self,
                    max_num,
                    default_ext='jpg',
                    queue_timeout=5,
                    req_timeout=5,
                    **kwargs):
        super(ImageDownloader, self).worker_exec(
            max_num, default_ext, queue_timeout, req_timeout, **kwargs)
<commit_msg>add option to skip download for existing files<commit_after># -*- coding: utf-8 -*-

from threading import current_thread

from PIL import Image
from six import BytesIO
from six.moves import queue
from six.moves.urllib.parse import urlparse
import os.path

from icrawler.utils import ThreadPool

class Downloader(ThreadPool):
    """Base class for downloader.

    A thread pool of downloader threads, in charge of downloading files and
    saving them in the corresponding paths.

    Attributes:
        task_queue (CachedQueue): A queue storing image downloading tasks,
            connecting :class:`Parser` and :class:`Downloader`.
        signal (Signal): A Signal object shared by all components.
        session (Session): A session object.
        logger: A logging.Logger object used for logging.
        workers (list): A list of downloader threads.
        thread_num (int): The number of downloader threads.
        lock (Lock): A threading.Lock object.
        storage (BaseStorage): storage backend.
    """

    def __init__(self, thread_num, signal, session, storage, skip_download_existing=True):
        """Init Parser with some shared variables."""
        super(Downloader, self).__init__(
            thread_num, out_queue=None, name='downloader')
        self.signal = signal
        self.session = session
        self.storage = storage
        self.file_idx_offset = 0
        self.clear_status()
        self.skip_download_existing = skip_download_existing

    def clear_status(self):
        """Reset fetched_num to 0."""
        self.fetched_num = 0

    def set_file_idx_offset(self, file_idx_offset=0):
        """Set offset of file index.

        Args:
            file_idx_offset: It can be either an integer or 'auto'. If set
                to an integer, the filename will start from
                ``file_idx_offset`` + 1. If set to ``'auto'``, the filename
                will start from existing max file index plus 1.
        """
        if isinstance(file_idx_offset, int):
            self.file_idx_offset = file_idx_offset
        elif file_idx_offset == 'auto':
            self.file_idx_offset = self.storage.max_file_idx()
        else:
            raise ValueError('"file_idx_offset" must be an integer or `auto`')

    def get_filename(self, task, default_ext):
        """Set the path where the image will be saved.

        The default strategy is to use an increasing 6-digit number as
        the filename. You can override this method if you want to set custom
        naming rules. The file extension is kept if it can be obtained from
        the url, otherwise ``default_ext`` is used as extension.

        Args:
            task (dict): The task dict got from ``task_queue``.

        Output:
            Filename with extension.
        """
        url_path = urlparse(task['file_url'])[2]
        extension = url_path.split('.')[-1] if '.' in url_path else default_ext
        file_idx = self.fetched_num + self.file_idx_offset
        return '{:06d}.{}'.format(file_idx, extension)

    def reach_max_num(self):
        """Check if downloaded images reached max num.

        Returns:
            bool: if downloaded images reached max num.
        """
        if self.signal.get('reach_max_num'):
            return True
        if self.max_num > 0 and self.fetched_num >= self.max_num:
            return True
        else:
            return False

    def keep_file(self, task, response, **kwargs):
        return True

    def download(self, task, default_ext, timeout=5, max_retry=3, **kwargs):
        """Download the image and save it to the corresponding path.

        Args:
            task (dict): The task dict got from ``task_queue``.
            timeout (int): Timeout of making requests for downloading images.
            max_retry (int): the max retry times if the request fails.
            **kwargs: reserved arguments for overriding.
        """
        file_url = task['file_url']
        task['success'] = False
        task['filename'] = None
        retry = max_retry
        if self.skip_download_existing and os.path.exists(self.get_filename(task, default_ext)):
            return

        while retry > 0 and not self.signal.get('reach_max_num'):
            try:
                response = self.session.get(file_url, timeout=timeout)
            except Exception as e:
                self.logger.error('Exception caught when downloading file %s, '
                                  'error: %s, remaining retry times: %d',
                                  file_url, e, retry - 1)
            else:
                if self.reach_max_num():
                    self.signal.set(reach_max_num=True)
                    break
                elif response.status_code != 200:
                    self.logger.error('Response status code %d, file %s',
                                      response.status_code, file_url)
                    break
                elif not self.keep_file(task, response, **kwargs):
                    break
                with self.lock:
                    self.fetched_num += 1
                    filename = self.get_filename(task, default_ext)
                self.logger.info('image #%s\t%s', self.fetched_num, file_url)
                self.storage.write(filename, response.content)
                task['success'] = True
                task['filename'] = filename
                break
            finally:
                retry -= 1

    def process_meta(self, task):
        """Process some meta data of the images.

        This method should be overridden by users if wanting to do more things
        other than just downloading the image, such as saving annotations.

        Args:
            task (dict): The task dict got from task_queue. This method will
                make use of fields other than ``file_url`` in the dict.
        """
        pass

    def start(self, file_idx_offset=0, *args, **kwargs):
        self.clear_status()
        self.set_file_idx_offset(file_idx_offset)
        self.init_workers(*args, **kwargs)
        for worker in self.workers:
            worker.start()
            self.logger.debug('thread %s started', worker.name)

    def worker_exec(self,
                    max_num,
                    default_ext='',
                    queue_timeout=5,
                    req_timeout=5,
                    **kwargs):
        """Target method of workers.

        Get task from ``task_queue`` and then download files and process meta
        data. A downloader thread will exit in either of the following cases:

        1. All parser threads have exited and the task_queue is empty.
        2. Downloaded image number has reached required number(max_num).

        Args:
            queue_timeout (int): Timeout of getting tasks from ``task_queue``.
            req_timeout (int): Timeout of making requests for downloading pages.
            **kwargs: Arguments passed to the :func:`download` method.
        """
        self.max_num = max_num
        while True:
            if self.signal.get('reach_max_num'):
                self.logger.info('downloaded images reach max num, thread %s'
                                 ' is ready to exit', current_thread().name)
                break
            try:
                task = self.in_queue.get(timeout=queue_timeout)
            except queue.Empty:
                if self.signal.get('parser_exited'):
                    self.logger.info('no more download task for thread %s',
                                     current_thread().name)
                    break
                else:
                    self.logger.info('%s is waiting for new download tasks',
                                     current_thread().name)
            except:
                self.logger.error('exception in thread %s',
                                  current_thread().name)
            else:
                self.download(task, default_ext, req_timeout, **kwargs)
                self.process_meta(task)
                self.in_queue.task_done()
        self.logger.info('thread {} exit'.format(current_thread().name))

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.logger.info('all downloader threads exited')

class ImageDownloader(Downloader):
    """Downloader specified for images.
    """

    def _size_lt(self, sz1, sz2):
        return max(sz1) <= max(sz2) and min(sz1) <= min(sz2)

    def _size_gt(self, sz1, sz2):
        return max(sz1) >= max(sz2) and min(sz1) >= min(sz2)

    def keep_file(self, task, response, min_size=None, max_size=None):
        """Decide whether to keep the image

        Compare image size with ``min_size`` and ``max_size`` to decide.

        Args:
            response (Response): response of requests.
            min_size (tuple or None): minimum size of required images.
            max_size (tuple or None): maximum size of required images.
        Returns:
            bool: whether to keep the image.
        """
        try:
            img = Image.open(BytesIO(response.content))
        except (IOError, OSError):
            return False
        task['img_size'] = img.size
        if min_size and not self._size_gt(img.size, min_size):
            return False
        if max_size and not self._size_lt(img.size, max_size):
            return False
        return True

    def get_filename(self, task, default_ext):
        url_path = urlparse(task['file_url'])[2]
        if '.' in url_path:
            extension = url_path.split('.')[-1]
            if extension.lower() not in [
                    'jpg', 'jpeg', 'png', 'bmp', 'tiff', 'gif', 'ppm', 'pgm'
            ]:
                extension = default_ext
        else:
            extension = default_ext
        file_idx = self.fetched_num + self.file_idx_offset
        return '{:06d}.{}'.format(file_idx, extension)

    def worker_exec(self,
                    max_num,
                    default_ext='jpg',
                    queue_timeout=5,
                    req_timeout=5,
                    **kwargs):
        super(ImageDownloader, self).worker_exec(
            max_num, default_ext, queue_timeout, req_timeout, **kwargs)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import unicode_literals

EXTENSIONS = {
    'apinotes': {'text', 'apinotes'},
    'asar': {'binary', 'asar'},
    'bash': {'text', 'shell', 'bash'},
    'bat': {'text', 'batch'},
    'bmp': {'binary', 'image', 'bitmap'},
    'bz2': {'binary', 'bzip2'},
    'c': {'text', 'c'},
    'cc': {'text', 'c++'},
    'cu': {'text', 'cuda'},
    'cfg': {'text'},
    'cmake': {'text', 'cmake'},
    'cnf': {'text'},
    'coffee': {'text', 'coffee'},
    'conf': {'text'},
    'cpp': {'text', 'c++'},
    'crt': {'text', 'pem'},
    'cs': {'text', 'c#'},
    'cson': {'text', 'cson'},
    'css': {'text', 'css'},
    'csv': {'text', 'csv'},
    'cxx': {'text', 'c++'},
    'def': {'text', 'def'},
    'dtd': {'text', 'dtd'},
    'ear': {'binary', 'zip', 'jar'},
    'ejs': {'text', 'ejs'},
    'eot': {'binary', 'eot'},
    'eps': {'binary', 'eps'},
    'erb': {'text', 'erb'},
    'exe': {'binary'},
    'eyaml': {'text', 'yaml'},
    'feature': {'text', 'gherkin'},
    'fish': {'text', 'fish'},
    'gemspec': {'text', 'ruby'},
    'gif': {'binary', 'image', 'gif'},
    'go': {'text', 'go'},
    'gotmpl': {'text', 'gotmpl'},
    'gpx': {'text', 'gpx', 'xml'},
    'gradle': {'text', 'groovy'},
    'groovy': {'text', 'groovy'},
    'gyb': {'text', 'gyb'},
    'gyp': {'text', 'gyp', 'python'},
    'gypi': {'text', 'gyp', 'python'},
    'gz': {'binary', 'gzip'},
    'h': {'text', 'header', 'c', 'c++'},
    'hpp': {'text', 'header', 'c++'},
    'htm': {'text', 'html'},
    'html': {'text', 'html'},
    'hxx': {'text', 'header', 'c++'},
    'icns': {'binary', 'icns'},
    'ico': {'binary', 'icon'},
    'idl': {'text', 'idl'},
    'inc': {'text', 'inc'},
    'ini': {'text', 'ini'},
    'j2': {'text', 'jinja'},
    'jade': {'text', 'jade'},
    'jar': {'binary', 'zip', 'jar'},
    'java': {'text', 'java'},
    'jenkinsfile': {'text', 'groovy'},
    'jinja': {'text', 'jinja'},
    'jinja2': {'text', 'jinja'},
    'jpeg': {'binary', 'image', 'jpeg'},
    'jpg': {'binary', 'image', 'jpeg'},
    'js': {'text', 'javascript'},
    'json': {'text', 'json'},
    'jsx': {'text', 'jsx'},
    'key': {'text', 'pem'},
    'kml': {'text', 'kml', 'xml'},
    'kt': {'text', 'kotlin'},
    'less': {'text', 'less'},
    'lua': {'text', 'lua'},
    'm': {'text', 'c', 'objective-c'},
    'manifest': {'text', 'manifest'},
    'map': {'text', 'map'},
    'markdown': {'text', 'markdown'},
    'md': {'text', 'markdown'},
    'mib': {'text', 'mib'},
    'mk': {'text', 'makefile'},
    'mm': {'text', 'c++', 'objective-c++'},
    'modulemap': {'text', 'modulemap'},
    'ngdoc': {'text', 'ngdoc'},
    'nim': {'text', 'nim'},
    'nims': {'text', 'nim'},
    'nimble': {'text', 'nimble'},
    'otf': {'binary', 'otf'},
    'p12': {'binary', 'p12'},
    'patch': {'text', 'diff'},
    'pdf': {'binary', 'pdf'},
    'pem': {'text', 'pem'},
    'php': {'text', 'php'},
    'php4': {'text', 'php'},
    'php5': {'text', 'php'},
    'phtml': {'text', 'php'},
    'pl': {'text', 'perl'},
    'plantuml': {'text', 'plantuml'},
    'png': {'binary', 'image', 'png'},
    'po': {'text', 'pofile'},
    'pp': {'text', 'puppet'},
    'proto': {'text', 'proto'},
    'purs': {'text', 'purescript'},
    'py': {'text', 'python'},
    'pyi': {'text', 'pyi'},
    'pyx': {'text', 'cython'},
    'pxd': {'text', 'cython'},
    'pxi': {'text', 'cython'},
    'r': {'text', 'r'},
    'rb': {'text', 'ruby'},
    'rs': {'text', 'rust'},
    'rst': {'text', 'rst'},
    's': {'text', 'asm'},
    'scss': {'text', 'scss'},
    'sh': {'text', 'shell'},
    'sls': {'text', 'salt'},
    'so': {'binary'},
    'sol': {'text', 'solidity'},
    'spec': {'text', 'spec'},
    'styl': {'text', 'stylus'},
    'sql': {'text', 'sql'},
    'svg': {'text', 'svg'},
    'swf': {'binary', 'swf'},
    'swift': {'text', 'swift'},
    'swiftdeps': {'text', 'swiftdeps'},
    'tac': {'text', 'twisted', 'python'},
    'tar': {'binary', 'tar'},
    'tgz': {'binary', 'gzip'},
    'thrift': {'text', 'thrift'},
    'tiff': {'binary', 'image', 'tiff'},
    'toml': {'text', 'toml'},
    'tf': {'text', 'terraform'},
    'ts': {'text', 'ts'},
    'tsx': {'text', 'tsx'},
    'ttf': {'binary', 'ttf'},
    'txt': {'text', 'plain-text'},
    'vdx': {'text', 'vdx'},
    'vim': {'text', 'vim'},
    'vue': {'text', 'vue'},
    'war': {'binary', 'zip', 'jar'},
    'wav': {'binary', 'audio', 'wav'},
    'wkt': {'text', 'wkt'},
    'whl': {'binary', 'wheel', 'zip'},
    'woff': {'binary', 'woff'},
    'woff2': {'binary', 'woff2'},
    'wsgi': {'text', 'wsgi', 'python'},
    'xml': {'text', 'xml'},
    'xq': {'text', 'xquery'},
    'xql': {'text', 'xquery'},
    'xqm': {'text', 'xquery'},
    'xqu': {'text', 'xquery'},
    'xquery': {'text', 'xquery'},
    'xqy': {'text', 'xquery'},
    'xsd': {'text', 'xml', 'xsd'},
    'xsl': {'text', 'xml', 'xsl'},
    'yaml': {'text', 'yaml'},
    'yang': {'text', 'yang'},
    'yin': {'text', 'xml', 'yin'},
    'yml': {'text', 'yaml'},
    'zip': {'binary', 'zip'},
    'zsh': {'text', 'shell', 'zsh'},
}
EXTENSIONS_NEED_BINARY_CHECK = {
    'plist': {'plist'},
}

NAMES = {
    '.babelrc': {'text', 'json', 'babelrc'},
    '.bowerrc': {'text', 'json', 'bowerrc'},
    '.coveragerc': {'text', 'ini', 'coveragerc'},
    '.dockerignore': {'text', 'dockerignore'},
    '.editorconfig': {'text', 'editorconfig'},
    '.gitattributes': {'text', 'gitattributes'},
    '.gitignore': {'text', 'gitignore'},
    '.gitmodules': {'text', 'gitmodules'},
    '.jshintrc': {'text', 'json', 'jshintrc'},
    '.mailmap': {'text', 'mailmap'},
    '.mention-bot': {'text', 'json', 'mention-bot'},
    '.npmignore': {'text', 'npmignore'},
    '.yamllint': {'text', 'yaml', 'yamllint'},
    'AUTHORS': EXTENSIONS['txt'],
    'COPYING': EXTENSIONS['txt'],
    'Dockerfile': {'text', 'dockerfile'},
    'Gemfile': EXTENSIONS['rb'],
    'Jenkinsfile': {'text', 'groovy'},
    'LICENSE': EXTENSIONS['txt'],
    'MAINTAINERS': EXTENSIONS['txt'],
    'Makefile': EXTENSIONS['mk'],
    'NOTICE': EXTENSIONS['txt'],
    'PATENTS': EXTENSIONS['txt'],
    'README': EXTENSIONS['txt'],
    'Rakefile': EXTENSIONS['rb'],
    'setup.cfg': EXTENSIONS['ini'],
}
<commit_msg>Know that Pipfile is a TOML file<commit_after># -*- coding: utf-8 -*-
from __future__ import absolute_import
from __future__ import unicode_literals

EXTENSIONS = {
    'apinotes': {'text', 'apinotes'},
    'asar': {'binary', 'asar'},
    'bash': {'text', 'shell', 'bash'},
    'bat': {'text', 'batch'},
    'bmp': {'binary', 'image', 'bitmap'},
    'bz2': {'binary', 'bzip2'},
    'c': {'text', 'c'},
    'cc': {'text', 'c++'},
    'cu': {'text', 'cuda'},
    'cfg': {'text'},
    'cmake': {'text', 'cmake'},
    'cnf': {'text'},
    'coffee': {'text', 'coffee'},
    'conf': {'text'},
    'cpp': {'text', 'c++'},
    'crt': {'text', 'pem'},
    'cs': {'text', 'c#'},
    'cson': {'text', 'cson'},
    'css': {'text', 'css'},
    'csv': {'text', 'csv'},
    'cxx': {'text', 'c++'},
    'def': {'text', 'def'},
    'dtd': {'text', 'dtd'},
    'ear': {'binary', 'zip', 'jar'},
    'ejs': {'text', 'ejs'},
    'eot': {'binary', 'eot'},
    'eps': {'binary', 'eps'},
    'erb': {'text', 'erb'},
    'exe': {'binary'},
    'eyaml': {'text', 'yaml'},
    'feature': {'text', 'gherkin'},
    'fish': {'text', 'fish'},
    'gemspec': {'text', 'ruby'},
    'gif': {'binary', 'image', 'gif'},
    'go': {'text', 'go'},
    'gotmpl': {'text', 'gotmpl'},
    'gpx': {'text', 'gpx', 'xml'},
    'gradle': {'text', 'groovy'},
    'groovy': {'text', 'groovy'},
    'gyb': {'text', 'gyb'},
    'gyp': {'text', 'gyp', 'python'},
    'gypi': {'text', 'gyp', 'python'},
    'gz': {'binary', 'gzip'},
    'h': {'text', 'header', 'c', 'c++'},
    'hpp': {'text', 'header', 'c++'},
    'htm': {'text', 'html'},
    'html': {'text', 'html'},
    'hxx': {'text', 'header', 'c++'},
    'icns': {'binary', 'icns'},
    'ico': {'binary', 'icon'},
    'idl': {'text', 'idl'},
    'inc': {'text', 'inc'},
    'ini': {'text', 'ini'},
    'j2': {'text', 'jinja'},
    'jade': {'text', 'jade'},
    'jar': {'binary', 'zip', 'jar'},
    'java': {'text', 'java'},
    'jenkinsfile': {'text', 'groovy'},
    'jinja': {'text', 'jinja'},
    'jinja2': {'text', 'jinja'},
    'jpeg': {'binary', 'image', 'jpeg'},
    'jpg': {'binary', 'image', 'jpeg'},
    'js': {'text', 'javascript'},
    'json': {'text', 'json'},
    'jsx': {'text', 'jsx'},
    'key': {'text', 'pem'},
    'kml': {'text', 'kml', 'xml'},
    'kt': {'text', 'kotlin'},
    'less': {'text', 'less'},
    'lua': {'text', 'lua'},
    'm': {'text', 'c', 'objective-c'},
    'manifest': {'text', 'manifest'},
    'map': {'text', 'map'},
    'markdown': {'text', 'markdown'},
    'md': {'text', 'markdown'},
    'mib': {'text', 'mib'},
    'mk': {'text', 'makefile'},
    'mm': {'text', 'c++', 'objective-c++'},
    'modulemap': {'text', 'modulemap'},
    'ngdoc': {'text', 'ngdoc'},
    'nim': {'text', 'nim'},
    'nims': {'text', 'nim'},
    'nimble': {'text', 'nimble'},
    'otf': {'binary', 'otf'},
    'p12': {'binary', 'p12'},
    'patch': {'text', 'diff'},
    'pdf': {'binary', 'pdf'},
    'pem': {'text', 'pem'},
    'php': {'text', 'php'},
    'php4': {'text', 'php'},
    'php5': {'text', 'php'},
    'phtml': {'text', 'php'},
    'pl': {'text', 'perl'},
    'plantuml': {'text', 'plantuml'},
    'png': {'binary', 'image', 'png'},
    'po': {'text', 'pofile'},
    'pp': {'text', 'puppet'},
    'proto': {'text', 'proto'},
    'purs': {'text', 'purescript'},
    'py': {'text', 'python'},
    'pyi': {'text', 'pyi'},
    'pyx': {'text', 'cython'},
    'pxd': {'text', 'cython'},
    'pxi': {'text', 'cython'},
    'r': {'text', 'r'},
    'rb': {'text', 'ruby'},
    'rs': {'text', 'rust'},
    'rst': {'text', 'rst'},
    's': {'text', 'asm'},
    'scss': {'text', 'scss'},
    'sh': {'text', 'shell'},
    'sls': {'text', 'salt'},
    'so': {'binary'},
    'sol': {'text', 'solidity'},
    'spec': {'text', 'spec'},
    'styl': {'text', 'stylus'},
    'sql': {'text', 'sql'},
    'svg': {'text', 'svg'},
    'swf': {'binary', 'swf'},
    'swift': {'text', 'swift'},
    'swiftdeps': {'text', 'swiftdeps'},
    'tac': {'text', 'twisted', 'python'},
    'tar': {'binary', 'tar'},
    'tgz': {'binary', 'gzip'},
    'thrift': {'text', 'thrift'},
    'tiff': {'binary', 'image', 'tiff'},
    'toml': {'text', 'toml'},
    'tf': {'text', 'terraform'},
    'ts': {'text', 'ts'},
    'tsx': {'text', 'tsx'},
    'ttf': {'binary', 'ttf'},
    'txt': {'text', 'plain-text'},
    'vdx': {'text', 'vdx'},
    'vim': {'text', 'vim'},
    'vue': {'text', 'vue'},
    'war': {'binary', 'zip', 'jar'},
    'wav': {'binary', 'audio', 'wav'},
    'wkt': {'text', 'wkt'},
    'whl': {'binary', 'wheel', 'zip'},
    'woff': {'binary', 'woff'},
    'woff2': {'binary', 'woff2'},
    'wsgi': {'text', 'wsgi', 'python'},
    'xml': {'text', 'xml'},
    'xq': {'text', 'xquery'},
    'xql': {'text', 'xquery'},
    'xqm': {'text', 'xquery'},
    'xqu': {'text', 'xquery'},
    'xquery': {'text', 'xquery'},
    'xqy': {'text', 'xquery'},
    'xsd': {'text', 'xml', 'xsd'},
    'xsl': {'text', 'xml', 'xsl'},
    'yaml': {'text', 'yaml'},
    'yang': {'text', 'yang'},
    'yin': {'text', 'xml', 'yin'},
    'yml': {'text', 'yaml'},
    'zip': {'binary', 'zip'},
    'zsh': {'text', 'shell', 'zsh'},
}
EXTENSIONS_NEED_BINARY_CHECK = {
    'plist': {'plist'},
}

NAMES = {
    '.babelrc': {'text', 'json', 'babelrc'},
    '.bowerrc': {'text', 'json', 'bowerrc'},
    '.coveragerc': {'text', 'ini', 'coveragerc'},
    '.dockerignore': {'text', 'dockerignore'},
    '.editorconfig': {'text', 'editorconfig'},
    '.gitattributes': {'text', 'gitattributes'},
    '.gitignore': {'text', 'gitignore'},
    '.gitmodules': {'text', 'gitmodules'},
    '.jshintrc': {'text', 'json', 'jshintrc'},
    '.mailmap': {'text', 'mailmap'},
    '.mention-bot': {'text', 'json', 'mention-bot'},
    '.npmignore': {'text', 'npmignore'},
    '.yamllint': {'text', 'yaml', 'yamllint'},
    'AUTHORS': EXTENSIONS['txt'],
    'COPYING': EXTENSIONS['txt'],
    'Dockerfile': {'text', 'dockerfile'},
    'Gemfile': EXTENSIONS['rb'],
    'Jenkinsfile': {'text', 'groovy'},
    'LICENSE': EXTENSIONS['txt'],
    'MAINTAINERS': EXTENSIONS['txt'],
    'Makefile': EXTENSIONS['mk'],
    'NOTICE': EXTENSIONS['txt'],
    'PATENTS': EXTENSIONS['txt'],
    'Pipfile': EXTENSIONS['toml'],
    'README': EXTENSIONS['txt'],
    'Rakefile': EXTENSIONS['rb'],
    'setup.cfg': EXTENSIONS['ini'],
}
<|endoftext|>


<commit_before>import ast
from myhdl._convutils import _makeAST
from myhdl._util import _flatten

class Data():
    pass

def _resolveRefs(symdict, arg):
    #temporary hack to prevent circular dependencies
    from myhdl._always_comb import _AlwaysComb
    from myhdl._always_seq import _AlwaysSeq
    from myhdl._always import _Always
    gens = _flatten(arg)
    data = Data()
    data.symdict = symdict
    v = _AttrRefTransformer(data)
    for gen in gens:
        if isinstance(gen, (_AlwaysComb, _AlwaysSeq, _Always)):
            func = gen.func
        else:
            func = gen.gen.gi_frame
        tree = _makeAST(func)
        v.visit(tree)
    return data.objlist

class _AttrRefTransformer(ast.NodeTransformer):
    def __init__(self, data):
        self.data = data
        self.data.objlist = []

    def visit_Attribute(self, node):
        reserved = ('next', 'posedge', 'negedge', 'max', 'min', 'val', 'signed')
        self.generic_visit(node)
        if node.attr in reserved:
            return node
        else:
            obj = self.data.symdict[node.value.id]
            attrobj = getattr(obj, node.attr)
            new_name = node.value.id+'.'+node.attr
            if new_name not in self.data.symdict:
                self.data.symdict[new_name] = attrobj
                self.data.objlist.append(new_name)
            else:
                pass
                #assert self.data.symdict[new_name] == attrobj
            new_node = ast.Name(id=new_name, ctx=node.value.ctx)
            return ast.copy_location(new_node, node)

    def visit_FunctionDef(self, node):
        nodes = _flatten(node.body, node.args)
        for n in nodes:
            self.visit(n)
        self.visit(node.args)
        return node
<commit_msg>Removed redundant line.<commit_after>import ast
from myhdl._convutils import _makeAST
from myhdl._util import _flatten

class Data():
    pass

def _resolveRefs(symdict, arg):
    #temporary hack to prevent circular dependencies
    from myhdl._always_comb import _AlwaysComb
    from myhdl._always_seq import _AlwaysSeq
    from myhdl._always import _Always
    gens = _flatten(arg)
    data = Data()
    data.symdict = symdict
    v = _AttrRefTransformer(data)
    for gen in gens:
        if isinstance(gen, (_AlwaysComb, _AlwaysSeq, _Always)):
            func = gen.func
        else:
            func = gen.gen.gi_frame
        tree = _makeAST(func)
        v.visit(tree)
    return data.objlist

class _AttrRefTransformer(ast.NodeTransformer):
    def __init__(self, data):
        self.data = data
        self.data.objlist = []

    def visit_Attribute(self, node):
        reserved = ('next', 'posedge', 'negedge', 'max', 'min', 'val', 'signed')
        self.generic_visit(node)
        if node.attr in reserved:
            return node
        else:
            obj = self.data.symdict[node.value.id]
            attrobj = getattr(obj, node.attr)
            new_name = node.value.id+'.'+node.attr
            if new_name not in self.data.symdict:
                self.data.symdict[new_name] = attrobj
                self.data.objlist.append(new_name)
            else:
                pass
                #assert self.data.symdict[new_name] == attrobj
            new_node = ast.Name(id=new_name, ctx=node.value.ctx)
            return ast.copy_location(new_node, node)

    def visit_FunctionDef(self, node):
        nodes = _flatten(node.body, node.args)
        for n in nodes:
            self.visit(n)
        return node
<|endoftext|>


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

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

# Modified based on Kafka's patch review tool

# Required Modules:
# - python-argparse

import argparse
import sys
import os
import time
import datetime
import tempfile
from jira.client import JIRA

def get_jira():
  options = {
    'server': 'https://issues.apache.org/jira'
  }
  # read the config file
  home=jira_home=os.getenv('HOME')
  home=home.rstrip('/')
  jira_config = dict(line.strip().split('=') for line in open(home + '/jira.ini'))
  jira = JIRA(options,basic_auth=(jira_config['user'], jira_config['password']))
  return jira

def main():
  ''' main(), shut up, pylint '''
  popt = argparse.ArgumentParser(description='Apache Lens patch review tool')
  popt.add_argument('-b', '--branch', action='store', dest='branch', required=True, help='Tracking branch to create diff against')
  popt.add_argument('-j', '--jira', action='store', dest='jira', required=True, help='JIRA corresponding to the reviewboard')
  popt.add_argument('-s', '--summary', action='store', dest='summary', required=False, help='Summary for the reviewboard')
  popt.add_argument('-d', '--description', action='store', dest='description', required=False, help='Description for reviewboard')
  popt.add_argument('-r', '--rb', action='store', dest='reviewboard', required=False, help='Review board that needs to be updated')
  popt.add_argument('-t', '--testing-done', action='store', dest='testing', required=False, help='Text for the Testing Done section of the reviewboard')
  popt.add_argument('-db', '--debug', action='store_true', required=False, help='Enable debug mode')
  popt.add_argument('-rbu', '--reviewboard-user', action='store', dest='reviewboard_user', required=True, help='Review board user name')
  popt.add_argument('-rbp', '--reviewboard-password', action='store', dest='reviewboard_password', required=True, help='Review board user password')
  popt.add_argument('-v', '--version', action='store', dest='version', required=False, help='Version for patch')
  opt = popt.parse_args()

  patch_file=tempfile.gettempdir() + "/" + opt.jira + "." + opt.version + ".patch"
  if opt.reviewboard:
    ts = time.time()
    st = datetime.datetime.fromtimestamp(ts).strftime('%Y-%m-%d_%H:%M:%S')
    patch_file=tempfile.gettempdir() + "/" + opt.jira + '_' + st + '.patch'

  git_remote_update="git remote update"
  print "Updating your remote branches to pull the latest changes"
  p=os.popen(git_remote_update)
  p.close()

  rb_command="post-review --publish --tracking-branch " + opt.branch + " --target-groups=lens --bugs-closed=" + opt.jira
  rb_command=rb_command + " --username " + opt.reviewboard_user + " --password " + opt.reviewboard_password

  if opt.debug:
    rb_command=rb_command + " --debug"
  summary="Patch for " + opt.jira
  if opt.summary:
    summary=opt.summary
  rb_command=rb_command + " --summary \"" + summary + "\""
  if opt.description:
    rb_command=rb_command + " --description \"" + opt.description + "\""
  if opt.reviewboard:
    rb_command=rb_command + " -r " + opt.reviewboard
  if opt.testing:
    rb_command=rb_command + " --testing-done=\"" + opt.testing + "\""
  if opt.debug:
    print rb_command
  p=os.popen(rb_command)
  rb_url=""
  for line in p:
    print line
    if line.startswith('http'):
      rb_url = line
    elif line.startswith("There don't seem to be any diffs"):
      print 'ERROR: Your reviewboard was not created/updated since there was no diff to upload. The reasons that can cause this issue are 1) Your diff is not checked into your local branch. Please check in the diff to the local branch and retry 2) You are not specifying the local branch name as part of the --branch option. Please specify the remote branch name obtained from git branch -r'
      p.close()
      sys.exit(1)
    elif line.startswith("Your review request still exists, but the diff is not attached") and not opt.debug:
      print 'ERROR: Your reviewboard was not created/updated. Please run the script with the --debug option to troubleshoot the problem'
      p.close()
      sys.exit(1)
  p.close()
  if opt.debug:
    print 'rb url=',rb_url

  git_command="git diff " + opt.branch + " > " + patch_file
  if opt.debug:
    print git_command
  p=os.popen(git_command)
  p.close()

  print 'Creating diff against', opt.branch, 'and uploading patch to JIRA',opt.jira
  jira=get_jira()
  issue = jira.issue(opt.jira)
  attachment=open(patch_file)
  jira.add_attachment(issue,attachment)
  attachment.close()

  comment="Created reviewboard "
  if not opt.reviewboard:
    print 'Created a new reviewboard ',rb_url
  else:
    print 'Updated reviewboard',opt.reviewboard
    comment="Updated reviewboard "

  comment = comment + rb_url
  jira.add_comment(opt.jira, comment)

if __name__ == '__main__':
  sys.exit(main())

<commit_msg>Revert "LENS-541 : Lens patch-review tool (Yash Sharma via Rajat Khandelwal)"<commit_after><|endoftext|>


<commit_before>#!/usr/bin/env python3
#
# Change json output to be human-readable
import argparse
import csv
import json
import re
import os
import sys

csv.field_size_limit(sys.maxsize)

def int_of_label(s):
    total = 0
    for c in s:
        total += ord(c) - ord('a') + 1
    return total

def label_of_int(i):
    s = ""
    maxint = ord('z') - ord('a') + 1
    while i > maxint:
        s += 'z'
        i -= maxint
    s += chr(i + ord('a') - 1)
    return s

def convert_dict(d):
    # for addresses, options, records, etc, just dereference
    if "type" in d and d["type"] in ["address", "option_or_ind", "record", "tuple", "Collection", "Map", "Seq", "Set", "VMap", "MapE", "SortedMapE", "SortedMap", "SortedSet", "Vector", "MultiIndexVMap"]:
        return convert_any(d["value"])
    # change record mapping back to tuple
    elif "ra" in d:
        res = []
        max = 0
        for key in d:
            k = int_of_label(key[1:])
            if k > max:
                max = k
        for i in range(max): # +1 for record fields
            res.append(convert_any(d["r" + label_of_int(i+1)]))
        return res
    elif "key" in d and "value" in d:
        res = []
        if d["key"] != "()":
            res.append(convert_any(d["key"]))
        res.append(convert_any(d["value"]))
        return res
    elif "elem" in d:
        res = []
        res.append(convert_any(d["elem"]))
        return res
    elif "addr" in d:
        res = []
        res.append(convert_any(d["addr"]))
        return res
    # standard record
    else:
        res = {}
        for key in d:
            res[key] = convert_any(d[key])
        return res

def convert_list(xs):
    res = []
    for x in xs:
        res.append(convert_any(x))
    return res

def convert_any(x):
    if type(x) is dict:
        return convert_dict(x)
    elif type(x) is list:
        return convert_list(x)
    else:
        return x

# Special conversion for first level of object
# So we have all array values
def convert_fst_level(x):
    if type(x) is dict:
        return convert_dict(x)
    elif type(x) is list:
        return convert_list(x)
    elif x == "()":
        return []
    else:
        return [x]

def convert_file(file_nm, writer, treat_as_ints):
    with open(file_nm, 'r', newline='') as csvfile:
        reader = csv.reader(csvfile, delimiter='|', quotechar="'")
        i = 0
        for row in reader:
            res = []
            j = 0
            for x in row:
                #print(x)
                try:
                    obj = json.loads(x)
                    # print("obj: " + str(obj))
                    # keep certain fields as raw integers rather than wrapping in array
                    if j in treat_as_ints:
                        new_obj = convert_any(obj)
                    else:
                        new_obj = convert_fst_level(obj)
                    #print(new_obj)
                    res += [json.dumps(new_obj)]
                except:
                    # if we can't convert, just add it as a string
                    res += [x]
                j += 1
            writer.writerow(res)
            i += 1

def process_files(files, prefix_path):
    with open(os.path.join(prefix_path, "messages.dsv"), "w", newline='') as msg_file:
        with open(os.path.join(prefix_path, "globals.dsv"), "w", newline='') as glb_file:
            msg_writer = csv.writer(msg_file, delimiter='|', quotechar="'")
            glb_writer = csv.writer(glb_file, delimiter='|', quotechar="'")
            for f in files:
                # is it a messages or a globals?
                if re.search("Globals", f):
                    convert_file(f, glb_writer, [0])
                else:
                    convert_file(f, msg_writer, [0, 5])

# Translate all files in their respective directories
def inplace_process(prefix):
    for root, dirs, files in os.walk(prefix):
        for f in files:
            name, ext = os.path.splitext(f)
            new_name = name + "_clean" + ext
            tests = [("Globals", [0]), ("Messages", [0, 5])]
            for regex, indices in tests:
                if re.search(regex, name):
                    with open(os.path.join(root, new_name), "w", newline='') as out_file:
                        writer = csv.writer(out_file, delimiter='|', quotechar="'")
                        convert_file(os.path.join(root, f), writer, indices)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("json_files", type=str, nargs='*', help="Specify path")
    parser.add_argument("--prefix_path", type=str, help="Specify prefix to output", default="")
    parser.add_argument("--inplace", type=str, help="Modify files in-place", default="")
    args = parser.parse_args()
    if args.inplace != "":
        inplace_process(args.inplace)
    else:
        process_files(args.json_files, args.prefix_path)

if __name__=='__main__':
    main ()
<commit_msg>mosaic/clean_json: don't create too many clean files<commit_after>#!/usr/bin/env python3
#
# Change json output to be human-readable
import argparse
import csv
import json
import re
import os
import sys

csv.field_size_limit(sys.maxsize)

def int_of_label(s):
    total = 0
    for c in s:
        total += ord(c) - ord('a') + 1
    return total

def label_of_int(i):
    s = ""
    maxint = ord('z') - ord('a') + 1
    while i > maxint:
        s += 'z'
        i -= maxint
    s += chr(i + ord('a') - 1)
    return s

def convert_dict(d):
    # for addresses, options, records, etc, just dereference
    if "type" in d and d["type"] in ["address", "option_or_ind", "record", "tuple", "Collection", "Map", "Seq", "Set", "VMap", "MapE", "SortedMapE", "SortedMap", "SortedSet", "Vector", "MultiIndexVMap"]:
        return convert_any(d["value"])
    # change record mapping back to tuple
    elif "ra" in d:
        res = []
        max = 0
        for key in d:
            k = int_of_label(key[1:])
            if k > max:
                max = k
        for i in range(max): # +1 for record fields
            res.append(convert_any(d["r" + label_of_int(i+1)]))
        return res
    elif "key" in d and "value" in d:
        res = []
        if d["key"] != "()":
            res.append(convert_any(d["key"]))
        res.append(convert_any(d["value"]))
        return res
    elif "elem" in d:
        res = []
        res.append(convert_any(d["elem"]))
        return res
    elif "addr" in d:
        res = []
        res.append(convert_any(d["addr"]))
        return res
    # standard record
    else:
        res = {}
        for key in d:
            res[key] = convert_any(d[key])
        return res

def convert_list(xs):
    res = []
    for x in xs:
        res.append(convert_any(x))
    return res

def convert_any(x):
    if type(x) is dict:
        return convert_dict(x)
    elif type(x) is list:
        return convert_list(x)
    else:
        return x

# Special conversion for first level of object
# So we have all array values
def convert_fst_level(x):
    if type(x) is dict:
        return convert_dict(x)
    elif type(x) is list:
        return convert_list(x)
    elif x == "()":
        return []
    else:
        return [x]

def convert_file(file_nm, writer, treat_as_ints):
    with open(file_nm, 'r', newline='') as csvfile:
        reader = csv.reader(csvfile, delimiter='|', quotechar="'")
        i = 0
        for row in reader:
            res = []
            j = 0
            for x in row:
                #print(x)
                try:
                    obj = json.loads(x)
                    # print("obj: " + str(obj))
                    # keep certain fields as raw integers rather than wrapping in array
                    if j in treat_as_ints:
                        new_obj = convert_any(obj)
                    else:
                        new_obj = convert_fst_level(obj)
                    #print(new_obj)
                    res += [json.dumps(new_obj)]
                except:
                    # if we can't convert, just add it as a string
                    res += [x]
                j += 1
            writer.writerow(res)
            i += 1

def process_files(files, prefix_path):
    with open(os.path.join(prefix_path, "messages.dsv"), "w", newline='') as msg_file:
        with open(os.path.join(prefix_path, "globals.dsv"), "w", newline='') as glb_file:
            msg_writer = csv.writer(msg_file, delimiter='|', quotechar="'")
            glb_writer = csv.writer(glb_file, delimiter='|', quotechar="'")
            for f in files:
                # is it a messages or a globals?
                if re.search("Globals", f):
                    convert_file(f, glb_writer, [0])
                else:
                    convert_file(f, msg_writer, [0, 5])

# Translate all files in their respective directories
def inplace_process(prefix):
    for root, dirs, files in os.walk(prefix):
        for f in files:
            name, ext = os.path.splitext(f)
            new_name = name + "_clean" + ext
            tests = [("Globals$", [0]), ("Messages$", [0, 5])]
            for regex, indices in tests:
                if re.search(regex, name):
                    with open(os.path.join(root, new_name), "w", newline='') as out_file:
                        writer = csv.writer(out_file, delimiter='|', quotechar="'")
                        convert_file(os.path.join(root, f), writer, indices)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("json_files", type=str, nargs='*', help="Specify path")
    parser.add_argument("--prefix_path", type=str, help="Specify prefix to output", default="")
    parser.add_argument("--inplace", type=str, help="Modify files in-place", default="")
    args = parser.parse_args()
    if args.inplace != "":
        inplace_process(args.inplace)
    else:
        process_files(args.json_files, args.prefix_path)

if __name__=='__main__':
    main ()
<|endoftext|>