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

import logging
import time
import argparse
import sys

from telegram.ext import Updater, CommandHandler, CallbackQueryHandler
from telegram import ReplyKeyboardMarkup, InlineKeyboardButton, InlineKeyboardMarkup
#aggiungee "Bot" alla linea sopra
#presentazione conclusa

<<<<<<< HEAD
from naoqi import *
=======
from naoqi import ALProxy, ALModule
>>>>>>> aggiunta funzione di lettura ddel file e chat_id

from settings import TOKEN, NAO_IP, NAO_PORT, chat_id

event_received = 0
sockinfo = None

filepath = '/home/remus/play-nao/LifeIsNAO-bot/abc.txt'

logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
                    level=logging.INFO)
logger = logging.getLogger(__name__)

def hello(bot, update):
    msg = 'Hello {},sei in chat {}'.format(
		update.message.from_user.first_name,
		update.message.chat.title)

    update.message.reply_text(msg)

    tts = ALProxy("ALTextToSpeech", sockinfo.ip, sockinfo.port)
    tts.say(msg)

def naosay(bot, update):
    txt = update.message.text[len('/naosay '):]
    update.message.reply_text(txt)

    tts = ALProxy("ALTextToSpeech", sockinfo.ip, sockinfo.port)
    tts.say(txt.encode("utf-8"))

def build_menu(buttons,
               n_cols,
               header_buttons=None,
               footer_buttons=None):
    menu = [buttons[i:i + n_cols] for i in range(0, len(buttons), n_cols)]
    if header_buttons:
        menu.insert(0, header_buttons)
    if footer_buttons:
        menu.append(footer_buttons)
    return menu

def present(bot, update):
  button_table = [
      [InlineKeyboardButton("/p Valerio", callback_data="1"),
      InlineKeyboardButton("/p Guramrit", callback_data="Guramrit")
      ],
      [
      InlineKeyboardButton("/p Gursharn", callback_data="Gursharn"),
      InlineKeyboardButton("/p Angelo", callback_data="Angelo")
      ],
      [
      InlineKeyboardButton("/p Ferdinando", callback_data="Ferdinando"),
      InlineKeyboardButton("/p Nicolas", callback_data="Nicolas")
      ],
      [
      InlineKeyboardButton("/p Marta", callback_data="Marta"),
      InlineKeyboardButton("/p Amira", callback_data="Amira")
      ],
      [
      InlineKeyboardButton("/p Jamila", callback_data="Jamila"),
      InlineKeyboardButton("/p Giacomo", callback_data="Giacomo")
      ]
  ]
  print(button_table)
  reply_markup = InlineKeyboardMarkup(build_menu(button_table, n_cols=2))
  reply_markup = ReplyKeyboardMarkup(button_table)
  print("bau")
  chat_id = update.message.chat.id
  print(chat_id)
  bot.send_message(chat_id, "Ecco qui il menu", reply_markup=reply_markup)

def p(bot, update):
    pst = 'Ora parla {}'.format(
		update.message.text[len('/p '):])

    update.message.reply_text(pst)

    tts = ALProxy("ALTextToSpeech", sockinfo.ip, sockinfo.port)
    tts.say(pst)

# def button(update, context):
#     print("Bottone")
#     query = update.callback_query
#     print ("bottone")
#     query.edit_message_text(text="Selected option: {}".format(query.data))

def error(update, context):
    """Log Errors caused by Updates."""
    logger.warning('Update "%s" caused error "%s"', update, context.error)

def fileread(bot, update):
    with open(filepath) as F:
      line = F.readline()
      if line != 'NODATA' and line != 'ERROR':
        value = float(line)/160

        update.message.reply_text(value)

        msg = 'Il valore è {}'.format(
        value
        ).replace('.', ',')

        tts = ALProxy("ALTextToSpeech", NAO_IP, NAO_PORT)
        tts.say(msg)

class HumanAnsweredQuestionModule(ALModule):
    """ A simple module able to react
    to facedetection events
    """
    def __init__(self, name):
        ALModule.__init__(self, name)
        # No need for IP and port here because
        # we have our Python broker connected to NAOqi broker

        # Create a proxy to ALTextToSpeech for later use
        self.tts = ALProxy("ALTextToSpeech")

        # Subscribe to the FaceDetected event:
        global memory
        memory = ALProxy("ALMemory")
        memory.subscribeToEvent("AnswerGiven",
            "HumanAnsweredQuestion",
            "onAnswerGiven")

        self.answers = []
        # Inserire qui le domande vere
        self.questions = [
            "Quanti hanni hai?",
            "Hai mangiato oggi?",
        ]

    def onAnswerGiven(self, *args):
        """ This will be called each time an answer is given
        """
        # OLD Unsubscribe to the event when talking,
        # OLD to avoid repetitions
        # OLD memory.unsubscribeToEvent("FaceDetected",
        # OLD     "HumanAnsweredQuestion")

        print("Rilevata risposta")
        msg = ""
        for i in range(len(self.questions)):
            try:
                question = self.questions[i-1]
                answer = memory.GetData("Domanda/%s" % i)
                msg += u"* Q:%s A:%s\n" % (question, answer)
            except Exception as e:
                # Qui ci va alla prima chiave che non esiste
                # Esce dal ciclo
                break

        # Leggere il valore istantaneo dal sensore
        freq_value = fileread(bot, update)
                     # TODO  scrivere funzione fileread_sensor()
        msg += " * FREQUENZA: {}".format(
        value
        ).replace('.', ',') % freq_value
        #msg += " * FREQUENZA: %s\n" % freq_value
        # Leggere gli ultimi 10 valori dal sensore
        #i = 0
        last_frequencies = "" # TODO scrivere funzione fileread_sensor_last_values(n=10)
        msg += " * ULTIME: %s\n" % last_frequencies

        # Inviare il messaggio al BOT
        # TODO bot.send_message(chat_id, msg)
        bot.send_message(chat_id, msg)

        # OLD Subscribe again to the event
        # OLD memory.subscribeToEvent("FaceDetected",
        # OLD     "HumanAnsweredQuestion",
        # OLD     "onFaceDetected")

def main():

    # Parse Args
    parser = argparse.ArgumentParser()
    parser.add_argument('--ip', default=NAO_IP)
    parser.add_argument('--port', default=NAO_PORT, type=int)
    # Enable or disable naoqi
    # parser.add_argument('--naoqi', action="store_true", default=False)
    global sockinfo
    sockinfo = parser.parse_args()

    # Start NAO Broker
    myBroker = ALBroker("myBroker", "0.0.0.0", 0, sockinfo.ip, sockinfo.port)
    # Warning: HumanAnsweredQuestion must be a global variable
    # The name given to the constructor must be the name of the
    # variable
    global HumanAnsweredQuestion
    HumanAnsweredQuestion = HumanAnsweredQuestionModule("HumanAnsweredQuestion")

    # Start BOT
    def terminate(self, signal_name):
        print("Interrupted bot and NAO module by user, shutdown")
        myBroker.shutdown()
        sys.exit(0)

#def Q&A():
#   risp = ALProxy("ALMemory", sockinfo.ip, sockinfo.port)

# bot = Bot(TOKEN)
# subscribe to event
# nella funzione o nel modulo chiamato come callback si fa
# bot.send_message(chat_id, msg)
# dove msg è il testo recuperato dalla memoria di NAO

class HumanAnsweredQuestionModule(ALModule):
    """ A simple module able to react
    to facedetection events

    """
    def __init__(self, name):
        ALModule.__init__(self, name)
        # No need for IP and port here because
        # we have our Python broker connected to NAOqi broker

        # Create a proxy to ALTextToSpeech for later use
        self.tts = ALProxy("ALTextToSpeech")

        # Subscribe to the FaceDetected event:
        global memory
        memory = ALProxy("ALMemory")
        memory.subscribeToEvent("AnswerGiven",
            "HumanAnsweredQuestion",
            "onAnswerGiven")

        self.answers = []
        # Inserire qui le domande vere
        self.questions = [
            "Quanti hanni hai?",
            "Hai mangiato oggi?",
        ]

    def onAnswerGiven(self, *args):
        """ This will be called each time an answer is given
        """
        # OLD Unsubscribe to the event when talking,
        # OLD to avoid repetitions
        # OLD memory.unsubscribeToEvent("FaceDetected",
        # OLD     "HumanAnsweredQuestion")

        print("Rilevata risposta")
        msg = ""
        for i in range(len(self.questions)):
            try:
                question = self.questions[i-1]
                answer = memory.GetData("Domanda/%s" % i)
                msg += u"* Q:%s A:%s\n" % (question, answer)
            except Exception as e:
                # Qui ci va alla prima chiave che non esiste
                # Esce dal ciclo
                break

        # Leggere il valore istantaneo dal sensore
        freq_value = ""  # TODO  scrivere funzione fileread_sensor()
        msg += " * FREQUENZA: %s\n" % freq_value

        # Leggere gli ultimi 10 valori dal sensore
        last_frequencies = ""  # TODO scrivere funzione fileread_sensor_last_values(n=10)
        msg += " * ULTIME: %s\n" % last_frequencies

        # Inviare il messaggio al BOT
        # TODO bot.send_message(chat_id, msg)

        # OLD Subscribe again to the event
        # OLD memory.subscribeToEvent("FaceDetected",
        # OLD     "HumanAnsweredQuestion",
        # OLD     "onFaceDetected")

def main():

    # Parse Args
    parser = argparse.ArgumentParser()
    parser.add_argument('--ip', default=NAO_IP)
    parser.add_argument('--port', default=NAO_PORT, type=int)
    # Enable or disable naoqi
    # parser.add_argument('--naoqi', action="store_true", default=False)
    global sockinfo
    sockinfo = parser.parse_args()

    # Start NAO Broker
    myBroker = ALBroker("myBroker", "0.0.0.0", 0, sockinfo.ip, sockinfo.port)
    # Warning: HumanAnsweredQuestion must be a global variable
    # The name given to the constructor must be the name of the
    # variable
    global HumanAnsweredQuestion
    HumanAnsweredQuestion = HumanAnsweredQuestionModule("HumanAnsweredQuestion")

    # Start BOT
    def terminate(self, signal_name):
        print("Interrupted bot and NAO module by user, shutdown")
        myBroker.shutdown()
        sys.exit(0)

    updater = Updater(TOKEN, user_sig_handler=terminate)
    updater.dispatcher.add_handler(CommandHandler('naohello', hello))
    updater.dispatcher.add_handler(CommandHandler('naosay', naosay))
    updater.dispatcher.add_handler(CommandHandler('naopresent', present))
    updater.dispatcher.add_handler(CommandHandler('p', p))
    updater.dispatcher.add_error_handler(error)

    updater.start_polling()
    updater.idle()

