<commit_before>import asyncio
import logging
import datetime
import sys
import copy
import time
import re
import os
import traceback
import functools
import textwrap

import discord # pip install git+https://github.com/Rapptz/discord.py@async
# pip install git+https://github.com/Julian/jsonschema

from . import utils, errors, clientextended

from .serverbotinstance import ServerBotInstance
from .messagecache import MessageCache

discord_logger = logging.getLogger('discord')
discord_logger.setLevel(logging.CRITICAL)
logger = logging.getLogger()
logger.setLevel(logging.CRITICAL)
handler = logging.FileHandler(filename='mentionbot.log', encoding='utf-8', mode='w')
handler.setFormatter(logging.Formatter('%(asctime)s:%(levelname)s:%(name)s: %(message)s'))
logger.addHandler(handler)

LOGIN_DETAILS_FILENAME = "bot_user_token" # This file is used to login. Only contains two lines. Line 1 is email, line 2 is password.

class MentionBot(clientextended.ClientExtended):
   BOTOWNER_ID = "119384097473822727"
   INITIAL_GAME_STATUS = "bot is running"
   INITIALIZING_GAME_STATUS = "bot is initializing"
   CACHE_DIRECTORY = "cache/" # This MUST end with a forward-slash. e.g. "cache/"

   def __init__(self, **kwargs):
      super(MentionBot, self).__init__(**kwargs)

      self.on_message_lock = asyncio.Lock()

      print("BOTOWNER_ID = '{}'".format(MentionBot.BOTOWNER_ID))
      print("INITIAL_GAME_STATUS = '{}'".format(MentionBot.INITIAL_GAME_STATUS))
      print("CACHE_DIRECTORY = '{}'".format(MentionBot.CACHE_DIRECTORY))

      self.botowner = None

      self._bot_instances = None
      return

   async def on_ready(self):
      try:
         await self.set_game_status(MentionBot.INITIALIZING_GAME_STATUS)
         self.botowner = self.search_for_user(MentionBot.BOTOWNER_ID)

         self._bot_instances = {}
         for server in self.servers:
            self._bot_instances[server] = await ServerBotInstance.get_instance(self, server)

         self.message_cache = await MessageCache.get_instance(self, self.CACHE_DIRECTORY)

         await self.set_game_status(MentionBot.INITIAL_GAME_STATUS)
         try:
            botowner = self.search_for_user(MentionBot.BOTOWNER_ID)
            print("Bot owner: " + self.botowner.name)
         except:
            print("Bot owner: (NOT FOUND. UNKNOWN ERROR.)")
            print("Bot owner ID: " + MentionBot.BOTOWNER_ID)
         print("Bot name: " + self.user.name)
         print("")
         self.on_message_lock.release()
         logging.info("MentionBot initialization complete.")
         print("Initialization complete.")
      except (SystemExit, KeyboardInterrupt):
         raise
      except BaseException as e:
         print(traceback.format_exc())
         sys.exit(1)
      return

   # General routine for uncaught exceptions from events (called by the API).
   # However, event handler routines should ideally implement their own.
   async def on_error(self, event, *args, **kwargs):
      buf = "MentionBot.on_error() called. \n\n"
      buf += "```\n" + traceback.format_exc() + "\n```"
      buf += "\nEVENT:\n" + str(event)
      buf += "\nARGS:\n"
      if len(args) == 0:
         buf += "(none)\n"
      else:
         for arg in args:
            buf += str(arg) + "\n"
      buf += "KWARGS:\n"
      if not kwargs:
         buf += "(none)\n"
      else:
         for (kwarg, item) in kwargs.items():
            buf += "{0}: {1}\n".format(str(kwarg), str(item))
      buf += "\n**Attempting to exit with code 1.**"
      print(buf, file=sys.stderr)
      try:
         await self.send_msg(self.botowner, buf)
      except:
         print("FAILED TO SEND BOTOWNER STACKTRACE.")
      logging.critical(buf)
      sys.exit(1)

   async def on_message(self, msg):
      await self.on_message_lock.acquire()
      try:
         await self._on_message(msg)
      finally:
         self.on_message_lock.release()
      return

   # TODO: Ensure this method actually lets in messages in a queued fashion...
   async def _on_message(self, msg):
      await self.message_cache.record_message(msg)
      if msg.author == self.user:
         return # Should no longer process own messages.

      # Common code for exception handling in _on_message().
      async def handle_general_error(e, msg, *, close_bot=True):
         # This is only for feedback. Exception will continue to propagate.
         print(traceback.format_exc())
         buf = "**EXCEPTION**"
         buf += "\n**From:** <#" + msg.channel.id + "> **in** " + msg.server.name
         buf += "\n**Command issued by:** <@" + msg.author.id + ">"
         buf += "\n**Full message:**"
         buf += "\n" + msg.content
         buf += "\n**Stack Trace:**"
         buf += "\n```" + traceback.format_exc() + "```"
         try:
            await self.send_msg(self.botowner, buf)
         except:
            print("FAILED TO SEND BOTOWNER STACKTRACE.")
         buf = "**EXCEPTION:** " + type(e).__name__
         buf += "\n" + str(e)
         buf += "\n<@" + MentionBot.BOTOWNER_ID + "> Check it out, will ya?"
         if close_bot:
            buf += "\n\n**THIS BOT WILL NOW RESTART.**"
         try:
            await self.send_msg(msg, buf)
         except:
            print("FAILED TO MESSAGE BOT TERMINATION BACK TO THE CHANNEL.")
         if close_bot:
            sys.exit(1)
         return

      try:
         text = msg.content.strip()
         if isinstance(msg.channel, discord.Channel):
            await self._bot_instances[msg.server].process_text(text, msg)
            try:
               print("msg rcv #" + msg.channel.name + ": " + utils.str_asciionly(text))
            except Exception:
               print("msg rcv (CAUGHT EXCEPTION; UNKNOWN DISPLAY ERROR)")
         else: # Assumed to be a private message.
            await self.send_msg(msg, "sry m8 im not programmed to do anything fancy with pms yet")
            try:
               print("private msg rcv from" + utils.str_asciionly(msg.author.name) + ": " + utils.str_asciionly(text))
            except Exception:
               print("private msg rcv (CAUGHT EXCEPTION; UNKNOWN DISPLAY ERROR)")

      except errors.SilentUnknownCommandError:
         print("Caught SilentUnknownCommandError.")
      except errors.UnknownCommandError:
         print("Caught UnknownCommandError.")
         await self.send_msg(msg, "Error: Unknown command.")
      except errors.InvalidCommandArgumentsError as e:
         print("Caught InvalidCommandArgumentsError.")
         if str(e) == "":
            buf = "Error: Invalid command arguments."
         else:
            buf = str(e)
         await self.send_msg(msg, buf)
      except errors.CommandPrivilegeError:
         print("Caught CommandPrivilegeError.")
         await self.send_msg(msg, "Error: Permission denied.")
      except errors.NoHelpContentExists:
         print("Caught NoHelpContentExists.")
         await self.send_msg(msg, "No help content exists.")
      except errors.OperationAborted:
         print("Caught OperationAborted.")
      except Exception as e:
         await handle_general_error(e, msg, close_bot=True)
      except (SystemExit, KeyboardInterrupt):
         raise
      except BaseException as e:
         await handle_general_error(e, msg, close_bot=True)
      return

   async def on_server_join(self, server):
      raise RuntimeError("Undefined behaviour on server join. Must restart.")

   async def on_server_remove(self, server):
      raise RuntimeError("Undefined behaviour on server leave. Must restart.")

   ##################
   # Other Services #
   ##################

   # Common code for handling an error.
   # Does:
   #  - Printing stacktrace to stderr
   #  - Logging the error report
   #  - Sending the error report to the bot owner
   #  - Give user feedback in the channel if a command caused the error.
   # And for clarity, this function does NOT:
   #  - Reraise errors
   #  - Exit the program
   #  - Kill bot modules
   #
   # PARAMETERS:
   #    e: Error event that caused the error.
   #    cmd_msg: If the error involved a command, this is the message object
   #             of the command that triggered this error.
   #             If this is supplied, a message will be automatically sent to
   #             the channel to inform the user of the error.
   #    extra_info: An additional message content (string) to attach to the
   #                error report. This is useful for reporting more information
   #                about an error.
   #                This only appears in the error report.
   #    final_info: A string to append to the end of the error report.
   #                This string will appear in both the error report and
   #                command feedback message.
   #                For example, "THIS BOT WILL NOW RESTART."
   #
   # RETURNS: A string which may be sent back to a channel for user feedback.
   async def report_exception(self, e, **kwargs):
      traceback_str = traceback.format_exc()
      print(traceback_str, file=sys.stderr)

      # Obtain keyword arguments
      cmd_msg = kwargs.get("cmd_msg", None) # Discord Message Object
      extra_info = str(kwargs.get("extra_info", None)) # String
      final_info = str(kwargs.get("final_info", None)) # String

      # STEP 1: Send the error report to the bot owner (via PM) and log it.

      buf = "**EXCEPTION**"
      if not cmd_msg is None:
         buf0 = textwrap.dedent("""
            **The error was triggered by a command.**
            **From:** <#{msg.channel.id}> **in** {msg.server.name}
            **Command issued by:** <@{msg.author.id}>
            **Full Message:**
            {msg.content}
            """)
         buf0 = buf0.format(msg=cmd_msg).strip()
         buf += "\n\n" + buf0

      if not extra_info is None:
         buf += "\n\n**Extra Info:**\n" + extra_info

      buf0 = textwrap.dedent("""\
         **Traceback:**
         ```
         {traceback_str}
         ```
         """)
      buf0 = buf0.format(traceback_str=traceback_str, final_info=final_info).strip()
      buf += "\n\n" + buf0

      if not final_info is None:
         buf += "\n\n" + final_info

      logging.critical(buf)
      try:
         await self.send_msg(self.botowner, buf)
      except:
         buf = "FAILED TO SEND BOTOWNER STACKTRACE."
         logging.critical(buf)
         print(buf, file=sys.stderr)

      # STEP 2: Create user feedback, and if a command message was supplied,
      #         send this message into that channel.

      buf = None
      if e is None:
         buf = "**AN UNKNOWN ERROR OCCURRED.**"
      else:
         buf = textwrap.dedent("""\
            **AN ERROR OCCURRED.**
            **EXCEPTION:** {e_name}
            {e_des}
            """)
         buf = buf.format(e_name=type(e).__name__, e_des=str(e)).strip()

      buf += "\n<@{}> Check it out, will ya?".format(MentionBot.BOTOWNER_ID)

      if not final_info is None:
         buf += "\n\n" + final_info

      if not cmd_msg is None:
         try:
            await self.send_msg(cmd_msg, buf)
         except:
            buf0 = "FAILED TO MESSAGE BOT TERMINATION BACK TO THE CHANNEL."
            logging.critical(buf0)
            print(buf0, file=sys.stderr)

      return buf

   def message_cache_read(self, server_id, ch_id):
      return self.message_cache.read_messages(server_id, ch_id)

   def message_cache_debug_str(self):
      return self.message_cache.get_debugging_info()

