<commit_before>import bpy
from bpy.app.handlers import persistent
from bpy.props import StringProperty
from bpy.utils import register_class
from bpy.utils import unregister_class

import json
import os
import sys
import time
import threading
from urllib import request
from zipfile import ZipFile
from configparser import ConfigParser
from subprocess import Popen, STDOUT, PIPE
from queue import Queue, Empty

__version__ = '1.0.1'

bl_info = \
    {
        "name":        "Wakatime plugin for Blender",
        "category":    "Development",
        "author":      "Allis Tauri <allista@gmail.com>",
        "version":     (1, 0, 1),
        "blender":     (2, 77, 0),
        "description": "Submits your working stats to the Wakatime time tracking service.",
        "warning":     "Beta",
        "tracker_url": "https://github.com/allista/WakatimeBlender/issues",
    }

# globals
HEARTBEAT_FREQUENCY = 2
_heartbeats = Queue()
_hb_processor = None
_last_hb = None
_filename = ''

REGISTERED = False
SHOW_KEY_DIALOG = False
PLUGIN_DIR = os.path.dirname(os.path.realpath(__file__))
RESOURCES_DIR = os.path.join(os.path.expanduser('~'), '.wakatime')
API_CLIENT_URL = 'https://github.com/wakatime/wakatime/archive/master.zip'
API_CLIENT = os.path.join(os.path.expanduser('~'), '.wakatime', 'wakatime-master', 'wakatime', 'cli.py')
SETTINGS_FILE = os.path.join(os.path.expanduser('~'), '.wakatime.cfg')
SETTINGS = None
settings = 'settings'

# Log Levels
DEBUG   = 'DEBUG'
INFO    = 'INFO'
WARNING = 'WARNING'
ERROR   = 'ERROR'

def u(text):
    if text is None: return None
    if isinstance(text, bytes):
        try: return text.decode('utf-8')
        except:
            try: return text.decode(sys.getdefaultencoding())
            except: pass
    try: return str(text)
    except: return text

def log(lvl, message, *args, **kwargs):
    if lvl == DEBUG and not SETTINGS.getboolean(settings, 'debug'): return
    msg = message
    if len(args) > 0: msg = message.format(*args)
    elif len(kwargs) > 0: msg = message.format(**kwargs)
    print('[WakaTime] [{lvl}] {msg}'.format(lvl=lvl, msg=msg))

class API_Key_Dialog(bpy.types.Operator):
    bl_idname = "ui.wakatime_api_key_dialog"
    bl_label  = "Enter WakaTime API Key"
    api_key   = StringProperty(name="API Key")
    is_shown  = False

    @classmethod
    def show(cls):
        global SHOW_KEY_DIALOG
        if not cls.is_shown and REGISTERED:
            cls.is_shown = True
            SHOW_KEY_DIALOG = False
            bpy.ops.ui.wakatime_api_key_dialog('INVOKE_DEFAULT')

    def execute(self, context):
        if self.api_key:
            SETTINGS.set(settings, 'api_key', u(self.api_key))
            save_settings()
        API_Key_Dialog.is_shown = False
        return {'FINISHED'}

    def invoke(self, context, event):
        self.api_key = SETTINGS.get(settings, 'api_key', fallback='')
        return context.window_manager.invoke_props_dialog(self)

class HeartbeatQueueProcessor(threading.Thread):
    def __init__(self, q):
        threading.Thread.__init__(self)
        self.daemon = True
        self._queue = q

    @classmethod
    def send(cls, heartbeat, extra_heartbeats=[]):
        global SHOW_KEY_DIALOG
        ua = 'blender/%s blender-wakatime/%s' % (bpy.app.version_string.split()[0], __version__)
        cmd = [
            bpy.app.binary_path_python,
            API_CLIENT,
            '--entity', heartbeat['entity'],
            '--time', str('%f' % heartbeat['timestamp']),
            '--plugin', ua,
        ]
        if heartbeat['is_write']:
            cmd.append('--write')
        for pattern in SETTINGS.get(settings, 'ignore', fallback=[]):  # or should it be blender-specific?
            cmd.extend(['--ignore', pattern])
        if SETTINGS.getboolean(settings, 'debug'):
            cmd.append('--verbose')
        if extra_heartbeats:
            cmd.append('--extra-heartbeats')
            extra_heartbeats = json.dumps(extra_heartbeats)
            stdin = PIPE
        else:
            extra_heartbeats = None
            stdin = None
        log(DEBUG, ' '.join(cmd))
        try:
            process = Popen(cmd, stdin=stdin, stdout=PIPE, stderr=STDOUT)
            inp = None
            if extra_heartbeats:
                inp = "{0}\n".format(extra_heartbeats)
                inp = inp.encode('utf-8')
            output, err = process.communicate(input=inp)
            output = u(output)
            retcode = process.poll()
            if (not retcode or retcode == 102) and not output:
                log(DEBUG, 'OK')
            elif retcode == 104:  # wrong API key
                log(ERROR, 'Wrong API key. Asking for a new one...')
                SETTINGS.set(settings, 'api_key', '')
                SHOW_KEY_DIALOG = True
            else:
                log(ERROR, 'Error')
            if retcode:
                log(DEBUG if retcode == 102 else ERROR, 'wakatime-core exited with status: {0}'.format(retcode))
            if output:
                log(ERROR, u('wakatime-core output: {0}').format(output))
        except:
            log(ERROR, u(sys.exc_info()[1]))

    def run(self):
        while True:
            time.sleep(1)
            if not SETTINGS.get(settings, 'api_key', fallback=''): continue
            try: heartbeat = self._queue.get_nowait()
            except Empty: continue
            if heartbeat is None: return
            extra_heartbeats = []
            try:
                while True:
                    extra_heartbeats.append(self._queue.get_nowait())
            except Empty: pass
            self.send(heartbeat, extra_heartbeats)

class DownloadWakatime(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)
        self.daemon = True

    def run(self):
        log(INFO, 'WakatimeBlender is registered')
        if not os.path.isdir(RESOURCES_DIR):
            try: os.mkdir(RESOURCES_DIR)
            except:
                log(ERROR, 'Unable to create directory:\n%s' % RESOURCES_DIR)
                return
        if not os.path.isfile(API_CLIENT):
            log(INFO, 'Downloading Wakatime...')
            zip_file = os.path.join(RESOURCES_DIR, 'wakatime-master.zip')
            request.urlretrieve(API_CLIENT_URL, zip_file)
            log(INFO, 'Extracting Wakatime...')
            with ZipFile(zip_file) as zf: zf.extractall(RESOURCES_DIR)
            try: os.remove(zip_file)
            except: pass
            log(INFO, 'Finished extracting Wakatime.')
        else: log(INFO, 'Found Wakatime client')

def save_settings():
    with open(SETTINGS_FILE, 'w') as out:
        SETTINGS.write(out)

def setup():
    global SETTINGS, _hb_processor
    download = DownloadWakatime()
    download.start()
    SETTINGS = ConfigParser()
    SETTINGS.read(SETTINGS_FILE)
    #common wakatime settings
    if not SETTINGS.has_section(settings):
        SETTINGS.add_section(settings)
    if not SETTINGS.has_option(settings, 'debug'):
        SETTINGS.set(settings, 'debug', str(False))
    _hb_processor = HeartbeatQueueProcessor(_heartbeats)
    _hb_processor.start()

@persistent
def load_handler(dummy):
    global _filename
    _filename = bpy.data.filepath
    handle_activity()

@persistent
def save_handler(dummy):
    global _filename
    _filename = bpy.data.filepath
    handle_activity(is_write=True)

@persistent
def activity_handler(dummy):
    handle_activity()

def enough_time_passed(now, is_write):
    return _last_hb is None or (now - _last_hb['timestamp'] > (2 if is_write else HEARTBEAT_FREQUENCY * 60))

def handle_activity(is_write=False):
    global _last_hb
    if SHOW_KEY_DIALOG or not SETTINGS.get(settings, 'api_key', fallback=''):
        API_Key_Dialog.show()
    timestamp = time.time()
    last_file = _last_hb['entity'] if _last_hb is not None else ''
    if _filename and (_filename != last_file or enough_time_passed(timestamp, is_write)):
        _last_hb = {'entity':_filename, 'timestamp':timestamp, 'is_write':is_write}
        _heartbeats.put_nowait(_last_hb)

classes = (
    API_Key_Dialog,
    HeartbeatQueueProcessor,
    DownloadWakatime,
)
register, unregister = bpy.utils.register_classes_factory(classes)

def register():
    global  REGISTERED
    log(INFO, 'Initializing WakaTime plugin v%s' % __version__)
    setup()

    for cls in classes:
        register_class(cls)

    bpy.app.handlers.load_post.append(load_handler)
    bpy.app.handlers.save_post.append(save_handler)
    bpy.app.handlers.scene_update_post.append(activity_handler)
    REGISTERED = True

def unregister():
    global REGISTERED
    log(INFO, 'Unregistering WakaTime plugin v%s' % __version__)
    save_settings()
    REGISTERED = False
    bpy.app.handlers.load_post.remove(load_handler)
    bpy.app.handlers.save_post.remove(save_handler)
    bpy.app.handlers.scene_update_post.remove(activity_handler)

    for cls in classes:
        unregister_class(cls)

    _heartbeats.put_nowait(None)
    _heartbeats.task_done()
    _hb_processor.join()

if __name__ == '__main__':
    register()
<commit_msg>Update the blender2.80 API, delete deprecated code<commit_after>import bpy
from bpy.app.handlers import persistent
from bpy.props import StringProperty

import json
import os
import sys
import time
import threading
from urllib import request
from zipfile import ZipFile
from configparser import ConfigParser
from subprocess import Popen, STDOUT, PIPE
from queue import Queue, Empty

__version__ = '1.0.1'

bl_info = \
    {
        "name":        "Wakatime plugin for Blender",
        "category":    "Development",
        "author":      "Allis Tauri <allista@gmail.com>",
        "version":     (1, 0, 1),
        "blender":     (2, 80, 0),
        "description": "Submits your working stats to the Wakatime time tracking service.",
        "warning":     "Beta",
        "tracker_url": "https://github.com/allista/WakatimeBlender/issues",
    }