if __name__ == '__main__':
    main()
<commit_msg>Aggiunyta funzione di letture e chat_id<commit_after>#-*- coding: utf-8 -*-

import logging
import time
import argparse
import sys

from telegram.ext import Updater, CommandHandler, CallbackQueryHandler
from telegram import ReplyKeyboardMarkup, InlineKeyboardButton, InlineKeyboardMarkup
#aggiungee "Bot" alla linea sopra
#presentazione conclusa

<<<<<<< HEAD
from naoqi import *
=======
from naoqi import ALProxy, ALModule
>>>>>>> aggiunta funzione di lettura ddel file e chat_id

from settings import TOKEN, NAO_IP, NAO_PORT, chat_id

event_received = 0
sockinfo = None

import time

import argparse

import sys
filepath = '/home/remus/play-nao/LifeIsNAO-bot/abc.txt'

logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
                    level=logging.INFO)
logger = logging.getLogger(__name__)

def hello(bot, update):
    msg = 'Hello {},sei in chat {}'.format(
		update.message.from_user.first_name,
		update.message.chat.title)

    update.message.reply_text(msg)

    tts = ALProxy("ALTextToSpeech", sockinfo.ip, sockinfo.port)
    tts.say(msg)

def naosay(bot, update):
    txt = update.message.text[len('/naosay '):]
    update.message.reply_text(txt)

    tts = ALProxy("ALTextToSpeech", sockinfo.ip, sockinfo.port)
    tts.say(txt.encode("utf-8"))

def build_menu(buttons,
               n_cols,
               header_buttons=None,
               footer_buttons=None):
    menu = [buttons[i:i + n_cols] for i in range(0, len(buttons), n_cols)]
    if header_buttons:
        menu.insert(0, header_buttons)
    if footer_buttons:
        menu.append(footer_buttons)
    return menu

def present(bot, update):
  button_table = [
      [InlineKeyboardButton("/p Valerio", callback_data="1"),
      InlineKeyboardButton("/p Guramrit", callback_data="Guramrit")
      ],
      [
      InlineKeyboardButton("/p Gursharn", callback_data="Gursharn"),
      InlineKeyboardButton("/p Angelo", callback_data="Angelo")
      ],
      [
      InlineKeyboardButton("/p Ferdinando", callback_data="Ferdinando"),
      InlineKeyboardButton("/p Nicolas", callback_data="Nicolas")
      ],
      [
      InlineKeyboardButton("/p Marta", callback_data="Marta"),
      InlineKeyboardButton("/p Amira", callback_data="Amira")
      ],
      [
      InlineKeyboardButton("/p Jamila", callback_data="Jamila"),
      InlineKeyboardButton("/p Giacomo", callback_data="Giacomo")
      ]
  ]
  print(button_table)
  reply_markup = InlineKeyboardMarkup(build_menu(button_table, n_cols=2))
  reply_markup = ReplyKeyboardMarkup(button_table)
  print("bau")
  chat_id = update.message.chat.id
  print(chat_id)
  bot.send_message(chat_id, "Ecco qui il menu", reply_markup=reply_markup)

def p(bot, update):
    pst = 'Ora parla {}'.format(
		update.message.text[len('/p '):])

    update.message.reply_text(pst)

    tts = ALProxy("ALTextToSpeech", sockinfo.ip, sockinfo.port)
    tts.say(pst)

# def button(update, context):
#     print("Bottone")
#     query = update.callback_query
#     print ("bottone")
#     query.edit_message_text(text="Selected option: {}".format(query.data))

def error(update, context):
    """Log Errors caused by Updates."""
    logger.warning('Update "%s" caused error "%s"', update, context.error)

def fileread(bot, update):
    with open(filepath) as F:
      line = F.readline()
      if line != 'NODATA' and line != 'ERROR':
        value = float(line)/160

        update.message.reply_text(value)

        msg = 'Il valore è {}'.format(
        value
        ).replace('.', ',')

        tts = ALProxy("ALTextToSpeech", NAO_IP, NAO_PORT)
        tts.say(msg)

class HumanAnsweredQuestionModule(ALModule):
    """ A simple module able to react
    to facedetection events
    """
    def __init__(self, name):
        ALModule.__init__(self, name)
        # No need for IP and port here because
        # we have our Python broker connected to NAOqi broker

        # Create a proxy to ALTextToSpeech for later use
        self.tts = ALProxy("ALTextToSpeech")

        # Subscribe to the FaceDetected event:
        global memory
        memory = ALProxy("ALMemory")
        memory.subscribeToEvent("AnswerGiven",
            "HumanAnsweredQuestion",
            "onAnswerGiven")

        self.answers = []
        # Inserire qui le domande vere
        self.questions = [
            "Quanti hanni hai?",
            "Hai mangiato oggi?",
        ]

    def onAnswerGiven(self, *args):
        """ This will be called each time an answer is given
        """
        # OLD Unsubscribe to the event when talking,
        # OLD to avoid repetitions
        # OLD memory.unsubscribeToEvent("FaceDetected",
        # OLD     "HumanAnsweredQuestion")

        print("Rilevata risposta")
        msg = ""
        for i in range(len(self.questions)):
            try:
                question = self.questions[i-1]
                answer = memory.GetData("Domanda/%s" % i)
                msg += u"* Q:%s A:%s\n" % (question, answer)
            except Exception as e:
                # Qui ci va alla prima chiave che non esiste
                # Esce dal ciclo
                break

        # Leggere il valore istantaneo dal sensore
        freq_value = fileread(bot, update)
                     # TODO  scrivere funzione fileread_sensor()
        msg += " * FREQUENZA: {}".format(
        value
        ).replace('.', ',') % freq_value
        #msg += " * FREQUENZA: %s\n" % freq_value
        # Leggere gli ultimi 10 valori dal sensore
        #i = 0
        last_frequencies = "" # TODO scrivere funzione fileread_sensor_last_values(n=10)
        msg += " * ULTIME: %s\n" % last_frequencies

        # Inviare il messaggio al BOT
        # TODO bot.send_message(chat_id, msg)
        bot.send_message(chat_id, msg)

        # OLD Subscribe again to the event
        # OLD memory.subscribeToEvent("FaceDetected",
        # OLD     "HumanAnsweredQuestion",
        # OLD     "onFaceDetected")

def main():

    # Parse Args
    parser = argparse.ArgumentParser()
    parser.add_argument('--ip', default=NAO_IP)
    parser.add_argument('--port', default=NAO_PORT, type=int)
    # Enable or disable naoqi
    # parser.add_argument('--naoqi', action="store_true", default=False)
    global sockinfo
    sockinfo = parser.parse_args()

    # Start NAO Broker
    myBroker = ALBroker("myBroker", "0.0.0.0", 0, sockinfo.ip, sockinfo.port)
    # Warning: HumanAnsweredQuestion must be a global variable
    # The name given to the constructor must be the name of the
    # variable
    global HumanAnsweredQuestion
    HumanAnsweredQuestion = HumanAnsweredQuestionModule("HumanAnsweredQuestion")

    # Start BOT
    def terminate(self, signal_name):
        print("Interrupted bot and NAO module by user, shutdown")
        myBroker.shutdown()
        sys.exit(0)

#def Q&A():
#   risp = ALProxy("ALMemory", sockinfo.ip, sockinfo.port)

# bot = Bot(TOKEN)
# subscribe to event
# nella funzione o nel modulo chiamato come callback si fa
# bot.send_message(chat_id, msg)
# dove msg è il testo recuperato dalla memoria di NAO

class HumanAnsweredQuestionModule(ALModule):
    """ A simple module able to react
    to facedetection events

    """
    def __init__(self, name):
        ALModule.__init__(self, name)
        # No need for IP and port here because
        # we have our Python broker connected to NAOqi broker

        # Create a proxy to ALTextToSpeech for later use
        self.tts = ALProxy("ALTextToSpeech")

        # Subscribe to the FaceDetected event:
        global memory
        memory = ALProxy("ALMemory")
        memory.subscribeToEvent("AnswerGiven",
            "HumanAnsweredQuestion",
            "onAnswerGiven")

        self.answers = []
        # Inserire qui le domande vere
        self.questions = [
            "Quanti hanni hai?",
            "Hai mangiato oggi?",
        ]

    def onAnswerGiven(self, *args):
        """ This will be called each time an answer is given
        """
        # OLD Unsubscribe to the event when talking,
        # OLD to avoid repetitions
        # OLD memory.unsubscribeToEvent("FaceDetected",
        # OLD     "HumanAnsweredQuestion")

        print("Rilevata risposta")
        msg = ""
        for i in range(len(self.questions)):
            try:
                question = self.questions[i-1]
                answer = memory.GetData("Domanda/%s" % i)
                msg += u"* Q:%s A:%s\n" % (question, answer)
            except Exception as e:
                # Qui ci va alla prima chiave che non esiste
                # Esce dal ciclo
                break

        # Leggere il valore istantaneo dal sensore
        freq_value = ""  # TODO  scrivere funzione fileread_sensor()
        msg += " * FREQUENZA: %s\n" % freq_value

        # Leggere gli ultimi 10 valori dal sensore
        last_frequencies = ""  # TODO scrivere funzione fileread_sensor_last_values(n=10)
        msg += " * ULTIME: %s\n" % last_frequencies

        # Inviare il messaggio al BOT
        # TODO bot.send_message(chat_id, msg)

        # OLD Subscribe again to the event
        # OLD memory.subscribeToEvent("FaceDetected",
        # OLD     "HumanAnsweredQuestion",
        # OLD     "onFaceDetected")