async def _client_login(client, token):
   await client.on_message_lock.acquire() # To be released when ready.
   # TODO: Acquire the lock elsewhere. It's a little out-of-place here...
   await client.login(token)
   await client.connect()
   return

def run():
   loop = asyncio.get_event_loop()

   # Log in to discord
   client = MentionBot()
   print("\nAttempting to log in using file '" + LOGIN_DETAILS_FILENAME + "'.")
   if not os.path.isfile(LOGIN_DETAILS_FILENAME):
      login_file = open(LOGIN_DETAILS_FILENAME, "w")
      login_file.write("TOKEN")
      login_file.close()
      print("File does not exist. Please edit the file {} with your login details.")
      sys.exit()
   login_file = open(LOGIN_DETAILS_FILENAME, "r")
   email = "(No email. This variable is depreciated.)"
   password = "(No password. This variable is depreciated.)"
   bot_user_token = login_file.readline().strip()
   login_file.close()
   print("Token found.")
   print("Logging in...") # print("Logging in...", end="")

   try:
      loop.run_until_complete(_client_login(client, bot_user_token))
   except Exception as e:
      loop.run_until_complete(client.logout())
      print("Error launching client!")
      print(traceback.format_exc())
   finally:
      try:
         loop.close()
      except:
         print(traceback.format_exc())
   sys.exit(1) # Should only return on error.

if __name__ == '__main__':
   run()
<commit_msg>Replaced some old redundant code with report_exception().<commit_after>import asyncio
import logging
import datetime
import sys
import copy
import time
import re
import os
import traceback
import functools
import textwrap

import discord # pip install git+https://github.com/Rapptz/discord.py@async
# pip install git+https://github.com/Julian/jsonschema

from . import utils, errors, clientextended

from .serverbotinstance import ServerBotInstance
from .messagecache import MessageCache

discord_logger = logging.getLogger('discord')
discord_logger.setLevel(logging.CRITICAL)
logger = logging.getLogger()
logger.setLevel(logging.CRITICAL)
handler = logging.FileHandler(filename='mentionbot.log', encoding='utf-8', mode='w')
handler.setFormatter(logging.Formatter('%(asctime)s:%(levelname)s:%(name)s: %(message)s'))
logger.addHandler(handler)

LOGIN_DETAILS_FILENAME = "bot_user_token" # This file is used to login. Only contains two lines. Line 1 is email, line 2 is password.

class MentionBot(clientextended.ClientExtended):
   BOTOWNER_ID = "119384097473822727"
   INITIAL_GAME_STATUS = "bot is running"
   INITIALIZING_GAME_STATUS = "bot is initializing"
   CACHE_DIRECTORY = "cache/" # This MUST end with a forward-slash. e.g. "cache/"

   def __init__(self, **kwargs):
      super(MentionBot, self).__init__(**kwargs)

      self.on_message_lock = asyncio.Lock()

      print("BOTOWNER_ID = '{}'".format(MentionBot.BOTOWNER_ID))
      print("INITIAL_GAME_STATUS = '{}'".format(MentionBot.INITIAL_GAME_STATUS))
      print("CACHE_DIRECTORY = '{}'".format(MentionBot.CACHE_DIRECTORY))

      self.botowner = None

      self._bot_instances = None
      return

   async def on_ready(self):
      try:
         await self.set_game_status(MentionBot.INITIALIZING_GAME_STATUS)
         self.botowner = self.search_for_user(MentionBot.BOTOWNER_ID)

         self._bot_instances = {}
         for server in self.servers:
            self._bot_instances[server] = await ServerBotInstance.get_instance(self, server)

         self.message_cache = await MessageCache.get_instance(self, self.CACHE_DIRECTORY)

         await self.set_game_status(MentionBot.INITIAL_GAME_STATUS)
         try:
            botowner = self.search_for_user(MentionBot.BOTOWNER_ID)
            print("Bot owner: " + self.botowner.name)
         except:
            print("Bot owner: (NOT FOUND. UNKNOWN ERROR.)")
            print("Bot owner ID: " + MentionBot.BOTOWNER_ID)
         print("Bot name: " + self.user.name)
         print("")
         self.on_message_lock.release()
         logging.info("MentionBot initialization complete.")
         print("Initialization complete.")
      except (SystemExit, KeyboardInterrupt):
         raise
      except BaseException as e:
         print(traceback.format_exc())
         sys.exit(1)
      return

   # General routine for uncaught exceptions from events (called by the API).
   # However, event handler routines should ideally implement their own.
   async def on_error(self, event, *args, **kwargs):
      try:
         buf_hb = "MentionBot.on_error()."
         buf_ei = "EVENT:\n" + str(event)
         buf_ei += "\nARGS:\n"
         if len(args) == 0:
            buf_ei += "(none)\n"
         else:
            for arg in args:
               buf_ei += str(arg) + "\n"
         buf_ei += "KWARGS:\n"
         if not kwargs:
            buf_ei += "(none)\n"
         else:
            for (kwarg, item) in kwargs.items():
               buf_ei += "{0}: {1}\n".format(str(kwarg), str(item))
         buf_fi = "Attempting to exit with code 1."
         await self.report_exception(event, handled_by=buf_hb, extra_info=buf_ei, final_info=buf_fi)
      except:
         buf = "ERROR WITHIN on_error()!\n" + traceback.format_exc()
         print(buf, file=sys.stderr)
         logging.critical(buf)
      sys.exit(1)

   async def on_message(self, msg):
      await self.on_message_lock.acquire()
      try:
         await self._on_message(msg)
      finally:
         self.on_message_lock.release()
      return

   # TODO: Ensure this method actually lets in messages in a queued fashion...
   async def _on_message(self, msg):
      await self.message_cache.record_message(msg)
      if msg.author == self.user:
         return # Should no longer process own messages.

      # Common code for exception handling in _on_message().
      async def handle_general_error(e, msg, *, close_bot=True):
         buf_hb = "MentionBot._on_message()."
         buf_fi = None
         if close_bot:
            buf_fi = "**THIS BOT WILL NOW RESTART.**"
         await self.report_exception(e, cmd_msg=msg, handled_by=buf_hb, final_info=buf_fi)
         if close_bot:
            sys.exit(1)
         return

      try:
         text = msg.content.strip()
         if isinstance(msg.channel, discord.Channel):
            await self._bot_instances[msg.server].process_text(text, msg)
            try:
               print("msg rcv #" + msg.channel.name + ": " + utils.str_asciionly(text))
            except Exception:
               print("msg rcv (CAUGHT EXCEPTION; UNKNOWN DISPLAY ERROR)")
         else: # Assumed to be a private message.
            await self.send_msg(msg, "sry m8 im not programmed to do anything fancy with pms yet")
            try:
               print("private msg rcv from" + utils.str_asciionly(msg.author.name) + ": " + utils.str_asciionly(text))
            except Exception:
               print("private msg rcv (CAUGHT EXCEPTION; UNKNOWN DISPLAY ERROR)")

      except errors.SilentUnknownCommandError:
         print("Caught SilentUnknownCommandError.")
      except errors.UnknownCommandError:
         print("Caught UnknownCommandError.")
         await self.send_msg(msg, "Error: Unknown command.")
      except errors.InvalidCommandArgumentsError as e:
         print("Caught InvalidCommandArgumentsError.")
         if str(e) == "":
            buf = "Error: Invalid command arguments."
         else:
            buf = str(e)
         await self.send_msg(msg, buf)
      except errors.CommandPrivilegeError:
         print("Caught CommandPrivilegeError.")
         await self.send_msg(msg, "Error: Permission denied.")
      except errors.NoHelpContentExists:
         print("Caught NoHelpContentExists.")
         await self.send_msg(msg, "No help content exists.")
      except errors.OperationAborted:
         print("Caught OperationAborted.")
      except Exception as e:
         await handle_general_error(e, msg, close_bot=True)
      except (SystemExit, KeyboardInterrupt):
         raise
      except BaseException as e:
         await handle_general_error(e, msg, close_bot=True)
      return

   async def on_server_join(self, server):
      raise RuntimeError("Undefined behaviour on server join. Must restart.")

   async def on_server_remove(self, server):
      raise RuntimeError("Undefined behaviour on server leave. Must restart.")

   ##################
   # Other Services #
   ##################

   # Common code for handling an error.
   # Does:
   #  - Printing stacktrace to stderr
   #  - Logging the error report
   #  - Sending the error report to the bot owner
   #  - Give user feedback in the channel if a command caused the error.
   # And for clarity, this function does NOT:
   #  - Reraise errors
   #  - Exit the program
   #  - Kill bot modules
   #
   # PARAMETERS:
   #    e: Error event that caused the error.
   #    cmd_msg: If the error involved a command, this is the message object
   #             of the command that triggered this error.
   #             If this is supplied, a message will be automatically sent to
   #             the channel to inform the user of the error.
   #    extra_info: An additional message content (string) to attach to the
   #                error report. This is useful for reporting more information
   #                about an error.
   #                This only appears in the error report.
   #    final_info: A string to append to the end of the error report.
   #                This string will appear in both the error report and
   #                command feedback message.
   #                For example, "THIS BOT WILL NOW RESTART."
   #
   # RETURNS: A string which may be sent back to a channel for user feedback.
   async def report_exception(self, e, **kwargs):
      try:
         traceback_str = traceback.format_exc()
         print(traceback_str, file=sys.stderr)

         # Obtain keyword arguments
         cmd_msg = kwargs.get("cmd_msg", None) # Discord Message Object
         handled_by = str(kwargs.get("handled_by", "")).strip()
         extra_info = str(kwargs.get("extra_info", "")).strip()
         final_info = str(kwargs.get("final_info", "")).strip()

         # STEP 1: Send the error report to the bot owner (via PM) and log it.

         buf = "**EXCEPTION**"

         if len(handled_by) != 0:
            buf += "\n\n**This error is being handled by:**\n" + handled_by

         if not cmd_msg is None:
            buf0 = textwrap.dedent("""
               **The error was triggered by a command.**
               **From:** <#{msg.channel.id}> **in** {msg.server.name}
               **Command issued by:** <@{msg.author.id}>
               **Full Message:**
               {msg.content}
               """)
            buf0 = buf0.format(msg=cmd_msg).strip()
            buf += "\n\n" + buf0

         if len(extra_info) != 0:
            buf += "\n\n**Extra Info:**\n" + extra_info

         buf0 = textwrap.dedent("""\
            **Traceback:**
            ```
            {traceback_str}
            ```
            """)
         buf0 = buf0.format(traceback_str=traceback_str).strip()
         buf += "\n\n" + buf0

         if len(final_info) != 0:
            buf += "\n\n" + final_info

         logging.critical(buf)
         try:
            await self.send_msg(self.botowner, buf)
         except:
            buf = "FAILED TO SEND BOTOWNER STACKTRACE."
            logging.critical(buf)
            print(buf, file=sys.stderr)

         # STEP 2: Create user feedback, and if a command message was supplied,
         #         send this message into that channel.

         buf = None
         if e is None:
            buf = "**AN UNKNOWN ERROR OCCURRED.**"
         else:
            buf = textwrap.dedent("""\
               **AN ERROR OCCURRED.**
               **EXCEPTION:** {e_name}
               {e_des}
               """)
            buf = buf.format(e_name=type(e).__name__, e_des=str(e)).strip()

         buf += "\n<@{}> Check it out, will ya?".format(MentionBot.BOTOWNER_ID)

         if len(final_info) != 0:
            buf += "\n\n" + final_info

         if not cmd_msg is None:
            try:
               await self.send_msg(cmd_msg, buf)
            except:
               buf0 = "FAILED TO MESSAGE BOT TERMINATION BACK TO THE CHANNEL."
               logging.critical(buf0)
               print(buf0, file=sys.stderr)
      except:
         buf = "ERROR WITHIN report_exception()!\n" + traceback.format_exc()
         buf += "\nThis error is completely unrecoverable."
         buf += "\nExiting with code 1."
         print(buf, file=sys.stderr)
         logging.critical(buf)
         sys.exit(1)
      return buf

   def message_cache_read(self, server_id, ch_id):
      return self.message_cache.read_messages(server_id, ch_id)

   def message_cache_debug_str(self):
      return self.message_cache.get_debugging_info()