# globals
HEARTBEAT_FREQUENCY = 2
_heartbeats = Queue()
_hb_processor = None
_last_hb = None
_filename = ''

REGISTERED = False
SHOW_KEY_DIALOG = False
PLUGIN_DIR = os.path.dirname(os.path.realpath(__file__))
RESOURCES_DIR = os.path.join(os.path.expanduser('~'), '.wakatime')
API_CLIENT_URL = 'https://github.com/wakatime/wakatime/archive/master.zip'
API_CLIENT = os.path.join(os.path.expanduser('~'), '.wakatime', 'wakatime-master', 'wakatime', 'cli.py')
SETTINGS_FILE = os.path.join(os.path.expanduser('~'), '.wakatime.cfg')
SETTINGS = None
settings = 'settings'

# Log Levels
DEBUG   = 'DEBUG'
INFO    = 'INFO'
WARNING = 'WARNING'
ERROR   = 'ERROR'

def u(text):
    if text is None: return None
    if isinstance(text, bytes):
        try: return text.decode('utf-8')
        except:
            try: return text.decode(sys.getdefaultencoding())
            except: pass
    try: return str(text)
    except: return text

def log(lvl, message, *args, **kwargs):
    if lvl == DEBUG and not SETTINGS.getboolean(settings, 'debug'): return
    msg = message
    if len(args) > 0: msg = message.format(*args)
    elif len(kwargs) > 0: msg = message.format(**kwargs)
    print('[WakaTime] [{lvl}] {msg}'.format(lvl=lvl, msg=msg))

class API_Key_Dialog(bpy.types.Operator):
    bl_idname = "ui.wakatime_api_key_dialog"
    bl_label  = "Enter WakaTime API Key"
    api_key   : StringProperty(name="API Key")
    is_shown  = False

    @classmethod
    def show(cls):
        global SHOW_KEY_DIALOG
        if not cls.is_shown and REGISTERED:
            cls.is_shown = True
            SHOW_KEY_DIALOG = False
            bpy.ops.ui.wakatime_api_key_dialog('INVOKE_DEFAULT')

    def execute(self, context):
        if self.api_key:
            SETTINGS.set(settings, 'api_key', u(self.api_key))
            save_settings()
        API_Key_Dialog.is_shown = False
        return {'FINISHED'}

    def invoke(self, context, event):
        self.api_key = SETTINGS.get(settings, 'api_key', fallback='')
        return context.window_manager.invoke_props_dialog(self)

class HeartbeatQueueProcessor(threading.Thread):
    def __init__(self, q):
        threading.Thread.__init__(self)
        self.daemon = True
        self._queue = q

    @classmethod
    def send(cls, heartbeat, extra_heartbeats=[]):
        global SHOW_KEY_DIALOG
        ua = 'blender/%s blender-wakatime/%s' % (bpy.app.version_string.split()[0], __version__)
        cmd = [
            bpy.app.binary_path_python,
            API_CLIENT,
            '--entity', heartbeat['entity'],
            '--time', str('%f' % heartbeat['timestamp']),
            '--plugin', ua,
        ]
        if heartbeat['is_write']:
            cmd.append('--write')
        for pattern in SETTINGS.get(settings, 'ignore', fallback=[]):  # or should it be blender-specific?
            cmd.extend(['--ignore', pattern])
        if SETTINGS.getboolean(settings, 'debug'):
            cmd.append('--verbose')
        if extra_heartbeats:
            cmd.append('--extra-heartbeats')
            extra_heartbeats = json.dumps(extra_heartbeats)
            stdin = PIPE
        else:
            extra_heartbeats = None
            stdin = None
        log(DEBUG, ' '.join(cmd))
        try:
            process = Popen(cmd, stdin=stdin, stdout=PIPE, stderr=STDOUT)
            inp = None
            if extra_heartbeats:
                inp = "{0}\n".format(extra_heartbeats)
                inp = inp.encode('utf-8')
            output, err = process.communicate(input=inp)
            output = u(output)
            retcode = process.poll()
            if (not retcode or retcode == 102) and not output:
                log(DEBUG, 'OK')
            elif retcode == 104:  # wrong API key
                log(ERROR, 'Wrong API key. Asking for a new one...')
                SETTINGS.set(settings, 'api_key', '')
                SHOW_KEY_DIALOG = True
            else:
                log(ERROR, 'Error')
            if retcode:
                log(DEBUG if retcode == 102 else ERROR, 'wakatime-core exited with status: {0}'.format(retcode))
            if output:
                log(ERROR, u('wakatime-core output: {0}').format(output))
        except:
            log(ERROR, u(sys.exc_info()[1]))

    def run(self):
        while True:
            time.sleep(1)
            if not SETTINGS.get(settings, 'api_key', fallback=''): continue
            try: heartbeat = self._queue.get_nowait()
            except Empty: continue
            if heartbeat is None: return
            extra_heartbeats = []
            try:
                while True:
                    extra_heartbeats.append(self._queue.get_nowait())
            except Empty: pass
            self.send(heartbeat, extra_heartbeats)

class DownloadWakatime(threading.Thread):
    def __init__(self):
        threading.Thread.__init__(self)
        self.daemon = True

    def run(self):
        log(INFO, 'WakatimeBlender is registered')
        if not os.path.isdir(RESOURCES_DIR):
            try: os.mkdir(RESOURCES_DIR)
            except:
                log(ERROR, 'Unable to create directory:\n%s' % RESOURCES_DIR)
                return
        if not os.path.isfile(API_CLIENT):
            log(INFO, 'Downloading Wakatime...')
            zip_file = os.path.join(RESOURCES_DIR, 'wakatime-master.zip')
            request.urlretrieve(API_CLIENT_URL, zip_file)
            log(INFO, 'Extracting Wakatime...')
            with ZipFile(zip_file) as zf: zf.extractall(RESOURCES_DIR)
            try: os.remove(zip_file)
            except: pass
            log(INFO, 'Finished extracting Wakatime.')
        else: log(INFO, 'Found Wakatime client')

def save_settings():
    with open(SETTINGS_FILE, 'w') as out:
        SETTINGS.write(out)

def setup():
    global SETTINGS, _hb_processor
    download = DownloadWakatime()
    download.start()
    SETTINGS = ConfigParser()
    SETTINGS.read(SETTINGS_FILE)
    #common wakatime settings
    if not SETTINGS.has_section(settings):
        SETTINGS.add_section(settings)
    if not SETTINGS.has_option(settings, 'debug'):
        SETTINGS.set(settings, 'debug', str(False))
    _hb_processor = HeartbeatQueueProcessor(_heartbeats)
    _hb_processor.start()

@persistent
def load_handler(dummy):
    global _filename
    _filename = bpy.data.filepath
    handle_activity()

@persistent
def save_handler(dummy):
    global _filename
    _filename = bpy.data.filepath
    handle_activity(is_write=True)

@persistent
def activity_handler(dummy):
    handle_activity()

def enough_time_passed(now, is_write):
    return _last_hb is None or (now - _last_hb['timestamp'] > (2 if is_write else HEARTBEAT_FREQUENCY * 60))

def handle_activity(is_write=False):
    global _last_hb
    if SHOW_KEY_DIALOG or not SETTINGS.get(settings, 'api_key', fallback=''):
        API_Key_Dialog.show()
    timestamp = time.time()
    last_file = _last_hb['entity'] if _last_hb is not None else ''
    if _filename and (_filename != last_file or enough_time_passed(timestamp, is_write)):
        _last_hb = {'entity':_filename, 'timestamp':timestamp, 'is_write':is_write}
        _heartbeats.put_nowait(_last_hb)

def register():
    global  REGISTERED
    log(INFO, 'Initializing WakaTime plugin v%s' % __version__)
    setup()
    bpy.utils.register_class(API_Key_Dialog)
    bpy.app.handlers.load_post.append(load_handler)
    bpy.app.handlers.save_post.append(save_handler)
    bpy.app.handlers.depsgraph_update_pre.append(activity_handler)
    REGISTERED = True

def unregister():
    global REGISTERED
    log(INFO, 'Unregistering WakaTime plugin v%s' % __version__)
    save_settings()
    REGISTERED = False
    bpy.app.handlers.load_post.remove(load_handler)
    bpy.app.handlers.save_post.remove(save_handler)
    bpy.app.handlers.depsgraph_update_pre.remove(activity_handler)
    bpy.utils.unregister_class(API_Key_Dialog)
    _heartbeats.put_nowait(None)
    _heartbeats.task_done()
    _hb_processor.join()

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


<commit_before>from asyncore import dispatcher
from asynchat import async_chat
import socket, asyncore, time
from django.db import models
from django.db import connection

from apps.config.models import CommandAlias
import scheduler
import functions_db
import functions_general
import defines_global as global_defines
import session_mgr
import gameconf
import settings

class Server(dispatcher):
   """
   The main server class from which everything branches.
   """
   def __init__(self):
      self.cmd_alias_list = {}
      self.game_running = True

      # Wipe our temporary flags on all of the objects.
      if settings.DATABASE_ENGINE == "sqlite3":
         self.sqlite3_prep()
      cursor = connection.cursor()
      cursor.execute("UPDATE objects_object SET nosave_flags=''")

      print '-'*50
      # Load stuff up into memory for easy/quick access.
      self.load_cmd_aliases()
      self.port = gameconf.get_configvalue('site_port')

      # Start accepting connections.
      dispatcher.__init__(self)
      self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
      self.set_reuse_addr()
      self.bind(('', int(self.port)))
      self.listen(100)
      self.start_time = time.time()

      print ' %s started on port %s.' % (gameconf.get_configvalue('site_name'), self.port,)
      print '-'*50

   """
   BEGIN SERVER STARTUP METHODS
   """

   def load_cmd_aliases(self):
      """
      Load up our command aliases.
      """
      alias_list = CommandAlias.objects.all()
      for alias in alias_list:
         self.cmd_alias_list[alias.user_input] = alias.equiv_command
      print ' Command Aliases Loaded: %i' % (len(self.cmd_alias_list),)

   def handle_accept(self):
      """
      What to do when we get a connection.
      """
      conn, addr = self.accept()
      session = session_mgr.new_session(self, conn, addr)
      session.game_connect_screen(session)
      print 'Connection:', str(session)
      print 'Sessions active:', len(session_mgr.get_session_list())

   def sqlite3_prep(self):
      """
      Optimize some SQLite stuff at startup since we can't save it to the
      database.
      """
      cursor = connection.cursor()
      cursor.execute("PRAGMA cache_size=10000")
      cursor.execute("PRAGMA synchronous=OFF")
      cursor.execute("PRAGMA count_changes=OFF")
      cursor.execute("PRAGMA temp_store=2")

   """
   BEGIN GENERAL METHODS
   """
   def shutdown(self, message='The server has been shutdown. Please check back soon.'):
      functions_general.announce_all(message)
      self.game_running = False

   """
   END Server CLASS
   """