def main():

    # Parse Args
    parser = argparse.ArgumentParser()
    parser.add_argument('--ip', default=NAO_IP)
    parser.add_argument('--port', default=NAO_PORT, type=int)
    # Enable or disable naoqi
    # parser.add_argument('--naoqi', action="store_true", default=False)
    global sockinfo
    sockinfo = parser.parse_args()

    # Start NAO Broker
    myBroker = ALBroker("myBroker", "0.0.0.0", 0, sockinfo.ip, sockinfo.port)
    # Warning: HumanAnsweredQuestion must be a global variable
    # The name given to the constructor must be the name of the
    # variable
    global HumanAnsweredQuestion
    HumanAnsweredQuestion = HumanAnsweredQuestionModule("HumanAnsweredQuestion")

    # Start BOT
    def terminate(self, signal_name):
        print("Interrupted bot and NAO module by user, shutdown")
        myBroker.shutdown()
        sys.exit(0)

    updater = Updater(TOKEN, user_sig_handler=terminate)
    updater.dispatcher.add_handler(CommandHandler('naohello', hello))
    updater.dispatcher.add_handler(CommandHandler('naosay', naosay))
    updater.dispatcher.add_handler(CommandHandler('naopresent', present))
    updater.dispatcher.add_handler(CommandHandler('p', p))
    updater.dispatcher.add_error_handler(error)

    updater.start_polling()
    updater.idle()

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


<commit_before>"""
Classes for datafile syncronization.

"""

import Aggregators
import Workers

class File(object):
    """File object class. Implements file locking and syncronization.

    """
    def __init__(self, filename, reader, writer):
        """Create File instance for interacting with file on disk.

        The File object keeps its own cached version of a data structure
        written to file in sync with the file itself. It does this by
        consistently applying locks to files before reading and writing.

        At all times, reading and modifying the data structure is the same as
        reading and writing to the file. Accessing an element of the data
        structure results in locking, reading, and unlocking the file each
        time. Modifying an element results in locking, reading, writing, and
        unlocking. All operations are performed atomically to ensure unintended
        overwrites are avoided.

        :Arguments:
           *filename*
              name of file on disk object synchronizes with
           *reader*
              function used to translate file into data structure
           *writer*
              function used to tranlate data structure into file

        """
        self.filename = filename
        self.lockname = "{}.lock".format(self.file.name)

        self.reader = reader
        self.writer = writer

        self.read(self.filename)

    def read(self):
        """Read contents of file into python representation.

        :Returns:
           *success*
              True if write successful
        """
        # keep attempting lock until successful (failsafe)
        while not self.lock():
            continue

        # keep reading until unlock gives success (only if synchronized)
        while not self.unlock()
            with open(self.filename, 'r') as f:
                self.data = self.reader(f)

        return True

    def write(self, writer):
        """Write python representation to file.

        :Returns:
           *success*
              True if write successful
        """
        # keep attempting lock until successful (failsafe)
        while not self.lock():
            continue

        # keep writing until unlock gives success (only if synchronized)
        while not self.unlock()
            with open(self.filename, 'w') as f:
                self.data = self.writer(self.data)

        return True

    def lock(self):
        """Get exclusive lock on file.

        The lock is just a symlink of the target file, since making a symlink
        appears to be an atomic operation on most platforms. This is important,
        since creating a symlink also serves to check if one is already present
        all in one operation.

        :Returns:
           *success*
              True if lockfile successfully created
        """
        # if lockfile already present, wait until it disappears; make lock
        while True:
            try:
                os.symlink(self.filename, self.lockname):
                break
            except OSError:
                time.sleep(1)

        # return lockfile existence
        return os.path.exists(self.lockname)

    def unlock(self):
        """Remove exclusive lock on file.

        Before removing the lock, checks that the python representation is
        the same as what is on file.

        :Returns:
           *success*
              True if comparison successful
        """
        # check that python representation matches file
        success = self.compare()

        if success:
            os.remove(self.lockname)

        return success

    def compare(self, reader):
        """Compare data structure with file contents.

        :Returns:
           *same*
              True if file synchronized with data structure
        """
        datatemp =

class ContainerFile(File):
    """Container file object; syncronized access to Container data.

    """

class OperatorFile(File):
    """Operator file object; syncronized access to Operator data.

    """

class DatabaseFile(File):
    """Database file object; syncronized access to Database data.

    """

class DataFile(object):
    """Universal datafile interface.

    Allows for safe reading and writing of datafiles, which can be of a wide
    array of formats. Handles the details of conversion from pythonic data
    structure to persistent file form.

    """

<commit_msg>File base class complete.<commit_after>"""
Classes for datafile syncronization.

"""

import Aggregators
import Workers

class File(object):
    """File object class. Implements file locking and syncronization.

    """
    def __init__(self, filename, reader, writer, logger=None):
        """Create File instance for interacting with file on disk.

        The File object keeps its own cached version of a data structure
        written to file in sync with the file itself. It does this by
        consistently applying locks to files before reading and writing.

        At all times, reading and modifying the data structure is the same as
        reading and writing to the file. Accessing an element of the data
        structure results in locking, reading, and unlocking the file each
        time. Modifying an element results in locking, reading, writing, and
        unlocking. All operations are performed atomically to ensure unintended
        overwrites are avoided.

        :Arguments:
           *filename*
              name of file on disk object synchronizes with
           *reader*
              function used to translate file into data structure
           *writer*
              function used to tranlate data structure into file
           *logger*
              logger to send warnings and errors to

        """
        self.filename = filename
        self.lockname = "{}.lock".format(self.file.name)

        self.reader = reader
        self.writer = writer

        self.read(self.filename)

    def read(self):
        """Read contents of file into python representation.

        :Returns:
           *success*
              True if write successful
        """
        # keep attempting lock until successful (failsafe)
        while not self.lock():
            continue

        # keep reading until unlock gives success (only if synchronized)
        while not self.unlock()
            with open(self.filename, 'r') as f:
                self.data = self.reader(f)

        return True

    def write(self):
        """Write python representation to file.

        :Returns:
           *success*
              True if write successful
        """
        # keep attempting lock until successful (failsafe)
        while not self.lock():
            continue

        # keep writing until unlock gives success (only if synchronized)
        while not self.unlock()
            with open(self.filename, 'w') as f:
                self.data = self.writer(self.data)

        return True

    def lock(self):
        """Get exclusive lock on file.

        The lock is just a symlink of the target file, since making a symlink
        appears to be an atomic operation on most platforms. This is important,
        since creating a symlink also serves to check if one is already present
        all in one operation.

        :Returns:
           *success*
              True if lockfile successfully created
        """
        # if lockfile already present, wait until it disappears; make lock
        while True:
            try:
                os.symlink(self.filename, self.lockname):
                break
            except OSError:
                time.sleep(1)

        # return lockfile existence
        return os.path.exists(self.lockname)

    def unlock(self):
        """Remove exclusive lock on file.

        Before removing the lock, checks that the python representation is
        the same as what is on file.

        :Returns:
           *success*
              True if comparison successful
        """
        # check that python representation matches file
        success = self.compare()

        if success:
            os.remove(self.lockname)

        return success

    def compare(self):
        """Compare data structure with file contents.

        :Returns:
           *same*
              True if file synchronized with data structure
        """
        with open(self.filename, 'r') as f:
            datatemp = self.reader(f)

        return self.data == datatemp

class ContainerFile(File):
    """Container file object; syncronized access to Container data.

    """

class OperatorFile(File):
    """Operator file object; syncronized access to Operator data.

    """

class DatabaseFile(File):
    """Database file object; syncronized access to Database data.

    """

class DataFile(object):
    """Universal datafile interface.

    Allows for safe reading and writing of datafiles, which can be of a wide
    array of formats. Handles the details of conversion from pythonic data
    structure to persistent file form.

    """

<|endoftext|>


<commit_before>import os
from pymongo import MongoClient

# from pymongo import MongoClient
from settings import TEST_MONGO_HOST, TEST_MONGO_PORT, TEST_MONGO_USERNAME, TEST_MONGO_PASSWORD
import pprint

class model_service (object):
    """Service which connects to Models via MongoDB"""

    def __init__(self, mongo_uri=None, db=None, worker_id=None, client=None):
        self.serviceURL = os.environ.get('SERVICE_URL', None)
        if client:
            self.client = client
        else:
            self.client = MongoClient(mongo_uri)
        self.db = self.client[db]

    def retrieveAllModels(self):
        """Returns all models for worker"""
        result = []
        for model in self.db.models.find():
            result.append(model)
        return result

    def getModelByID(self, identifier):
        return self.db.models.find_one({'_id': identifier})

    def removeModel(self, filter):
        """Removes all models for filter"""
        self.db.models.remove(filter)

    def updateModel(self, model, set_obj):
        """Updates the given model"""
        return self.db.models.update_one(
            {'_id': ObjectId(model["_id"])}, set_obj,
            upsert=False)

    def retrieveRequestedDeploy(self):
        """Claims a requested deploy from the DB
            returns the model object, once claim verified"""
        if self.serviceURL is None:
            print ("No serviceURL!")
            return None
        result = None
        for model in self.db.models.find({'deployRequested':True}):
            self.db.models.update_one(
                {'_id': model["_id"]},
                {'$set': {'serviceURL': self.serviceURL, 'deployRequested': False}},
                upsert=False)
            result = self.db.models.find_one({'_id': model["_id"]})
            if result.get("deployRequested") == False and result.get("serviceURL") == self.serviceURL:
                break
        return result

    def insertModel(self, model):
        """Inserts the given model"""
        return self.db.models.insert(model)
        <commit_msg>Prediction feature<commit_after>import os
from pymongo import MongoClient
from bson import ObjectId

# from pymongo import MongoClient
from settings import TEST_MONGO_HOST, TEST_MONGO_PORT, TEST_MONGO_USERNAME, TEST_MONGO_PASSWORD
import pprint