async def _client_login(client, token):
   await client.on_message_lock.acquire() # To be released when ready.
   # TODO: Acquire the lock elsewhere. It's a little out-of-place here...
   await client.login(token)
   await client.connect()
   return

def run():
   loop = asyncio.get_event_loop()

   # Log in to discord
   client = MentionBot()
   print("\nAttempting to log in using file '" + LOGIN_DETAILS_FILENAME + "'.")
   if not os.path.isfile(LOGIN_DETAILS_FILENAME):
      login_file = open(LOGIN_DETAILS_FILENAME, "w")
      login_file.write("TOKEN")
      login_file.close()
      print("File does not exist. Please edit the file {} with your login details.")
      sys.exit()
   login_file = open(LOGIN_DETAILS_FILENAME, "r")
   email = "(No email. This variable is depreciated.)"
   password = "(No password. This variable is depreciated.)"
   bot_user_token = login_file.readline().strip()
   login_file.close()
   print("Token found.")
   print("Logging in...") # print("Logging in...", end="")

   try:
      loop.run_until_complete(_client_login(client, bot_user_token))
   except Exception as e:
      loop.run_until_complete(client.logout())
      print("Error launching client!")
      print(traceback.format_exc())
   finally:
      try:
         loop.close()
      except:
         print(traceback.format_exc())
   sys.exit(1) # Should only return on error.

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


<commit_before>
import dbus.glib
import gobject
import sys

from account import read_account, connect

from telepathy.client.channel import Channel
from telepathy.constants import (
    CONNECTION_HANDLE_TYPE_NONE, CONNECTION_HANDLE_TYPE_CONTACT,
    CONNECTION_STATUS_CONNECTED, CONNECTION_STATUS_DISCONNECTED,
    MEDIA_STREAM_TYPE_AUDIO, MEDIA_STREAM_TYPE_VIDEO)
from telepathy.interfaces import (
    CHANNEL_INTERFACE, CHANNEL_INTERFACE_GROUP, CHANNEL_TYPE_STREAMED_MEDIA,
    CONN_INTERFACE, CONN_INTERFACE_CAPABILITIES)

def get_stream_engine():
    bus = dbus.Bus()
    return bus.get_object(
        'org.freedesktop.Telepathy.StreamEngine',
        '/org/freedesktop/Telepathy/StreamEngine')

class Call:
    def __init__(self, conn):
        self.conn = conn
        self.channel = None

        conn[CONN_INTERFACE].connect_to_signal('StatusChanged',
            self.status_changed_cb)
        conn[CONN_INTERFACE].connect_to_signal('NewChannel',
            self.new_channel_cb)

    def run(self):
        self.loop = gobject.MainLoop()

        try:
            self.loop.run()
        except KeyboardInterrupt:
            print "killed"

            if self.channel:
                print "closing channel"
                self.channel[CHANNEL_INTERFACE].Close()

    def quit(self):
        if self.loop:
            self.loop.quit()
            self.loop = None

    def status_changed_cb(self, state, reason):
        if state == CONNECTION_STATUS_DISCONNECTED:
            print 'connection closed'
            self.quit()

    def request_channel_error_cb(self, exception):
        print 'error:', exception
        self.quit()

    def new_channel_cb(self, object_path, channel_type, handle_type, handle,
            suppress_handler):
        if channel_type != CHANNEL_TYPE_STREAMED_MEDIA:
            return

        self.chan_handle_type = handle_type
        self.chan_handle = handle

        print "new streamed media channel"
        Channel(self.conn._dbus_object._named_service, object_path,
                ready_handler=self.channel_ready_cb)

    def channel_ready_cb(self, channel):
        print "channel ready"
        channel[CHANNEL_INTERFACE].connect_to_signal('Closed', self.closed_cb)
        channel[CHANNEL_INTERFACE_GROUP].connect_to_signal('MembersChanged',
            self.members_changed_cb)

        stream_engine = get_stream_engine()
        handler = dbus.Interface(stream_engine,
            'org.freedesktop.Telepathy.ChannelHandler')
        handler.HandleChannel(
            self.conn._dbus_object._named_service,
            self.conn._dbus_object._object_path,
            CHANNEL_TYPE_STREAMED_MEDIA,
            channel._dbus_object._object_path,
            self.chan_handle_type,
            self.chan_handle)

        self.channel = channel

    def closed_cb(self):
        print "channel closed"
        self.quit()

    def members_changed_cb(self, message, added, removed, local_pending,
            remote_pending, actor, reason):
        print 'MembersChanged', (
            added, removed, local_pending, remote_pending, actor, reason)

class OutgoingCall(Call):
    def __init__(self, conn, contact):
        Call.__init__(self, conn)
        self.contact = contact

    def status_changed_cb (self, state, reason):
        if state == CONNECTION_STATUS_CONNECTED:
            handle = self.conn[CONN_INTERFACE].RequestHandles(
                CONNECTION_HANDLE_TYPE_CONTACT, [self.contact])[0]
            self.handle = handle

            print 'got handle %d for %s' % (handle, self.contact)

            # hack
            import time
            time.sleep(5)

            self.conn[CONN_INTERFACE].RequestChannel(
                CHANNEL_TYPE_STREAMED_MEDIA, CONNECTION_HANDLE_TYPE_NONE,
                0, True,
                reply_handler=lambda *stuff: None,
                error_handler=self.request_channel_error_cb)

        Call.status_changed_cb(self, state, reason)

    def channel_ready_cb(self, channel):
        Call.channel_ready_cb(self, channel)

        channel[CHANNEL_INTERFACE_GROUP].AddMembers([self.handle], "")

        print "requesting audio/video streams"

        try:
            channel[CHANNEL_TYPE_STREAMED_MEDIA].RequestStreams(
                self.handle,
                [MEDIA_STREAM_TYPE_AUDIO, MEDIA_STREAM_TYPE_VIDEO]);
        except dbus.DBusException, e:
            print "failed:", e
            print "requesting audio stream"

            try:
                channel[CHANNEL_TYPE_STREAMED_MEDIA].RequestStreams(
                    self.handle, [MEDIA_STREAM_TYPE_AUDIO]);
            except dbus.DBusException, e:
                print "failed:", e
                print "giving up"
                self.quit()

class IncomingCall(Call):
    def status_changed_cb (self, state, reason):
        if state == CONNECTION_STATUS_CONNECTED:
            print "waiting for incoming call"

            self.conn._valid_interfaces.add(CONN_INTERFACE_CAPABILITIES)
            self.conn[CONN_INTERFACE_CAPABILITIES].AdvertiseCapabilities(
                [(CHANNEL_TYPE_STREAMED_MEDIA, 3)], [])

        Call.status_changed_cb(self, state, reason)

    def channel_ready_cb(self, channel):
        Call.channel_ready_cb(self, channel)

        print "accepting incoming call"
        pending = channel[CHANNEL_INTERFACE_GROUP].GetLocalPendingMembers()
        channel[CHANNEL_INTERFACE_GROUP].AddMembers(pending, "")

if __name__ == '__main__':
    assert len(sys.argv) in (2, 3)
    account_file = sys.argv[1]

    manager, protocol, account = read_account(account_file)
    conn = connect(manager, protocol, account)

    if len(sys.argv) > 2:
        contact = sys.argv[2]
        call = OutgoingCall(conn, sys.argv[2])
    else:
        call = IncomingCall(conn)

    print "connecting"
    conn[CONN_INTERFACE].Connect()
    call.run()

    try:
        print "disconnecting"
        conn[CONN_INTERFACE].Disconnect()
    except dbus.DBusException:
        pass

<commit_msg>examples/call.py: use connection_from_file<commit_after>
import dbus.glib
import gobject
import sys

from account import connection_from_file

from telepathy.client.channel import Channel
from telepathy.constants import (
    CONNECTION_HANDLE_TYPE_NONE, CONNECTION_HANDLE_TYPE_CONTACT,
    CONNECTION_STATUS_CONNECTED, CONNECTION_STATUS_DISCONNECTED,
    MEDIA_STREAM_TYPE_AUDIO, MEDIA_STREAM_TYPE_VIDEO)
from telepathy.interfaces import (
    CHANNEL_INTERFACE, CHANNEL_INTERFACE_GROUP, CHANNEL_TYPE_STREAMED_MEDIA,
    CONN_INTERFACE, CONN_INTERFACE_CAPABILITIES)

def get_stream_engine():
    bus = dbus.Bus()
    return bus.get_object(
        'org.freedesktop.Telepathy.StreamEngine',
        '/org/freedesktop/Telepathy/StreamEngine')