"""
BEGIN MAIN APPLICATION LOGIC
"""
if __name__ == '__main__':
   server = Server()

   try:
      while server.game_running:
         asyncore.loop(timeout=5, count=1)
         scheduler.heartbeat()

   except KeyboardInterrupt:
      server.shutdown()
      print '--> Server killed by keystroke.'
<commit_msg>A few minor cleaning things.<commit_after>from asyncore import dispatcher
from asynchat import async_chat
import socket, asyncore, time
from django.db import models
from django.db import connection

from apps.config.models import CommandAlias
import scheduler
import functions_general
import session_mgr
import gameconf
import settings

class Server(dispatcher):
   """
   The main server class from which everything branches.
   """
   def __init__(self):
      self.cmd_alias_list = {}
      self.game_running = True

      # Database-specific startup optimizations.
      if settings.DATABASE_ENGINE == "sqlite3":
         self.sqlite3_prep()

      # Wipe our temporary flags on all of the objects.
      cursor = connection.cursor()
      cursor.execute("UPDATE objects_object SET nosave_flags=''")

      print '-'*50
      # Load command aliases into memory for easy/quick access.
      self.load_cmd_aliases()
      self.port = gameconf.get_configvalue('site_port')

      # Start accepting connections.
      dispatcher.__init__(self)
      self.create_socket(socket.AF_INET, socket.SOCK_STREAM)
      self.set_reuse_addr()
      self.bind(('', int(self.port)))
      self.listen(100)
      self.start_time = time.time()

      print ' %s started on port %s.' % (gameconf.get_configvalue('site_name'), self.port,)
      print '-'*50

   """
   BEGIN SERVER STARTUP METHODS
   """

   def load_cmd_aliases(self):
      """
      Load up our command aliases.
      """
      alias_list = CommandAlias.objects.all()
      for alias in alias_list:
         self.cmd_alias_list[alias.user_input] = alias.equiv_command
      print ' Command Aliases Loaded: %i' % (len(self.cmd_alias_list),)

   def handle_accept(self):
      """
      What to do when we get a connection.
      """
      conn, addr = self.accept()
      session = session_mgr.new_session(self, conn, addr)
      session.game_connect_screen(session)
      print 'Connection:', str(session)
      print 'Sessions active:', len(session_mgr.get_session_list())

   def sqlite3_prep(self):
      """
      Optimize some SQLite stuff at startup since we can't save it to the
      database.
      """
      cursor = connection.cursor()
      cursor.execute("PRAGMA cache_size=10000")
      cursor.execute("PRAGMA synchronous=OFF")
      cursor.execute("PRAGMA count_changes=OFF")
      cursor.execute("PRAGMA temp_store=2")

   """
   BEGIN GENERAL METHODS
   """
   def shutdown(self, message='The server has been shutdown. Please check back soon.'):
      functions_general.announce_all(message)
      self.game_running = False

   """
   END Server CLASS
   """

"""
BEGIN MAIN APPLICATION LOGIC
"""
if __name__ == '__main__':
   server = Server()

   try:
      while server.game_running:
         asyncore.loop(timeout=5, count=1)
         scheduler.heartbeat()

   except KeyboardInterrupt:
      server.shutdown()
      print '--> Server killed by keystroke.'

   except:
      server.shutdown(message="The server has encountered an error and has been shut down. Please check back soon.")
      print '-!> Server crashed.'
<|endoftext|>


<commit_before>import os
import codecs
import re
import time
import pandas as pd
import six

from urbanaccess import config
from urbanaccess.utils import log
from urbanaccess.gtfs.gtfsfeeds_dataframe import gtfsfeeds_dfs
from urbanaccess.gtfs import utils_validation
from urbanaccess.gtfs import utils_format

def _standardize_txt(csv_rootpath=os.path.join(config.settings.data_folder,
                                               'gtfsfeed_text')):
    """
    Standardize all text files inside a GTFS feed for machine readability

    Parameters
    ----------
    csv_rootpath : str, optional
        root path where all gtfs feeds that make up a contiguous metropolitan
        area are stored

    Returns
    -------
    None
    """

    gtfsfiles_to_use = ['stops.txt', 'routes.txt', 'trips.txt',
                        'stop_times.txt', 'calendar.txt',
                        'agency.txt', 'calendar_dates.txt']

    if six.PY2:
        _txt_encoder_check(gtfsfiles_to_use, csv_rootpath)
    _txt_header_whitespace_check(gtfsfiles_to_use, csv_rootpath)

def _txt_encoder_check(gtfsfiles_to_use,
                       csv_rootpath=os.path.join(
                           config.settings.data_folder,
                           'gtfsfeed_text')):
    """
    Standardize all text files inside a GTFS feed for encoding problems.
    Has not been updated for Python 3.

    Parameters
    ----------
    gtfsfiles_to_use : list
        list of gtfs feed txt files to utilize
    csv_rootpath : str, optional
        root path where all gtfs feeds that make up a contiguous metropolitan
        area are stored

    Returns
    -------
    None
    """
    # UnicodeDecodeError
    start_time = time.time()

    folderlist = [foldername for foldername in os.listdir(csv_rootpath) if
                  os.path.isdir(os.path.join(csv_rootpath, foldername))]

    if not folderlist:
        folderlist = [csv_rootpath]

    for folder in folderlist:
        textfilelist = [textfilename for textfilename in
                        os.listdir(os.path.join(csv_rootpath, folder)) if
                        textfilename.endswith(".txt")]

        for textfile in textfilelist:
            if textfile in gtfsfiles_to_use:
                # Read from file
                file_open = open(os.path.join(csv_rootpath, folder, textfile))
                raw = file_open.read()
                file_open.close()
                if raw.startswith(codecs.BOM_UTF8):
                    raw = raw.replace(codecs.BOM_UTF8, '', 1)
                    # Write to file
                    file_open = open(
                        os.path.join(csv_rootpath, folder, textfile), 'w')
                    file_open.write(raw)
                    file_open.close()

    log('GTFS text file encoding check completed. Took {:,.2f} seconds'.format(
        time.time() - start_time))

def _txt_header_whitespace_check(gtfsfiles_to_use,
                                 csv_rootpath=os.path.join(
                                     config.settings.data_folder,
                                     'gtfsfeed_text')):
    """
    Standardize all text files inside a GTFS feed to remove whitespace
    in headers

    Parameters
    ----------
    gtfsfiles_to_use : list
        list of gtfs feed txt files to utilize
    csv_rootpath : str, optional
        root path where all gtfs feeds that make up a contiguous metropolitan
        area are stored

    Returns
    -------
    None
    """
    start_time = time.time()

    folderlist = [foldername for foldername in os.listdir(csv_rootpath) if
                  os.path.isdir(os.path.join(csv_rootpath, foldername))]

    if not folderlist:
        folderlist = [csv_rootpath]

    for folder in folderlist:
        textfilelist = [textfilename for textfilename in
                        os.listdir(os.path.join(csv_rootpath, folder)) if
                        textfilename.endswith(".txt")]

        for textfile in textfilelist:
            if textfile in gtfsfiles_to_use:
                # Read from file
                with open(os.path.join(csv_rootpath, folder, textfile)) as f:
                    lines = f.readlines()
                lines[0] = re.sub(r'\s+', '', lines[0]) + '\n'
                # Write to file
                try:
                    with open(os.path.join(csv_rootpath, folder, textfile),
                              'w') as f:
                        f.writelines(lines)
                except Exception:
                    log('Unable to read {}. Check that file is not currently'
                        'being read or is not already in memory as this is '
                        'likely the cause of the error.'
                        ''.format(os.path.join(csv_rootpath,
                                               folder, textfile)))
    log(
        'GTFS text file header whitespace check completed. Took {:,'
        '.2f} seconds'.format(
            time.time() - start_time))