class model_service (object):
    """Service which connects to Models via MongoDB"""

    def __init__(self, mongo_uri=None, db=None, worker_id=None, client=None):
        self.serviceURL = os.environ.get('SERVICE_URL', None)
        if client:
            self.client = client
        else:
            self.client = MongoClient(mongo_uri)
        self.db = self.client[db]

    def retrieveAllModels(self):
        """Returns all models for worker"""
        result = []
        for model in self.db.models.find():
            result.append(model)
        return result

    def getModelByID(self, identifier):
        return self.db.models.find_one({'_id': identifier})

    def removeModel(self, filter):
        """Removes all models for filter"""
        self.db.models.remove(filter)

    def updateModel(self, model, set_obj):
        """Updates the given model"""
        return self.db.models.update_one(
            {'_id': ObjectId(model["_id"])}, set_obj,
            upsert=False)

    def retrieveRequestedDeploy(self):
        """Claims a requested deploy from the DB
            returns the model object, once claim verified"""
        if self.serviceURL is None:
            print ("No serviceURL!")
            return None
        result = None
        for model in self.db.models.find({'deployRequested':True}):
            self.db.models.update_one(
                {'_id': model["_id"]},
                {'$set': {'serviceURL': self.serviceURL, 'deployRequested': False}},
                upsert=False)
            result = self.db.models.find_one({'_id': model["_id"]})
            if result.get("deployRequested") == False and result.get("serviceURL") == self.serviceURL:
                break
        return result

    def insertModel(self, model):
        """Inserts the given model"""
        return self.db.models.insert(model)
        <|endoftext|>


<commit_before># Rigid bodies and joints, Compliant-style.
#
# Authors: maxime.tournier@inria.fr, ... ?
#
# The basic philosophy is that python provides nice data structures to
# define scene component *semantically*. Once this is done, the actual
# scene graph must be generated through the 'insert' methods. Please
# refer to the python examples in Compliant for more information.
#

import Sofa

from subprocess import Popen, PIPE

from numpy import *
import SofaPython.Quaternion as quat
import Vec as vec

import Tools
from Tools import cat as concat

import os

class Frame:
        # a rigid frame, group operations are available.

        # TODO kwargs
        def __init__(self, value = None):
                if value != None:
                        self.translation = value[:3]
                        self.rotation = value[3:]
                else:
                        self.translation = [0, 0, 0]
                        self.rotation = [0, 0, 0, 1]

        def insert(self, parent, **args):
                return parent.createObject('MechanicalObject',
                                           template = 'Rigid',
                                           position = str(self),
                                           **args)

        def __str__(self):
                return concat(self.translation) + ' ' + concat(self.rotation)

        def copy(self):
                # TODO better !
                return Rigid.Frame().read( str(self) )

        def setValue(self, v):
             self.translation = v[:3]
             self.rotation = v[3:]
             return self

        def read(self, str):
                num = map(float, str.split())
                self.translation = num[:3]
                self.rotation = num[3:]
                return self

        def offset(self):
            return hstack((self.translation, self.rotation))

	def __mul__(self, other):
		res = Frame()
		res.translation = vec.sum(self.translation,
					  quat.rotate(self.rotation,
						      other.translation))
		res.rotation = quat.prod( self.rotation,
					  other.rotation)

		return res

	def inv(self):
		res = Frame()
		res.rotation = quat.conj( self.rotation )
		res.translation = vec.minus( quat.rotate(res.rotation,
							 self.translation) )
		return res

        def set(self, **kwargs):
                for k in kwargs:
                        setattr(self, k, kwargs[k])

                return self

        # TODO more: apply( vec3 ), wrench/twist frame change.

class MassInfo:
        pass

# front-end to sofa GenerateRigid tool. density unit is kg/m^3
def generate_rigid(filename, density = 1000.0, scale=[1,1,1]):

        # TODO bind GenerateRigid
        # - faster than writing in a file
        # - more robust (if several processes try to work in the same file)

        tmpfilename = Tools.path( __file__ ) +"/tmp.rigid"

        cmd = [ Sofa.build_dir() + '/bin/GenerateRigid', filename, tmpfilename, str(density), str(scale[0]), str(scale[1]), str(scale[2]) ]

#        print cmd

        try:

            output = Popen(cmd, stdout=PIPE)

        except OSError:
            # try the debug version
            cmd += 'd'

            try:
                    output = Popen(cmd, stdout=PIPE)
            except OSError:
                    print 'error when calling GenerateRigid, do you have GenerateRigid built in SOFA?'
                    raise

        output.communicate() # wait until Popen command id finished!!!

        # GenerateRigid output is stored in the file tmpfilename
        rigidFile = open( tmpfilename, "r" )
        line = list( rigidFile )
        rigidFile.close()

#        for i in xrange(len(line)):
#            print str(i) + str(line[i])

        start = 1

        res = MassInfo()

        res.mass = float( line[start].split(' ')[1] )
        #volm = float( line[start + 1].split(' ')[1] )
        res.inertia = map(float, line[start + 2].split(' ')[1:] ); res.inertia = [res.mass * x for x in res.inertia]
        res.com = map(float, line[start + 3].split(' ')[1:] )

        # TODO extract principal axes basis if needed
        # or at least say that we screwd up

#        print "generate_rigid " + str(res.mass) + " " + str( res.inertia )

        return res

class Body:
        # generic rigid body

        def __init__(self, name = "unnamed"):
                self.name = name         # node name
                self.collision = None # collision mesh
                self.visual = None    # visual mesh
                self.dofs = Frame()   # initial dofs
                self.mass = 1         # mass
                self.inertia = [1, 1, 1] # inertia tensor
                self.color = [1, 1, 1]   # not sure this is used
                self.offset = None       # rigid offset for com/inertia axes
                self.inertia_forces = False # compute inertia forces flag
                self.group = None
                self.mu = 0           # friction coefficient
                self.scale = [1, 1, 1]

                # TODO more if needed (scale, color)

        def mass_from_mesh(self, name, density = 1000.0):
                info = generate_rigid(name, density)

                self.mass = info.mass

                # TODO svd inertia tensor, extract rotation quaternion

                self.inertia = [info.inertia[0],
                                info.inertia[3 + 1],
                                info.inertia[6 + 2]]

                self.offset = Frame()
                self.offset.translation = info.com

                # TODO handle principal axes

        def insert(self, node):
                res = node.createChild( self.name )

                # mass offset, if any
                off = Frame()
                if self.offset != None:
                        off = self.offset

                # kinematic dofs
                frame = self.dofs * off
                dofs = frame.insert(res, name = 'dofs' )

                # dofs are now located at the mass frame, good
                mass = res.createObject('RigidMass',
                                        template = 'Rigid',
                                        name = 'mass',
                                        mass = self.mass,
                                        inertia = concat(self.inertia),
                                        inertia_forces = self.inertia_forces )

                # user node i.e. the one the user provided
                user = res.createChild( self.name + '-user' )
                off.inv().insert(user, name = 'dofs')
                user.createObject('AssembledRigidRigidMapping',
                                  template = 'Rigid,Rigid',
                                  source = '0 ' + str(off.inv()) )

                # visual model
                if self.visual != None:
                        visual_template = 'ExtVec3f'

                        visual = user.createChild( 'visual' )
                        ogl = visual.createObject('OglModel',
                                                  template = visual_template,
                                                  name = 'mesh',
                                                  fileMesh = self.visual,
                                                  color = concat(self.color),
                                                  scale3d = concat(self.scale))

                        visual_map = visual.createObject('RigidMapping',
                                                         template = 'Rigid' + ', ' + visual_template,
                                                         input = '@../')
                # collision model
                if self.collision != None:
                        collision = user.createChild('collision')

                        collision.createObject("MeshObjLoader",
					       name = 'loader',
					       filename = self.collision,
                                               scale3d = concat(self.scale) )

                        collision.createObject('MeshTopology',
                                               name = 'topology',
                                               triangles = '@loader.triangles')

			collision.createObject('MechanicalObject',
                                               name = 'dofs',
                                               position = '@loader.position')

			model = collision.createObject('TriangleModel',
                                               name = 'model',
                                               template = 'Vec3d',
                                               contactFriction = self.mu)
                        if self.group != None:
                                model.group = self.group

			collision.createObject('RigidMapping',
                                               template = 'Rigid,Vec3d',
                                               input = '@../',
                                               output = '@./')

                self.node = res
                self.user = user

                return res