class Call:
    def __init__(self, conn):
        self.conn = conn
        self.channel = None

        conn[CONN_INTERFACE].connect_to_signal('StatusChanged',
            self.status_changed_cb)
        conn[CONN_INTERFACE].connect_to_signal('NewChannel',
            self.new_channel_cb)

    def run(self):
        self.loop = gobject.MainLoop()

        try:
            self.loop.run()
        except KeyboardInterrupt:
            print "killed"

            if self.channel:
                print "closing channel"
                self.channel[CHANNEL_INTERFACE].Close()

    def quit(self):
        if self.loop:
            self.loop.quit()
            self.loop = None

    def status_changed_cb(self, state, reason):
        if state == CONNECTION_STATUS_DISCONNECTED:
            print 'connection closed'
            self.quit()

    def request_channel_error_cb(self, exception):
        print 'error:', exception
        self.quit()

    def new_channel_cb(self, object_path, channel_type, handle_type, handle,
            suppress_handler):
        if channel_type != CHANNEL_TYPE_STREAMED_MEDIA:
            return

        self.chan_handle_type = handle_type
        self.chan_handle = handle

        print "new streamed media channel"
        Channel(self.conn._dbus_object._named_service, object_path,
                ready_handler=self.channel_ready_cb)

    def channel_ready_cb(self, channel):
        print "channel ready"
        channel[CHANNEL_INTERFACE].connect_to_signal('Closed', self.closed_cb)
        channel[CHANNEL_INTERFACE_GROUP].connect_to_signal('MembersChanged',
            self.members_changed_cb)

        stream_engine = get_stream_engine()
        handler = dbus.Interface(stream_engine,
            'org.freedesktop.Telepathy.ChannelHandler')
        handler.HandleChannel(
            self.conn._dbus_object._named_service,
            self.conn._dbus_object._object_path,
            CHANNEL_TYPE_STREAMED_MEDIA,
            channel._dbus_object._object_path,
            self.chan_handle_type,
            self.chan_handle)

        self.channel = channel

    def closed_cb(self):
        print "channel closed"
        self.quit()

    def members_changed_cb(self, message, added, removed, local_pending,
            remote_pending, actor, reason):
        print 'MembersChanged', (
            added, removed, local_pending, remote_pending, actor, reason)

class OutgoingCall(Call):
    def __init__(self, conn, contact):
        Call.__init__(self, conn)
        self.contact = contact

    def status_changed_cb (self, state, reason):
        if state == CONNECTION_STATUS_CONNECTED:
            handle = self.conn[CONN_INTERFACE].RequestHandles(
                CONNECTION_HANDLE_TYPE_CONTACT, [self.contact])[0]
            self.handle = handle

            print 'got handle %d for %s' % (handle, self.contact)

            # hack
            import time
            time.sleep(5)

            self.conn[CONN_INTERFACE].RequestChannel(
                CHANNEL_TYPE_STREAMED_MEDIA, CONNECTION_HANDLE_TYPE_NONE,
                0, True,
                reply_handler=lambda *stuff: None,
                error_handler=self.request_channel_error_cb)

        Call.status_changed_cb(self, state, reason)

    def channel_ready_cb(self, channel):
        Call.channel_ready_cb(self, channel)

        channel[CHANNEL_INTERFACE_GROUP].AddMembers([self.handle], "")

        print "requesting audio/video streams"

        try:
            channel[CHANNEL_TYPE_STREAMED_MEDIA].RequestStreams(
                self.handle,
                [MEDIA_STREAM_TYPE_AUDIO, MEDIA_STREAM_TYPE_VIDEO]);
        except dbus.DBusException, e:
            print "failed:", e
            print "requesting audio stream"

            try:
                channel[CHANNEL_TYPE_STREAMED_MEDIA].RequestStreams(
                    self.handle, [MEDIA_STREAM_TYPE_AUDIO]);
            except dbus.DBusException, e:
                print "failed:", e
                print "giving up"
                self.quit()

class IncomingCall(Call):
    def status_changed_cb (self, state, reason):
        if state == CONNECTION_STATUS_CONNECTED:
            print "waiting for incoming call"

            self.conn._valid_interfaces.add(CONN_INTERFACE_CAPABILITIES)
            self.conn[CONN_INTERFACE_CAPABILITIES].AdvertiseCapabilities(
                [(CHANNEL_TYPE_STREAMED_MEDIA, 3)], [])

        Call.status_changed_cb(self, state, reason)

    def channel_ready_cb(self, channel):
        Call.channel_ready_cb(self, channel)

        print "accepting incoming call"
        pending = channel[CHANNEL_INTERFACE_GROUP].GetLocalPendingMembers()
        channel[CHANNEL_INTERFACE_GROUP].AddMembers(pending, "")

if __name__ == '__main__':
    assert len(sys.argv) in (2, 3)
    conn = connection_from_file(sys.argv[1])

    if len(sys.argv) > 2:
        contact = sys.argv[2]
        call = OutgoingCall(conn, sys.argv[2])
    else:
        call = IncomingCall(conn)

    print "connecting"
    conn[CONN_INTERFACE].Connect()
    call.run()

    try:
        print "disconnecting"
        conn[CONN_INTERFACE].Disconnect()
    except dbus.DBusException:
        pass

<|endoftext|>


<commit_before>from __future__ import absolute_import, division, print_function

import os
import sys
import shutil
import tarfile
import zipfile
import subprocess
from os.path import dirname, getmtime, getsize, isdir, join

from conda.utils import md5_file

from conda_build import external

# Backwards compatibility import. Do not remove.
from conda.install import rm_rf

def copy_into(src, dst):
    "Copy all the files and directories in src to the directory dst"

    tocopy = os.listdir(src)
    for afile in tocopy:
        srcname = os.path.join(src, afile)
        dstname = os.path.join(dst, afile)

        if os.path.isdir(srcname):
            shutil.copytree(srcname, dstname)
        else:
            shutil.copy2(srcname, dstname)

def relative(f, d='lib'):
    assert not f.startswith('/'), f
    assert not d.startswith('/'), d
    d = d.strip('/').split('/')
    if d == ['.']:
        d = []
    f = dirname(f).split('/')
    if f == ['']:
        f = []
    while d and f and d[0] == f[0]:
        d.pop(0)
        f.pop(0)
    return '/'.join(((['..'] * len(f)) if f else ['.']) + d)

def _check_call(args, **kwargs):
    try:
        subprocess.check_call(args, **kwargs)
    except subprocess.CalledProcessError:
        sys.exit('Command failed: %s' % ' '.join(args))

def tar_xf(tarball, dir_path, mode='r:*'):
    if tarball.endswith('.tar.xz'):
        unxz = external.find_executable('unxz')
        if not unxz:
            sys.exit("""\
unxz is required to unarchive .xz source files.
""")

        subprocess.check_call([unxz, '-f', '-k', tarball])
        tarball = tarball[:-3]
    t = tarfile.open(tarball, mode)
    t.extractall(path=dir_path)
    t.close()

def unzip(zip_path, dir_path):
    z = zipfile.ZipFile(zip_path)
    for name in z.namelist():
        if name.endswith('/'):
            continue
        path = join(dir_path, *name.split('/'))
        dp = dirname(path)
        if not isdir(dp):
            os.makedirs(dp)
        with open(path, 'wb') as fo:
            fo.write(z.read(name))
    z.close()

def file_info(path):
    return {'size': getsize(path),
            'md5': md5_file(path),
            'mtime': getmtime(path)}
<commit_msg>Python tarfile supports xz in Python 3<commit_after>from __future__ import absolute_import, division, print_function

import os
import sys
import shutil
import tarfile
import zipfile
import subprocess
from os.path import dirname, getmtime, getsize, isdir, join

from conda.utils import md5_file
from conda.compat import PY3

from conda_build import external

# Backwards compatibility import. Do not remove.
from conda.install import rm_rf

def copy_into(src, dst):
    "Copy all the files and directories in src to the directory dst"

    tocopy = os.listdir(src)
    for afile in tocopy:
        srcname = os.path.join(src, afile)
        dstname = os.path.join(dst, afile)

        if os.path.isdir(srcname):
            shutil.copytree(srcname, dstname)
        else:
            shutil.copy2(srcname, dstname)

def relative(f, d='lib'):
    assert not f.startswith('/'), f
    assert not d.startswith('/'), d
    d = d.strip('/').split('/')
    if d == ['.']:
        d = []
    f = dirname(f).split('/')
    if f == ['']:
        f = []
    while d and f and d[0] == f[0]:
        d.pop(0)
        f.pop(0)
    return '/'.join(((['..'] * len(f)) if f else ['.']) + d)

def _check_call(args, **kwargs):
    try:
        subprocess.check_call(args, **kwargs)
    except subprocess.CalledProcessError:
        sys.exit('Command failed: %s' % ' '.join(args))

def tar_xf(tarball, dir_path, mode='r:*'):
    if not PY3 and tarball.endswith('.tar.xz'):
        unxz = external.find_executable('unxz')
        if not unxz:
            sys.exit("""\
unxz is required to unarchive .xz source files.
""")

        subprocess.check_call([unxz, '-f', '-k', tarball])
        tarball = tarball[:-3]
    t = tarfile.open(tarball, mode)
    t.extractall(path=dir_path)
    t.close()

def unzip(zip_path, dir_path):
    z = zipfile.ZipFile(zip_path)
    for name in z.namelist():
        if name.endswith('/'):
            continue
        path = join(dir_path, *name.split('/'))
        dp = dirname(path)
        if not isdir(dp):
            os.makedirs(dp)
        with open(path, 'wb') as fo:
            fo.write(z.read(name))
    z.close()

def file_info(path):
    return {'size': getsize(path),
            'md5': md5_file(path),
            'mtime': getmtime(path)}
<|endoftext|>


<commit_before>"""
This test file is meant for developing purposes. Providing an easy method to
test the functioning of PASTA during development.
"""

from pasta import *
from pasta.read.menydata import MenyData

fname = 'data/MenyanthesTest.men'
meny = MenyData(fname)

# Round the time of observations to days (using ceil, so the observation is never before the stress)
meny.H[0].series.index = meny.H[0].series.index.ceil('d')

# Create the time series model
ml = Model(meny.H[0].series)

# Add drainage level
d = Constant(value=meny.H[0].series.min())
ml.add_tseries(d)

# Add precipitation
IN = next(x for x in meny.IN if x.name == 'Neerslag (Hoogerheide)')
# round to days (precipitation is measured at 9:00)
IN.series = IN.series.resample('d').bfill()
IN.name = IN.name.replace('(','')
IN.name = IN.name.replace(')','')
IN.name = IN.name.replace(' ','_')
ts = Tseries(IN.series, Gamma, IN.name)
ml.add_tseries(ts)

# Add well extraction at Ossendrecht
IN = next(x for x in meny.IN if x.name == 'Onttrekking (Ossendrecht)')
# extraction amound counts for the previeous month
IN.series = IN.series.resample('d').bfill()
IN.name = IN.name.replace('(','')
IN.name = IN.name.replace(')','')
IN.name = IN.name.replace(' ','_')
ts = Tseries(IN.series, Gamma, IN.name)
ml.add_tseries(ts)

# Add well extraction at Huijbergen
IN = next(x for x in meny.IN if x.name == 'Onttrekking (Huijbergen)')
# extraction amound counts for the previeous month
IN.series = IN.series.resample('d').bfill()
IN.name = IN.name.replace('(','')
IN.name = IN.name.replace(')','')
IN.name = IN.name.replace(' ','_')
ts = Tseries(IN.series, Gamma, IN.name)
ml.add_tseries(ts)