def gtfsfeed_to_df(gtfsfeed_path=None, validation=False, verbose=True,
                   bbox=None, remove_stops_outsidebbox=None,
                   append_definitions=False):
    """
    Read all GTFS feed components as a dataframe in a gtfsfeeds_dfs object and
    merge all individual GTFS feeds into a regional metropolitan data table.
    Optionally, data can also be validated before its use.

    Parameters
    ----------
    gtfsfeed_path : str, optional
        root path where all gtfs feeds that make up a contiguous metropolitan
        area are stored
    validation : bool
        if true, the validation check on stops checking for stops outside
        of a bounding box and stop coordinate
        hemisphere will be run. this is required to remove stops outside of
        a bbox
    verbose : bool
        if true and stops are found outside of the bbox, the stops that are
        outside will be printed for your reference
    bbox : tuple
        Bounding box formatted as a 4 element tuple:
        (lng_max, lat_min, lng_min, lat_max)
        example: (-122.304611,37.798933,-122.263412,37.822802)
        a bbox can be extracted for an area using: the CSV format bbox
        from http://boundingbox.klokantech.com/
    remove_stops_outsidebbox : bool
        if true stops that are outside the bbox will be removed
    append_definitions : bool
        if true, columns that use the GTFS data schema for their attribute
        codes will have the corresponding GTFS definition information of
        that code appended to the resulting dataframes for reference

    Returns
    -------
    gtfsfeeds_dfs : object
        processed dataframes of corresponding GTFS feed text files
    gtfsfeeds_dfs.stops : pandas.DataFrame
    gtfsfeeds_dfs.routes : pandas.DataFrame
    gtfsfeeds_dfs.trips : pandas.DataFrame
    gtfsfeeds_dfs.stop_times : pandas.DataFrame
    gtfsfeeds_dfs.calendar : pandas.DataFrame
    gtfsfeeds_dfs.calendar_dates : pandas.DataFrame
    """

    merged_stops_df = pd.DataFrame()
    merged_routes_df = pd.DataFrame()
    merged_trips_df = pd.DataFrame()
    merged_stop_times_df = pd.DataFrame()
    merged_calendar_df = pd.DataFrame()
    merged_calendar_dates_df = pd.DataFrame()

    start_time = time.time()

    if gtfsfeed_path is None:
        gtfsfeed_path = os.path.join(config.settings.data_folder,
                                     'gtfsfeed_text')
        if not os.path.exists(gtfsfeed_path):
            raise ValueError('{} does not exist'.format(gtfsfeed_path))
    else:
        if not os.path.exists(gtfsfeed_path):
            raise ValueError('{} does not exist'.format(gtfsfeed_path))
    if not isinstance(gtfsfeed_path, str):
        raise ValueError('gtfsfeed_path must be a string')

    if validation:
        if bbox is None or remove_stops_outsidebbox is None or verbose is \
                None:
            raise ValueError(
                'Attempted to run validation but bbox, verbose, and or '
                'remove_stops_outsidebbox were set to None. These parameters '
                'must be specified for validation.')

    _standardize_txt(csv_rootpath=gtfsfeed_path)

    folderlist = [foldername for foldername in os.listdir(gtfsfeed_path) if
                  os.path.isdir(os.path.join(gtfsfeed_path, foldername))]
    if not folderlist:
        folderlist = [gtfsfeed_path]

    for index, folder in enumerate(folderlist):

        # print break to visually separate each gtfs feed log
        log('--------------------------------')
        log('Processing GTFS feed: {!s}'.format(os.path.split(folder)[1]))

        textfilelist = [textfilename for textfilename in
                        os.listdir(os.path.join(gtfsfeed_path, folder)) if
                        textfilename.endswith(".txt")]
        required_gtfsfiles = ['stops.txt', 'routes.txt', 'trips.txt',
                              'stop_times.txt']
        optional_gtfsfiles = ['agency.txt']
        # either calendar or calendar_dates is required
        calendar_gtfsfiles = ['calendar.txt', 'calendar_dates.txt']

        calendar_files = [i for i in calendar_gtfsfiles if i in textfilelist]
        if len(calendar_files) == 0:
            error_msg = (
                'at least one of `calendar.txt` or `calendar_dates.txt` is '
                'required to complete a GTFS dataset but neither was found in '
                'folder {}')
            raise ValueError(error_msg.format(os.path.join(
                gtfsfeed_path, folder)))

        for required_file in required_gtfsfiles:
            if required_file not in textfilelist:
                raise ValueError(
                    '{} is a required GTFS text file and was not found in '
                    'folder {}'.format(
                        required_file,
                        os.path.join(gtfsfeed_path, folder)))

        for textfile in required_gtfsfiles:
            if textfile == 'stops.txt':
                stops_df = utils_format._read_gtfs_stops(
                    textfile_path=os.path.join(gtfsfeed_path, folder),
                    textfile=textfile)
            if textfile == 'routes.txt':
                routes_df = utils_format._read_gtfs_routes(
                    textfile_path=os.path.join(gtfsfeed_path, folder),
                    textfile=textfile)
            if textfile == 'trips.txt':
                trips_df = utils_format._read_gtfs_trips(
                    textfile_path=os.path.join(gtfsfeed_path, folder),
                    textfile=textfile)
            if textfile == 'stop_times.txt':
                stop_times_df = utils_format._read_gtfs_stop_times(
                    textfile_path=os.path.join(gtfsfeed_path, folder),
                    textfile=textfile)

        for textfile in calendar_files:
            # use both calendar and calendar_dates if they exist, otherwise
            # if only one of them exists use the one that exists and set the
            # other one that does not exist to a blank df
            if textfile == 'calendar.txt':
                calendar_df = utils_format._read_gtfs_calendar(
                    textfile_path=os.path.join(gtfsfeed_path, folder),
                    textfile=textfile)
                # if only calendar, set calendar_dates as blank
                # with default required columns
                if len(calendar_files) == 1:
                    calendar_dates_df = pd.DataFrame(
                        columns=['service_id', 'dates', 'exception_type'])
            else:
                calendar_dates_df = utils_format._read_gtfs_calendar_dates(
                    textfile_path=os.path.join(gtfsfeed_path, folder),
                    textfile=textfile)
                # if only calendar_files, set calendar as blank
                # with default required columns
                if len(calendar_files) == 1:
                    calendar_df = pd.DataFrame(
                        columns=['service_id', 'monday',
                                 'tuesday', 'wednesday', 'thursday',
                                 'friday', 'saturday', 'sunday',
                                 'start_date', 'end_date'])

        for textfile in optional_gtfsfiles:
            if textfile == 'agency.txt':
                if textfile in textfilelist:
                    agency_df = utils_format._read_gtfs_agency(
                        textfile_path=os.path.join(gtfsfeed_path, folder),
                        textfile=textfile)
                else:
                    agency_df = pd.DataFrame()

        stops_df, routes_df, trips_df, stop_times_df, calendar_df, \
            calendar_dates_df = (utils_format
                                 ._add_unique_agencyid(
                                    agency_df=agency_df,
                                    stops_df=stops_df,
                                    routes_df=routes_df,
                                    trips_df=trips_df,
                                    stop_times_df=stop_times_df,
                                    calendar_df=calendar_df,
                                    calendar_dates_df=calendar_dates_df,
                                    nulls_as_folder=True,
                                    feed_folder=os.path.join(
                                        gtfsfeed_path, folder)))

        stops_df, routes_df, trips_df, stop_times_df, calendar_df, \
            calendar_dates_df = (utils_format
                                 ._add_unique_gtfsfeed_id(
                                    stops_df=stops_df,
                                    routes_df=routes_df,
                                    trips_df=trips_df,
                                    stop_times_df=stop_times_df,
                                    calendar_df=calendar_df,
                                    calendar_dates_df=calendar_dates_df,
                                    feed_folder=folder,
                                    feed_number=index+1))

        if validation:
            stops_df = (utils_validation
                        ._validate_gtfs(
                            stops_df=stops_df,
                            feed_folder=os.path.join(gtfsfeed_path, folder),
                            verbose=verbose,
                            bbox=bbox,
                            remove_stops_outsidebbox=remove_stops_outsidebbox))
            if remove_stops_outsidebbox:
                stops_inside_bbox = list(stops_df['stop_id'])
                stop_times_df = stop_times_df[stop_times_df['stop_id'].isin(
                    stops_inside_bbox)]

        stops_df = (utils_format
                    ._append_route_type(
                        stops_df=stops_df,
                        stop_times_df=stop_times_df,
                        routes_df=routes_df[['route_id', 'route_type']],
                        trips_df=trips_df[['trip_id', 'route_id']],
                        info_to_append='route_type_to_stops'))
        stop_times_df = (utils_format
                         ._append_route_type(
                            stops_df=stops_df,
                            stop_times_df=stop_times_df,
                            routes_df=routes_df[['route_id', 'route_type']],
                            trips_df=trips_df[['trip_id', 'route_id']],
                            info_to_append='route_type_to_stop_times'))

        merged_stops_df = merged_stops_df.append(stops_df, ignore_index=True)
        merged_routes_df = merged_routes_df.append(routes_df,
                                                   ignore_index=True)
        merged_trips_df = merged_trips_df.append(trips_df, ignore_index=True)
        merged_stop_times_df = merged_stop_times_df.append(stop_times_df,
                                                           ignore_index=True)
        merged_calendar_df = merged_calendar_df.append(calendar_df,
                                                       ignore_index=True)
        merged_calendar_dates_df = merged_calendar_dates_df.append(
            calendar_dates_df, ignore_index=True)

        # print break to visually separate each gtfs feed log
        log('--------------------------------')

    if append_definitions:
        merged_stops_df, merged_routes_df, merged_stop_times_df, \
            merged_trips_df = (utils_format
                               ._add_txt_definitions(
                                    stops_df=merged_stops_df,
                                    routes_df=merged_routes_df,
                                    stop_times_df=merged_stop_times_df,
                                    trips_df=merged_trips_df))

    merged_stop_times_df = utils_format._timetoseconds(
        df=merged_stop_times_df, time_cols=['departure_time'])

    # set gtfsfeeds_dfs object to merged GTFS dfs
    gtfsfeeds_dfs.stops = merged_stops_df
    gtfsfeeds_dfs.routes = merged_routes_df
    gtfsfeeds_dfs.trips = merged_trips_df
    gtfsfeeds_dfs.stop_times = merged_stop_times_df
    gtfsfeeds_dfs.calendar = merged_calendar_df
    gtfsfeeds_dfs.calendar_dates = merged_calendar_dates_df

    # TODO: add to print the list of gtfs feed txt files read in for each feed
    log('{:,} GTFS feed file(s) successfully read as dataframes:'.format(
        len(folderlist)))
    for folder in folderlist:
        log('     {}'.format(os.path.split(folder)[1]))
    log('     Took {:,.2f} seconds'.format(time.time() - start_time))

    return gtfsfeeds_dfs
<commit_msg>update note<commit_after>import os
import codecs
import re
import time
import pandas as pd
import six

from urbanaccess import config
from urbanaccess.utils import log
from urbanaccess.gtfs.gtfsfeeds_dataframe import gtfsfeeds_dfs
from urbanaccess.gtfs import utils_validation
from urbanaccess.gtfs import utils_format