class Joint:
        # generic rigid joint

        def __init__(self, name = 'joint'):
                self.dofs = [0] * 6
                self.body = []
                self.offset = []
                self.name = name
                self.pathToBodies = '' # to be used when bodies are not inserted in rootNode

                # link constraints compliance
                self.compliance = 0

                # TODO if you're looking for damping/stiffness, you
                # now should do it yourself, directly on joint dofs

        def append(self, node, offset = None):
                self.body.append(node)
                self.offset.append(offset)
                self.name = self.name + '-' + node.name

        # convenience: define joint using absolute frame and vararg nodes
        def absolute(self, frame, *nodes):
                for n in nodes:
                        pos = n.getObject('dofs').position
                        s = concat(pos[0])
                        local = Frame().read( s )
                        self.append(n, local.inv() * frame)

        # joint dimension
        def dim(self):
                return sum( self.dofs )

        class Node:
                pass

        def insert(self, parent):
                self.node = parent.createChild(self.name)

                # build input data for multimapping
                input = []
                for b, o in zip(self.body, self.offset):
                        if o is None:
                                input.append( '@' + Tools.node_path_rel(self.node, b) + '/dofs' )
                        else:
                                joint = b.createChild( self.name + '-offset' )

                                joint.createObject('MechanicalObject',
                                                   template = 'Rigid',
                                                   name = 'dofs' )

                                joint.createObject('AssembledRigidRigidMapping',
                                                   template = "Rigid,Rigid",
                                                   source = '0 ' + str( o ) )

                                input.append( '@' + Tools.node_path_rel(self.node, b) + '/' + joint.name + '/dofs' )

                if len(input) == 0:
                        print 'warning: empty joint'
                        return None

                dofs = self.node.createObject('MechanicalObject',
                                         template = 'Vec6d',
                                         name = 'dofs',
                                         position = '0 0 0 0 0 0' )

                map = self.node.createObject('RigidJointMultiMapping',
                                        name = 'mapping',
                                        template = 'Rigid,Vec6d',
                                        input = concat(input),
                                        output = '@dofs',
                                        pairs = "0 0")

		sub = self.node.createChild("constraints")

		sub.createObject('MechanicalObject',
				 template = 'Vec1d',
				 name = 'dofs')

		mask = [ (1 - d) for d in self.dofs ]

		map = sub.createObject('MaskMapping',
				       name = 'mapping',
				       template = 'Vec6d,Vec1d',
				       input = '@../',
				       output = '@dofs',
				       dofs = concat(mask) )

                compliance = sub.createObject('UniformCompliance',
					      name = 'compliance',
					      template = 'Vec1d',
					      compliance = self.compliance)

                stab = sub.createObject('Stabilization')

		return self.node

        def setTargetPose(self, targetPose, compliance="1e-3", damping="1e3"):
            """ Set the target pose of the joint -
            targetPose vector is filtered out according to free dofs of the joint
            """
            target = self.node.createChild("target")

            maskedTargetPose=[]
            print "dofs", self.dofs
            for i,d in enumerate(self.dofs):
                if d is 1:
                    maskedTargetPose.append(targetPose[i])
            print "maskedTargetPose", maskedTargetPose

            target.createObject('MechanicalObject', template = 'Vec1d', name = 'dofs')
            target.createObject('MaskMapping', name = 'mapping', template = 'Vec6d,Vec1d', input = '@../', output = '@dofs', dofs = "0 0 0 1 1 1" )

            target_constraint = target.createChild("target_constraint")
            target_constraint.createObject('MechanicalObject', template = 'Vec1d', name = 'dofs')
            target_constraint.createObject('OffsetMapping', name = 'mapping', template = 'Vec1d,Vec1d', input = '@../', output = '@dofs', offsets = concat(maskedTargetPose) )
            target_constraint.createObject('UniformCompliance', name = 'compliance', template = 'Vec1d', compliance = compliance, damping=damping)

# and now for more specific joints:

class SphericalJoint(Joint):

        def __init__(self, **args):
                Joint.__init__(self)
                self.dofs = [0, 0, 0, 1, 1, 1]
                self.name = 'spherical'

                for k in args:
                        setattr(self, k, args[k])

# this one has limits \o/
class RevoluteJoint(Joint):

        # TODO make this 'x', 'y', 'z' instead
        def __init__(self, axis, **args):
                Joint.__init__(self)
                self.dofs[3 + axis] = 1
                self.name = 'revolute'
                self.lower_limit = None
                self.upper_limit = None

                for k in args:
                        setattr(self, k, args[k])

        def insert(self, parent):
                res = Joint.insert(self, parent)

                if self.lower_limit == None and self.upper_limit == None:
                        return res

                limit = res.createChild('limit')

                dofs = limit.createObject('MechanicalObject', template = 'Vec1d')
                map = limit.createObject('ProjectionMapping', template = 'Vec6d, Vec1d' )

                limit.createObject('UniformCompliance', template = 'Vec1d', compliance = '0' )
                limit.createObject('UnilateralConstraint');

                # don't stabilize as we need to detect violated
                # constraints first
                # limit.createObject('Stabilization');

                set = []
                position = []
                offset = []

                if self.lower_limit != None:
                        set = set + [0] + self.dofs
                        position.append(0)
                        offset.append(self.lower_limit)

                if self.upper_limit != None:
                        set = set + [0] + vec.minus(self.dofs)
                        position.append(0)
                        offset.append(- self.upper_limit)

                map.set = concat(set)
                map.offset = concat(offset)
                dofs.position = concat(position)

                return res

class CylindricalJoint(Joint):

        def __init__(self, axis ):
                Joint.__init__(self)
                self.dofs[0 + axis] = 1
                self.dofs[3 + axis] = 1
                self.name = 'cylindrical'

class PrismaticJoint(Joint):

        def __init__(self, axis):
                Joint.__init__(self)
                self.dofs[0 + axis] = 1
                self.name = 'prismatic'

class PlanarJoint(Joint):

        def __init__(self, normal):
                Joint.__init__(self)
                self.dofs = [
                        int( (i != normal) if i < 3 else (i - 3 == normal) )
                        for i in xrange(6)
                ]
                self.name = 'planar'
<commit_msg>r11334/sofa : Compliant: fixing python generate_rigid in debug<commit_after># Rigid bodies and joints, Compliant-style.
#
# Authors: maxime.tournier@inria.fr, ... ?
#
# The basic philosophy is that python provides nice data structures to
# define scene component *semantically*. Once this is done, the actual
# scene graph must be generated through the 'insert' methods. Please
# refer to the python examples in Compliant for more information.
#

import Sofa

from subprocess import Popen, PIPE

from numpy import *
import SofaPython.Quaternion as quat
import Vec as vec

import Tools
from Tools import cat as concat

import os

class Frame:
        # a rigid frame, group operations are available.

        # TODO kwargs
        def __init__(self, value = None):
                if value != None:
                        self.translation = value[:3]
                        self.rotation = value[3:]
                else:
                        self.translation = [0, 0, 0]
                        self.rotation = [0, 0, 0, 1]

        def insert(self, parent, **args):
                return parent.createObject('MechanicalObject',
                                           template = 'Rigid',
                                           position = str(self),
                                           **args)

        def __str__(self):
                return concat(self.translation) + ' ' + concat(self.rotation)

        def copy(self):
                # TODO better !
                return Rigid.Frame().read( str(self) )

        def setValue(self, v):
             self.translation = v[:3]
             self.rotation = v[3:]
             return self

        def read(self, str):
                num = map(float, str.split())
                self.translation = num[:3]
                self.rotation = num[3:]
                return self

        def offset(self):
            return hstack((self.translation, self.rotation))

	def __mul__(self, other):
		res = Frame()
		res.translation = vec.sum(self.translation,
					  quat.rotate(self.rotation,
						      other.translation))
		res.rotation = quat.prod( self.rotation,
					  other.rotation)

		return res

	def inv(self):
		res = Frame()
		res.rotation = quat.conj( self.rotation )
		res.translation = vec.minus( quat.rotate(res.rotation,
							 self.translation) )
		return res

        def set(self, **kwargs):
                for k in kwargs:
                        setattr(self, k, kwargs[k])

                return self

        # TODO more: apply( vec3 ), wrench/twist frame change.

class MassInfo:
        pass

# front-end to sofa GenerateRigid tool. density unit is kg/m^3
def generate_rigid(filename, density = 1000.0, scale=[1,1,1]):

        # TODO bind GenerateRigid
        # - faster than writing in a file
        # - more robust (if several processes try to work in the same file)

        tmpfilename = Tools.path( __file__ ) +"/tmp.rigid"

        cmd = [ Sofa.build_dir() + '/bin/GenerateRigid', filename, tmpfilename, str(density), str(scale[0]), str(scale[1]), str(scale[2]) ]

#        print cmd

        try:

            output = Popen(cmd, stdout=PIPE)

        except OSError:
            # try the debug version
            cmd[0] += 'd'

            try:
                    output = Popen(cmd, stdout=PIPE)
            except OSError:
                    print 'error when calling GenerateRigid, do you have GenerateRigid built in SOFA?'
                    raise

        output.communicate() # wait until Popen command is finished!!!

        # GenerateRigid output is stored in the file tmpfilename
        rigidFile = open( tmpfilename, "r" )
        line = list( rigidFile )
        rigidFile.close()

#        for i in xrange(len(line)):
#            print str(i) + str(line[i])

        start = 1

        res = MassInfo()

        res.mass = float( line[start].split(' ')[1] )
        #volm = float( line[start + 1].split(' ')[1] )
        res.inertia = map(float, line[start + 2].split(' ')[1:] ); res.inertia = [res.mass * x for x in res.inertia]
        res.com = map(float, line[start + 3].split(' ')[1:] )

        # TODO extract principal axes basis if needed
        # or at least say that we screwd up

#        print "generate_rigid " + str(res.mass) + " " + str( res.inertia )

        return res

class Body:
        # generic rigid body

        def __init__(self, name = "unnamed"):
                self.name = name         # node name
                self.collision = None # collision mesh
                self.visual = None    # visual mesh
                self.dofs = Frame()   # initial dofs
                self.mass = 1         # mass
                self.inertia = [1, 1, 1] # inertia tensor
                self.color = [1, 1, 1]   # not sure this is used
                self.offset = None       # rigid offset for com/inertia axes
                self.inertia_forces = False # compute inertia forces flag
                self.group = None
                self.mu = 0           # friction coefficient
                self.scale = [1, 1, 1]

                # TODO more if needed (scale, color)

        def mass_from_mesh(self, name, density = 1000.0):
                info = generate_rigid(name, density)

                self.mass = info.mass

                # TODO svd inertia tensor, extract rotation quaternion

                self.inertia = [info.inertia[0],
                                info.inertia[3 + 1],
                                info.inertia[6 + 2]]

                self.offset = Frame()
                self.offset.translation = info.com

                # TODO handle principal axes

        def insert(self, node):
                res = node.createChild( self.name )

                # mass offset, if any
                off = Frame()
                if self.offset != None:
                        off = self.offset

                # kinematic dofs
                frame = self.dofs * off
                dofs = frame.insert(res, name = 'dofs' )

                # dofs are now located at the mass frame, good
                mass = res.createObject('RigidMass',
                                        template = 'Rigid',
                                        name = 'mass',
                                        mass = self.mass,
                                        inertia = concat(self.inertia),
                                        inertia_forces = self.inertia_forces )

                # user node i.e. the one the user provided
                user = res.createChild( self.name + '-user' )
                off.inv().insert(user, name = 'dofs')
                user.createObject('AssembledRigidRigidMapping',
                                  template = 'Rigid,Rigid',
                                  source = '0 ' + str(off.inv()) )

                # visual model
                if self.visual != None:
                        visual_template = 'ExtVec3f'

                        visual = user.createChild( 'visual' )
                        ogl = visual.createObject('OglModel',
                                                  template = visual_template,
                                                  name = 'mesh',
                                                  fileMesh = self.visual,
                                                  color = concat(self.color),
                                                  scale3d = concat(self.scale))

                        visual_map = visual.createObject('RigidMapping',
                                                         template = 'Rigid' + ', ' + visual_template,
                                                         input = '@../')
                # collision model
                if self.collision != None:
                        collision = user.createChild('collision')

                        collision.createObject("MeshObjLoader",
					       name = 'loader',
					       filename = self.collision,
                                               scale3d = concat(self.scale) )

                        collision.createObject('MeshTopology',
                                               name = 'topology',
                                               triangles = '@loader.triangles')

			collision.createObject('MechanicalObject',
                                               name = 'dofs',
                                               position = '@loader.position')

			model = collision.createObject('TriangleModel',
                                               name = 'model',
                                               template = 'Vec3d',
                                               contactFriction = self.mu)
                        if self.group != None:
                                model.group = self.group

			collision.createObject('RigidMapping',
                                               template = 'Rigid,Vec3d',
                                               input = '@../',
                                               output = '@./')

                self.node = res
                self.user = user

                return res