# Add well extraction at Essen
IN = next(x for x in meny.IN if x.name == 'Onttrekking (Essen)')
# extraction amound counts for the previeous month
IN.series = IN.series.resample('d').bfill()
IN.name = IN.name.replace('(','')
IN.name = IN.name.replace(')','')
IN.name = IN.name.replace(' ','_')
ts = Tseries(IN.series, Gamma, IN.name)
ml.add_tseries(ts)

# Add noise model
n = NoiseModel()
ml.add_noisemodel(n)

# Solve
ml.solve()
ml.plot_decomposition()
<commit_msg>Changed menyanthes-example, so fit is now better<commit_after>"""
This test file is meant for developing purposes. Providing an easy method to
test the functioning of PASTA during development.
"""

from pasta import *
from pasta.read.menydata import MenyData

fname = 'data/MenyanthesTest.men'
meny = MenyData(fname)

# Round the time of observations to days (using ceil, so the observation is never before the stress)
meny.H[0].series.index = meny.H[0].series.index.ceil('d')

# Create the time series model
ml = Model(meny.H[0].series)

# Add drainage level
d = Constant(value=meny.H[0].series.min())
ml.add_tseries(d)

# Add precipitation
IN = next(x for x in meny.IN if x.name == 'Neerslag (Hoogerheide)')
# round to days (precipitation is measured at 9:00)
IN.series = IN.series.resample('d').bfill()
IN.name = IN.name.replace('(','')
IN.name = IN.name.replace(')','')
IN.name = IN.name.replace(' ','_')
if False:
    ts = Tseries(IN.series, Gamma, IN.name)
else:
    IN2 = next(x for x in meny.IN if x.name == 'Verdamping (GILZE-RIJEN)')
    IN2.series = IN2.series.resample('d').bfill()
    IN2.name = IN2.name.replace('(', '')
    IN2.name = IN2.name.replace(')', '')
    IN2.name = IN2.name.replace(' ', '_')
    ts = Tseries2(IN.series, IN2.series, Gamma, 'Recharge')
ml.add_tseries(ts)

# Add well extraction at Ossendrecht
IN = next(x for x in meny.IN if x.name == 'Onttrekking (Ossendrecht)')
# extraction amount counts for the previous month
IN.series = IN.series.resample('d').bfill()
IN.name = IN.name.replace('(','')
IN.name = IN.name.replace(')','')
IN.name = IN.name.replace(' ','_')
ts = Tseries(-IN.series/1000, Gamma, IN.name)
ml.add_tseries(ts)

# Add well extraction at Huijbergen
IN = next(x for x in meny.IN if x.name == 'Onttrekking (Huijbergen)')
# extraction amount counts for the previous month
IN.series = IN.series.resample('d').bfill()
IN.name = IN.name.replace('(','')
IN.name = IN.name.replace(')','')
IN.name = IN.name.replace(' ','_')
ts = Tseries(-IN.series/1000, Gamma, IN.name)
ml.add_tseries(ts)

# Add well extraction at Essen
IN = next(x for x in meny.IN if x.name == 'Onttrekking (Essen)')
# extraction amount counts for the previous month
IN.series = IN.series.resample('d').bfill()
IN.name = IN.name.replace('(','')
IN.name = IN.name.replace(')','')
IN.name = IN.name.replace(' ','_')
ts = Tseries(-IN.series/1000, Gamma, IN.name)
ml.add_tseries(ts)

# Add noise model
n = NoiseModel()
ml.add_noisemodel(n)

# Solve
ml.solve()
ml.plot_decomposition()
<|endoftext|>


<commit_before># swift_build_support/shell.py ----------------------------------*- python -*-
#
# This source file is part of the Swift.org open source project
#
# Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
# Licensed under Apache License v2.0 with Runtime Library Exception
#
# See https://swift.org/LICENSE.txt for license information
# See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
# ----------------------------------------------------------------------------
"""
Centralized command line and file system interface for the build script.
"""
# ----------------------------------------------------------------------------

from __future__ import print_function

import os
import pipes
import shutil
import subprocess
import sys
from contextlib import contextmanager
from multiprocessing import Lock, Pool, cpu_count

from . import diagnostics

DEVNULL = getattr(subprocess, 'DEVNULL', subprocess.PIPE)

dry_run = False

def _quote(arg):
    return pipes.quote(str(arg))

def quote_command(args):
    """
    quote_command(args) -> str

    Quote the command for passing to a shell.
    """
    return ' '.join([_quote(a) for a in args])

def _coerce_dry_run(dry_run_override):
    if dry_run_override is None:
        return dry_run
    else:
        return dry_run_override

def _echo_command(dry_run, command, env=None, prompt="+ "):
    output = []
    if env is not None:
        output += ['env'] + [_quote("%s=%s" % (k, v))
                             for (k, v) in sorted(env.items())]
    output += [_quote(arg) for arg in command]
    file = sys.stderr
    if dry_run:
        file = sys.stdout
    print(prompt + ' '.join(output), file=file)
    file.flush()

def call(command, stderr=None, env=None, dry_run=None, echo=True):
    """
    call(command, ...) -> str

    Execute the given command.

    This function will raise an exception on any command failure.
    """
    dry_run = _coerce_dry_run(dry_run)
    if dry_run or echo:
        _echo_command(dry_run, command, env=env)
    if dry_run:
        return
    _env = None
    if env is not None:
        _env = dict(os.environ)
        _env.update(env)
    try:
        subprocess.check_call(command, env=_env, stderr=stderr)
    except subprocess.CalledProcessError as e:
        diagnostics.fatal(
            "command terminated with a non-zero exit status " +
            str(e.returncode) + ", aborting")
    except OSError as e:
        diagnostics.fatal(
            "could not execute '" + quote_command(command) +
            "': " + e.strerror)

def capture(command, stderr=None, env=None, dry_run=None, echo=True,
            optional=False, allow_non_zero_exit=False):
    """
    capture(command, ...) -> str

    Execute the given command and return the standard output.

    This function will raise an exception on any command failure.
    """
    dry_run = _coerce_dry_run(dry_run)
    if dry_run or echo:
        _echo_command(dry_run, command, env=env)
    if dry_run:
        return
    _env = None
    if env is not None:
        _env = dict(os.environ)
        _env.update(env)
    try:
        out = subprocess.check_output(command, env=_env, stderr=stderr)
        # Coerce to `str` hack. not py3 `byte`, not py2 `unicode`.
        return str(out.decode())
    except subprocess.CalledProcessError as e:
        if allow_non_zero_exit:
            return e.output
        if optional:
            return None
        diagnostics.fatal(
            "command terminated with a non-zero exit status " +
            str(e.returncode) + ", aborting")
    except OSError as e:
        if optional:
            return None
        diagnostics.fatal(
            "could not execute '" + quote_command(command) +
            "': " + e.strerror)

@contextmanager
def pushd(path, dry_run=None, echo=True):
    dry_run = _coerce_dry_run(dry_run)
    old_dir = os.getcwd()
    if dry_run or echo:
        _echo_command(dry_run, ["pushd", path])
    if not dry_run:
        os.chdir(path)
    yield
    if dry_run or echo:
        _echo_command(dry_run, ["popd"])
    if not dry_run:
        os.chdir(old_dir)

def makedirs(path, dry_run=None, echo=True):
    dry_run = _coerce_dry_run(dry_run)
    if dry_run or echo:
        _echo_command(dry_run, ['mkdir', '-p', path])
    if dry_run:
        return
    if not os.path.isdir(path):
        os.makedirs(path)

def rmtree(path, dry_run=None, echo=True):
    dry_run = _coerce_dry_run(dry_run)
    if dry_run or echo:
        _echo_command(dry_run, ['rm', '-rf', path])
    if dry_run:
        return
    if os.path.exists(path):
        shutil.rmtree(path)

def copytree(src, dest, dry_run=None, echo=True):
    dry_run = _coerce_dry_run(dry_run)
    if dry_run or echo:
        _echo_command(dry_run, ['cp', '-r', src, dest])
    if dry_run:
        return
    shutil.copytree(src, dest)

# Initialized later
lock = None

def run(*args, **kwargs):
    repo_path = os.getcwd()
    echo_output = kwargs.pop('echo', False)
    dry_run = kwargs.pop('dry_run', False)
    env = kwargs.pop('env', None)
    if dry_run:
        _echo_command(dry_run, *args, env=env)
        return(None, 0, args)

    my_pipe = subprocess.Popen(
        *args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, **kwargs)
    (stdout, stderr) = my_pipe.communicate()
    ret = my_pipe.wait()

    if lock:
        lock.acquire()
    if echo_output:
        print(repo_path)
        _echo_command(dry_run, *args, env=env)
        if stdout:
            print(stdout, end="")
        if stderr:
            print(stderr, end="")
        print()
    if lock:
        lock.release()

    if ret != 0:
        eout = Exception()
        eout.ret = ret
        eout.args = args
        eout.repo_path = repo_path
        eout.stderr = stderr
        raise eout
    return (stdout, 0, args)

def run_parallel(fn, pool_args, n_processes=0):
    def init(l):
        global lock
        lock = l

    if n_processes == 0:
        n_processes = cpu_count() * 2

    l = Lock()
    print("Running ``%s`` with up to %d processes." %
          (fn.__name__, n_processes))
    pool = Pool(processes=n_processes, initializer=init, initargs=(l,))
    results = pool.map_async(func=fn, iterable=pool_args).get(9999999)
    pool.close()
    pool.join()
    return results

def check_parallel_results(results, op):
    fail_count = 0
    if results is None:
        return 0
    for r in results:
        if r is not None:
            if fail_count == 0:
                print("======%s FAILURES======" % op)
            print("%s failed (ret=%d): %s" % (r.repo_path, r.ret, r))
            fail_count += 1
            if r.stderr:
                print(r.stderr)
    return fail_count
<commit_msg>Fix build-script shell to work on Windows<commit_after># swift_build_support/shell.py ----------------------------------*- python -*-
#
# This source file is part of the Swift.org open source project
#
# Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
# Licensed under Apache License v2.0 with Runtime Library Exception
#
# See https://swift.org/LICENSE.txt for license information
# See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
# ----------------------------------------------------------------------------
"""
Centralized command line and file system interface for the build script.
"""
# ----------------------------------------------------------------------------

from __future__ import print_function

import os
import pipes
import shutil
import subprocess
import sys
from contextlib import contextmanager
from multiprocessing import Lock, Pool, cpu_count

from . import diagnostics

DEVNULL = getattr(subprocess, 'DEVNULL', subprocess.PIPE)

dry_run = False

def _quote(arg):
    return pipes.quote(str(arg))

def quote_command(args):
    """
    quote_command(args) -> str

    Quote the command for passing to a shell.
    """
    return ' '.join([_quote(a) for a in args])