def _standardize_txt(csv_rootpath=os.path.join(config.settings.data_folder,
                                               'gtfsfeed_text')):
    """
    Standardize all text files inside a GTFS feed for machine readability

    Parameters
    ----------
    csv_rootpath : str, optional
        root path where all gtfs feeds that make up a contiguous metropolitan
        area are stored

    Returns
    -------
    None
    """

    gtfsfiles_to_use = ['stops.txt', 'routes.txt', 'trips.txt',
                        'stop_times.txt', 'calendar.txt',
                        'agency.txt', 'calendar_dates.txt']

    if six.PY2:
        _txt_encoder_check(gtfsfiles_to_use, csv_rootpath)
    _txt_header_whitespace_check(gtfsfiles_to_use, csv_rootpath)

def _txt_encoder_check(gtfsfiles_to_use,
                       csv_rootpath=os.path.join(
                           config.settings.data_folder,
                           'gtfsfeed_text')):
    """
    Standardize all text files inside a GTFS feed for encoding problems.
    Has not been updated for Python 3.

    Parameters
    ----------
    gtfsfiles_to_use : list
        list of gtfs feed txt files to utilize
    csv_rootpath : str, optional
        root path where all gtfs feeds that make up a contiguous metropolitan
        area are stored

    Returns
    -------
    None
    """
    # UnicodeDecodeError
    start_time = time.time()

    folderlist = [foldername for foldername in os.listdir(csv_rootpath) if
                  os.path.isdir(os.path.join(csv_rootpath, foldername))]

    if not folderlist:
        folderlist = [csv_rootpath]

    for folder in folderlist:
        textfilelist = [textfilename for textfilename in
                        os.listdir(os.path.join(csv_rootpath, folder)) if
                        textfilename.endswith(".txt")]

        for textfile in textfilelist:
            if textfile in gtfsfiles_to_use:
                # Read from file
                file_open = open(os.path.join(csv_rootpath, folder, textfile))
                raw = file_open.read()
                file_open.close()
                if raw.startswith(codecs.BOM_UTF8):
                    raw = raw.replace(codecs.BOM_UTF8, '', 1)
                    # Write to file
                    file_open = open(
                        os.path.join(csv_rootpath, folder, textfile), 'w')
                    file_open.write(raw)
                    file_open.close()

    log('GTFS text file encoding check completed. Took {:,.2f} seconds'.format(
        time.time() - start_time))

def _txt_header_whitespace_check(gtfsfiles_to_use,
                                 csv_rootpath=os.path.join(
                                     config.settings.data_folder,
                                     'gtfsfeed_text')):
    """
    Standardize all text files inside a GTFS feed to remove whitespace
    in headers

    Parameters
    ----------
    gtfsfiles_to_use : list
        list of gtfs feed txt files to utilize
    csv_rootpath : str, optional
        root path where all gtfs feeds that make up a contiguous metropolitan
        area are stored

    Returns
    -------
    None
    """
    start_time = time.time()

    folderlist = [foldername for foldername in os.listdir(csv_rootpath) if
                  os.path.isdir(os.path.join(csv_rootpath, foldername))]

    if not folderlist:
        folderlist = [csv_rootpath]

    for folder in folderlist:
        textfilelist = [textfilename for textfilename in
                        os.listdir(os.path.join(csv_rootpath, folder)) if
                        textfilename.endswith(".txt")]

        for textfile in textfilelist:
            if textfile in gtfsfiles_to_use:
                # Read from file
                with open(os.path.join(csv_rootpath, folder, textfile)) as f:
                    lines = f.readlines()
                lines[0] = re.sub(r'\s+', '', lines[0]) + '\n'
                # Write to file
                try:
                    with open(os.path.join(csv_rootpath, folder, textfile),
                              'w') as f:
                        f.writelines(lines)
                except Exception:
                    log('Unable to read {}. Check that file is not currently'
                        'being read or is not already in memory as this is '
                        'likely the cause of the error.'
                        ''.format(os.path.join(csv_rootpath,
                                               folder, textfile)))
    log(
        'GTFS text file header whitespace check completed. Took {:,'
        '.2f} seconds'.format(
            time.time() - start_time))

def gtfsfeed_to_df(gtfsfeed_path=None, validation=False, verbose=True,
                   bbox=None, remove_stops_outsidebbox=None,
                   append_definitions=False):
    """
    Read all GTFS feed components as a dataframe in a gtfsfeeds_dfs object and
    merge all individual GTFS feeds into a regional metropolitan data table.
    Optionally, data can also be validated before its use.

    Parameters
    ----------
    gtfsfeed_path : str, optional
        root path where all gtfs feeds that make up a contiguous metropolitan
        area are stored
    validation : bool
        if true, the validation check on stops checking for stops outside
        of a bounding box and stop coordinate
        hemisphere will be run. this is required to remove stops outside of
        a bbox
    verbose : bool
        if true and stops are found outside of the bbox, the stops that are
        outside will be printed for your reference
    bbox : tuple
        Bounding box formatted as a 4 element tuple:
        (lng_max, lat_min, lng_min, lat_max)
        example: (-122.304611,37.798933,-122.263412,37.822802)
        a bbox can be extracted for an area using: the CSV format bbox
        from http://boundingbox.klokantech.com/
    remove_stops_outsidebbox : bool
        if true stops that are outside the bbox will be removed
    append_definitions : bool
        if true, columns that use the GTFS data schema for their attribute
        codes will have the corresponding GTFS definition information of
        that code appended to the resulting dataframes for reference

    Returns
    -------
    gtfsfeeds_dfs : object
        processed dataframes of corresponding GTFS feed text files
    gtfsfeeds_dfs.stops : pandas.DataFrame
    gtfsfeeds_dfs.routes : pandas.DataFrame
    gtfsfeeds_dfs.trips : pandas.DataFrame
    gtfsfeeds_dfs.stop_times : pandas.DataFrame
    gtfsfeeds_dfs.calendar : pandas.DataFrame
    gtfsfeeds_dfs.calendar_dates : pandas.DataFrame
    """

    merged_stops_df = pd.DataFrame()
    merged_routes_df = pd.DataFrame()
    merged_trips_df = pd.DataFrame()
    merged_stop_times_df = pd.DataFrame()
    merged_calendar_df = pd.DataFrame()
    merged_calendar_dates_df = pd.DataFrame()

    start_time = time.time()

    if gtfsfeed_path is None:
        gtfsfeed_path = os.path.join(config.settings.data_folder,
                                     'gtfsfeed_text')
        if not os.path.exists(gtfsfeed_path):
            raise ValueError('{} does not exist'.format(gtfsfeed_path))
    else:
        if not os.path.exists(gtfsfeed_path):
            raise ValueError('{} does not exist'.format(gtfsfeed_path))
    if not isinstance(gtfsfeed_path, str):
        raise ValueError('gtfsfeed_path must be a string')

    if validation:
        if bbox is None or remove_stops_outsidebbox is None or verbose is \
                None:
            raise ValueError(
                'Attempted to run validation but bbox, verbose, and or '
                'remove_stops_outsidebbox were set to None. These parameters '
                'must be specified for validation.')

    _standardize_txt(csv_rootpath=gtfsfeed_path)

    folderlist = [foldername for foldername in os.listdir(gtfsfeed_path) if
                  os.path.isdir(os.path.join(gtfsfeed_path, foldername))]
    if not folderlist:
        folderlist = [gtfsfeed_path]

    for index, folder in enumerate(folderlist):

        # print break to visually separate each gtfs feed log
        log('--------------------------------')
        log('Processing GTFS feed: {!s}'.format(os.path.split(folder)[1]))

        textfilelist = [textfilename for textfilename in
                        os.listdir(os.path.join(gtfsfeed_path, folder)) if
                        textfilename.endswith(".txt")]
        required_gtfsfiles = ['stops.txt', 'routes.txt', 'trips.txt',
                              'stop_times.txt']
        optional_gtfsfiles = ['agency.txt']
        # either calendar or calendar_dates is required
        calendar_gtfsfiles = ['calendar.txt', 'calendar_dates.txt']

        calendar_files = [i for i in calendar_gtfsfiles if i in textfilelist]
        if len(calendar_files) == 0:
            error_msg = (
                'at least one of `calendar.txt` or `calendar_dates.txt` is '
                'required to complete a GTFS dataset but neither was found in '
                'folder {}')
            raise ValueError(error_msg.format(os.path.join(
                gtfsfeed_path, folder)))

        for required_file in required_gtfsfiles:
            if required_file not in textfilelist:
                raise ValueError(
                    '{} is a required GTFS text file and was not found in '
                    'folder {}'.format(
                        required_file,
                        os.path.join(gtfsfeed_path, folder)))

        for textfile in required_gtfsfiles:
            if textfile == 'stops.txt':
                stops_df = utils_format._read_gtfs_stops(
                    textfile_path=os.path.join(gtfsfeed_path, folder),
                    textfile=textfile)
            if textfile == 'routes.txt':
                routes_df = utils_format._read_gtfs_routes(
                    textfile_path=os.path.join(gtfsfeed_path, folder),
                    textfile=textfile)
            if textfile == 'trips.txt':
                trips_df = utils_format._read_gtfs_trips(
                    textfile_path=os.path.join(gtfsfeed_path, folder),
                    textfile=textfile)
            if textfile == 'stop_times.txt':
                stop_times_df = utils_format._read_gtfs_stop_times(
                    textfile_path=os.path.join(gtfsfeed_path, folder),
                    textfile=textfile)

        for textfile in calendar_files:
            # use both calendar and calendar_dates if they exist, otherwise
            # if only one of them exists use the one that exists and set the
            # other one that does not exist to a blank df
            if textfile == 'calendar.txt':
                calendar_df = utils_format._read_gtfs_calendar(
                    textfile_path=os.path.join(gtfsfeed_path, folder),
                    textfile=textfile)
                # if only calendar, set calendar_dates as blank
                # with default required columns
                if len(calendar_files) == 1:
                    calendar_dates_df = pd.DataFrame(
                        columns=['service_id', 'dates', 'exception_type'])
            else:
                calendar_dates_df = utils_format._read_gtfs_calendar_dates(
                    textfile_path=os.path.join(gtfsfeed_path, folder),
                    textfile=textfile)
                # if only calendar_dates, set calendar as blank
                # with default required columns
                if len(calendar_files) == 1:
                    calendar_df = pd.DataFrame(
                        columns=['service_id', 'monday',
                                 'tuesday', 'wednesday', 'thursday',
                                 'friday', 'saturday', 'sunday',
                                 'start_date', 'end_date'])

        for textfile in optional_gtfsfiles:
            if textfile == 'agency.txt':
                if textfile in textfilelist:
                    agency_df = utils_format._read_gtfs_agency(
                        textfile_path=os.path.join(gtfsfeed_path, folder),
                        textfile=textfile)
                else:
                    agency_df = pd.DataFrame()

        stops_df, routes_df, trips_df, stop_times_df, calendar_df, \
            calendar_dates_df = (utils_format
                                 ._add_unique_agencyid(
                                    agency_df=agency_df,
                                    stops_df=stops_df,
                                    routes_df=routes_df,
                                    trips_df=trips_df,
                                    stop_times_df=stop_times_df,
                                    calendar_df=calendar_df,
                                    calendar_dates_df=calendar_dates_df,
                                    nulls_as_folder=True,
                                    feed_folder=os.path.join(
                                        gtfsfeed_path, folder)))

        stops_df, routes_df, trips_df, stop_times_df, calendar_df, \
            calendar_dates_df = (utils_format
                                 ._add_unique_gtfsfeed_id(
                                    stops_df=stops_df,
                                    routes_df=routes_df,
                                    trips_df=trips_df,
                                    stop_times_df=stop_times_df,
                                    calendar_df=calendar_df,
                                    calendar_dates_df=calendar_dates_df,
                                    feed_folder=folder,
                                    feed_number=index+1))

        if validation:
            stops_df = (utils_validation
                        ._validate_gtfs(
                            stops_df=stops_df,
                            feed_folder=os.path.join(gtfsfeed_path, folder),
                            verbose=verbose,
                            bbox=bbox,
                            remove_stops_outsidebbox=remove_stops_outsidebbox))
            if remove_stops_outsidebbox:
                stops_inside_bbox = list(stops_df['stop_id'])
                stop_times_df = stop_times_df[stop_times_df['stop_id'].isin(
                    stops_inside_bbox)]

        stops_df = (utils_format
                    ._append_route_type(
                        stops_df=stops_df,
                        stop_times_df=stop_times_df,
                        routes_df=routes_df[['route_id', 'route_type']],
                        trips_df=trips_df[['trip_id', 'route_id']],
                        info_to_append='route_type_to_stops'))
        stop_times_df = (utils_format
                         ._append_route_type(
                            stops_df=stops_df,
                            stop_times_df=stop_times_df,
                            routes_df=routes_df[['route_id', 'route_type']],
                            trips_df=trips_df[['trip_id', 'route_id']],
                            info_to_append='route_type_to_stop_times'))

        merged_stops_df = merged_stops_df.append(stops_df, ignore_index=True)
        merged_routes_df = merged_routes_df.append(routes_df,
                                                   ignore_index=True)
        merged_trips_df = merged_trips_df.append(trips_df, ignore_index=True)
        merged_stop_times_df = merged_stop_times_df.append(stop_times_df,
                                                           ignore_index=True)
        merged_calendar_df = merged_calendar_df.append(calendar_df,
                                                       ignore_index=True)
        merged_calendar_dates_df = merged_calendar_dates_df.append(
            calendar_dates_df, ignore_index=True)

        # print break to visually separate each gtfs feed log
        log('--------------------------------')

    if append_definitions:
        merged_stops_df, merged_routes_df, merged_stop_times_df, \
            merged_trips_df = (utils_format
                               ._add_txt_definitions(
                                    stops_df=merged_stops_df,
                                    routes_df=merged_routes_df,
                                    stop_times_df=merged_stop_times_df,
                                    trips_df=merged_trips_df))

    merged_stop_times_df = utils_format._timetoseconds(
        df=merged_stop_times_df, time_cols=['departure_time'])

    # set gtfsfeeds_dfs object to merged GTFS dfs
    gtfsfeeds_dfs.stops = merged_stops_df
    gtfsfeeds_dfs.routes = merged_routes_df
    gtfsfeeds_dfs.trips = merged_trips_df
    gtfsfeeds_dfs.stop_times = merged_stop_times_df
    gtfsfeeds_dfs.calendar = merged_calendar_df
    gtfsfeeds_dfs.calendar_dates = merged_calendar_dates_df

    # TODO: add to print the list of gtfs feed txt files read in for each feed
    log('{:,} GTFS feed file(s) successfully read as dataframes:'.format(
        len(folderlist)))
    for folder in folderlist:
        log('     {}'.format(os.path.split(folder)[1]))
    log('     Took {:,.2f} seconds'.format(time.time() - start_time))

    return gtfsfeeds_dfs