class Joint:
        # generic rigid joint

        def __init__(self, name = 'joint'):
                self.dofs = [0] * 6
                self.body = []
                self.offset = []
                self.name = name
                self.pathToBodies = '' # to be used when bodies are not inserted in rootNode

                # link constraints compliance
                self.compliance = 0

                # TODO if you're looking for damping/stiffness, you
                # now should do it yourself, directly on joint dofs

        def append(self, node, offset = None):
                self.body.append(node)
                self.offset.append(offset)
                self.name = self.name + '-' + node.name

        # convenience: define joint using absolute frame and vararg nodes
        def absolute(self, frame, *nodes):
                for n in nodes:
                        pos = n.getObject('dofs').position
                        s = concat(pos[0])
                        local = Frame().read( s )
                        self.append(n, local.inv() * frame)

        # joint dimension
        def dim(self):
                return sum( self.dofs )

        class Node:
                pass

        def insert(self, parent):
                self.node = parent.createChild(self.name)

                # build input data for multimapping
                input = []
                for b, o in zip(self.body, self.offset):
                        if o is None:
                                input.append( '@' + Tools.node_path_rel(self.node, b) + '/dofs' )
                        else:
                                joint = b.createChild( self.name + '-offset' )

                                joint.createObject('MechanicalObject',
                                                   template = 'Rigid',
                                                   name = 'dofs' )

                                joint.createObject('AssembledRigidRigidMapping',
                                                   template = "Rigid,Rigid",
                                                   source = '0 ' + str( o ) )

                                input.append( '@' + Tools.node_path_rel(self.node, b) + '/' + joint.name + '/dofs' )

                if len(input) == 0:
                        print 'warning: empty joint'
                        return None

                dofs = self.node.createObject('MechanicalObject',
                                         template = 'Vec6d',
                                         name = 'dofs',
                                         position = '0 0 0 0 0 0' )

                map = self.node.createObject('RigidJointMultiMapping',
                                        name = 'mapping',
                                        template = 'Rigid,Vec6d',
                                        input = concat(input),
                                        output = '@dofs',
                                        pairs = "0 0")

		sub = self.node.createChild("constraints")

		sub.createObject('MechanicalObject',
				 template = 'Vec1d',
				 name = 'dofs')

		mask = [ (1 - d) for d in self.dofs ]

		map = sub.createObject('MaskMapping',
				       name = 'mapping',
				       template = 'Vec6d,Vec1d',
				       input = '@../',
				       output = '@dofs',
				       dofs = concat(mask) )

                compliance = sub.createObject('UniformCompliance',
					      name = 'compliance',
					      template = 'Vec1d',
					      compliance = self.compliance)

                stab = sub.createObject('Stabilization')

		return self.node

        def setTargetPose(self, targetPose, compliance="1e-3", damping="1e3"):
            """ Set the target pose of the joint -
            targetPose vector is filtered out according to free dofs of the joint
            """
            target = self.node.createChild("target")

            maskedTargetPose=[]
            print "dofs", self.dofs
            for i,d in enumerate(self.dofs):
                if d is 1:
                    maskedTargetPose.append(targetPose[i])
            print "maskedTargetPose", maskedTargetPose

            target.createObject('MechanicalObject', template = 'Vec1d', name = 'dofs')
            target.createObject('MaskMapping', name = 'mapping', template = 'Vec6d,Vec1d', input = '@../', output = '@dofs', dofs = "0 0 0 1 1 1" )

            target_constraint = target.createChild("target_constraint")
            target_constraint.createObject('MechanicalObject', template = 'Vec1d', name = 'dofs')
            target_constraint.createObject('OffsetMapping', name = 'mapping', template = 'Vec1d,Vec1d', input = '@../', output = '@dofs', offsets = concat(maskedTargetPose) )
            target_constraint.createObject('UniformCompliance', name = 'compliance', template = 'Vec1d', compliance = compliance, damping=damping)

# and now for more specific joints:

class SphericalJoint(Joint):

        def __init__(self, **args):
                Joint.__init__(self)
                self.dofs = [0, 0, 0, 1, 1, 1]
                self.name = 'spherical'

                for k in args:
                        setattr(self, k, args[k])

# this one has limits \o/
class RevoluteJoint(Joint):

        # TODO make this 'x', 'y', 'z' instead
        def __init__(self, axis, **args):
                Joint.__init__(self)
                self.dofs[3 + axis] = 1
                self.name = 'revolute'
                self.lower_limit = None
                self.upper_limit = None

                for k in args:
                        setattr(self, k, args[k])

        def insert(self, parent):
                res = Joint.insert(self, parent)

                if self.lower_limit == None and self.upper_limit == None:
                        return res

                limit = res.createChild('limit')

                dofs = limit.createObject('MechanicalObject', template = 'Vec1d')
                map = limit.createObject('ProjectionMapping', template = 'Vec6d, Vec1d' )

                limit.createObject('UniformCompliance', template = 'Vec1d', compliance = '0' )
                limit.createObject('UnilateralConstraint');

                # don't stabilize as we need to detect violated
                # constraints first
                # limit.createObject('Stabilization');

                set = []
                position = []
                offset = []

                if self.lower_limit != None:
                        set = set + [0] + self.dofs
                        position.append(0)
                        offset.append(self.lower_limit)

                if self.upper_limit != None:
                        set = set + [0] + vec.minus(self.dofs)
                        position.append(0)
                        offset.append(- self.upper_limit)

                map.set = concat(set)
                map.offset = concat(offset)
                dofs.position = concat(position)

                return res

class CylindricalJoint(Joint):

        def __init__(self, axis ):
                Joint.__init__(self)
                self.dofs[0 + axis] = 1
                self.dofs[3 + axis] = 1
                self.name = 'cylindrical'

class PrismaticJoint(Joint):

        def __init__(self, axis):
                Joint.__init__(self)
                self.dofs[0 + axis] = 1
                self.name = 'prismatic'

class PlanarJoint(Joint):

        def __init__(self, normal):
                Joint.__init__(self)
                self.dofs = [
                        int( (i != normal) if i < 3 else (i - 3 == normal) )
                        for i in xrange(6)
                ]
                self.name = 'planar'
<|endoftext|>


<commit_before>"""
Test lldb-vscode setBreakpoints request
"""

from __future__ import print_function

import unittest2
import vscode
from lldbsuite.test.decorators import *
from lldbsuite.test.lldbtest import *
from lldbsuite.test import lldbutil
import lldbvscode_testcase
import shutil
import subprocess
import tempfile
import threading
import time

def spawn_and_wait(program, delay):
    if delay:
        time.sleep(delay)
    process = subprocess.Popen([program],
                               stdin=subprocess.PIPE,
                               stdout=subprocess.PIPE,
                               stderr=subprocess.PIPE)
    process.wait()