def _coerce_dry_run(dry_run_override):
    if dry_run_override is None:
        return dry_run
    else:
        return dry_run_override

def _echo_command(dry_run, command, env=None, prompt="+ "):
    output = []
    if env is not None:
        output += ['env'] + [_quote("%s=%s" % (k, v))
                             for (k, v) in sorted(env.items())]
    output += [_quote(arg) for arg in command]
    file = sys.stderr
    if dry_run:
        file = sys.stdout
    print(prompt + ' '.join(output), file=file)
    file.flush()

def call(command, stderr=None, env=None, dry_run=None, echo=True):
    """
    call(command, ...) -> str

    Execute the given command.

    This function will raise an exception on any command failure.
    """
    dry_run = _coerce_dry_run(dry_run)
    if dry_run or echo:
        _echo_command(dry_run, command, env=env)
    if dry_run:
        return
    _env = None
    if env is not None:
        _env = dict(os.environ)
        _env.update(env)
    try:
        subprocess.check_call(command, env=_env, stderr=stderr)
    except subprocess.CalledProcessError as e:
        diagnostics.fatal(
            "command terminated with a non-zero exit status " +
            str(e.returncode) + ", aborting")
    except OSError as e:
        diagnostics.fatal(
            "could not execute '" + quote_command(command) +
            "': " + e.strerror)

def capture(command, stderr=None, env=None, dry_run=None, echo=True,
            optional=False, allow_non_zero_exit=False):
    """
    capture(command, ...) -> str

    Execute the given command and return the standard output.

    This function will raise an exception on any command failure.
    """
    dry_run = _coerce_dry_run(dry_run)
    if dry_run or echo:
        _echo_command(dry_run, command, env=env)
    if dry_run:
        return
    _env = None
    if env is not None:
        _env = dict(os.environ)
        _env.update(env)
    try:
        out = subprocess.check_output(command, env=_env, stderr=stderr)
        # Coerce to `str` hack. not py3 `byte`, not py2 `unicode`.
        return str(out.decode())
    except subprocess.CalledProcessError as e:
        if allow_non_zero_exit:
            return e.output
        if optional:
            return None
        diagnostics.fatal(
            "command terminated with a non-zero exit status " +
            str(e.returncode) + ", aborting")
    except OSError as e:
        if optional:
            return None
        diagnostics.fatal(
            "could not execute '" + quote_command(command) +
            "': " + e.strerror)

@contextmanager
def pushd(path, dry_run=None, echo=True):
    dry_run = _coerce_dry_run(dry_run)
    old_dir = os.getcwd()
    if dry_run or echo:
        _echo_command(dry_run, ["pushd", path])
    if not dry_run:
        os.chdir(path)
    yield
    if dry_run or echo:
        _echo_command(dry_run, ["popd"])
    if not dry_run:
        os.chdir(old_dir)

def makedirs(path, dry_run=None, echo=True):
    dry_run = _coerce_dry_run(dry_run)
    if dry_run or echo:
        _echo_command(dry_run, ['mkdir', '-p', path])
    if dry_run:
        return
    if not os.path.isdir(path):
        os.makedirs(path)

def rmtree(path, dry_run=None, echo=True):
    dry_run = _coerce_dry_run(dry_run)
    if dry_run or echo:
        _echo_command(dry_run, ['rm', '-rf', path])
    if dry_run:
        return
    if os.path.exists(path):
        shutil.rmtree(path)

def copytree(src, dest, dry_run=None, echo=True):
    dry_run = _coerce_dry_run(dry_run)
    if dry_run or echo:
        _echo_command(dry_run, ['cp', '-r', src, dest])
    if dry_run:
        return
    shutil.copytree(src, dest)

# Initialized later
lock = None

def run(*args, **kwargs):
    repo_path = os.getcwd()
    echo_output = kwargs.pop('echo', False)
    dry_run = kwargs.pop('dry_run', False)
    env = kwargs.pop('env', None)
    if dry_run:
        _echo_command(dry_run, *args, env=env)
        return(None, 0, args)

    my_pipe = subprocess.Popen(
        *args, stdout=subprocess.PIPE, stderr=subprocess.PIPE, **kwargs)
    (stdout, stderr) = my_pipe.communicate()
    ret = my_pipe.wait()

    if lock:
        lock.acquire()
    if echo_output:
        print(repo_path)
        _echo_command(dry_run, *args, env=env)
        if stdout:
            print(stdout, end="")
        if stderr:
            print(stderr, end="")
        print()
    if lock:
        lock.release()

    if ret != 0:
        eout = Exception()
        eout.ret = ret
        eout.args = args
        eout.repo_path = repo_path
        eout.stderr = stderr
        raise eout
    return (stdout, 0, args)

def init(l):
    global lock
    lock = l

def run_parallel(fn, pool_args, n_processes=0):
    if n_processes == 0:
        n_processes = cpu_count() * 2

    l = Lock()
    print("Running ``%s`` with up to %d processes." %
          (fn.__name__, n_processes))
    pool = Pool(processes=n_processes, initializer=init, initargs=(l,))
    results = pool.map_async(func=fn, iterable=pool_args).get(9999999)
    pool.close()
    pool.join()
    return results

def check_parallel_results(results, op):
    fail_count = 0
    if results is None:
        return 0
    for r in results:
        if r is not None:
            if fail_count == 0:
                print("======%s FAILURES======" % op)
            print("%s failed (ret=%d): %s" % (r.repo_path, r.ret, r))
            fail_count += 1
            if r.stderr:
                print(r.stderr)
    return fail_count
<|endoftext|>


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

from collections import namedtuple
from sortedcontainers import SortedList
import numpy as np

Point = namedtuple('Point','x y z')

class Section(object):

    def __init__(self, pos: Point, chord: float, alpha: float, airfoil: str):
        self.pos = pos
        self.chord = chord
        self.airfoil = airfoil

    def __lt__(self, other) -> bool:
        return self.pos.x < other.pos.x

    def __eq__(self, other):
        return self.pos.x == other.pos.x

class Wing(object):
    """describes symmetric wing"""

    def __init__(self, pos):
        self.sections = SortedList()
        self.pos = pos

    @classmethod
    def fast_creation(cls, span_positions: list, chord_lengths: list, offsets:list, twists:list):
        """Generates a Wing object faster than through adding sections manually"""

        pos1 = Point(0, 0, 0)
        wing1 = cls(pos1)

        data = zip(span_positions, chord_lengths, offsets, twists)

        for span_position, chord_length, offset, twist in data:
            pos = Point(offset, span_position, 0)
            section = Section(pos, chord_length, twist, None)

            wing1.add_section(section)

        return wing1

    def add_section(self, section: Section):

        self.sections.add(section)

    def area(self) -> float:
        """Calculate area of the wing."""

        span_positions = [section.pos.y for section in self.sections]
        chord_lengths = [section.chord for section in self.sections]

        area = np.trapz(chord_lengths, span_positions)

        return area

    def mac(self) -> Section:
        """Calculate the mean aerodynamic chord of wing."""

        def generate_segment_macs(x_positions, y_positions, chord_lenghts):
            """Generate the segments mean aerodynamic chord values."""

            last = None

            for x, y, chord in zip(x_positions, y_positions, chord_lenghts):

                if last is not None:
                    x_old, y_old, chord_old = last
                    span = y - y_old
                    area = (chord_old+chord) * span / 2
                    taper = chord / chord_old
                    taper1 = taper + 1
                    fraction = (taper + taper1) / (3 * taper1)

                    mac = chord_old * (taper**2 + taper + 1) / (1.5 * taper1)
                    mac_x = x_old + fraction * (x - x_old)
                    mac_y = y_old + fraction * span

                    yield (mac_x*area, mac_y*area, mac*area, area)

                last = (x, y, chord)

        x_positions = [section.pos.x for section in self.sections]
        y_positions = [section.pos.y for section in self.sections]
        chord_lengths = [section.chord for section in self.sections]

        res = np.array(list(generate_segment_macs(x_positions, y_positions, chord_lengths)))

        print(res)

        area = self.area()

        mac_x = sum(res[:, 0])/area
        mac_y = sum(res[:, 1])/area
        mac = sum(res[:, 2])/area

        return Section(Point(mac_x, mac_y, 0),  mac, 0, airfoil=None)

    def span_width(self) -> float:
        """Calculate the span width of wing."""
        return 2 * max((section.pos.y for section in self.sections))

    def chord_at(self, y:float) -> float:
        """Calculates the chord depth at given span position"""

        y_positions = [section.pos.y for section in self.sections]
        chord_lengths = [section.chord for section in self.sections]

        return np.interp(y, y_positions, chord_lengths)

    def airfoil_at(self, y: float) -> str:
        """Lookup airfoil at given span position"""

        y_positions = [section.pos.y for section in self.sections]

        airfoil_index = np.argmin(np.array(y_positions)-y)

        return self.sections[airfoil_index].airfoil

    def positions_x(self):
        return [section.pos.x for section in self.sections]

    def positions_y(self):
        return [section.pos.y for section in self.sections]

    def _repr_html_(self):
        return """

Wing

        <svg>
        </svg>

spanwidth:{0}

wing area:{1}

mac{2}
""".format(self.span_width(), self.area(), self.mac().chord)
        # TODO: create nice representation with svg

class Plane:
    def __init__(self, name, wing, hlw):
        self.name = name
        self.wing = wing
        self.hlw = hlw

    @classmethod
    def load_from_file(cls, filename):
        import yaml, io

        # Lade Daten
        with io.open(filename, encoding='utf8') as f:
            content = f.read()
            data = yaml.load(content)

        # Erstelle Flugzeug Objekt
        plane = cls.generate(data)

        return plane

    @classmethod
    def generate(cls, data):

        # Name einlesen
        name = data['Name']

        # Flügel einlesen
        fluegel = Wing.generate(data[u'Flügel'])

        # Höhenleitwerk einlesen
        Hoehenleitwerk = namedtuple('Hoehenleitwerk','pos c4tel ')

        hlw_data = data[u'Höhenleitwerk']

        hlw_pos = Point(x=hlw_data['pos']['x'], y=0, z=hlw_data['pos']['z'])
        hlw = Hoehenleitwerk(pos=hlw_pos, c4tel=hlw_data['c4tel'])

        # Flugzeug Objekt erstellen
        plane = Plane(name, wing=fluegel, hlw=hlw)

        return plane

    def plot(self):
        from matplotlib import pyplot as plt

        fig = plt.figure(figsize=(10,6))

        # 3D Plot
        ax = fig.add_subplot(221, projection='3d')

        xle = self.wing.xle
        yle = self.wing.yle
        zle = self.wing.zle
        c = self.wing.c
        kt = self.wing.kt

        ax.plot(xle, yle, zle)

        # top view
        plt.subplot(2,2,2)

        plt.plot(yle, xle, 'black')      # leading edge
        plt.plot(yle, xle+c, 'black')    # trailing edge
        plt.plot(yle, xle+(1-kt)*c, '*') # flap edge
        plt.axis('equal')

        # side view
        plt.subplot(2,2,3)<commit_msg>Adds first svg output for Wing representation<commit_after># -*- coding: utf-8 -*-