<|endoftext|>


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

# ==============================================================================
# MODULE DOCSTRING
# ==============================================================================

"""Utilities to run on MPI.

Provide functions and decorators that simplify running the same code on
multiple nodes. One benefit is that serial and parallel code is exactly
the same.

Global variables
----------------
disable_mpi : bool
    Set this to True to force running serially.

Routines
--------
get_mpicomm
    Automatically detect and configure MPI execution and return an
    MPI communicator.
run_single_node
    Run a task on a single node.
on_single_node
    Decorator version of run_single_node.
distribute
    Map a task on a sequence of arguments on all the nodes.
delay_termination
    A context manager to delay the response to termination signals.
delayed_termination
    A decorator version of delay_termination.

"""

# ==============================================================================
# GLOBAL IMPORTS
# ==============================================================================

import os
import sys
import signal
import logging
from contextlib import contextmanager

# TODO drop this when we drop Python 2 support
from openmoltools.utils import wraps_py2

logger = logging.getLogger(__name__)

# ==============================================================================
# GLOBAL VARIABLES
# ==============================================================================

disable_mpi = False

# ==============================================================================
# MAIN FUNCTIONS
# ==============================================================================

def get_mpicomm():
    """Retrieve the MPI communicator for this execution.

    The function automatically detects if the program runs on MPI by checking
    specific environment variables set by various MPI implementations. On
    first execution, it modifies sys.excepthook and register a handler for
    SIGINT, SIGTERM, SIGABRT to call Abort() to correctly terminate all
    processes.

    Returns
    -------
    mpicomm : mpi4py communicator or None
        The communicator for this node, None if the program doesn't run
        with MPI.

    """
    # If MPI execution is forcefully disabled, return None.
    if disable_mpi:
        return None

    # If we have already initialized MPI, return the cached MPI communicator.
    if get_mpicomm._is_initialized:
        return get_mpicomm._mpicomm

    # Check for environment variables set by mpirun. Variables are from
    # http://docs.roguewave.com/threadspotter/2012.1/linux/manual_html/apas03.html
    variables = ['PMI_RANK', 'OMPI_COMM_WORLD_RANK', 'OMPI_MCA_ns_nds_vpid',
                 'PMI_ID', 'SLURM_PROCID', 'LAMRANK', 'MPI_RANKID',
                 'MP_CHILD', 'MP_RANK', 'MPIRUN_RANK']
    use_mpi = False
    for var in variables:
        if var in os.environ:
            use_mpi = True
            break

    # Return None if we are not running on MPI.
    if not use_mpi:
        logger.debug('Cannot find MPI environment. MPI disabled.')
        get_mpicomm._mpicomm = None
        return get_mpicomm._mpicomm

    # Initialize MPI
    from mpi4py import MPI
    MPI.COMM_WORLD.barrier()
    mpicomm = MPI.COMM_WORLD

    # Override sys.excepthook to abort MPI on exception
    def mpi_excepthook(type, value, traceback):
        sys.__excepthook__(type, value, traceback)
        sys.stdout.flush()
        sys.stderr.flush()
        if mpicomm.size > 1:
            mpicomm.Abort(1)
    # Use our eception handler
    sys.excepthook = mpi_excepthook

    # Catch sigterm signals
    def handle_signal(signal, frame):
        if mpicomm.size > 1:
            mpicomm.Abort(1)
    for sig in [signal.SIGINT, signal.SIGTERM, signal.SIGABRT]:
        signal.signal(sig, handle_signal)

    # Cache and return the MPI communicator.
    get_mpicomm._is_initialized = True
    get_mpicomm._mpicomm = mpicomm

    # Report initialization
    logger.debug("MPI initialized on node {}/{}".format(mpicomm.rank+1, mpicomm.size))

    return mpicomm

get_mpicomm._is_initialized = False  # Static variable

def run_single_node(rank, task, *args, **kwargs):
    """Run task on a single node.

    If MPI is not activated, this simply runs locally.

    Parameters
    ----------
    task : callable
        The task to run on node rank.
    rank : int
        The rank of the MPI communicator that must execute the task.
    broadcast_result : bool, optional
        If True, the result is broadcasted to all nodes. If False,
        only the node executing the task will receive the return
        value of the task, and all other nodes will receive None
        (default is False).
    sync_nodes : bool, optional
        If True, the nodes will be synchronized at the end of the
        execution (i.e. the task will be blocking) even if the
        result is not broadcasted  (default is False).

    Other Parameters
    ----------------
    *args
        The ordered arguments to pass to task.
    *kwargs
        The keyword arguments to pass to task.

    Returns
    -------
    result
        The return value of the task. This will be None on all nodes
        that is not the rank unless broadcast_result is set to True.

    Examples
    --------
    >>> def add(a, b):
    ...     return a + b
    >>> # Run 3+4 on node 0.
    >>> run_single_node(0, task=add, a=3, b=4, broadcast_result=True)
    7

    """
    broadcast_result = kwargs.pop('broadcast_result', False)
    sync_nodes = kwargs.pop('sync_nodes', False)
    result = None
    mpicomm = get_mpicomm()

    if mpicomm is not None:
        node_name = 'Node {}/{}'.format(mpicomm.rank+1, mpicomm.size)
    else:
        node_name = 'Single node'

    # Execute the task only on the specified node.
    if mpicomm is None or mpicomm.rank == rank:
        logger.debug('{}: executing {}'.format(node_name, task))
        result = task(*args, **kwargs)

    # Broadcast the result if required.
    if mpicomm is not None:
        if broadcast_result is True:
            logger.debug('{}: waiting for broadcast of {}'.format(node_name, task))
            result = mpicomm.bcast(result, root=rank)
        elif sync_nodes is True:
            logger.debug('{}: waiting for barrier after {}'.format(node_name, task))
            mpicomm.barrier()

    # Return result.
    return result