class TestVSCode_attach(lldbvscode_testcase.VSCodeTestCaseBase):

    mydir = TestBase.compute_mydir(__file__)

    def set_and_hit_breakpoint(self, continueToExit=True):
        source = 'main.c'
        breakpoint1_line = line_number(source, '// breakpoint 1')
        lines = [breakpoint1_line]
        # Set breakoint in the thread function so we can step the threads
        breakpoint_ids = self.set_source_breakpoints(source, lines)
        self.assertTrue(len(breakpoint_ids) == len(lines),
                        "expect correct number of breakpoints")
        self.continue_to_breakpoints(breakpoint_ids)
        if continueToExit:
            self.continue_to_exit()

    @skipIfWindows
    @no_debug_info_test
    def test_by_pid(self):
        '''
            Tests attaching to a process by process ID.
        '''
        self.build_and_create_debug_adaptor()
        program = self.getBuildArtifact("a.out")
        self.process = subprocess.Popen([program],
                                        stdin=subprocess.PIPE,
                                        stdout=subprocess.PIPE,
                                        stderr=subprocess.PIPE)
        self.attach(pid=self.process.pid)
        self.set_and_hit_breakpoint(continueToExit=True)

    @skipIfWindows
    @no_debug_info_test
    def test_by_name(self):
        '''
            Tests attaching to a process by process name.
        '''
        self.build_and_create_debug_adaptor()
        orig_program = self.getBuildArtifact("a.out")
        # Since we are going to attach by process name, we need a unique
        # process name that has minimal chance to match a process that is
        # already running. To do this we use tempfile.mktemp() to give us a
        # full path to a location where we can copy our executable. We then
        # run this copy to ensure we don't get the error "more that one
        # process matches 'a.out'".
        program = tempfile.mktemp()
        shutil.copyfile(orig_program, program)
        shutil.copymode(orig_program, program)
        self.process = subprocess.Popen([program],
                                        stdin=subprocess.PIPE,
                                        stdout=subprocess.PIPE,
                                        stderr=subprocess.PIPE)
        # Wait for a bit to ensure the process is launched
        time.sleep(1)
        self.attach(program=program)
        self.set_and_hit_breakpoint(continueToExit=True)

    @skipUnlessDarwin
    @no_debug_info_test
    def test_by_name_waitFor(self):
        '''
            Tests attaching to a process by process name and waiting for the
            next instance of a process to be launched, ingoring all current
            ones.
        '''
        self.build_and_create_debug_adaptor()
        program = self.getBuildArtifact("a.out")
        self.spawn_thread = threading.Thread(target=spawn_and_wait,
                                             args=(program, 1.0,))
        self.spawn_thread.start()
        self.attach(program=program, waitFor=True)
        self.set_and_hit_breakpoint(continueToExit=True)

    @skipIfWindows
    @no_debug_info_test
    def test_commands(self):
        '''
            Tests the "initCommands", "preRunCommands", "stopCommands",
            "exitCommands", and "attachCommands" that can be passed during
            attach.

            "initCommands" are a list of LLDB commands that get executed
            before the targt is created.
            "preRunCommands" are a list of LLDB commands that get executed
            after the target has been created and before the launch.
            "stopCommands" are a list of LLDB commands that get executed each
            time the program stops.
            "exitCommands" are a list of LLDB commands that get executed when
            the process exits
            "attachCommands" are a list of LLDB commands that get executed and
            must have a valid process in the selected target in LLDB after
            they are done executing. This allows custom commands to create any
            kind of debug session.
        '''
        self.build_and_create_debug_adaptor()
        program = self.getBuildArtifact("a.out")
        # Here we just create a target and launch the process as a way to test
        # if we are able to use attach commands to create any kind of a target
        # and use it for debugging
        attachCommands = [
            'target create -d "%s"' % (program),
            'process launch -- arg1'
        ]
        initCommands = ['target list', 'platform list']
        preRunCommands = ['image list a.out', 'image dump sections a.out']
        stopCommands = ['frame variable', 'bt']
        exitCommands = ['expr 2+3', 'expr 3+4']
        self.attach(program=program,
                    attachCommands=attachCommands,
                    initCommands=initCommands,
                    preRunCommands=preRunCommands,
                    stopCommands=stopCommands,
                    exitCommands=exitCommands)

        # Get output from the console. This should contain both the
        # "initCommands" and the "preRunCommands".
        output = self.get_console()
        # Verify all "initCommands" were found in console output
        self.verify_commands('initCommands', output, initCommands)
        # Verify all "preRunCommands" were found in console output
        self.verify_commands('preRunCommands', output, preRunCommands)

        functions = ['main']
        breakpoint_ids = self.set_function_breakpoints(functions)
        self.assertTrue(len(breakpoint_ids) == len(functions),
                        "expect one breakpoint")
        self.continue_to_breakpoints(breakpoint_ids)
        output = self.get_console(timeout=1.0)
        self.verify_commands('stopCommands', output, stopCommands)

        # Continue after launch and hit the "pause()" call and stop the target.
        # Get output from the console. This should contain both the
        # "stopCommands" that were run after we stop.
        self.vscode.request_continue()
        time.sleep(0.5)
        self.vscode.request_pause()
        self.vscode.wait_for_stopped()
        output = self.get_console(timeout=1.0)
        self.verify_commands('stopCommands', output, stopCommands)

        # Continue until the program exits
        self.continue_to_exit()
        # Get output from the console. This should contain both the
        # "exitCommands" that were run after the second breakpoint was hit
        output = self.get_console(timeout=1.0)
        self.verify_commands('exitCommands', output, exitCommands)
<commit_msg>Change the attach test case to use self.assertEqual so we can get more info on what is going wrong on test bots.<commit_after>"""
Test lldb-vscode setBreakpoints request
"""

from __future__ import print_function

import unittest2
import vscode
from lldbsuite.test.decorators import *
from lldbsuite.test.lldbtest import *
from lldbsuite.test import lldbutil
import lldbvscode_testcase
import os
import shutil
import subprocess
import tempfile
import threading
import time

def spawn_and_wait(program, delay):
    if delay:
        time.sleep(delay)
    process = subprocess.Popen([program],
                               stdin=subprocess.PIPE,
                               stdout=subprocess.PIPE,
                               stderr=subprocess.PIPE)
    process.wait()