from collections import namedtuple
from sortedcontainers import SortedList
import numpy as np

Point = namedtuple('Point','x y z')

class Section(object):

    def __init__(self, pos: Point, chord: float, alpha: float, airfoil: str):
        self.pos = pos
        self.chord = chord
        self.airfoil = airfoil

    def __lt__(self, other) -> bool:
        return self.pos.x < other.pos.x

    def __eq__(self, other):
        return self.pos.x == other.pos.x

class Wing(object):
    """describes symmetric wing"""

    def __init__(self, pos):
        self.sections = SortedList()
        self.pos = pos

    @classmethod
    def fast_creation(cls, span_positions: list, chord_lengths: list, offsets:list, twists:list):
        """Generates a Wing object faster than through adding sections manually"""

        pos1 = Point(0, 0, 0)
        wing1 = cls(pos1)

        data = zip(span_positions, chord_lengths, offsets, twists)

        for span_position, chord_length, offset, twist in data:
            pos = Point(offset, span_position, 0)
            section = Section(pos, chord_length, twist, None)

            wing1.add_section(section)

        return wing1

    def add_section(self, section: Section):

        self.sections.add(section)

    def area(self) -> float:
        """Calculate area of the wing."""

        span_positions = [section.pos.y for section in self.sections]
        chord_lengths = [section.chord for section in self.sections]

        area = np.trapz(chord_lengths, span_positions)

        return area

    def mac(self) -> Section:
        """Calculate the mean aerodynamic chord of wing."""

        def generate_segment_macs(x_positions, y_positions, chord_lenghts):
            """Generate the segments mean aerodynamic chord values."""

            last = None

            for x, y, chord in zip(x_positions, y_positions, chord_lenghts):

                if last is not None:
                    x_old, y_old, chord_old = last
                    span = y - y_old
                    area = (chord_old+chord) * span / 2
                    taper = chord / chord_old
                    taper1 = taper + 1
                    fraction = (taper + taper1) / (3 * taper1)

                    mac = chord_old * (taper**2 + taper + 1) / (1.5 * taper1)
                    mac_x = x_old + fraction * (x - x_old)
                    mac_y = y_old + fraction * span

                    yield (mac_x*area, mac_y*area, mac*area, area)

                last = (x, y, chord)

        x_positions = [section.pos.x for section in self.sections]
        y_positions = [section.pos.y for section in self.sections]
        chord_lengths = [section.chord for section in self.sections]

        res = np.array(list(generate_segment_macs(x_positions, y_positions, chord_lengths)))

        print(res)

        area = self.area()

        mac_x = sum(res[:, 0])/area
        mac_y = sum(res[:, 1])/area
        mac = sum(res[:, 2])/area

        return Section(Point(mac_x, mac_y, 0),  mac, 0, airfoil=None)

    def span_width(self) -> float:
        """Calculate the span width of wing."""
        return 2 * max((section.pos.y for section in self.sections))

    def chord_at(self, y:float) -> float:
        """Calculates the chord depth at given span position"""

        y_positions = [section.pos.y for section in self.sections]
        chord_lengths = [section.chord for section in self.sections]

        return np.interp(y, y_positions, chord_lengths)

    def airfoil_at(self, y: float) -> str:
        """Lookup airfoil at given span position"""

        y_positions = [section.pos.y for section in self.sections]

        airfoil_index = np.argmin(np.array(y_positions)-y)

        return self.sections[airfoil_index].airfoil

    def _svg_(self):
        import svgwrite

        dwg = svgwrite.Drawing()

        # draw contourline
        x_positions = [section.pos.x for section in self.sections]
        y_positions = [section.pos.y for section in self.sections]
        chord_lengths = [section.chord for section in self.sections]

        leading_edge = [(x*100, y*100) for x, y in zip(x_positions, y_positions)]
        trailing_edge = [((x+c)*100, y*100) for x, y, c in zip(x_positions, y_positions, chord_lengths)]

        all_pts = leading_edge+trailing_edge[::-1]

        contour = dwg.add(dwg.g(id='contour', stroke='black'))

        contour.add(dwg.polyline(all_pts))

        return dwg.tostring()

        # TODO: fix this function

    def _repr_html_(self):
        return """

Wing

        {3}

spanwidth:{0}

wing area:{1}

mac{2}
""".format(self.span_width(), self.area(), self.mac().chord, self._svg_())

class Plane:
    def __init__(self, name, wing, hlw):
        self.name = name
        self.wing = wing
        self.hlw = hlw

    @classmethod
    def load_from_file(cls, filename):
        import yaml, io

        # Lade Daten
        with io.open(filename, encoding='utf8') as f:
            content = f.read()
            data = yaml.load(content)

        # Erstelle Flugzeug Objekt
        plane = cls.generate(data)

        return plane

    @classmethod
    def generate(cls, data):

        # Name einlesen
        name = data['Name']

        # Flügel einlesen
        fluegel = Wing.generate(data[u'Flügel'])

        # Höhenleitwerk einlesen
        Hoehenleitwerk = namedtuple('Hoehenleitwerk','pos c4tel ')

        hlw_data = data[u'Höhenleitwerk']

        hlw_pos = Point(x=hlw_data['pos']['x'], y=0, z=hlw_data['pos']['z'])
        hlw = Hoehenleitwerk(pos=hlw_pos, c4tel=hlw_data['c4tel'])

        # Flugzeug Objekt erstellen
        plane = Plane(name, wing=fluegel, hlw=hlw)

        return plane

    def plot(self):
        from matplotlib import pyplot as plt

        fig = plt.figure(figsize=(10,6))

        # 3D Plot
        ax = fig.add_subplot(221, projection='3d')

        xle = self.wing.xle
        yle = self.wing.yle
        zle = self.wing.zle
        c = self.wing.c
        kt = self.wing.kt

        ax.plot(xle, yle, zle)

        # top view
        plt.subplot(2,2,2)

        plt.plot(yle, xle, 'black')      # leading edge
        plt.plot(yle, xle+c, 'black')    # trailing edge
        plt.plot(yle, xle+(1-kt)*c, '*') # flap edge
        plt.axis('equal')

        # side view
        plt.subplot(2,2,3)<|endoftext|>


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

import numpy as np
import argparse
import os
import time
import sys

from SULI import which
from SULI.execute_command import execute_command
from SULI.work_within_directory import work_within_directory
from astropy.io import fits

if __name__ == "__main__":

    parser = argparse.ArgumentParser('Submit transient search to the farm at Stanford')

    # add the arguments needed to the parser
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('--date', help='date specifying file to load')
    group.add_argument("--src_dir", help="Directory containing input data to be searched", type=str)

    parser.add_argument("--irf", help="Instrument response function name to be used", type=str, required=True)
    parser.add_argument("--probability", help="Probability of null hypothesis", type=float, required=True)
    parser.add_argument("--min_dist", help="Distance above which regions are not considered to overlap", type=float,
                        required=True)

    parser.add_argument("--res_dir", help="Directory where to put the results and logs for the search",
                        required=False, type=str, default=os.getcwd())
    parser.add_argument("--job_size", help="Number of jobs to submit at a time", required=False, type=int, default=20)

    parser.add_argument('--test', dest='test_run', action='store_true')
    parser.set_defaults(test_run=False)

    # parse the arguments
    args = parser.parse_args()

    # get output directory from parser
    res_dir = os.path.abspath(os.path.expandvars(os.path.expanduser(args.res_dir)))

    # Check that the output directory exists
    if not os.path.exists(res_dir):

        raise RuntimeError("Directory %s does not exist" % res_dir)

    # get src directory from parser
    src_dir = os.path.abspath(os.path.expandvars(os.path.expanduser(args.src_dir)))

    # Check that the simulation directory exists
    if not os.path.exists(src_dir):

        raise RuntimeError("Directory %s does not exist" % src_dir)

    # Go to output directory
    with work_within_directory(res_dir):

        # Create logs directory if it does not exist
        if not os.path.exists('logs'):

            os.mkdir('logs')

        # Create generated_data directory if it does not exist, and get number of files in the directory
        if not os.path.exists('generated_data'):

            os.mkdir('generated_data')

        DIR = './generated_data'
        num_res_files = len([results for results in os.listdir(DIR) if os.path.isfile(os.path.join(DIR, results))])

        # Generate universal command line parameters
        log_path = os.path.abspath('logs')
        out_path = os.path.abspath('generated_data')
        exe_path = which.which('search_on_farm.py')

        # if using simulated data:
        if args.src_dir:

            # get list of ft1 files
            ft1_files = [f for f in os.listdir(src_dir) if (str(os.path.join(src_dir, f)).endswith(('ft1.fits',
                                                                                'ft1.fit')))]
            # (is basically glob, should be changed to glob at some point)

            # get list of ft2 files
            ft2_files = [f for f in os.listdir(src_dir) if (str(os.path.join(src_dir, f)).endswith(('ft2.fits',
                                                                                'ft2.fit')))]

            # sort them
            def ft_sort(in_list):

                out_list = in_list.split("_")[1]

                return float(out_list)

            ft1_files = sorted(ft1_files, key=ft_sort)
            ft2_files = sorted(ft2_files, key=ft_sort)

            print '\nFound %s ft1 files\nFound %s ft2 files\n' % (len(ft1_files), len(ft2_files))

            # make sure each ft1/ft2 is part of a pair
            if len(ft1_files) != len(ft2_files):

                # determine which type there is more of for error msg
                if len(ft1_files) > len(ft2_files):

                    x = 'ft1 files'
                    y = 'ft2 files'

                else:

                    x = 'ft2 files'
                    y = 'ft1 files'

                raise RuntimeError('There are more %s than %s' % (x, y))

            # make sure pairs match
            for i in range(len(ft1_files)):

                with fits.open(os.path.join(src_dir, ft1_files[i])) as fits_file:

                    # Check the start and stop in the binary table
                    ft1_times = fits_file['EVENTS'].data.field("TIME")
                    ft1_starts = fits_file['GTI'].data.field("START")
                    ft1_stops = fits_file['GTI'].data.field("STOP")

                with fits.open(os.path.join(src_dir, ft2_files[i])) as fits_file:

                    # Check the start and stop in the binary table
                    ft2_starts = fits_file['SC_DATA'].data.field("START")
                    ft2_stops = fits_file['SC_DATA'].data.field("STOP")

                if ft2_starts.min() - min(ft1_starts.min(), ft1_times.min()) > 0:

                    raise RuntimeError("Mismatch in ft pair %s (FT2 file starts after the start of the FT1 file)" % i)

                if ft2_stops.max() - max(ft1_stops.max(), ft1_stops.max()) < 0:

                    raise RuntimeError("Mismatch in ft pair %s (FT2 file stops before the end of the FT1 file)" % i)

            # generate command line
            def sim_cmd_line(ft1, ft2, jobid):

                this_cmd_line = "qsub -l vmem=30gb -o %s/%s.out -e %s/%s.err -V -F '--inp_fts %s,%s --irf %s " \
                                "--probability %s --min_dist %s --out_dir %s' %s" % (log_path, jobid, log_path,
                                                                                jobid, ft1, ft2, args.irf,
                                                                                args.probability, args.min_dist,
                                                                                out_path, exe_path)
                return this_cmd_line

            # iterate over input directory, calling search on each pair of fits
            for i in range(len(ft1_files)):

                this_ft1 = src_dir + '/' + ft1_files[i]
                this_ft2 = src_dir + '/' + ft2_files[i]
                this_id = ft1_files[i]

                cmd_line = sim_cmd_line(this_ft1, this_ft2, this_id)

                if not args.test_run:

                    execute_command(cmd_line)

                # don't spam the farm; if more than [jobsize] jobs have been submitted,
                # wait until they finish to submit more
                if (i + 1) % args.job_size == 0:

                    # check res_dir every 10s for new results

                    # while the current number of results is not i+1 more than the initial number
                    # i.e., while the number of new files hasn't caught up to i + 1
                    num_fin = len([results for results in os.listdir(DIR) if os.path.isfile(os.path.join(DIR,
                                                                                results))])
                    sleep_count = 0
                    while num_fin - num_res_files != i + 1:

                        # sleep for 10s
                        time.sleep(10)
                        sleep_count += 1

                        # update num_fin for any finished jobs
                        num_fin = len([results for results in os.listdir(DIR) if os.path.isfile(os.path.join(DIR,
                                                                                results))])
                        print "%s ouf of %s jobs in this pass finished." % (num_fin - num_res_files, args.job_size)

                        # if it is taking to long, prompt to continue
                        if sleep_count >= 60:

                            def continue_prompt(inp_string):

                                sys.stdout.write(inp_string)
                                sys.stdout.write("Job is taking longer than usual. Abort? [y/n]")
                                inp = raw_input().lower()

                                if inp == 'y':

                                    raise RuntimeError("Farm took too long to respond. Aborting search.")

                                elif inp =='n':

                                    return 0

                                else:

                                    continue_prompt('Invalid prompt\n')

                            sleep_count = continue_prompt('\n')

        else:

            def rl_cmd_line(start):

                this_cmd_line = "qsub -l vmem=10gb -o %s/%s.out -e %s/%s.err -V -F '--date %s --irf %s " \
                                "--probability %s --min_dist %s --out_dir %s' %s" % (log_path, start, log_path,
                                                                                start, start, args.irf,
                                                                                args.probability, args.min_dist,
                                                                                out_path, exe_path)
                return this_cmd_line

            # A year of Fermi data

            dates = np.arange(args.date, args.date + (365.0 * 86400.0), 86400.0)

            for this_tstart in dates:

                cmd_line = rl_cmd_line(this_tstart)

                if not args.test_run:

                    execute_command(cmd_line)