def on_single_node(rank, broadcast_result=False, sync_nodes=False):
    """A decorator version of run_single_node.

    Decorates a function to be always executed with run_single_node.

    Parameters
    ----------
    rank : int
        The rank of the MPI communicator that must execute the task.
    broadcast_result : bool, optional
        If True the result is broadcasted to all nodes. If False,
        only the node executing the function will receive its return
        value, and all other nodes will receive None (default is False).
    sync_nodes : bool, optional
        If True, the nodes will be synchronized at the end of the
        execution (i.e. the task will be blocking) even if the
        result is not broadcasted (default is False).

    See Also
    --------
    run_single_node

    Examples
    --------
    >>> @on_single_node(rank=0, broadcast_result=True)
    ... def add(a, b):
    ...     return a + b
    >>> add(3, 4)
    7

    """
    def _on_single_node(task):
        @wraps_py2(task)
        def _wrapper(*args, **kwargs):
            kwargs['broadcast_result'] = broadcast_result
            kwargs['sync_nodes'] = sync_nodes
            return run_single_node(rank, task, *args, **kwargs)
        return _wrapper
    return _on_single_node

def distribute(task, distributed_args, *other_args, **kwargs):
    """Map the task on a sequence of arguments to be executed on different nodes.

    If MPI is not activated, this simply runs serially on this node. The
    algorithm guarantees that each node will be assigned to the same job_id
    (i.e. the index of the argument in distributed_args) every time.

    Parameters
    ----------
    task : callable
        The task to be distributed among nodes. The task will be called as
        task(distributed_args[job_id], *other_args, **kwargs), so the parameter
        to be distributed must the the first one.
    distributed_args : iterable
        The sequence of the parameters to distribute among nodes.
    send_results_to : int or 'all', optional
        If the string 'all', the result will be sent to all nodes. If an
        int, the result will be send only to the node with rank send_results_to.
        The return value of distribute depends on the value of this parameter
        (default is None).
    sync_nodes : bool, optional
        If True, the nodes will be synchronized at the end of the
        execution (i.e. the task will be blocking) even if the
        result is not shared (default is False).

    Other Parameters
    ----------------
    *other_args
        Other parameters to pass to task beside the assigned distributed
        parameters.
    **kwargs
        Keyword arguments to pass to task beside the assigned distributed
        parameters.

    Returns
    -------
    all_results : list
        All the return values for all the arguments if the results where sent
        to the node, or only the return values of the arguments processed by
        this node otherwise.
    arg_indices : list of int, optional
        This is returned as part of a tuple (all_results, job_indices) only
        if send_results_to is set to an int or None. In this case all_results[i]
        is the return value of task(all_args[arg_indices[i]]).

    Examples
    --------
    >>> def square(x):
    ...     return x**2
    >>> distribute(square, [1, 2, 3, 4], send_results_to='all')
    [1, 4, 9, 16]

    When send_results_to is not set to `all`, the return value include also
    the indices of the arguments associated to the result.

    >>> distribute(square, [1, 2, 3, 4], send_results_to=0)
    ([1, 4, 9, 16], [0, 1, 2, 3])

    """
    send_results_to = kwargs.pop('send_results_to', None)
    sync_nodes = kwargs.pop('sync_nodes', False)
    mpicomm = get_mpicomm()
    n_jobs = len(distributed_args)

    # If MPI is not activated, just run serially.
    if mpicomm is None:
        logger.debug('Running {} serially.'.format(task.__name__))
        all_results = [task(job_args, *other_args, **kwargs) for job_args in distributed_args]
        if send_results_to == 'all':
            return all_results
        else:
            return all_results, list(range(n_jobs))

    node_job_ids = range(mpicomm.rank, n_jobs, mpicomm.size)

    # Compute all the results assigned to this node.
    results = []
    node_name = 'Node {}/{}'.format(mpicomm.rank+1, mpicomm.size)
    for job_id in node_job_ids:
        distributed_arg = distributed_args[job_id]
        logger.debug('{}: execute {}({})'.format(node_name, task.__name__, distributed_arg))
        results.append(task(distributed_arg, *other_args, **kwargs))

    # Share result as specified.
    if send_results_to == 'all':
        logger.debug('{}: allgather results of {}'.format(node_name, task.__name__))
        all_results = mpicomm.allgather(results)
    elif isinstance(send_results_to, int):
        logger.debug('{}: gather results of {}'.format(node_name, task.__name__))
        all_results = mpicomm.gather(results, root=send_results_to)

        # If this is not the receiving node, we can safely return.
        if mpicomm.rank != send_results_to:
            return results, list(node_job_ids)
    else:
        assert send_results_to is None  # Safety check.
        if sync_nodes is True:
            logger.debug('{}: waiting for barrier after {}'.format(node_name, task.__name__))
            mpicomm.barrier()
        return results, list(node_job_ids)

    # all_results is a list of list of results. The internal lists of
    # results are ordered by rank. We need to reorder the results as a
    # flat list or results ordered by job_id.

    # job_indices[job_id] is the tuple of indices (rank, i). The result
    # of job_id is stored in all_results[rank][i].
    job_indices = []
    max_jobs_per_node = max([len(r) for r in all_results])
    for i in range(max_jobs_per_node):
        for rank in range(mpicomm.size):
            # Not all nodes have executed max_jobs_per_node tasks.
            if len(all_results[rank]) > i:
                job_indices.append((rank, i))

    # Reorder the results.
    all_results = [all_results[rank][i] for rank, i in job_indices]

    # Return result.
    if send_results_to == 'all':
        return all_results
    else:
        return all_results, list(range(n_jobs))

@contextmanager
def delay_termination():
    """Context manager to delay handling of termination signals.

    This allows to avoid interrupting tasks such as writing to the file
    system, which could result in the corruption of the file.

    """
    signals_to_catch = [signal.SIGINT, signal.SIGTERM, signal.SIGABRT]
    old_handlers = {signum: signal.getsignal(signum) for signum in signals_to_catch}
    signals_received = {signum: None for signum in signals_to_catch}

    def delay_handler(signum, frame):
        signals_received[signum] = (signum, frame)

    # Set handlers fot delay
    for signum in signals_to_catch:
        signal.signal(signum, delay_handler)

    yield  # Resume program

    # Restore old handlers
    for signum, handler in old_handlers.items():
        signal.signal(signum, handler)

    # Fire delayed signals
    for signum, s in signals_received.items():
        if s is not None:
            old_handlers[signum](*s)

def delayed_termination(func):
    """Decorator that runs the function with delay_termination()."""
    @wraps_py2(func)
    def _delayed_termination(*args, **kwargs):
        with delay_termination():
            return func(*args, **kwargs)
    return _delayed_termination

# ==============================================================================
# MAIN AND TESTS
# ==============================================================================

if __name__ == "__main__":
    import doctest
    doctest.testmod()
<commit_msg>Fix get_mpicomm._is_initialized static variable in serial case.<commit_after>#!/usr/bin/env python

# ==============================================================================
# MODULE DOCSTRING
# ==============================================================================

"""Utilities to run on MPI.

Provide functions and decorators that simplify running the same code on
multiple nodes. One benefit is that serial and parallel code is exactly
the same.

Global variables
----------------
disable_mpi : bool
    Set this to True to force running serially.

Routines
--------
get_mpicomm
    Automatically detect and configure MPI execution and return an
    MPI communicator.
run_single_node
    Run a task on a single node.
on_single_node
    Decorator version of run_single_node.
distribute
    Map a task on a sequence of arguments on all the nodes.
delay_termination
    A context manager to delay the response to termination signals.
delayed_termination
    A decorator version of delay_termination.

"""

# ==============================================================================
# GLOBAL IMPORTS
# ==============================================================================

import os
import sys
import signal
import logging
from contextlib import contextmanager

# TODO drop this when we drop Python 2 support
from openmoltools.utils import wraps_py2

logger = logging.getLogger(__name__)

# ==============================================================================
# GLOBAL VARIABLES
# ==============================================================================

disable_mpi = False

# ==============================================================================
# MAIN FUNCTIONS
# ==============================================================================

def get_mpicomm():
    """Retrieve the MPI communicator for this execution.

    The function automatically detects if the program runs on MPI by checking
    specific environment variables set by various MPI implementations. On
    first execution, it modifies sys.excepthook and register a handler for
    SIGINT, SIGTERM, SIGABRT to call Abort() to correctly terminate all
    processes.

    Returns
    -------
    mpicomm : mpi4py communicator or None
        The communicator for this node, None if the program doesn't run
        with MPI.

    """
    # If MPI execution is forcefully disabled, return None.
    if disable_mpi:
        return None

    # If we have already initialized MPI, return the cached MPI communicator.
    if get_mpicomm._is_initialized:
        return get_mpicomm._mpicomm

    # Check for environment variables set by mpirun. Variables are from
    # http://docs.roguewave.com/threadspotter/2012.1/linux/manual_html/apas03.html
    variables = ['PMI_RANK', 'OMPI_COMM_WORLD_RANK', 'OMPI_MCA_ns_nds_vpid',
                 'PMI_ID', 'SLURM_PROCID', 'LAMRANK', 'MPI_RANKID',
                 'MP_CHILD', 'MP_RANK', 'MPIRUN_RANK']
    use_mpi = False
    for var in variables:
        if var in os.environ:
            use_mpi = True
            break

    # Return None if we are not running on MPI.
    if not use_mpi:
        logger.debug('Cannot find MPI environment. MPI disabled.')
        get_mpicomm._mpicomm = None
        get_mpicomm._is_initialized = True
        return get_mpicomm._mpicomm

    # Initialize MPI
    from mpi4py import MPI
    mpicomm = MPI.COMM_WORLD

    # Override sys.excepthook to abort MPI on exception
    def mpi_excepthook(type, value, traceback):
        sys.__excepthook__(type, value, traceback)
        sys.stdout.flush()
        sys.stderr.flush()
        if mpicomm.size > 1:
            mpicomm.Abort(1)
    # Use our eception handler
    sys.excepthook = mpi_excepthook

    # Catch sigterm signals
    def handle_signal(signal, frame):
        if mpicomm.size > 1:
            mpicomm.Abort(1)
    for sig in [signal.SIGINT, signal.SIGTERM, signal.SIGABRT]:
        signal.signal(sig, handle_signal)

    # Cache and return the MPI communicator.
    get_mpicomm._is_initialized = True
    get_mpicomm._mpicomm = mpicomm

    # Report initialization
    logger.debug("MPI initialized on node {}/{}".format(mpicomm.rank+1, mpicomm.size))

    return mpicomm