class TestVSCode_attach(lldbvscode_testcase.VSCodeTestCaseBase):

    mydir = TestBase.compute_mydir(__file__)

    def set_and_hit_breakpoint(self, continueToExit=True):
        source = 'main.c'
        breakpoint1_line = line_number(source, '// breakpoint 1')
        lines = [breakpoint1_line]
        # Set breakoint in the thread function so we can step the threads
        breakpoint_ids = self.set_source_breakpoints(source, lines)
        self.assertEqual(len(breakpoint_ids), len(lines),
                         "expect correct number of breakpoints")
        self.continue_to_breakpoints(breakpoint_ids)
        if continueToExit:
            self.continue_to_exit()

    @skipIfWindows
    @no_debug_info_test
    def test_by_pid(self):
        '''
            Tests attaching to a process by process ID.
        '''
        self.build_and_create_debug_adaptor()
        program = self.getBuildArtifact("a.out")
        self.process = subprocess.Popen([program],
                                        stdin=subprocess.PIPE,
                                        stdout=subprocess.PIPE,
                                        stderr=subprocess.PIPE)
        self.attach(pid=self.process.pid)
        self.set_and_hit_breakpoint(continueToExit=True)

    @skipIfWindows
    @no_debug_info_test
    def test_by_name(self):
        '''
            Tests attaching to a process by process name.
        '''
        self.build_and_create_debug_adaptor()
        orig_program = self.getBuildArtifact("a.out")
        # Since we are going to attach by process name, we need a unique
        # process name that has minimal chance to match a process that is
        # already running. To do this we use tempfile.mktemp() to give us a
        # full path to a location where we can copy our executable. We then
        # run this copy to ensure we don't get the error "more that one
        # process matches 'a.out'".
        program = tempfile.mktemp()
        shutil.copyfile(orig_program, program)
        shutil.copymode(orig_program, program)

        def cleanup():
            if os.path.exists(program):
                os.unlink(program)
        # Execute the cleanup function during test case tear down.
        self.addTearDownHook(cleanup)

        self.process = subprocess.Popen([program],
                                        stdin=subprocess.PIPE,
                                        stdout=subprocess.PIPE,
                                        stderr=subprocess.PIPE)
        # Wait for a bit to ensure the process is launched
        time.sleep(1)
        self.attach(program=program)
        self.set_and_hit_breakpoint(continueToExit=True)

    @skipUnlessDarwin
    @no_debug_info_test
    def test_by_name_waitFor(self):
        '''
            Tests attaching to a process by process name and waiting for the
            next instance of a process to be launched, ingoring all current
            ones.
        '''
        self.build_and_create_debug_adaptor()
        program = self.getBuildArtifact("a.out")
        self.spawn_thread = threading.Thread(target=spawn_and_wait,
                                             args=(program, 1.0,))
        self.spawn_thread.start()
        self.attach(program=program, waitFor=True)
        self.set_and_hit_breakpoint(continueToExit=True)

    @skipIfWindows
    @no_debug_info_test
    def test_commands(self):
        '''
            Tests the "initCommands", "preRunCommands", "stopCommands",
            "exitCommands", and "attachCommands" that can be passed during
            attach.

            "initCommands" are a list of LLDB commands that get executed
            before the targt is created.
            "preRunCommands" are a list of LLDB commands that get executed
            after the target has been created and before the launch.
            "stopCommands" are a list of LLDB commands that get executed each
            time the program stops.
            "exitCommands" are a list of LLDB commands that get executed when
            the process exits
            "attachCommands" are a list of LLDB commands that get executed and
            must have a valid process in the selected target in LLDB after
            they are done executing. This allows custom commands to create any
            kind of debug session.
        '''
        self.build_and_create_debug_adaptor()
        program = self.getBuildArtifact("a.out")
        # Here we just create a target and launch the process as a way to test
        # if we are able to use attach commands to create any kind of a target
        # and use it for debugging
        attachCommands = [
            'target create -d "%s"' % (program),
            'process launch -- arg1'
        ]
        initCommands = ['target list', 'platform list']
        preRunCommands = ['image list a.out', 'image dump sections a.out']
        stopCommands = ['frame variable', 'bt']
        exitCommands = ['expr 2+3', 'expr 3+4']
        self.attach(program=program,
                    attachCommands=attachCommands,
                    initCommands=initCommands,
                    preRunCommands=preRunCommands,
                    stopCommands=stopCommands,
                    exitCommands=exitCommands)

        # Get output from the console. This should contain both the
        # "initCommands" and the "preRunCommands".
        output = self.get_console()
        # Verify all "initCommands" were found in console output
        self.verify_commands('initCommands', output, initCommands)
        # Verify all "preRunCommands" were found in console output
        self.verify_commands('preRunCommands', output, preRunCommands)

        functions = ['main']
        breakpoint_ids = self.set_function_breakpoints(functions)
        self.assertTrue(len(breakpoint_ids) == len(functions),
                        "expect one breakpoint")
        self.continue_to_breakpoints(breakpoint_ids)
        output = self.get_console(timeout=1.0)
        self.verify_commands('stopCommands', output, stopCommands)

        # Continue after launch and hit the "pause()" call and stop the target.
        # Get output from the console. This should contain both the
        # "stopCommands" that were run after we stop.
        self.vscode.request_continue()
        time.sleep(0.5)
        self.vscode.request_pause()
        self.vscode.wait_for_stopped()
        output = self.get_console(timeout=1.0)
        self.verify_commands('stopCommands', output, stopCommands)

        # Continue until the program exits
        self.continue_to_exit()
        # Get output from the console. This should contain both the
        # "exitCommands" that were run after the second breakpoint was hit
        output = self.get_console(timeout=1.0)
        self.verify_commands('exitCommands', output, exitCommands)
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""Databases.py: Functions to load MINE databases from and dump compounds into common cheminformatics formats"""

__author__ = 'JGJeffryes'

from pymongo import DESCENDING
from rdkit.Chem import AllChem
from databases import MINE
import utils
import sys
import hashlib
import os
import datetime
import csv

def export_sdf(mine_db, dir_path, max_compounds=None):
    """
    Exports compounds from the database as an MDL SDF file
    :param mine_db: a MINE object
    :param path: directory for files
    :param max_compounds: maximum number of compounds per file (defaults to 100000)
    :return:
    """
    if not mine_db.compounds.find_one({"Product_of": {'$exists': 1}}):
        mine_db.add_rxn_pointers()

    print("Exporting %s compounds from %s as an SDF file" %(mine_db.compounds.count(), mine_db.name))
    target = utils.prevent_overwrite(os.path.join(dir_path, mine_db.name)+"_1.sdf")
    w = AllChem.SDWriter(target)
    w.SetKekulize(True)
    n_files = 1
    for compound in mine_db.compounds.find():
        mol = AllChem.MolFromSmiles(compound['SMILES'], True, {'CoA': '*', 'R': "*"})
        if mol:
            mol.SetProp('_id', compound['_id'])
            mol.SetProp('Generation', str(compound['Generation']))
            if 'Reactant_in' in compound:
                mol.SetProp('Reactant_in', str(compound['Reactant_in']))
            if 'Product_of' in compound:
                mol.SetProp('Product_of', str(compound['Product_of']))
            w.write(mol)
            if max_compounds and (w.NumMols() >= max_compounds):
                n_files += 1
                target = utils.prevent_overwrite(os.path.join(dir_path, mine_db.name) + "_%s.sdf" % n_files)
                w = AllChem.SmilesWriter(target)
    w.close()

def export_smiles(mine_db, dir_path, max_compounds=None):
    """
    Exports compounds from the database as an SMILES file
    :param mine_db: a MINE object
    :param dir_path: directory for files
    :param max_compounds: maximum number of compounds per file (defaults to unlimited)
    :return:
    """
    header = ['SMILES', "_id", "Generation", 'Reactant_in', 'Product_of']
    if not mine_db.compounds.find_one({"Product_of": {'$exists': 1}}):
        mine_db.add_rxn_pointers()

    print("Exporting %s compounds from %s as SMILES file" % (mine_db.compounds.count(), mine_db.name))
    target = open(utils.prevent_overwrite(os.path.join(dir_path, mine_db.name)+"_1.smiles"), 'w')

    w = csv.DictWriter(target, fieldnames=header, dialect='excel-tab')
    n_files = 1
    i = 0
    for compound in mine_db.compounds.find({}, dict([(x, 1) for x in header])):
        w.writerow(compound)
        i += 1
        if max_compounds and not (i % max_compounds):
            n_files += 1
            target = utils.prevent_overwrite(os.path.join(dir_path, mine_db.name)+"_%s.smiles" % n_files)
            w = csv.DictWriter(target, fieldnames=header, dialect='excel-tab')

def export_mol(mine_db, target, name_field='_id'):
    """
    Exports compounds from the database as MDL molfiles
    :param mine_db: The database to export
    :type mine_db: a MINE object
    :param target: a directory in which to place the files
    :type target: string
    :param name_field: the field to provide names for the mol files. Must be unique & universal
    :type name_field: string
    :return:
    :rtype:
    """
    if not os.path.exists(target):
        os.mkdir(target)

    if mine_db.compounds.find().count() != mine_db.compounds.find({name_field: {'$exists': 1}}).count():
        raise ValueError('%s does not exist for every compound in the database' % name_field)

    for compound in mine_db.compounds.find({'_id': {'$regex': '^C'}}):
        mol = AllChem.MolFromSmiles(compound['SMILES'], True, {'CoA': '*', 'R': "*"})
        if "." in name_field:
            compound[name_field] = utils.get_dotted_field(compound, name_field)
        if isinstance(compound[name_field], list):
            compound[name_field] = ','.join(compound[name_field])
        AllChem.MolToMolFile(mol, os.path.join(target, compound[name_field]+'.mol'))

def import_sdf(mine_db, target,):
    """
    Imports a SDF file as a MINE database
    :param mine_db: a MINE object, the database to insert the compound into
    :param target: a path, the SDF file to be loaded
    :return:
    """
    sdf_gen = AllChem.SDMolSupplier(target)
    for mol in sdf_gen:
        mine_db.insert_compound(mol, compound_dict=mol.GetPropsAsDict(), pubchem_db=None, kegg_db=None, modelseed_db=None)
    mine_db.meta_data.insert({"Timestamp": datetime.datetime.now(), "Action": "SDF Imported", "Filepath": target})

if __name__ == '__main__':
    task = sys.argv[1]
    db_name = sys.argv[2]
    path = sys.argv[3]
    database = MINE(db_name)
    if task == 'export-sdf':
        export_sdf(database, path, max_compounds=10000)
    if task == 'export-smi':
        export_smiles(database, path)
    if task == 'export-mol':
        if len(sys.argv) == 5:
            export_mol(database, path, sys.argv[4])
        else:
            export_mol(database, path)
    if task == 'import-sdf':
        import_sdf(database, path)
<commit_msg>file size limit support<commit_after>#!/usr/bin/env python
"""Databases.py: Functions to load MINE databases from and dump compounds into common cheminformatics formats"""

__author__ = 'JGJeffryes'

from pymongo import DESCENDING
from rdkit.Chem import AllChem
from databases import MINE
import utils
import sys
import hashlib
import os
import datetime
import csv

def export_sdf(mine_db, dir_path, max_compounds=None):
    """
    Exports compounds from the database as an MDL SDF file
    :param mine_db: a MINE object
    :param path: directory for files
    :param max_compounds: maximum number of compounds per file (defaults to 100000)
    :return:
    """
    if not mine_db.compounds.find_one({"Product_of": {'$exists': 1}}):
        mine_db.add_rxn_pointers()

    print("Exporting %s compounds from %s as an SDF file" %(mine_db.compounds.count(), mine_db.name))
    target = utils.prevent_overwrite(os.path.join(dir_path, mine_db.name)+"_1.sdf")
    w = AllChem.SDWriter(target)
    w.SetKekulize(True)
    n_files = 1
    for compound in mine_db.compounds.find():
        mol = AllChem.MolFromSmiles(compound['SMILES'], True, {'CoA': '*', 'R': "*"})
        if mol:
            mol.SetProp('_id', compound['_id'])
            mol.SetProp('Generation', str(compound['Generation']))
            if 'Reactant_in' in compound:
                mol.SetProp('Reactant_in', str(compound['Reactant_in']))
            if 'Product_of' in compound:
                mol.SetProp('Product_of', str(compound['Product_of']))
            w.write(mol)
            if max_compounds and (w.NumMols() >= max_compounds):
                n_files += 1
                target = utils.prevent_overwrite(os.path.join(dir_path, mine_db.name) + "_%s.sdf" % n_files)
                w = AllChem.SmilesWriter(target)
    w.close()

def export_smiles(mine_db, dir_path, max_compounds=None):
    """
    Exports compounds from the database as an SMILES file
    :param mine_db: a MINE object
    :param dir_path: directory for files
    :param max_compounds: maximum number of compounds per file (defaults to unlimited)
    :return:
    """
    header = ['SMILES', "_id", "Generation", 'Reactant_in', 'Product_of']
    if not mine_db.compounds.find_one({"Product_of": {'$exists': 1}}):
        mine_db.add_rxn_pointers()

    print("Exporting %s compounds from %s as SMILES file" % (mine_db.compounds.count(), mine_db.name))
    target = open(utils.prevent_overwrite(os.path.join(dir_path, mine_db.name)+"_1.smiles"), 'w')

    w = csv.DictWriter(target, fieldnames=header, dialect='excel-tab')
    n_files = 1
    i = 0
    for compound in mine_db.compounds.find({}, dict([(x, 1) for x in header])):
        w.writerow(compound)
        i += 1
        if max_compounds and not (i % max_compounds):
            n_files += 1
            target = open(utils.prevent_overwrite(os.path.join(dir_path, mine_db.name)+"_%s.smiles" % n_files), 'w')
            w = csv.DictWriter(target, fieldnames=header, dialect='excel-tab')

def export_mol(mine_db, target, name_field='_id'):
    """
    Exports compounds from the database as MDL molfiles
    :param mine_db: The database to export
    :type mine_db: a MINE object
    :param target: a directory in which to place the files
    :type target: string
    :param name_field: the field to provide names for the mol files. Must be unique & universal
    :type name_field: string
    :return:
    :rtype:
    """
    if not os.path.exists(target):
        os.mkdir(target)

    if mine_db.compounds.find().count() != mine_db.compounds.find({name_field: {'$exists': 1}}).count():
        raise ValueError('%s does not exist for every compound in the database' % name_field)

    for compound in mine_db.compounds.find({'_id': {'$regex': '^C'}}):
        mol = AllChem.MolFromSmiles(compound['SMILES'], True, {'CoA': '*', 'R': "*"})
        if "." in name_field:
            compound[name_field] = utils.get_dotted_field(compound, name_field)
        if isinstance(compound[name_field], list):
            compound[name_field] = ','.join(compound[name_field])
        AllChem.MolToMolFile(mol, os.path.join(target, compound[name_field]+'.mol'))

def import_sdf(mine_db, target,):
    """
    Imports a SDF file as a MINE database
    :param mine_db: a MINE object, the database to insert the compound into
    :param target: a path, the SDF file to be loaded
    :return:
    """
    sdf_gen = AllChem.SDMolSupplier(target)
    for mol in sdf_gen:
        mine_db.insert_compound(mol, compound_dict=mol.GetPropsAsDict(), pubchem_db=None, kegg_db=None, modelseed_db=None)
    mine_db.meta_data.insert({"Timestamp": datetime.datetime.now(), "Action": "SDF Imported", "Filepath": target})

if __name__ == '__main__':
    task = sys.argv[1]
    db_name = sys.argv[2]
    path = sys.argv[3]
    database = MINE(db_name)
    if task == 'export-sdf':
        if len(sys.argv) == 5:
            export_sdf(database, path, int(sys.argv[4]))
        else:
            export_sdf(database, path)
    if task == 'export-smi':
        if len(sys.argv) == 5:
            export_smiles(database, path, int(sys.argv[4]))
        else:
            export_smiles(database, path)
    if task == 'export-mol':
        if len(sys.argv) == 5:
            export_mol(database, path, sys.argv[4])
        else:
            export_mol(database, path)
    if task == 'import-sdf':
        import_sdf(database, path)
<|endoftext|>