<commit_msg>stagger jobs to farm<commit_after>#!/usr/bin/env python

import numpy as np
import argparse
import os
import time
import sys

from SULI import which
from SULI.execute_command import execute_command
from SULI.work_within_directory import work_within_directory
from astropy.io import fits

if __name__ == "__main__":

    parser = argparse.ArgumentParser('Submit transient search to the farm at Stanford')

    # add the arguments needed to the parser
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument('--date', help='date specifying file to load')
    group.add_argument("--src_dir", help="Directory containing input data to be searched", type=str)

    parser.add_argument("--irf", help="Instrument response function name to be used", type=str, required=True)
    parser.add_argument("--probability", help="Probability of null hypothesis", type=float, required=True)
    parser.add_argument("--min_dist", help="Distance above which regions are not considered to overlap", type=float,
                        required=True)

    parser.add_argument("--res_dir", help="Directory where to put the results and logs for the search",
                        required=False, type=str, default=os.getcwd())
    parser.add_argument("--job_size", help="Number of jobs to submit at a time", required=False, type=int, default=20)

    parser.add_argument('--test', dest='test_run', action='store_true')
    parser.set_defaults(test_run=False)

    # parse the arguments
    args = parser.parse_args()

    # get output directory from parser
    res_dir = os.path.abspath(os.path.expandvars(os.path.expanduser(args.res_dir)))

    # Check that the output directory exists
    if not os.path.exists(res_dir):

        raise RuntimeError("Directory %s does not exist" % res_dir)

    # get src directory from parser
    src_dir = os.path.abspath(os.path.expandvars(os.path.expanduser(args.src_dir)))

    # Check that the simulation directory exists
    if not os.path.exists(src_dir):

        raise RuntimeError("Directory %s does not exist" % src_dir)

    # Go to output directory
    with work_within_directory(res_dir):

        # Create logs directory if it does not exist
        if not os.path.exists('logs'):

            os.mkdir('logs')

        # Create generated_data directory if it does not exist, and get number of files in the directory
        if not os.path.exists('generated_data'):

            os.mkdir('generated_data')

        DIR = './generated_data'
        num_res_files = len([results for results in os.listdir(DIR) if os.path.isfile(os.path.join(DIR, results))])

        # Generate universal command line parameters
        log_path = os.path.abspath('logs')
        out_path = os.path.abspath('generated_data')
        exe_path = which.which('search_on_farm.py')

        # if using simulated data:
        if args.src_dir:

            # get list of ft1 files
            ft1_files = [f for f in os.listdir(src_dir) if (str(os.path.join(src_dir, f)).endswith(('ft1.fits',
                                                                                'ft1.fit')))]
            # (is basically glob, should be changed to glob at some point)

            # get list of ft2 files
            ft2_files = [f for f in os.listdir(src_dir) if (str(os.path.join(src_dir, f)).endswith(('ft2.fits',
                                                                                'ft2.fit')))]

            # sort them
            def ft_sort(in_list):

                out_list = in_list.split("_")[1]

                return float(out_list)

            ft1_files = sorted(ft1_files, key=ft_sort)
            ft2_files = sorted(ft2_files, key=ft_sort)

            print '\nFound %s ft1 files\nFound %s ft2 files\n' % (len(ft1_files), len(ft2_files))

            # make sure each ft1/ft2 is part of a pair
            if len(ft1_files) != len(ft2_files):

                # determine which type there is more of for error msg
                if len(ft1_files) > len(ft2_files):

                    x = 'ft1 files'
                    y = 'ft2 files'

                else:

                    x = 'ft2 files'
                    y = 'ft1 files'

                raise RuntimeError('There are more %s than %s' % (x, y))

            # make sure pairs match
            for i in range(len(ft1_files)):

                with fits.open(os.path.join(src_dir, ft1_files[i])) as fits_file:

                    # Check the start and stop in the binary table
                    ft1_times = fits_file['EVENTS'].data.field("TIME")
                    ft1_starts = fits_file['GTI'].data.field("START")
                    ft1_stops = fits_file['GTI'].data.field("STOP")

                with fits.open(os.path.join(src_dir, ft2_files[i])) as fits_file:

                    # Check the start and stop in the binary table
                    ft2_starts = fits_file['SC_DATA'].data.field("START")
                    ft2_stops = fits_file['SC_DATA'].data.field("STOP")

                if ft2_starts.min() - min(ft1_starts.min(), ft1_times.min()) > 0:

                    raise RuntimeError("Mismatch in ft pair %s (FT2 file starts after the start of the FT1 file)" % i)

                if ft2_stops.max() - max(ft1_stops.max(), ft1_stops.max()) < 0:

                    raise RuntimeError("Mismatch in ft pair %s (FT2 file stops before the end of the FT1 file)" % i)

            # generate command line
            def sim_cmd_line(ft1, ft2, jobid):

                this_cmd_line = "qsub -l vmem=30gb -o %s/%s.out -e %s/%s.err -V -F '--inp_fts %s,%s --irf %s " \
                                "--probability %s --min_dist %s --out_dir %s' %s" % (log_path, jobid, log_path,
                                                                                jobid, ft1, ft2, args.irf,
                                                                                args.probability, args.min_dist,
                                                                                out_path, exe_path)
                return this_cmd_line

            # iterate over input directory, calling search on each pair of fits
            for i in range(len(ft1_files)):

                this_ft1 = src_dir + '/' + ft1_files[i]
                this_ft2 = src_dir + '/' + ft2_files[i]
                this_id = ft1_files[i]

                cmd_line = sim_cmd_line(this_ft1, this_ft2, this_id)

                if not args.test_run:

                    execute_command(cmd_line)

                # don't spam the farm; if more than [jobsize] jobs have been submitted,
                # wait until they finish to submit more
                if (i + 1) % args.job_size == 0:

                    # check res_dir every 10s for new results

                    # while the current number of results is not i+1 more than the initial number
                    # i.e., while the number of new files hasn't caught up to i + 1

                    num_fin = len([results for results in os.listdir(DIR) if os.path.isfile(os.path.join(DIR,
                                                                                results))])
                    sleep_count = 0
                    while num_fin - num_res_files != i + 1:

                        # sleep for 10s
                        time.sleep(10)
                        sleep_count += 1

                        # update num_fin for any finished jobs
                        num_fin = len([results for results in os.listdir(DIR) if os.path.isfile(os.path.join(DIR,
                                                                                results))])
                        print "%s ouf of %s jobs in this pass finished." % (num_fin - num_res_files, args.job_size)

                        # if it is taking to long, prompt to continue
                        if sleep_count >= 60:

                            def continue_prompt(inp_string):

                                sys.stdout.write(inp_string)
                                sys.stdout.write("Job is taking longer than usual. Abort? [y/n]")
                                inp = raw_input().lower()

                                if inp == 'y':

                                    raise RuntimeError("Farm took too long to respond. Aborting search.")

                                elif inp =='n':

                                    return 0

                                else:

                                    continue_prompt('Invalid prompt\n')

                            sleep_count = continue_prompt('\n')

        else:

            def rl_cmd_line(start):

                this_cmd_line = "qsub -l vmem=10gb -o %s/%s.out -e %s/%s.err -V -F '--date %s --irf %s " \
                                "--probability %s --min_dist %s --out_dir %s' %s" % (log_path, start, log_path,
                                                                                start, start, args.irf,
                                                                                args.probability, args.min_dist,
                                                                                out_path, exe_path)
                return this_cmd_line

            # A year of Fermi data

            dates = np.arange(args.date, args.date + (365.0 * 86400.0), 86400.0)

            for this_tstart in dates:

                cmd_line = rl_cmd_line(this_tstart)

                if not args.test_run:

                    execute_command(cmd_line)
<|endoftext|>