get_mpicomm._is_initialized = False  # Static variable

def run_single_node(rank, task, *args, **kwargs):
    """Run task on a single node.

    If MPI is not activated, this simply runs locally.

    Parameters
    ----------
    task : callable
        The task to run on node rank.
    rank : int
        The rank of the MPI communicator that must execute the task.
    broadcast_result : bool, optional
        If True, the result is broadcasted to all nodes. If False,
        only the node executing the task will receive the return
        value of the task, and all other nodes will receive None
        (default is False).
    sync_nodes : bool, optional
        If True, the nodes will be synchronized at the end of the
        execution (i.e. the task will be blocking) even if the
        result is not broadcasted  (default is False).

    Other Parameters
    ----------------
    *args
        The ordered arguments to pass to task.
    **kwargs
        The keyword arguments to pass to task.

    Returns
    -------
    result
        The return value of the task. This will be None on all nodes
        that is not the rank unless broadcast_result is set to True.

    Examples
    --------
    >>> def add(a, b):
    ...     return a + b
    >>> # Run 3+4 on node 0.
    >>> run_single_node(0, task=add, a=3, b=4, broadcast_result=True)
    7

    """
    broadcast_result = kwargs.pop('broadcast_result', False)
    sync_nodes = kwargs.pop('sync_nodes', False)
    result = None
    mpicomm = get_mpicomm()

    if mpicomm is not None:
        node_name = 'Node {}/{}'.format(mpicomm.rank+1, mpicomm.size)
    else:
        node_name = 'Single node'

    # Execute the task only on the specified node.
    if mpicomm is None or mpicomm.rank == rank:
        logger.debug('{}: executing {}'.format(node_name, task))
        result = task(*args, **kwargs)

    # Broadcast the result if required.
    if mpicomm is not None:
        if broadcast_result is True:
            logger.debug('{}: waiting for broadcast of {}'.format(node_name, task))
            result = mpicomm.bcast(result, root=rank)
        elif sync_nodes is True:
            logger.debug('{}: waiting for barrier after {}'.format(node_name, task))
            mpicomm.barrier()

    # Return result.
    return result

def on_single_node(rank, broadcast_result=False, sync_nodes=False):
    """A decorator version of run_single_node.

    Decorates a function to be always executed with run_single_node.

    Parameters
    ----------
    rank : int
        The rank of the MPI communicator that must execute the task.
    broadcast_result : bool, optional
        If True the result is broadcasted to all nodes. If False,
        only the node executing the function will receive its return
        value, and all other nodes will receive None (default is False).
    sync_nodes : bool, optional
        If True, the nodes will be synchronized at the end of the
        execution (i.e. the task will be blocking) even if the
        result is not broadcasted (default is False).

    See Also
    --------
    run_single_node

    Examples
    --------
    >>> @on_single_node(rank=0, broadcast_result=True)
    ... def add(a, b):
    ...     return a + b
    >>> add(3, 4)
    7

    """
    def _on_single_node(task):
        @wraps_py2(task)
        def _wrapper(*args, **kwargs):
            kwargs['broadcast_result'] = broadcast_result
            kwargs['sync_nodes'] = sync_nodes
            return run_single_node(rank, task, *args, **kwargs)
        return _wrapper
    return _on_single_node

def distribute(task, distributed_args, *other_args, **kwargs):
    """Map the task on a sequence of arguments to be executed on different nodes.

    If MPI is not activated, this simply runs serially on this node. The
    algorithm guarantees that each node will be assigned to the same job_id
    (i.e. the index of the argument in distributed_args) every time.

    Parameters
    ----------
    task : callable
        The task to be distributed among nodes. The task will be called as
        task(distributed_args[job_id], *other_args, **kwargs), so the parameter
        to be distributed must the the first one.
    distributed_args : iterable
        The sequence of the parameters to distribute among nodes.
    send_results_to : int or 'all', optional
        If the string 'all', the result will be sent to all nodes. If an
        int, the result will be send only to the node with rank send_results_to.
        The return value of distribute depends on the value of this parameter
        (default is None).
    sync_nodes : bool, optional
        If True, the nodes will be synchronized at the end of the
        execution (i.e. the task will be blocking) even if the
        result is not shared (default is False).

    Other Parameters
    ----------------
    *other_args
        Other parameters to pass to task beside the assigned distributed
        parameters.
    **kwargs
        Keyword arguments to pass to task beside the assigned distributed
        parameters.

    Returns
    -------
    all_results : list
        All the return values for all the arguments if the results where sent
        to the node, or only the return values of the arguments processed by
        this node otherwise.
    arg_indices : list of int, optional
        This is returned as part of a tuple (all_results, job_indices) only
        if send_results_to is set to an int or None. In this case all_results[i]
        is the return value of task(all_args[arg_indices[i]]).

    Examples
    --------
    >>> def square(x):
    ...     return x**2
    >>> distribute(square, [1, 2, 3, 4], send_results_to='all')
    [1, 4, 9, 16]

    When send_results_to is not set to `all`, the return value include also
    the indices of the arguments associated to the result.

    >>> distribute(square, [1, 2, 3, 4], send_results_to=0)
    ([1, 4, 9, 16], [0, 1, 2, 3])

    """
    send_results_to = kwargs.pop('send_results_to', None)
    sync_nodes = kwargs.pop('sync_nodes', False)
    mpicomm = get_mpicomm()
    n_jobs = len(distributed_args)

    # If MPI is not activated, just run serially.
    if mpicomm is None:
        logger.debug('Running {} serially.'.format(task.__name__))
        all_results = [task(job_args, *other_args, **kwargs) for job_args in distributed_args]
        if send_results_to == 'all':
            return all_results
        else:
            return all_results, list(range(n_jobs))

    node_job_ids = range(mpicomm.rank, n_jobs, mpicomm.size)

    # Compute all the results assigned to this node.
    results = []
    node_name = 'Node {}/{}'.format(mpicomm.rank+1, mpicomm.size)
    for job_id in node_job_ids:
        distributed_arg = distributed_args[job_id]
        logger.debug('{}: execute {}({})'.format(node_name, task.__name__, distributed_arg))
        results.append(task(distributed_arg, *other_args, **kwargs))

    # Share result as specified.
    if send_results_to == 'all':
        logger.debug('{}: allgather results of {}'.format(node_name, task.__name__))
        all_results = mpicomm.allgather(results)
    elif isinstance(send_results_to, int):
        logger.debug('{}: gather results of {}'.format(node_name, task.__name__))
        all_results = mpicomm.gather(results, root=send_results_to)

        # If this is not the receiving node, we can safely return.
        if mpicomm.rank != send_results_to:
            return results, list(node_job_ids)
    else:
        assert send_results_to is None  # Safety check.
        if sync_nodes is True:
            logger.debug('{}: waiting for barrier after {}'.format(node_name, task.__name__))
            mpicomm.barrier()
        return results, list(node_job_ids)

    # all_results is a list of list of results. The internal lists of
    # results are ordered by rank. We need to reorder the results as a
    # flat list or results ordered by job_id.

    # job_indices[job_id] is the tuple of indices (rank, i). The result
    # of job_id is stored in all_results[rank][i].
    job_indices = []
    max_jobs_per_node = max([len(r) for r in all_results])
    for i in range(max_jobs_per_node):
        for rank in range(mpicomm.size):
            # Not all nodes have executed max_jobs_per_node tasks.
            if len(all_results[rank]) > i:
                job_indices.append((rank, i))

    # Reorder the results.
    all_results = [all_results[rank][i] for rank, i in job_indices]

    # Return result.
    if send_results_to == 'all':
        return all_results
    else:
        return all_results, list(range(n_jobs))

@contextmanager
def delay_termination():
    """Context manager to delay handling of termination signals.

    This allows to avoid interrupting tasks such as writing to the file
    system, which could result in the corruption of the file.

    """
    signals_to_catch = [signal.SIGINT, signal.SIGTERM, signal.SIGABRT]
    old_handlers = {signum: signal.getsignal(signum) for signum in signals_to_catch}
    signals_received = {signum: None for signum in signals_to_catch}

    def delay_handler(signum, frame):
        signals_received[signum] = (signum, frame)

    # Set handlers fot delay
    for signum in signals_to_catch:
        signal.signal(signum, delay_handler)

    yield  # Resume program

    # Restore old handlers
    for signum, handler in old_handlers.items():
        signal.signal(signum, handler)

    # Fire delayed signals
    for signum, s in signals_received.items():
        if s is not None:
            old_handlers[signum](*s)

def delayed_termination(func):
    """Decorator that runs the function with delay_termination()."""
    @wraps_py2(func)
    def _delayed_termination(*args, **kwargs):
        with delay_termination():
            return func(*args, **kwargs)
    return _delayed_termination

# ==============================================================================
# MAIN AND TESTS
# ==============================================================================

if __name__ == "__main__":
    import doctest
    doctest.testmod()
<|endoftext|>