<commit_before>"""
IDAPython script providing an user interface for setting the configuration for
annotating database files with MSDN information.

Authors: Moritz Raabe, William Ballenthin
Copyright 2014 Mandiant, A FireEye Company

Mandiant licenses this file to you under the Apache License, Version
2.0 (the "License"); you may not use this file except in compliance with the
License. You may obtain a copy of the License at:

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied. See the License for the specific language governing
permissions and limitations under the License.
"""

import sys
import logging
import traceback
import ConfigParser
from ConfigParser import SafeConfigParser
from PyQt5 import QtWidgets
from PyQt5 import QtCore
from PyQt5.QtCore import Qt

idaapi.require("IDB_MSDN_Annotator")

CONFIG_FILE = 'MSDN_annotations.cfg'

g_logger = logging.getLogger(__name__)

def getDefaultMsdnDataDir():
    return os.path.abspath(os.path.join(idaapi.get_user_idadir(), 'MSDN_data'))

class MSDNAnnotationDialog(QtWidgets.QDialog):

    def read_config(self):
        config = {}
        if not self.config_parser.has_section('Functions') or \
           not self.config_parser.has_section('Arguments') or \
           not self.config_parser.has_section('Constants'):
            # Create default
            self.config_parser.add_section('Functions')
            self.config_parser.add_section('Arguments')
            self.config_parser.add_section('Constants')
            config['functions_annotate'] = True
            config['functions_repeatable_comment'] = False
            config['arguments_annotate'] = True
            config['constants_import'] = True
            config['msdn_data_dir'] = getDefaultMsdnDataDir()

        else:
            # Read existing
            config['functions_annotate'] = self.config_parser.getboolean('Functions', 'annotate')
            config['functions_repeatable_comment'] = self.config_parser.getboolean('Functions', 'repeatable_comment')
            config['arguments_annotate'] = self.config_parser.getboolean('Arguments', 'annotate')
            config['constants_import'] = self.config_parser.getboolean('Constants', 'import')
            try:
                config['msdn_data_dir'] = self.config_parser.get('Constants', 'msdn_data_dir')
            except ConfigParser.NoOptionError:
                config['msdn_data_dir'] = getDefaultMsdnDataDir()

        return config

    def save_config(self):
        self.config_parser.set('Functions', 'annotate', str(self.chkFunctionsAnnotate.isChecked()))
        self.config_parser.set('Functions', 'repeatable_comment', str(self.chkFuntcsRepeatable.isChecked()))
        self.config_parser.set('Arguments', 'annotate', str(self.chkArgumentsAnnotate.isChecked()))
        self.config_parser.set('Constants', 'import', str(self.chkConstantsImport.isChecked()))
        self.config_parser.set('Constants', 'msdn_data_dir', str(self.dirText.text()) )

        with open(self.file_path, 'wb') as conffile:
            self.config_parser.write(conffile)

    def change_image(self):
        funct = self.chkFunctionsAnnotate.isChecked() and \
            self.chkFuntcsRepeatable.isChecked()
        image = "{}-{}-{}.png".format(int(funct),
                                          int(self.chkArgumentsAnnotate
                                              .isChecked()),
                                          int(self.chkConstantsImport
                                              .isChecked()))
        img_path = os.path.abspath(os.path.join(os.path.dirname(__file__), 'IDB_MSDN_Annotator', 'img'))
        self.pic.setPixmap(QtWidgets.QPixmap(os.path.join(img_path, image)))

    def on_select_dir(self):
        msdnDir = QtWidgets.QFileDialog.getExistingDirectory(caption='Select directory containing MSDN XML Database')
        if len(msdnDir) != 0:
            self.dirText.setText(msdnDir)

    def toggle_option(self):
        disable = not self.chkFunctionsAnnotate.isChecked()
        self.chkFuntcsRepeatable.setDisabled(disable)
        self.change_image()

    def on_ok_button(self):
        #test the msdn data dir

        msdnpath = os.path.join(self.dirText.text(), IDB_MSDN_Annotator.MSDN_INFO_FILE)
        if not os.path.exists(msdnpath):
            g_logger.info('Error - no msdn info file: %s', msdnpath)
            ret = QtWidgets.QMessageBox.warning(self, 'MSDN Info Not Found', 'The file %s was not found in the specified MSDN Data Directory' % IDB_MSDN_Annotator.MSDN_INFO_FILE, QtWidgets.QMessageBox.Ok)
            #self.done(QtWidgets.QDialog.Rejected)
            return

        self.done(QtWidgets.QDialog.Accepted)
        g_logger.info('Saving config')
        self.save_config()
        config = self.read_config()
        idaapi.set_script_timeout(1)
        IDB_MSDN_Annotator.main(config)
        idaapi.set_script_timeout(0)

    def set_form_values(self):
        # Set values according to configuration file
        if self.config['functions_annotate']:
            self.chkFunctionsAnnotate.setCheckState(QtCore.Qt.Checked)
            if self.config['functions_repeatable_comment']:
                self.chkFuntcsRepeatable.setCheckState(QtCore.Qt.Checked)
        else:
            self.chkFuntcsRepeatable.setDisabled(True)
            self.chkFuntcsRepeatable.setCheckState(QtCore.Qt.Unchecked)
        if self.config['arguments_annotate']:
            self.chkArgumentsAnnotate.setCheckState(QtCore.Qt.Checked)
        if self.config['constants_import']:
            self.chkConstantsImport.setCheckState(QtCore.Qt.Checked)
        self.dirText.setText(self.config['msdn_data_dir'])

    def populate_form(self):
        layout = QtWidgets.QVBoxLayout()

        # Functions
        layout1 = QtWidgets.QVBoxLayout()
        groupBox = QtWidgets.QGroupBox('Markup Options')
        self.chkFunctionsAnnotate = QtWidgets.QCheckBox("Annotate function names"
                                                    " (see note)")
        layout1.addWidget(self.chkFunctionsAnnotate)
        self.chkFuntcsRepeatable = QtWidgets.QCheckBox("Use repeatable comments "
                                                   "for function name "
                                                   "annotations")
        layout1.addWidget(self.chkFuntcsRepeatable)

        # Arguments
        self.chkArgumentsAnnotate = QtWidgets.QCheckBox("Annotate function "
                                                    "arguments (see note)")
        layout1.addWidget(self.chkArgumentsAnnotate)

        # Constants
        self.chkConstantsImport = QtWidgets.QCheckBox("Rename constants")
        layout1.addWidget(self.chkConstantsImport)

        groupBox.setLayout(layout1)
        layout.addWidget(groupBox)

        #MSDN data dir
        hlayout = QtWidgets.QHBoxLayout()
        self.selectDirButton = QtWidgets.QPushButton('...')
        self.selectDirButton.clicked.connect(self.on_select_dir)
        hlayout.addWidget(self.selectDirButton)
        self.dirText = QtWidgets.QLineEdit()
        self.dirText.setReadOnly(True)
        hlayout.addWidget(self.dirText)
        groupBox = QtWidgets.QGroupBox('MSDN Data Directory')
        groupBox.setLayout(hlayout)
        layout.addWidget(groupBox)

        # Toggle
        self.chkFunctionsAnnotate.clicked.connect(self.toggle_option)
        self.chkFuntcsRepeatable.clicked.connect(self.change_image)
        self.chkArgumentsAnnotate.clicked.connect(self.change_image)
        self.chkConstantsImport.clicked.connect(self.change_image)

        self.set_form_values()

        info_string = "Note: Annotating functions and/or arguments allows " \
                      "you to hover\nthe respective element in order to " \
                      "show its description."
        layout.addWidget(QtWidgets.QLabel(info_string))

        # Buttons
        button_ok = QtWidgets.QPushButton('&OK')
        button_ok.setDefault(True)
        button_ok.clicked.connect(self.on_ok_button)
        #button_ok.clicked.connect(self.close)
        layout.addWidget(button_ok)
        button_cancel = QtWidgets.QPushButton('&Cancel')
        button_cancel.clicked.connect(self.close)
        layout.addWidget(button_cancel)

        # Image
        self.pic = QtWidgets.QLabel()
        self.pic.setGeometry(0, 0, 663, 203)
        self.change_image()

        # Layout right
        layout_r = QtWidgets.QVBoxLayout()
        #layout_r.addWidget(QtWidgets.QLabel("Annotation preview"))
        layout_r.addWidget(self.pic)
        groupBox = QtWidgets.QGroupBox('Annotation preview')
        groupBox.setLayout(layout_r)

        # Setup layouts
        h_layout = QtWidgets.QHBoxLayout()
        h_layout.addLayout(layout)
        #h_layout.addLayout(layout_r)
        h_layout.addWidget(groupBox)
        self.setLayout(h_layout)

    def __init__(self, parent=None):
        self._logger = logging.getLogger(__name__ + '.' +
                                         self.__class__.__name__)
        self._logger.debug('Starting UI')
        QtWidgets.QDialog.__init__(self, parent, QtCore.Qt.WindowSystemMenuHint |
                               QtCore.Qt.WindowTitleHint)
        self.setWindowTitle("MSDN Annotations Configuration")

        # Parse configuration file to dictionary
        self.file_path = os.path.abspath(os.path.join(idaapi.get_user_idadir(), CONFIG_FILE))
        self.config_parser = SafeConfigParser()
        self.config_parser.read(self.file_path)
        self.config = self.read_config()

        self.populate_form()

if __name__ == '__main__':
    #logging.basicConfig(level=logging.DEBUG)
    logging.basicConfig(level=logging.INFO)

    dlg = MSDNAnnotationDialog()
    # Disable script timeout -> otherwise cancel script dialog pops up
    oldTo = idaapi.set_script_timeout(0)
    dlg.exec_()
    # Restore the timeout
    idaapi.set_script_timeout(oldTo)
    g_logger.debug('UI closed')
<commit_msg>Fixed QtGui<commit_after>"""
IDAPython script providing an user interface for setting the configuration for
annotating database files with MSDN information.

Authors: Moritz Raabe, William Ballenthin
Copyright 2014 Mandiant, A FireEye Company

Mandiant licenses this file to you under the Apache License, Version
2.0 (the "License"); you may not use this file except in compliance with the
License. You may obtain a copy of the License at:

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied. See the License for the specific language governing
permissions and limitations under the License.
"""

import sys
import logging
import traceback
import ConfigParser
from ConfigParser import SafeConfigParser
from PyQt5 import QtWidgets
from PyQt5 import QtGui
from PyQt5 import QtCore
from PyQt5.QtCore import Qt

idaapi.require("IDB_MSDN_Annotator")

CONFIG_FILE = 'MSDN_annotations.cfg'

g_logger = logging.getLogger(__name__)

def getDefaultMsdnDataDir():
    return os.path.abspath(os.path.join(idaapi.get_user_idadir(), 'MSDN_data'))

class MSDNAnnotationDialog(QtWidgets.QDialog):

    def read_config(self):
        config = {}
        if not self.config_parser.has_section('Functions') or \
           not self.config_parser.has_section('Arguments') or \
           not self.config_parser.has_section('Constants'):
            # Create default
            self.config_parser.add_section('Functions')
            self.config_parser.add_section('Arguments')
            self.config_parser.add_section('Constants')
            config['functions_annotate'] = True
            config['functions_repeatable_comment'] = False
            config['arguments_annotate'] = True
            config['constants_import'] = True
            config['msdn_data_dir'] = getDefaultMsdnDataDir()

        else:
            # Read existing
            config['functions_annotate'] = self.config_parser.getboolean('Functions', 'annotate')
            config['functions_repeatable_comment'] = self.config_parser.getboolean('Functions', 'repeatable_comment')
            config['arguments_annotate'] = self.config_parser.getboolean('Arguments', 'annotate')
            config['constants_import'] = self.config_parser.getboolean('Constants', 'import')
            try:
                config['msdn_data_dir'] = self.config_parser.get('Constants', 'msdn_data_dir')
            except ConfigParser.NoOptionError:
                config['msdn_data_dir'] = getDefaultMsdnDataDir()

        return config

    def save_config(self):
        self.config_parser.set('Functions', 'annotate', str(self.chkFunctionsAnnotate.isChecked()))
        self.config_parser.set('Functions', 'repeatable_comment', str(self.chkFuntcsRepeatable.isChecked()))
        self.config_parser.set('Arguments', 'annotate', str(self.chkArgumentsAnnotate.isChecked()))
        self.config_parser.set('Constants', 'import', str(self.chkConstantsImport.isChecked()))
        self.config_parser.set('Constants', 'msdn_data_dir', str(self.dirText.text()) )

        with open(self.file_path, 'wb') as conffile:
            self.config_parser.write(conffile)

    def change_image(self):
        funct = self.chkFunctionsAnnotate.isChecked() and \
            self.chkFuntcsRepeatable.isChecked()
        image = "{}-{}-{}.png".format(int(funct),
                                          int(self.chkArgumentsAnnotate
                                              .isChecked()),
                                          int(self.chkConstantsImport
                                              .isChecked()))
        img_path = os.path.abspath(os.path.join(os.path.dirname(__file__), 'IDB_MSDN_Annotator', 'img'))
        self.pic.setPixmap(QtGui.QPixmap(os.path.join(img_path, image)))

    def on_select_dir(self):
        msdnDir = QtWidgets.QFileDialog.getExistingDirectory(caption='Select directory containing MSDN XML Database')
        if len(msdnDir) != 0:
            self.dirText.setText(msdnDir)

    def toggle_option(self):
        disable = not self.chkFunctionsAnnotate.isChecked()
        self.chkFuntcsRepeatable.setDisabled(disable)
        self.change_image()

    def on_ok_button(self):
        #test the msdn data dir

        msdnpath = os.path.join(self.dirText.text(), IDB_MSDN_Annotator.MSDN_INFO_FILE)
        if not os.path.exists(msdnpath):
            g_logger.info('Error - no msdn info file: %s', msdnpath)
            ret = QtWidgets.QMessageBox.warning(self, 'MSDN Info Not Found', 'The file %s was not found in the specified MSDN Data Directory' % IDB_MSDN_Annotator.MSDN_INFO_FILE, QtWidgets.QMessageBox.Ok)
            #self.done(QtWidgets.QDialog.Rejected)
            return

        self.done(QtWidgets.QDialog.Accepted)
        g_logger.info('Saving config')
        self.save_config()
        config = self.read_config()
        idaapi.set_script_timeout(1)
        IDB_MSDN_Annotator.main(config)
        idaapi.set_script_timeout(0)

    def set_form_values(self):
        # Set values according to configuration file
        if self.config['functions_annotate']:
            self.chkFunctionsAnnotate.setCheckState(QtCore.Qt.Checked)
            if self.config['functions_repeatable_comment']:
                self.chkFuntcsRepeatable.setCheckState(QtCore.Qt.Checked)
        else:
            self.chkFuntcsRepeatable.setDisabled(True)
            self.chkFuntcsRepeatable.setCheckState(QtCore.Qt.Unchecked)
        if self.config['arguments_annotate']:
            self.chkArgumentsAnnotate.setCheckState(QtCore.Qt.Checked)
        if self.config['constants_import']:
            self.chkConstantsImport.setCheckState(QtCore.Qt.Checked)
        self.dirText.setText(self.config['msdn_data_dir'])

    def populate_form(self):
        layout = QtWidgets.QVBoxLayout()

        # Functions
        layout1 = QtWidgets.QVBoxLayout()
        groupBox = QtWidgets.QGroupBox('Markup Options')
        self.chkFunctionsAnnotate = QtWidgets.QCheckBox("Annotate function names"
                                                    " (see note)")
        layout1.addWidget(self.chkFunctionsAnnotate)
        self.chkFuntcsRepeatable = QtWidgets.QCheckBox("Use repeatable comments "
                                                   "for function name "
                                                   "annotations")
        layout1.addWidget(self.chkFuntcsRepeatable)

        # Arguments
        self.chkArgumentsAnnotate = QtWidgets.QCheckBox("Annotate function "
                                                    "arguments (see note)")
        layout1.addWidget(self.chkArgumentsAnnotate)

        # Constants
        self.chkConstantsImport = QtWidgets.QCheckBox("Rename constants")
        layout1.addWidget(self.chkConstantsImport)

        groupBox.setLayout(layout1)
        layout.addWidget(groupBox)

        #MSDN data dir
        hlayout = QtWidgets.QHBoxLayout()
        self.selectDirButton = QtWidgets.QPushButton('...')
        self.selectDirButton.clicked.connect(self.on_select_dir)
        hlayout.addWidget(self.selectDirButton)
        self.dirText = QtWidgets.QLineEdit()
        self.dirText.setReadOnly(True)
        hlayout.addWidget(self.dirText)
        groupBox = QtWidgets.QGroupBox('MSDN Data Directory')
        groupBox.setLayout(hlayout)
        layout.addWidget(groupBox)

        # Toggle
        self.chkFunctionsAnnotate.clicked.connect(self.toggle_option)
        self.chkFuntcsRepeatable.clicked.connect(self.change_image)
        self.chkArgumentsAnnotate.clicked.connect(self.change_image)
        self.chkConstantsImport.clicked.connect(self.change_image)

        self.set_form_values()

        info_string = "Note: Annotating functions and/or arguments allows " \
                      "you to hover\nthe respective element in order to " \
                      "show its description."
        layout.addWidget(QtWidgets.QLabel(info_string))

        # Buttons
        button_ok = QtWidgets.QPushButton('&OK')
        button_ok.setDefault(True)
        button_ok.clicked.connect(self.on_ok_button)
        #button_ok.clicked.connect(self.close)
        layout.addWidget(button_ok)
        button_cancel = QtWidgets.QPushButton('&Cancel')
        button_cancel.clicked.connect(self.close)
        layout.addWidget(button_cancel)

        # Image
        self.pic = QtWidgets.QLabel()
        self.pic.setGeometry(0, 0, 663, 203)
        self.change_image()

        # Layout right
        layout_r = QtWidgets.QVBoxLayout()
        #layout_r.addWidget(QtWidgets.QLabel("Annotation preview"))
        layout_r.addWidget(self.pic)
        groupBox = QtWidgets.QGroupBox('Annotation preview')
        groupBox.setLayout(layout_r)

        # Setup layouts
        h_layout = QtWidgets.QHBoxLayout()
        h_layout.addLayout(layout)
        #h_layout.addLayout(layout_r)
        h_layout.addWidget(groupBox)
        self.setLayout(h_layout)

    def __init__(self, parent=None):
        self._logger = logging.getLogger(__name__ + '.' +
                                         self.__class__.__name__)
        self._logger.debug('Starting UI')
        QtWidgets.QDialog.__init__(self, parent, QtCore.Qt.WindowSystemMenuHint |
                               QtCore.Qt.WindowTitleHint)
        self.setWindowTitle("MSDN Annotations Configuration")

        # Parse configuration file to dictionary
        self.file_path = os.path.abspath(os.path.join(idaapi.get_user_idadir(), CONFIG_FILE))
        self.config_parser = SafeConfigParser()
        self.config_parser.read(self.file_path)
        self.config = self.read_config()

        self.populate_form()

if __name__ == '__main__':
    #logging.basicConfig(level=logging.DEBUG)
    logging.basicConfig(level=logging.INFO)

    dlg = MSDNAnnotationDialog()
    # Disable script timeout -> otherwise cancel script dialog pops up
    oldTo = idaapi.set_script_timeout(0)
    dlg.exec_()
    # Restore the timeout
    idaapi.set_script_timeout(oldTo)
    g_logger.debug('UI closed')
<|endoftext|>


<commit_before>from templates import TextTemplate, add_quick_reply
import requests
import config
import os
import json
from app import usersdb

ACCESS_TOKEN = os.environ.get('ACCESS_TOKEN', config.ACCESS_TOKEN)

def send_interest_menu(sender):
    out = {"keyword":"interest"}
    message = TextTemplate(text="Whom do you want to chat with?").get_message()
    out["interest"] = "male"
    message = add_quick_reply(message=message, title="Men", payload=json.dumps(out))
    out["interest"] = "female"
    message = add_quick_reply(message=message, title="Women", payload=json.dumps(out))
    out["interest"] = "random"
    message = add_quick_reply(message=message, title="Random", payload=json.dumps(out))

    send_message(message, sender)

def send_newchat_prompt(id):
    payload = {"keyword": "newchat"}
    payload["ans"] = "y"
    message = TextTemplate(text="Are you ready to start a new chat").get_message()
    message = add_quick_reply(message=message, title="Oh! Yes!", payload=json.dumps(payload))
    payload["ans"] = "n"
    message = add_quick_reply(message=message, title="No. Later", payload=json.dumps(payload))
    send_message(message, id)

def send_message(message, id, pause_check=False):

    if isinstance(message, dict) == False:
        message = message.get_message()

    if pause_check and usersdb.getPauseStatus(id):
        print("USER PAUSED, MESSAGE STORED")
        usersdb.addMessage(id=id, message=json.dumps(message))
        return

    if 'quick_replies' in message:
        usersdb.setPauseStatus(id=id, status=True)

    payload = {
        'recipient': {
            'id': id
        },
        'message': message
    }
    r = requests.post('https://graph.facebook.com/v2.6/me/messages', params={'access_token': ACCESS_TOKEN},
                      json=payload)

def send_help(sender):
    helptext = "BlindChat allows you to chat with people without revealing your identity. "+\
        "The bot will match you with strangers all over the world. You can choose to share your profile with the other person after ending the chat.\n"+\
        "\nAvailable commands:\n1. quit/exit: quits from the active chat or from the waitlist"+\
        "\n2. help: view the help menu\n4. start: starts to look for a new chat\n5. restart: restart the bot."

    message = TextTemplate(text=helptext)
    send_message(message.get_message(), sender)

def show_typing(id, duration):
    from time import sleep
    payload = {
        'recipient': {
            'id': id
        },
        "sender_action":"typing_on"
    }
    r = requests.post('https://graph.facebook.com/v2.6/me/messages', params={'access_token': ACCESS_TOKEN},
                      json=payload)

    sleep(duration)
    payload["sender_action"] = "typing_off"
    r = requests.post('https://graph.facebook.com/v2.6/me/messages', params={'access_token': ACCESS_TOKEN},
                      json=payload)

def send_paused_messages(id):
    if usersdb.getPauseStatus(id) == False:
        m_list = usersdb.getMessages(id)
        for m in m_list:
            send_message(message=json.loads(m), id=id)
<commit_msg>bug fix: pause conversation at quick replies<commit_after>from templates import TextTemplate, add_quick_reply
import requests
import config
import os
import json
from app import usersdb

ACCESS_TOKEN = os.environ.get('ACCESS_TOKEN', config.ACCESS_TOKEN)

def send_interest_menu(sender):
    out = {"keyword":"interest"}
    message = TextTemplate(text="Whom do you want to chat with?").get_message()
    out["interest"] = "male"
    message = add_quick_reply(message=message, title="Men", payload=json.dumps(out))
    out["interest"] = "female"
    message = add_quick_reply(message=message, title="Women", payload=json.dumps(out))
    out["interest"] = "random"
    message = add_quick_reply(message=message, title="Random", payload=json.dumps(out))

    send_message(message, sender)

def send_newchat_prompt(id):
    payload = {"keyword": "newchat"}
    payload["ans"] = "y"
    message = TextTemplate(text="Are you ready to start a new chat").get_message()
    message = add_quick_reply(message=message, title="Oh! Yes!", payload=json.dumps(payload))
    payload["ans"] = "n"
    message = add_quick_reply(message=message, title="No. Later", payload=json.dumps(payload))
    send_message(message, id)

def send_message(message, id, pause_check=False):

    if isinstance(message, dict) == False:
        message = message.get_message()

    if pause_check and usersdb.getPauseStatus(id):
        print("USER PAUSED, MESSAGE STORED")
        usersdb.addMessage(id=id, message=json.dumps(message))
        return

    if 'quick_replies' in message:
        usersdb.setPauseStatus(id=id, status=True)

    payload = {
        'recipient': {
            'id': id
        },
        'message': message
    }
    r = requests.post('https://graph.facebook.com/v2.6/me/messages', params={'access_token': ACCESS_TOKEN},
                      json=payload)

def send_help(sender):
    helptext = "BlindChat allows you to chat with people without revealing your identity. "+\
        "The bot will match you with strangers all over the world. You can choose to share your profile with the other person after ending the chat.\n"+\
        "\nAvailable commands:\n1. quit/exit: quits from the active chat or from the waitlist"+\
        "\n2. help: view the help menu\n4. start: starts to look for a new chat\n5. restart: restart the bot."

    message = TextTemplate(text=helptext)
    send_message(message.get_message(), sender)

def show_typing(id, duration):
    from time import sleep
    payload = {
        'recipient': {
            'id': id
        },
        "sender_action":"typing_on"
    }
    r = requests.post('https://graph.facebook.com/v2.6/me/messages', params={'access_token': ACCESS_TOKEN},
                      json=payload)

    sleep(duration)
    payload["sender_action"] = "typing_off"
    r = requests.post('https://graph.facebook.com/v2.6/me/messages', params={'access_token': ACCESS_TOKEN},
                      json=payload)

def send_paused_messages(id):
    if usersdb.getPauseStatus(id) == False:
        m_list = usersdb.getMessages(id)
        if len(m_list[0]) == 0:
            return
        for m in m_list:
            send_message(message=json.loads(m), id=id)
<|endoftext|>


<commit_before># -- start license --
# Highly Optimized Object-oriented Many-particle Dynamics -- Blue Edition
# (HOOMD-blue) Open Source Software License Copyright 2009-2014 The Regents of
# the University of Michigan All rights reserved.

# HOOMD-blue may contain modifications ("Contributions") provided, and to which
# copyright is held, by various Contributors who have granted The Regents of the
# University of Michigan the right to modify and/or distribute such Contributions.

# You may redistribute, use, and create derivate works of HOOMD-blue, in source
# and binary forms, provided you abide by the following conditions:

# * Redistributions of source code must retain the above copyright notice, this
# list of conditions, and the following disclaimer both in the code and
# prominently in any materials provided with the distribution.

# * Redistributions in binary form must reproduce the above copyright notice, this
# list of conditions, and the following disclaimer in the documentation and/or
# other materials provided with the distribution.

# * All publications and presentations based on HOOMD-blue, including any reports
# or published results obtained, in whole or in part, with HOOMD-blue, will
# acknowledge its use according to the terms posted at the time of submission on:
# http://codeblue.umich.edu/hoomd-blue/citations.html

# * Any electronic documents citing HOOMD-Blue will link to the HOOMD-Blue website:
# http://codeblue.umich.edu/hoomd-blue/

# * Apart from the above required attributions, neither the name of the copyright
# holder nor the names of HOOMD-blue's contributors may be used to endorse or
# promote products derived from this software without specific prior written
# permission.

# Disclaimer

# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDER AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, AND/OR ANY
# WARRANTIES THAT THIS SOFTWARE IS FREE OF INFRINGEMENT ARE DISCLAIMED.

# IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
# -- end license --

# Maintainer: joaander / All Developers are free to add commands for new features

## \package hoomd_script.meta
# \brief Write out simulation and environment context metadata
#
# Metadata is stored in form of key-value pairs in a JSON file and used
# to summarize the per-run simulation parameters so that they can be easily
# taken up by other scripts and stored in a database.
#
# The metadata is returned by dump_metadata() and can optionally be written to a file.
# Provide user defined metadata as mapping type.
#
# \code
# metadata = meta.dump_metadata()
# meta.dump_metadata(filename = "metadata.json", overwrite = False, extra = {'debug': True})
# \endcode

import hoomd;
import hoomd_script;
import json, collections;
import time
import datetime

from collections import OrderedDict

## \internal
# \brief A Mixin to facilitate storage of simulation metadata
class _metadata:
    def __init__(self):
        # No metadata provided per default
        self.metadata_fields = []

    ## \internal
    # \brief Return the metadata
    def get_metadata(self):
        data = OrderedDict()

        for m in self.metadata_fields:
            data[m] = getattr(self, m)

        return data

class _metadata_from_dict:
    def __init__(self, d):
        self.d = d;

    def get_metadata(self):
        data = OrderedDict()

        for m in self.d.keys():
            data[m] = self.d[m];

        return data

# Writes simulation metadata into a file
#
# When called, this function will query all registered forces, updaters etc.
# and ask them to provide metadata. E.g. a pair potential will return
# information about parameters, the Logger will output the filename it is
# logging to, etc.
#
# Custom metadata can be provided as a dictionary.
#
# The output is aggregated into a dictionary and written to a
# JSON file, together with a timestamp.
#
# \param filename The name of the file to write JSON metadata to (optional)
# \param user Additional metadata, needs to be a mapping type, such as a dict()
# \param overwrite If true, overwrite output file if it already exists
# \param indent The json indentation size
#
# \returns metadata as a dictionary
def dump_metadata(filename=None,user=None,overwrite=False,indent=4):
    hoomd_script.util.print_status_line();

    if not hoomd_script.init.is_initialized():
        hoomd_script.globals.msg.error("Need to initialize system first.\n")
        raise RuntimeError("Error writing out metadata.")

    metadata = []
    obj = OrderedDict()

    if user is not None:
        if not isinstance(user, collections.Mapping):
            globals.msg.warning("Extra meta data needs to be a mapping type. Ignoring.\n")
        else:
            obj['user'] = _metadata_from_dict(user);

    if not overwrite and filename is not None:
        try:
            with open(filename) as f:
                metadata = json.load(f)
                globals.msg.notice(2,"Appending to file {1}." % filename)
        except Exception:
            pass

    # Generate time stamp
    ts = time.time()
    st = datetime.datetime.fromtimestamp(ts).strftime('%Y-%m-%d %H:%M:%S')
    obj['timestamp'] = st
    obj['context'] = hoomd_script.context.ExecutionContext()
    obj['hoomd'] = hoomd_script.context.HOOMDContext()

    global_objs = [hoomd_script.data.system_data(hoomd_script.globals.system_definition)];
    global_objs += hoomd_script.globals.forces;
    global_objs += hoomd_script.globals.constraint_forces;
    global_objs += [hoomd_script.globals.integrator];
    global_objs += hoomd_script.globals.integration_methods;
    global_objs += hoomd_script.globals.forces
    global_objs += hoomd_script.globals.analyzers;
    global_objs += hoomd_script.globals.updaters;

    # add list of objects to JSON
    for o in global_objs:
        if o is not None:
            name = o.__module__+'.'+o.__class__.__name__;
            if len(name) > 13 and name[:13] == 'hoomd_script.':
                name = name[13:];
            obj[name] = o

    metadata.append(obj)

    # handler for unknown objects
    def default_handler(obj):
        try:
            return obj.get_metadata()
        except (AttributeError, NotImplementedError):
            return None

    # dump to JSON
    meta_str = json.dumps(metadata,default=default_handler,indent=indent)
    if filename is not None:
        with open(filename, 'w') as file:
            file.write(meta_str)
    return json.loads(meta_str)
<commit_msg>improve hoomd_script.meta doxygen docs<commit_after># -- start license --
# Highly Optimized Object-oriented Many-particle Dynamics -- Blue Edition
# (HOOMD-blue) Open Source Software License Copyright 2009-2014 The Regents of
# the University of Michigan All rights reserved.

# HOOMD-blue may contain modifications ("Contributions") provided, and to which
# copyright is held, by various Contributors who have granted The Regents of the
# University of Michigan the right to modify and/or distribute such Contributions.

# You may redistribute, use, and create derivate works of HOOMD-blue, in source
# and binary forms, provided you abide by the following conditions:

# * Redistributions of source code must retain the above copyright notice, this
# list of conditions, and the following disclaimer both in the code and
# prominently in any materials provided with the distribution.

# * Redistributions in binary form must reproduce the above copyright notice, this
# list of conditions, and the following disclaimer in the documentation and/or
# other materials provided with the distribution.

# * All publications and presentations based on HOOMD-blue, including any reports
# or published results obtained, in whole or in part, with HOOMD-blue, will
# acknowledge its use according to the terms posted at the time of submission on:
# http://codeblue.umich.edu/hoomd-blue/citations.html

# * Any electronic documents citing HOOMD-Blue will link to the HOOMD-Blue website:
# http://codeblue.umich.edu/hoomd-blue/

# * Apart from the above required attributions, neither the name of the copyright
# holder nor the names of HOOMD-blue's contributors may be used to endorse or
# promote products derived from this software without specific prior written
# permission.

# Disclaimer

# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDER AND CONTRIBUTORS ``AS IS'' AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, AND/OR ANY
# WARRANTIES THAT THIS SOFTWARE IS FREE OF INFRINGEMENT ARE DISCLAIMED.

# IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
# BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
# LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE
# OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
# -- end license --

# Maintainer: joaander / All Developers are free to add commands for new features

## \package hoomd_script.meta
# \brief Write out simulation and environment context metadata
#
# Metadata is stored in form of key-value pairs in a JSON file and used
# to summarize the per-run simulation parameters so that they can be easily
# taken up by other scripts and stored in a database.
#
# The metadata is returned by dump_metadata() and can optionally be written to a file.
# Provide user defined metadata as mapping type.
#
# \code
# metadata = meta.dump_metadata()
# meta.dump_metadata(filename = "metadata.json", overwrite = False, user = {'debug': True}, indent=2)
# \endcode

import hoomd;
import hoomd_script;
import json, collections;
import time
import datetime

from collections import OrderedDict

## \internal
# \brief A Mixin to facilitate storage of simulation metadata
class _metadata:
    def __init__(self):
        # No metadata provided per default
        self.metadata_fields = []

    ## \internal
    # \brief Return the metadata
    def get_metadata(self):
        data = OrderedDict()

        for m in self.metadata_fields:
            data[m] = getattr(self, m)

        return data

class _metadata_from_dict:
    def __init__(self, d):
        self.d = d;

    def get_metadata(self):
        data = OrderedDict()

        for m in self.d.keys():
            data[m] = self.d[m];

        return data

## Writes simulation metadata into a file
#
# When called, this function will query all registered forces, updaters etc.
# and ask them to provide metadata. E.g. a pair potential will return
# information about parameters, the Logger will output the filename it is
# logging to, etc.
#
# Custom metadata can be provided as a dictionary.
#
# The output is aggregated into a dictionary and written to a
# JSON file, together with a timestamp.
#
# \param filename The name of the file to write JSON metadata to (optional)
# \param user Additional metadata, needs to be a mapping type, such as a dict()
# \param overwrite If true, overwrite output file if it already exists
# \param indent The json indentation size
#
# \returns metadata as a dictionary
def dump_metadata(filename=None,user=None,overwrite=False,indent=4):
    hoomd_script.util.print_status_line();

    if not hoomd_script.init.is_initialized():
        hoomd_script.globals.msg.error("Need to initialize system first.\n")
        raise RuntimeError("Error writing out metadata.")

    metadata = []
    obj = OrderedDict()

    if user is not None:
        if not isinstance(user, collections.Mapping):
            globals.msg.warning("Extra meta data needs to be a mapping type. Ignoring.\n")
        else:
            obj['user'] = _metadata_from_dict(user);

    if not overwrite and filename is not None:
        try:
            with open(filename) as f:
                metadata = json.load(f)
                globals.msg.notice(2,"Appending to file {1}." % filename)
        except Exception:
            pass

    # Generate time stamp
    ts = time.time()
    st = datetime.datetime.fromtimestamp(ts).strftime('%Y-%m-%d %H:%M:%S')
    obj['timestamp'] = st
    obj['context'] = hoomd_script.context.ExecutionContext()
    obj['hoomd'] = hoomd_script.context.HOOMDContext()

    global_objs = [hoomd_script.data.system_data(hoomd_script.globals.system_definition)];
    global_objs += hoomd_script.globals.forces;
    global_objs += hoomd_script.globals.constraint_forces;
    global_objs += [hoomd_script.globals.integrator];
    global_objs += hoomd_script.globals.integration_methods;
    global_objs += hoomd_script.globals.forces
    global_objs += hoomd_script.globals.analyzers;
    global_objs += hoomd_script.globals.updaters;

    # add list of objects to JSON
    for o in global_objs:
        if o is not None:
            name = o.__module__+'.'+o.__class__.__name__;
            if len(name) > 13 and name[:13] == 'hoomd_script.':
                name = name[13:];
            obj[name] = o

    metadata.append(obj)

    # handler for unknown objects
    def default_handler(obj):
        try:
            return obj.get_metadata()
        except (AttributeError, NotImplementedError):
            return None

    # dump to JSON
    meta_str = json.dumps(metadata,default=default_handler,indent=indent)
    if filename is not None:
        with open(filename, 'w') as file:
            file.write(meta_str)
    return json.loads(meta_str)
<|endoftext|>


<commit_before>#!/usr/bin/env python3
"""
make_a_star_cluster.py creates a model star cluster,
which can then be used in N-body simulations or for other purposes.

It requires AMUSE, which can be downloaded from http://amusecode.org or
https://github.com/amusecode/amuse.

Currently not feature-complete yet, and function/argument names are
subject to change.

-- Steven Rieder steven at rieder punt nl
"""

import os
import argparse

import numpy as np

from amuse.io import write_set_to_file, read_set_from_file
from amuse.datamodel import Particles
from amuse.units import units
from amuse.support.console import set_printing_strategy
from amuse.ext.masc import new_star_cluster, new_stars_from_sink

def new_argument_parser():
    "Parse command line arguments"
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--sinks',
        dest='sinks',
        default=None,
        help="Generate stars from sink particles in this file (EXPERIMENTAL)",
    )
    parser.add_argument(
        '-o',
        dest='clustername',
        default="auto",
        help="Outputfile name [auto]",
    )
    parser.add_argument(
        '-type',
        dest='filetype',
        default="amuse",
        help="Output file type ([amuse]/txt/starlab/nemo)",
    )
    parser.add_argument(
        '-N',
        dest='number_of_stars',
        default=1024,
        type=int,
        help='Number of stars [1024]',
    )
    parser.add_argument(
        '-M',
        dest='cluster_mass',
        type=float,
        default=0,
        help='Cluster mass (in MSun) (takes precedence over number of stars)',
    )
    parser.add_argument(
        '-dist',
        dest='star_distribution',
        default="plummer",
        help="Star distribution ([plummer]/king/fractal)",
    )
    parser.add_argument(
        '-gasdist',
        dest='gas_distribution',
        default="none",
        help="Gas distribution ([none]/plummer/king/fractal) \
                NOT IMPLEMENTED YET",
    )
    parser.add_argument(
        '-imf',
        dest='initial_mass_function',
        default="kroupa",
        help="IMF ([kroupa]/salpeter/fixed)",
    )
    parser.add_argument(
        '-modelnr',
        dest='cluster_model_number',
        type=int,
        default=30483,
        help='Model number',
    )
    parser.add_argument(
        '-kingw',
        dest='king_parameter_w0',
        type=float,
        default=7.0,
        help='King W0 parameter [7.0]',
    )
    parser.add_argument(
        '-fractald',
        dest='fractal_parameter_fd',
        type=float,
        default=2.0,
        help='Fractal distribution [2.0]',
    )
    parser.add_argument(
        '-uml',
        dest='upper_mass_limit',
        type=float,
        default=125.,
        help='Upper mass limit for stars (in MSun) [125.]',
    )
    parser.add_argument(
        '-lml',
        dest='lower_mass_limit',
        type=float,
        default=0.1,
        help='Lower mass limit for stars (in MSun) [0.1]',
    )
    parser.add_argument(
        '-Z',
        dest='metallicity',
        type=float,
        default=0.01,
        help='Metallicity [0.01]',
    )
    parser.add_argument(
        '-ibf',
        dest='initial_binary_fraction',
        type=float,
        default=0.,
        help='Initial binary fraction (0-1) [0.] NOT IMPLEMENTED YET',
    )
    parser.add_argument(
        '-R',
        dest='effective_radius',
        type=float,
        default=3.0,
        help=(
            "Cluster radius / default radius when forming stars from sinks"
            " [3.0] (in parsec)"
        ),
    )
    parser.add_argument(
        '--vel',
        dest='velocity_dispersion',
        type=float,
        default=1.0,
        help=(
            "Default velocity dispersion when forming stars from sinks"
            " [1.0] (in km/s)"
        ),
    )
    # parser.add_argument(
    #     '-vr',
    #     dest='virial_ratio',
    #     type=float,
    #     default=0.5,
    #     help="Virial ration [0.5], 0.5=stable, 0.75=just expelled gas, \
    #             0.1=collapsing",
    # )
    args = parser.parse_args()
    return args

def main():
    "Make a star cluster"
    set_printing_strategy(
        "custom",
        preferred_units=[units.MSun, units.parsec, units.yr, units.kms],
        precision=5,
    )
    clustertemplate = "TESTCluster_%08i"

    args = new_argument_parser()
    sinks = args.sinks
    cluster_model_number = args.cluster_model_number
    star_distribution = args.star_distribution
    # gas_distribution = args.gas_distribution
    king_parameter_w0 = args.king_parameter_w0
    fractal_parameter_fd = args.fractal_parameter_fd
    initial_mass_function = args.initial_mass_function.lower()

    number_of_stars = args.number_of_stars
    if args.cluster_mass != 0:
        cluster_mass = args.cluster_mass | units.MSun
    else:
        cluster_mass = False
    upper_mass_limit = args.upper_mass_limit | units.MSun
    lower_mass_limit = args.lower_mass_limit | units.MSun
    effective_radius = args.effective_radius | units.parsec
    metallicity = args.metallicity
    # virial_ratio = args.virial_ratio
    filetype = args.filetype

    # not implemented yet
    # initial_binary_fraction = args.initial_binary_fraction

    np.random.seed(cluster_model_number)

    if not (number_of_stars or cluster_mass or sinks):
        print(
            "no number of stars, cluster mass or origin sinks given, exiting"
        )
        exit()

    if sinks is not None:
        sinks = read_set_from_file(sinks, "amuse")
        stars = Particles()
        for sink in sinks:
            try:
                velocity_dispersion = sink.u.sqrt()
            except AttributeError:
                velocity_dispersion = args.velocity_dispersion | units.kms
            new_stars = new_stars_from_sink(
                sink,
                upper_mass_limit=upper_mass_limit,
                lower_mass_limit=lower_mass_limit,
                default_radius=effective_radius,
                velocity_dispersion=velocity_dispersion,
                initial_mass_function=initial_mass_function,
                # logger=logger,
            )
            stars.add_particles(
                new_stars
            )
    else:
        stars = new_star_cluster(
            stellar_mass=cluster_mass,
            initial_mass_function=initial_mass_function,
            upper_mass_limit=upper_mass_limit,
            lower_mass_limit=lower_mass_limit,
            number_of_stars=number_of_stars,
            effective_radius=effective_radius,
            star_distribution=star_distribution,
            star_distribution_w0=king_parameter_w0,
            star_distribution_fd=fractal_parameter_fd,
            star_metallicity=metallicity,
        )

    print(
        "%i stars generated (%s)"
        % (len(stars), stars.total_mass().in_(units.MSun))
    )

    if args.clustername != "auto":
        clustertemplate = args.clustername + "%s"

    stars_file_exists = True
    sinks_file_exists = True
    N = -1
    while (stars_file_exists or sinks_file_exists):
        N += 1
        starsfilename = (
            clustertemplate % N
            + "-stars." + filetype
        )
        stars_file_exists = os.path.isfile(starsfilename)
        sinksfilename = (
            clustertemplate % N
            + "-sinks." + filetype
        )
        sinks_file_exists = os.path.isfile(sinksfilename)

    write_set_to_file(stars, starsfilename, filetype)
    if sinks is not None:
        write_set_to_file(sinks, sinksfilename, filetype)

if __name__ in ["__main__"]:
    main()
<commit_msg>small improvements<commit_after>#!/usr/bin/env python3
"""
make_a_star_cluster.py creates a model star cluster,
which can then be used in N-body simulations or for other purposes.

It requires AMUSE, which can be downloaded from http://amusecode.org or
https://github.com/amusecode/amuse.

Currently not feature-complete yet, and function/argument names are
subject to change.

-- Steven Rieder steven at rieder punt nl
"""

import sys
import os
import argparse

import numpy as np

from amuse.io import write_set_to_file, read_set_from_file
from amuse.datamodel import Particles
from amuse.units import units
from amuse.support.console import set_printing_strategy
from amuse.ext.masc import new_star_cluster, new_stars_from_sink

def new_argument_parser():
    "Parse command line arguments"
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--sinks',
        dest='sinks',
        default=None,
        help="Generate stars from sink particles in this file (EXPERIMENTAL)",
    )
    parser.add_argument(
        '-o',
        dest='clustername',
        default="auto",
        help="Outputfile name [auto]",
    )
    parser.add_argument(
        '-type',
        dest='filetype',
        default="amuse",
        help="Output file type ([amuse]/txt/starlab/nemo)",
    )
    parser.add_argument(
        '-N',
        dest='number_of_stars',
        default=1024,
        type=int,
        help='Number of stars [1024]',
    )
    parser.add_argument(
        '-M',
        dest='cluster_mass',
        type=float,
        default=0,
        help='Cluster mass (in MSun) (takes precedence over number of stars)',
    )
    parser.add_argument(
        '-dist',
        dest='star_distribution',
        default="plummer",
        help="Star distribution ([plummer]/king/fractal)",
    )
    parser.add_argument(
        '-gasdist',
        dest='gas_distribution',
        default="none",
        help="Gas distribution ([none]/plummer/king/fractal) \
                NOT IMPLEMENTED YET",
    )
    parser.add_argument(
        '-imf',
        dest='initial_mass_function',
        default="kroupa",
        help="IMF ([kroupa]/salpeter/fixed)",
    )
    parser.add_argument(
        '-modelnr',
        dest='cluster_model_number',
        type=int,
        default=30483,
        help='Model number',
    )
    parser.add_argument(
        '-kingw',
        dest='king_parameter_w0',
        type=float,
        default=7.0,
        help='King W0 parameter [7.0]',
    )
    parser.add_argument(
        '-fractald',
        dest='fractal_parameter_fd',
        type=float,
        default=2.0,
        help='Fractal distribution [2.0]',
    )
    parser.add_argument(
        '-uml',
        dest='upper_mass_limit',
        type=float,
        default=125.,
        help='Upper mass limit for stars (in MSun) [125.]',
    )
    parser.add_argument(
        '-lml',
        dest='lower_mass_limit',
        type=float,
        default=0.1,
        help='Lower mass limit for stars (in MSun) [0.1]',
    )
    parser.add_argument(
        '-Z',
        dest='metallicity',
        type=float,
        default=0.01,
        help='Metallicity [0.01]',
    )
    parser.add_argument(
        '-ibf',
        dest='initial_binary_fraction',
        type=float,
        default=0.,
        help='Initial binary fraction (0-1) [0.] NOT IMPLEMENTED YET',
    )
    parser.add_argument(
        '-R',
        dest='effective_radius',
        type=float,
        default=3.0,
        help=(
            "Cluster radius / default radius when forming stars from sinks"
            " [3.0] (in parsec)"
        ),
    )
    parser.add_argument(
        '--vel',
        dest='velocity_dispersion',
        type=float,
        default=1.0,
        help=(
            "Default velocity dispersion when forming stars from sinks"
            " [1.0] (in km/s)"
        ),
    )
    # parser.add_argument(
    #     '-vr',
    #     dest='virial_ratio',
    #     type=float,
    #     default=0.5,
    #     help="Virial ration [0.5], 0.5=stable, 0.75=just expelled gas, \
    #             0.1=collapsing",
    # )
    args = parser.parse_args()
    return args

def main():
    "Make a star cluster"
    set_printing_strategy(
        "custom",
        preferred_units=[units.MSun, units.parsec, units.yr, units.kms],
        precision=5,
    )
    clustertemplate = "TESTCluster_%08i"

    args = new_argument_parser()
    sinks = args.sinks
    cluster_model_number = args.cluster_model_number
    star_distribution = args.star_distribution
    # gas_distribution = args.gas_distribution
    king_parameter_w0 = args.king_parameter_w0
    fractal_parameter_fd = args.fractal_parameter_fd
    initial_mass_function = args.initial_mass_function.lower()

    number_of_stars = args.number_of_stars
    if args.cluster_mass != 0:
        cluster_mass = args.cluster_mass | units.MSun
    else:
        cluster_mass = False
    upper_mass_limit = args.upper_mass_limit | units.MSun
    lower_mass_limit = args.lower_mass_limit | units.MSun
    effective_radius = args.effective_radius | units.parsec
    metallicity = args.metallicity
    # virial_ratio = args.virial_ratio
    filetype = args.filetype

    # not implemented yet
    # initial_binary_fraction = args.initial_binary_fraction

    np.random.seed(cluster_model_number)

    if not (number_of_stars or cluster_mass or sinks):
        print(
            "no number of stars, cluster mass or origin sinks given, exiting"
        )
        sys.exit()

    if sinks is not None:
        sinks = read_set_from_file(sinks, "amuse")
        stars = Particles()
        for sink in sinks:
            try:
                velocity_dispersion = sink.u.sqrt()
            except AttributeError:
                velocity_dispersion = args.velocity_dispersion | units.kms
            new_stars = new_stars_from_sink(
                sink,
                upper_mass_limit=upper_mass_limit,
                lower_mass_limit=lower_mass_limit,
                default_radius=effective_radius,
                velocity_dispersion=velocity_dispersion,
                initial_mass_function=initial_mass_function,
                # logger=logger,
            )
            stars.add_particles(
                new_stars
            )
    else:
        stars = new_star_cluster(
            stellar_mass=cluster_mass,
            initial_mass_function=initial_mass_function,
            upper_mass_limit=upper_mass_limit,
            lower_mass_limit=lower_mass_limit,
            number_of_stars=number_of_stars,
            effective_radius=effective_radius,
            star_distribution=star_distribution,
            star_distribution_w0=king_parameter_w0,
            star_distribution_fd=fractal_parameter_fd,
            star_metallicity=metallicity,
        )

    print(
        f"{len(stars)} stars generated "
        f"({stars.total_mass().in_(units.MSun)})"
    )

    if args.clustername != "auto":
        clustertemplate = args.clustername + "%s"

    stars_file_exists = True
    sinks_file_exists = True
    filenumber = -1
    while (stars_file_exists or sinks_file_exists):
        filenumber += 1
        starsfilename = (
            clustertemplate % filenumber
            + "-stars." + filetype
        )
        stars_file_exists = os.path.isfile(starsfilename)
        sinksfilename = (
            clustertemplate % filenumber
            + "-sinks." + filetype
        )
        sinks_file_exists = os.path.isfile(sinksfilename)

    write_set_to_file(stars, starsfilename, filetype)
    if sinks is not None:
        write_set_to_file(sinks, sinksfilename, filetype)

if __name__ in ["__main__"]:
    main()
<|endoftext|>


<commit_before>##########################################################################
#
#  Copyright (c) 2011-2012, John Haddon. All rights reserved.
#  Copyright (c) 2011-2013, Image Engine Design Inc. All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are
#  met:
#
#      * Redistributions of source code must retain the above
#        copyright notice, this list of conditions and the following
#        disclaimer.
#
#      * Redistributions in binary form must reproduce the above
#        copyright notice, this list of conditions and the following
#        disclaimer in the documentation and/or other materials provided with
#        the distribution.
#
#      * Neither the name of John Haddon nor the names of
#        any other contributors to this software may be used to endorse or
#        promote products derived from this software without specific prior
#        written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
#  IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
#  THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
#  PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
#  CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
#  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
#  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
#  PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
#  LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
#  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
##########################################################################

import re
import fnmatch

import IECore

import Gaffer
import GafferUI

## \todo Potential optimisation.
# Currently every PlugValueWidget instance connects to the various plug signals
# on the node, and tests to see if the plug is of interest when signalled. When many
# PlugValueWidgets are instantiated for a node this might prove a little slow. In this
# eventuality perhaps we can optimise things by having CompoundPlugValueWidget be
# responsible for updating only the correct child, so the children don't need to be
# connected themselves. PlugValueWidget would need to maintain the ability to do things
# itself when used alone, but this might give a good speedup for the most common case.
class PlugValueWidget( GafferUI.Widget ) :

	def __init__( self, topLevelWidget, plug, **kw ) :

		GafferUI.Widget.__init__( self, topLevelWidget, **kw )

		# we don't want to call _updateFromPlug yet because the derived
		# classes haven't constructed yet. they can call it themselves
		# upon completing construction.
		self.__setPlugInternal( plug, callUpdateFromPlug=False )

		self.__popupMenuConnections = []
		self.__readOnly = False

	def setPlug( self, plug ) :

		self.__setPlugInternal( plug, callUpdateFromPlug=True )

	def getPlug( self ) :

		return self.__plug

	## By default, PlugValueWidgets operate in the main context held by the script node
	# for the script the plug belongs to. This function allows an alternative context
	# to be provided, making it possible to view a plug at a custom frame (or with any
	# other context modification).
	def setContext( self, context ) :

		assert( isinstance( context, Gaffer.Context ) )
		if context is self.__context :
			return

		self.__context = context
		self.__updateContextConnection()
		self._updateFromPlug()

	def getContext( self ) :

		return self.__context

	## This method allows editing of the plug value
	# to be disabled for this ui. Note that even when getReadOnly()
	# is False, the ui may not allow editing due to the plug
	# itself being read only for other reasons.
	def setReadOnly( self, readOnly ) :

		assert( isinstance( readOnly, bool ) )
		if readOnly == self.__readOnly :
			return

		self.__readOnly = readOnly
		self._updateFromPlug()

	def getReadOnly( self ) :

		return self.__readOnly

	## Should be reimplemented to return True if this widget includes
	# some sort of labelling for the plug. This is used to prevent
	# extra labels being created in the NodeUI when they're not necessary.
	def hasLabel( self ) :

		return False

	## Must be implemented by subclasses so that the widget reflects the current
	# status of the plug. To temporarily suspend calls to this function, use
	# Gaffer.BlockedConnection( self._plugConnections() ).
	def _updateFromPlug( self ) :

		raise NotImplementedError

	def _plugConnections( self ) :

		return [ self.__plugSetConnection,
			self.__plugDirtiedConnection,
			self.__plugInputChangedConnection,
			self.__plugFlagsChangedConnection
		]

	## Returns True if the plug value is editable as far as this ui is concerned
	# - that plug.settable() is True and self.getReadOnly() is False.
	def _editable( self ) :

		plug = self.getPlug()

		if plug is None :
			return False

		if hasattr(plug, 'settable') and not plug.settable():
			return False

		if self.__readOnly :
			return False

		return True

	## Adds a useful popup menu to the specified widget, providing useful functions that
	# operate on the plug. The menu is populated with the result of _popupMenuDefinition(),
	# and may also be customised by external code using the popupMenuSignal().
	def _addPopupMenu( self, widget = None, buttons = GafferUI.ButtonEvent.Buttons.Right ) :

		if widget is None :
			widget = self

		# it's unclear under what circumstances we get given a right-click vs a context menu event,
		# but we try to cover all our bases by connecting to both.

		self.__popupMenuConnections.append(
			widget.buttonPressSignal().connect( IECore.curry( Gaffer.WeakMethod( self.__buttonPress ), buttonMask = buttons ) )
		)

		if buttons & GafferUI.ButtonEvent.Buttons.Right :
			self.__popupMenuConnections.append(
				widget.contextMenuSignal().connect( IECore.curry( Gaffer.WeakMethod( self.__contextMenu ) ) )
			)

	## Returns a definition for the popup menu - this is called each time the menu is displayed
	# to allow for dynamic menus. Subclasses may override this method to customise the menu, but
	# should call the base class implementation first.
	def _popupMenuDefinition( self ) :

		menuDefinition = IECore.MenuDefinition()

		if self.getPlug().getInput() is not None :
			menuDefinition.append( "/Edit input...", { "command" : Gaffer.WeakMethod( self.__editInput ) } )
			menuDefinition.append( "/EditInputDivider", { "divider" : True } )
			menuDefinition.append( "/Remove input", { "command" : Gaffer.WeakMethod( self.__removeInput ) } )
		if hasattr( self.getPlug(), "defaultValue" ) :
			menuDefinition.append(
				"/Default", {
					"command" : IECore.curry( Gaffer.WeakMethod( self.__setValue ), self.getPlug().defaultValue() ),
					"active" : self._editable()
				}
			)

		self.popupMenuSignal()( menuDefinition, self )

		return menuDefinition

	__popupMenuSignal = Gaffer.Signal2()
	## This signal is emitted whenever a popup menu for a plug is about
	# to be shown. This provides an opportunity to customise the menu from
	# external code. The signature for slots is ( menuDefinition, plugValueWidget ),
	# and slots should just modify the menu definition in place.
	@classmethod
	def popupMenuSignal( cls ) :

		return cls.__popupMenuSignal

	## Returns a PlugValueWidget suitable for representing the specified plug. If
	# useTypeOnly is True, then custom registrations made by registerCreator() will
	# be ignored and only the plug type will be taken into account in creating a
	# PlugValueWidget.
	@classmethod
	def create( cls, plug, useTypeOnly=False ) :

		# first try to create one using a creator registered for the specific plug
		if not useTypeOnly :
			node = plug.node()
			if node is not None :
				plugPath = plug.relativeName( node )
				nodeHierarchy = IECore.RunTimeTyped.baseTypeIds( node.typeId() )
				for nodeTypeId in [ node.typeId() ] + nodeHierarchy :
					creators = cls.__nodeTypesToCreators.get( nodeTypeId, None )
					if creators :
						for creator in creators :
							if creator.plugPathMatcher.match( plugPath ) :
								if creator.creator is not None :
									return creator.creator( plug, **(creator.creatorKeywordArgs) )
								else :
									return None

		# if that failed, then just create something based on the type of the plug
		typeId = plug.typeId()
		for plugTypeId in [ plug.typeId() ] + IECore.RunTimeTyped.baseTypeIds( plug.typeId() ) :
			if plugTypeId in cls.__plugTypesToCreators :
				creator = cls.__plugTypesToCreators[plugTypeId]
				if creator is not None :
					return creator( plug )
				else :
					return None

		return None

	## Registers a PlugValueWidget type for a specific Plug type. Note
	# that the registerCreator function below provides the
	# opportunity to further customise the type of Widget used for specific
	# plug instances based on the node type and plug name.
	@classmethod
	def registerType( cls, plugTypeId, creator ) :

		cls.__plugTypesToCreators[plugTypeId] = creator

	## Registers a function to create a PlugWidget. None may be passed as creator, to
	# disable the creation of uis for specific plugs.
	@classmethod
	def registerCreator( cls, nodeTypeId, plugPath, creator, **creatorKeywordArgs ) :

		if isinstance( plugPath, basestring ) :
			plugPath = re.compile( fnmatch.translate( plugPath ) )
		else :
			assert( type( plugPath ) is type( re.compile( "" ) ) )

		creators = cls.__nodeTypesToCreators.setdefault( nodeTypeId, [] )

		creator = IECore.Struct(
			plugPathMatcher = plugPath,
			creator = creator,
			creatorKeywordArgs = creatorKeywordArgs,
		)

		creators.insert( 0, creator )

	__plugTypesToCreators = {}
	__nodeTypesToCreators = {}

	def __plugSet( self, plug ) :

		if plug.isSame( self.__plug ) :

			self._updateFromPlug()

	def __plugDirtied( self, plug ) :

		if plug.isSame( self.__plug ) :

			self._updateFromPlug()

	def __plugInputChanged( self, plug ) :

		if plug.isSame( self.__plug ) :
			self.__updateContextConnection()
			self._updateFromPlug()

	def __plugFlagsChanged( self, plug ) :

		if plug.isSame( self.__plug ) :
			self._updateFromPlug()

	def __contextChanged( self, context, key ) :

		self._updateFromPlug()

	def __setPlugInternal( self, plug, callUpdateFromPlug ) :

		self.__plug = plug

		context = self.__fallbackContext

		if self.__plug is not None :
			self.__plugSetConnection = plug.node().plugSetSignal().connect( Gaffer.WeakMethod( self.__plugSet ) )
			if hasattr( plug.node(), "plugDirtiedSignal" ) :
				# only DependencyNodes have the plugDirtiedSignal
				self.__plugDirtiedConnection = plug.node().plugDirtiedSignal().connect( Gaffer.WeakMethod( self.__plugDirtied ) )
			else :
				self.__plugDirtiedConnection = None
			self.__plugInputChangedConnection = plug.node().plugInputChangedSignal().connect( Gaffer.WeakMethod( self.__plugInputChanged ) )
			self.__plugFlagsChangedConnection = plug.node().plugFlagsChangedSignal().connect( Gaffer.WeakMethod( self.__plugFlagsChanged ) )
			scriptNode = self.__plug.ancestor( Gaffer.ScriptNode.staticTypeId() )
			if scriptNode is not None :
				context = scriptNode.context()
		else :
			self.__plugSetConnection = None
			self.__plugDirtiedConnection = None
			self.__plugInputChangedConnection = None
			self.__plugFlagsChangedConnection = None

		self.__context = context
		self.__updateContextConnection()

		if callUpdateFromPlug :
			self._updateFromPlug()

	def __updateContextConnection( self ) :

		# we only want to be notified of context changes if we have a plug and that
		# plug has an incoming connection. otherwise context changes are irrelevant
		# and we'd just be slowing things down by asking for notifications.

		context = self.__context
		plug = self.getPlug()
		if plug is None or plug.getInput() is None :
			context = None

		if context is not None :
			self.__contextChangedConnection = context.changedSignal().connect( Gaffer.WeakMethod( self.__contextChanged ) )
		else :
			self.__contextChangedConnection = None

	# we use this when the plug being viewed doesn't have a ScriptNode ancestor
	# to provide a context.
	__fallbackContext = Gaffer.Context()

	def __buttonPress( self, widget, event, buttonMask ) :

		if event.buttons & buttonMask :
			return self.__contextMenu()

		return False

	def __contextMenu( self, *unused ) :

		menuDefinition = self._popupMenuDefinition()
		if not len( menuDefinition.items() ) :
			return False

		self.__popupMenu = GafferUI.Menu( menuDefinition )
		self.__popupMenu.popup()

		return True

	def __setValue( self, value ) :

		with Gaffer.UndoContext( self.getPlug().ancestor( Gaffer.ScriptNode.staticTypeId() ) ) :
			self.getPlug().setValue( value )

	def __editInput( self ) :

		GafferUI.NodeEditor.acquire( self.getPlug().getInput().node() )

	def __removeInput( self ) :

		with Gaffer.UndoContext( self.getPlug().ancestor( Gaffer.ScriptNode.staticTypeId() ) ) :
			self.getPlug().setInput( None )

<commit_msg>Added title to PlugValueWidget popup menus.<commit_after>##########################################################################
#
#  Copyright (c) 2011-2012, John Haddon. All rights reserved.
#  Copyright (c) 2011-2013, Image Engine Design Inc. All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are
#  met:
#
#      * Redistributions of source code must retain the above
#        copyright notice, this list of conditions and the following
#        disclaimer.
#
#      * Redistributions in binary form must reproduce the above
#        copyright notice, this list of conditions and the following
#        disclaimer in the documentation and/or other materials provided with
#        the distribution.
#
#      * Neither the name of John Haddon nor the names of
#        any other contributors to this software may be used to endorse or
#        promote products derived from this software without specific prior
#        written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
#  IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
#  THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
#  PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
#  CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
#  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
#  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
#  PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
#  LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
#  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
##########################################################################

import re
import fnmatch

import IECore

import Gaffer
import GafferUI

## \todo Potential optimisation.
# Currently every PlugValueWidget instance connects to the various plug signals
# on the node, and tests to see if the plug is of interest when signalled. When many
# PlugValueWidgets are instantiated for a node this might prove a little slow. In this
# eventuality perhaps we can optimise things by having CompoundPlugValueWidget be
# responsible for updating only the correct child, so the children don't need to be
# connected themselves. PlugValueWidget would need to maintain the ability to do things
# itself when used alone, but this might give a good speedup for the most common case.
class PlugValueWidget( GafferUI.Widget ) :

	def __init__( self, topLevelWidget, plug, **kw ) :

		GafferUI.Widget.__init__( self, topLevelWidget, **kw )

		# we don't want to call _updateFromPlug yet because the derived
		# classes haven't constructed yet. they can call it themselves
		# upon completing construction.
		self.__setPlugInternal( plug, callUpdateFromPlug=False )

		self.__popupMenuConnections = []
		self.__readOnly = False

	def setPlug( self, plug ) :

		self.__setPlugInternal( plug, callUpdateFromPlug=True )

	def getPlug( self ) :

		return self.__plug

	## By default, PlugValueWidgets operate in the main context held by the script node
	# for the script the plug belongs to. This function allows an alternative context
	# to be provided, making it possible to view a plug at a custom frame (or with any
	# other context modification).
	def setContext( self, context ) :

		assert( isinstance( context, Gaffer.Context ) )
		if context is self.__context :
			return

		self.__context = context
		self.__updateContextConnection()
		self._updateFromPlug()

	def getContext( self ) :

		return self.__context

	## This method allows editing of the plug value
	# to be disabled for this ui. Note that even when getReadOnly()
	# is False, the ui may not allow editing due to the plug
	# itself being read only for other reasons.
	def setReadOnly( self, readOnly ) :

		assert( isinstance( readOnly, bool ) )
		if readOnly == self.__readOnly :
			return

		self.__readOnly = readOnly
		self._updateFromPlug()

	def getReadOnly( self ) :

		return self.__readOnly

	## Should be reimplemented to return True if this widget includes
	# some sort of labelling for the plug. This is used to prevent
	# extra labels being created in the NodeUI when they're not necessary.
	def hasLabel( self ) :

		return False

	## Must be implemented by subclasses so that the widget reflects the current
	# status of the plug. To temporarily suspend calls to this function, use
	# Gaffer.BlockedConnection( self._plugConnections() ).
	def _updateFromPlug( self ) :

		raise NotImplementedError

	def _plugConnections( self ) :

		return [ self.__plugSetConnection,
			self.__plugDirtiedConnection,
			self.__plugInputChangedConnection,
			self.__plugFlagsChangedConnection
		]

	## Returns True if the plug value is editable as far as this ui is concerned
	# - that plug.settable() is True and self.getReadOnly() is False.
	def _editable( self ) :

		plug = self.getPlug()

		if plug is None :
			return False

		if hasattr(plug, 'settable') and not plug.settable():
			return False

		if self.__readOnly :
			return False

		return True

	## Adds a useful popup menu to the specified widget, providing useful functions that
	# operate on the plug. The menu is populated with the result of _popupMenuDefinition(),
	# and may also be customised by external code using the popupMenuSignal().
	def _addPopupMenu( self, widget = None, buttons = GafferUI.ButtonEvent.Buttons.Right ) :

		if widget is None :
			widget = self

		# it's unclear under what circumstances we get given a right-click vs a context menu event,
		# but we try to cover all our bases by connecting to both.

		self.__popupMenuConnections.append(
			widget.buttonPressSignal().connect( IECore.curry( Gaffer.WeakMethod( self.__buttonPress ), buttonMask = buttons ) )
		)

		if buttons & GafferUI.ButtonEvent.Buttons.Right :
			self.__popupMenuConnections.append(
				widget.contextMenuSignal().connect( IECore.curry( Gaffer.WeakMethod( self.__contextMenu ) ) )
			)

	## Returns a definition for the popup menu - this is called each time the menu is displayed
	# to allow for dynamic menus. Subclasses may override this method to customise the menu, but
	# should call the base class implementation first.
	def _popupMenuDefinition( self ) :

		menuDefinition = IECore.MenuDefinition()

		if self.getPlug().getInput() is not None :
			menuDefinition.append( "/Edit input...", { "command" : Gaffer.WeakMethod( self.__editInput ) } )
			menuDefinition.append( "/EditInputDivider", { "divider" : True } )
			menuDefinition.append( "/Remove input", { "command" : Gaffer.WeakMethod( self.__removeInput ) } )
		if hasattr( self.getPlug(), "defaultValue" ) and self.getPlug().direction() == Gaffer.Plug.Direction.In :
			menuDefinition.append(
				"/Default", {
					"command" : IECore.curry( Gaffer.WeakMethod( self.__setValue ), self.getPlug().defaultValue() ),
					"active" : self._editable()
				}
			)

		self.popupMenuSignal()( menuDefinition, self )

		return menuDefinition

	__popupMenuSignal = Gaffer.Signal2()
	## This signal is emitted whenever a popup menu for a plug is about
	# to be shown. This provides an opportunity to customise the menu from
	# external code. The signature for slots is ( menuDefinition, plugValueWidget ),
	# and slots should just modify the menu definition in place.
	@classmethod
	def popupMenuSignal( cls ) :

		return cls.__popupMenuSignal

	## Returns a PlugValueWidget suitable for representing the specified plug. If
	# useTypeOnly is True, then custom registrations made by registerCreator() will
	# be ignored and only the plug type will be taken into account in creating a
	# PlugValueWidget.
	@classmethod
	def create( cls, plug, useTypeOnly=False ) :

		# first try to create one using a creator registered for the specific plug
		if not useTypeOnly :
			node = plug.node()
			if node is not None :
				plugPath = plug.relativeName( node )
				nodeHierarchy = IECore.RunTimeTyped.baseTypeIds( node.typeId() )
				for nodeTypeId in [ node.typeId() ] + nodeHierarchy :
					creators = cls.__nodeTypesToCreators.get( nodeTypeId, None )
					if creators :
						for creator in creators :
							if creator.plugPathMatcher.match( plugPath ) :
								if creator.creator is not None :
									return creator.creator( plug, **(creator.creatorKeywordArgs) )
								else :
									return None

		# if that failed, then just create something based on the type of the plug
		typeId = plug.typeId()
		for plugTypeId in [ plug.typeId() ] + IECore.RunTimeTyped.baseTypeIds( plug.typeId() ) :
			if plugTypeId in cls.__plugTypesToCreators :
				creator = cls.__plugTypesToCreators[plugTypeId]
				if creator is not None :
					return creator( plug )
				else :
					return None

		return None

	## Registers a PlugValueWidget type for a specific Plug type. Note
	# that the registerCreator function below provides the
	# opportunity to further customise the type of Widget used for specific
	# plug instances based on the node type and plug name.
	@classmethod
	def registerType( cls, plugTypeId, creator ) :

		cls.__plugTypesToCreators[plugTypeId] = creator

	## Registers a function to create a PlugWidget. None may be passed as creator, to
	# disable the creation of uis for specific plugs.
	@classmethod
	def registerCreator( cls, nodeTypeId, plugPath, creator, **creatorKeywordArgs ) :

		if isinstance( plugPath, basestring ) :
			plugPath = re.compile( fnmatch.translate( plugPath ) )
		else :
			assert( type( plugPath ) is type( re.compile( "" ) ) )

		creators = cls.__nodeTypesToCreators.setdefault( nodeTypeId, [] )

		creator = IECore.Struct(
			plugPathMatcher = plugPath,
			creator = creator,
			creatorKeywordArgs = creatorKeywordArgs,
		)

		creators.insert( 0, creator )

	__plugTypesToCreators = {}
	__nodeTypesToCreators = {}

	def __plugSet( self, plug ) :

		if plug.isSame( self.__plug ) :

			self._updateFromPlug()

	def __plugDirtied( self, plug ) :

		if plug.isSame( self.__plug ) :

			self._updateFromPlug()

	def __plugInputChanged( self, plug ) :

		if plug.isSame( self.__plug ) :
			self.__updateContextConnection()
			self._updateFromPlug()

	def __plugFlagsChanged( self, plug ) :

		if plug.isSame( self.__plug ) :
			self._updateFromPlug()

	def __contextChanged( self, context, key ) :

		self._updateFromPlug()

	def __setPlugInternal( self, plug, callUpdateFromPlug ) :

		self.__plug = plug

		context = self.__fallbackContext

		if self.__plug is not None :
			self.__plugSetConnection = plug.node().plugSetSignal().connect( Gaffer.WeakMethod( self.__plugSet ) )
			if hasattr( plug.node(), "plugDirtiedSignal" ) :
				# only DependencyNodes have the plugDirtiedSignal
				self.__plugDirtiedConnection = plug.node().plugDirtiedSignal().connect( Gaffer.WeakMethod( self.__plugDirtied ) )
			else :
				self.__plugDirtiedConnection = None
			self.__plugInputChangedConnection = plug.node().plugInputChangedSignal().connect( Gaffer.WeakMethod( self.__plugInputChanged ) )
			self.__plugFlagsChangedConnection = plug.node().plugFlagsChangedSignal().connect( Gaffer.WeakMethod( self.__plugFlagsChanged ) )
			scriptNode = self.__plug.ancestor( Gaffer.ScriptNode.staticTypeId() )
			if scriptNode is not None :
				context = scriptNode.context()
		else :
			self.__plugSetConnection = None
			self.__plugDirtiedConnection = None
			self.__plugInputChangedConnection = None
			self.__plugFlagsChangedConnection = None

		self.__context = context
		self.__updateContextConnection()

		if callUpdateFromPlug :
			self._updateFromPlug()

	def __updateContextConnection( self ) :

		# we only want to be notified of context changes if we have a plug and that
		# plug has an incoming connection. otherwise context changes are irrelevant
		# and we'd just be slowing things down by asking for notifications.

		context = self.__context
		plug = self.getPlug()
		if plug is None or plug.getInput() is None :
			context = None

		if context is not None :
			self.__contextChangedConnection = context.changedSignal().connect( Gaffer.WeakMethod( self.__contextChanged ) )
		else :
			self.__contextChangedConnection = None

	# we use this when the plug being viewed doesn't have a ScriptNode ancestor
	# to provide a context.
	__fallbackContext = Gaffer.Context()

	def __buttonPress( self, widget, event, buttonMask ) :

		if event.buttons & buttonMask :
			return self.__contextMenu()

		return False

	def __contextMenu( self, *unused ) :

		menuDefinition = self._popupMenuDefinition()
		if not len( menuDefinition.items() ) :
			return False

		title = self.getPlug().relativeName( self.getPlug().node() )
		title = ".".join( [ IECore.CamelCase.join( IECore.CamelCase.split( x ) ) for x in title.split( "." ) ] )

		self.__popupMenu = GafferUI.Menu( menuDefinition, title = title )
		self.__popupMenu.popup()

		return True

	def __setValue( self, value ) :

		with Gaffer.UndoContext( self.getPlug().ancestor( Gaffer.ScriptNode.staticTypeId() ) ) :
			self.getPlug().setValue( value )

	def __editInput( self ) :

		GafferUI.NodeEditor.acquire( self.getPlug().getInput().node() )

	def __removeInput( self ) :

		with Gaffer.UndoContext( self.getPlug().ancestor( Gaffer.ScriptNode.staticTypeId() ) ) :
			self.getPlug().setInput( None )

<|endoftext|>


<commit_before>import ast

from taichi.lang import impl
from taichi.lang.ast_resolver import ASTResolver
from taichi.lang.exception import TaichiSyntaxError

import taichi as ti

# Total transform
class ASTTransformerTotal(object):
    def __init__(self,
                 func=None,
                 excluded_parameters=(),
                 is_kernel=True,
                 arg_features=None):
        self.func = func
        self.excluded_parameters = excluded_parameters
        self.is_kernel = is_kernel
        self.arg_features = arg_features
        self.pass_Checks = ASTTransformerChecks(func=func)

    @staticmethod
    def print_ast(tree, title=None):
        if not impl.get_runtime().print_preprocessed:
            return
        if title is not None:
            print(f'{title}:')
        import astor
        print(astor.to_source(tree.body[0], indent_with='    '))

    def visit(self, tree):
        from taichi.lang.ast_builder_utils import BuilderContext
        from taichi.lang.stmt_builder import build_stmt
        self.print_ast(tree, 'Initial AST')
        ctx = BuilderContext(func=self.func,
                             excluded_parameters=self.excluded_parameters,
                             is_kernel=self.is_kernel,
                             arg_features=self.arg_features)
        # Convert Python AST to Python code that generates Taichi C++ AST.
        tree = build_stmt(ctx, tree)
        ast.fix_missing_locations(tree)
        self.print_ast(tree, 'Preprocessed')
        self.pass_Checks.visit(tree)  # does not modify the AST

class ASTTransformerBase(ast.NodeTransformer):
    def __init__(self, func):
        super().__init__()
        self.func = func

    @staticmethod
    def get_decorator(node):
        if not isinstance(node, ast.Call):
            return ''
        for wanted, name in [
            (ti.static, 'static'),
            (ti.grouped, 'grouped'),
            (ti.ndrange, 'ndrange'),
        ]:
            if ASTResolver.resolve_to(node.func, wanted, globals()):
                return name
        return ''

# Performs checks at the Python AST level. Does not modify the AST.
class ASTTransformerChecks(ASTTransformerBase):
    def __init__(self, func):
        super().__init__(func)
        self.has_return = False
        self.in_static_if = False

    def visit_If(self, node):
        node.test = self.visit(node.test)

        old_in_static_if = self.in_static_if
        self.in_static_if = self.get_decorator(node.test) == 'static'

        node.body = list(map(self.visit, node.body))
        if node.orelse is not None:
            node.orelse = list(map(self.visit, node.orelse))

        self.in_static_if = old_in_static_if

        return node

    def visit_Return(self, node):
        if self.in_static_if:  # we can have multiple return in static-if branches
            return node

        if not self.has_return:
            self.has_return = True
        else:
            raise TaichiSyntaxError(
                'Taichi functions/kernels cannot have multiple returns!'
                ' Consider using a local variable to walk around.')

        return node
<commit_msg>[Lang] Fix python AST print format issues in python/taichi/lang/transformer.py (#3061)<commit_after>import ast

from taichi.lang import impl
from taichi.lang.ast_resolver import ASTResolver
from taichi.lang.exception import TaichiSyntaxError

import taichi as ti

# Total transform
class ASTTransformerTotal(object):
    def __init__(self,
                 func=None,
                 excluded_parameters=(),
                 is_kernel=True,
                 arg_features=None):
        self.func = func
        self.excluded_parameters = excluded_parameters
        self.is_kernel = is_kernel
        self.arg_features = arg_features
        self.pass_Checks = ASTTransformerChecks(func=func)

    @staticmethod
    def print_ast(tree, title=None):
        if not impl.get_runtime().print_preprocessed:
            return
        if title is not None:
            ti.info(f'{title}:')
        import astor
        print(astor.to_source(tree.body[0], indent_with='    '))

    def visit(self, tree):
        from taichi.lang.ast_builder_utils import BuilderContext
        from taichi.lang.stmt_builder import build_stmt
        self.print_ast(tree, 'Initial AST')
        ctx = BuilderContext(func=self.func,
                             excluded_parameters=self.excluded_parameters,
                             is_kernel=self.is_kernel,
                             arg_features=self.arg_features)
        # Convert Python AST to Python code that generates Taichi C++ AST.
        tree = build_stmt(ctx, tree)
        ast.fix_missing_locations(tree)
        self.print_ast(tree, 'Preprocessed')
        self.pass_Checks.visit(tree)  # does not modify the AST

class ASTTransformerBase(ast.NodeTransformer):
    def __init__(self, func):
        super().__init__()
        self.func = func

    @staticmethod
    def get_decorator(node):
        if not isinstance(node, ast.Call):
            return ''
        for wanted, name in [
            (ti.static, 'static'),
            (ti.grouped, 'grouped'),
            (ti.ndrange, 'ndrange'),
        ]:
            if ASTResolver.resolve_to(node.func, wanted, globals()):
                return name
        return ''

# Performs checks at the Python AST level. Does not modify the AST.
class ASTTransformerChecks(ASTTransformerBase):
    def __init__(self, func):
        super().__init__(func)
        self.has_return = False
        self.in_static_if = False

    def visit_If(self, node):
        node.test = self.visit(node.test)

        old_in_static_if = self.in_static_if
        self.in_static_if = self.get_decorator(node.test) == 'static'

        node.body = list(map(self.visit, node.body))
        if node.orelse is not None:
            node.orelse = list(map(self.visit, node.orelse))

        self.in_static_if = old_in_static_if

        return node

    def visit_Return(self, node):
        if self.in_static_if:  # we can have multiple return in static-if branches
            return node

        if not self.has_return:
            self.has_return = True
        else:
            raise TaichiSyntaxError(
                'Taichi functions/kernels cannot have multiple returns!'
                ' Consider using a local variable to walk around.')

        return node
<|endoftext|>


<commit_before>#
# Copyright 2016 The BigDL Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import sys
import os
import glob
import warnings

def exist_pyspark():
    # check whether pyspark package exists
    try:
        import pyspark
        return True
    except ImportError:
        return False

def check_spark_source_conflict(spark_home, pyspark_path):
    # check if both spark_home env var and pyspark package exist
    # trigger a warning if two spark sources don't match
    if spark_home and not pyspark_path.startswith(spark_home):
        warning_msg = "Find both SPARK_HOME and pyspark. You may need to check whether they " + \
                      "match with each other. SPARK_HOME environment variable is set to: " + spark_home + \
                      ", and pyspark is found in: " + pyspark_path + ". If they are unmatched, " + \
                      "please use one source only to avoid conflict. " + \
                      "For example, you can unset SPARK_HOME and use pyspark only."
        warnings.warn(warning_msg)

def __prepare_spark_env():
    spark_home = os.environ.get('SPARK_HOME', None)
    if exist_pyspark():
        # use pyspark as the spark source
        import pyspark
        check_spark_source_conflict(spark_home, pyspark.__file__)
    else:
        # use SPARK_HOME as the spark source
        if not spark_home:
            raise ValueError(
                """Could not find Spark. Please make sure SPARK_HOME env is set:
                   export SPARK_HOME=path to your spark home directory.""")
        print("Using %s" % spark_home)
        py4j = glob.glob(os.path.join(spark_home, 'python/lib', 'py4j-*.zip'))[0]
        pyspark = glob.glob(os.path.join(spark_home, 'python/lib', 'pyspark*.zip'))[0]
        sys.path.insert(0, py4j)
        sys.path.insert(0, pyspark)

def __prepare_bigdl_env():
    jar_dir = os.path.abspath(__file__ + "/../../")
    conf_paths = glob.glob(os.path.join(jar_dir, "share/conf/*.conf"))
    bigdl_classpath = get_bigdl_classpath()

    def append_path(env_var_name, path):
        try:
            print("Adding %s to %s" % (path, env_var_name))
            os.environ[env_var_name] = path + ":" + os.environ[env_var_name]  # noqa
        except KeyError:
            os.environ[env_var_name] = path

    if bigdl_classpath:
        append_path("BIGDL_JARS", bigdl_classpath)

    if conf_paths:
        assert len(conf_paths) == 1, "Expecting one conf: %s" % len(conf_paths)
        print("Prepending %s to sys.path" % conf_paths[0])
        sys.path.insert(0, conf_paths[0])

    if os.environ.get("BIGDL_JARS", None) and is_spark_below_2_2():
        for jar in os.environ["BIGDL_JARS"].split(":"):
            append_path("SPARK_CLASSPATH", jar)

    if os.environ.get("BIGDL_PACKAGES", None):
        for package in os.environ["BIGDL_PACKAGES"].split(":"):
            sys.path.insert(0, package)

def get_bigdl_classpath():
    """
    Get and return the jar path for bigdl if exists.
    """
    if os.getenv("BIGDL_CLASSPATH"):
        return os.environ["BIGDL_CLASSPATH"]
    jar_dir = os.path.abspath(__file__ + "/../../")
    jar_paths = glob.glob(os.path.join(jar_dir, "share/lib/*.jar"))
    if jar_paths:
        assert len(jar_paths) == 1, "Expecting one jar: %s" % len(jar_paths)
        return jar_paths[0]
    return ""

def is_spark_below_2_2():
    """
    Check if spark version is below 2.2
    """
    import pyspark
    if(hasattr(pyspark,"version")):
        full_version = pyspark.version.__version__
        # We only need the general spark version (eg, 1.6, 2.2).
        parts = full_version.split(".")
        spark_version = parts[0] + "." + parts[1]
        if(compare_version(spark_version, "2.2")>=0):
            return False
    return True

def compare_version(version1, version2):
    """
    Compare version strings.
    :param version1;
    :param version2;
    :return: 1 if version1 is after version2; -1 if version1 is before version2; 0 if two versions are the same.
    """
    v1Arr = version1.split(".")
    v2Arr = version2.split(".")
    len1 = len(v1Arr)
    len2 = len(v2Arr)
    lenMax = max(len1, len2)
    for x in range(lenMax):
        v1Token = 0
        if x < len1:
            v1Token = int(v1Arr[x])
        v2Token = 0
        if x < len2:
            v2Token = int(v2Arr[x])
        if v1Token < v2Token:
            return -1
        if v1Token > v2Token:
            return 1
    return 0

def prepare_env():
    __prepare_spark_env()
    __prepare_bigdl_env()<commit_msg>quick fix for 2539 (#2603)<commit_after>#
# Copyright 2016 The BigDL Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

import sys
import os
import glob
import warnings

def exist_pyspark():
    # check whether pyspark package exists
    try:
        import pyspark
        return True
    except ImportError:
        return False

def check_spark_source_conflict(spark_home, pyspark_path):
    # check if both spark_home env var and pyspark package exist
    # trigger a warning if two spark sources don't match
    if spark_home and not pyspark_path.startswith(spark_home):
        warning_msg = "Find both SPARK_HOME and pyspark. You may need to check whether they " + \
                      "match with each other. SPARK_HOME environment variable is set to: " + spark_home + \
                      ", and pyspark is found in: " + pyspark_path + ". If they are unmatched, " + \
                      "please use one source only to avoid conflict. " + \
                      "For example, you can unset SPARK_HOME and use pyspark only."
        warnings.warn(warning_msg)

def __sys_path_insert(file_path):
    if file_path not in sys.path:
        print("Prepending %s to sys.path" % file_path)
        sys.path.insert(0, file_path)

def __prepare_spark_env():
    spark_home = os.environ.get('SPARK_HOME', None)
    if exist_pyspark():
        # use pyspark as the spark source
        import pyspark
        check_spark_source_conflict(spark_home, pyspark.__file__)
    else:
        # use SPARK_HOME as the spark source
        if not spark_home:
            raise ValueError(
                """Could not find Spark. Please make sure SPARK_HOME env is set:
                   export SPARK_HOME=path to your spark home directory.""")
        print("Using %s" % spark_home)
        py4j = glob.glob(os.path.join(spark_home, 'python/lib', 'py4j-*.zip'))[0]
        pyspark = glob.glob(os.path.join(spark_home, 'python/lib', 'pyspark*.zip'))[0]
        __sys_path_insert(py4j)
        __sys_path_insert(pyspark)

def __prepare_bigdl_env():
    jar_dir = os.path.abspath(__file__ + "/../../")
    conf_paths = glob.glob(os.path.join(jar_dir, "share/conf/*.conf"))
    bigdl_classpath = get_bigdl_classpath()

    def append_path(env_var_name, jar_path):
        try:
            if jar_path not in os.environ[env_var_name].split(":"):
	            print("Adding %s to %s" % (jar_path, env_var_name))
	            os.environ[env_var_name] = jar_path + ":" + os.environ[env_var_name]  # noqa
        except KeyError:
            os.environ[env_var_name] = jar_path

    if bigdl_classpath:
        append_path("BIGDL_JARS", bigdl_classpath)

    if conf_paths:
        assert len(conf_paths) == 1, "Expecting one conf: %s" % len(conf_paths)
        __sys_path_insert(conf_paths[0])

    if os.environ.get("BIGDL_JARS", None) and is_spark_below_2_2():
        for jar in os.environ["BIGDL_JARS"].split(":"):
            append_path("SPARK_CLASSPATH", jar)

    if os.environ.get("BIGDL_PACKAGES", None):
        for package in os.environ["BIGDL_PACKAGES"].split(":"):
            __sys_path_insert(package)

def get_bigdl_classpath():
    """
    Get and return the jar path for bigdl if exists.
    """
    if os.getenv("BIGDL_CLASSPATH"):
        return os.environ["BIGDL_CLASSPATH"]
    jar_dir = os.path.abspath(__file__ + "/../../")
    jar_paths = glob.glob(os.path.join(jar_dir, "share/lib/*.jar"))
    if jar_paths:
        assert len(jar_paths) == 1, "Expecting one jar: %s" % len(jar_paths)
        return jar_paths[0]
    return ""

def is_spark_below_2_2():
    """
    Check if spark version is below 2.2
    """
    import pyspark
    if(hasattr(pyspark,"version")):
        full_version = pyspark.version.__version__
        # We only need the general spark version (eg, 1.6, 2.2).
        parts = full_version.split(".")
        spark_version = parts[0] + "." + parts[1]
        if(compare_version(spark_version, "2.2")>=0):
            return False
    return True

def compare_version(version1, version2):
    """
    Compare version strings.
    :param version1;
    :param version2;
    :return: 1 if version1 is after version2; -1 if version1 is before version2; 0 if two versions are the same.
    """
    v1Arr = version1.split(".")
    v2Arr = version2.split(".")
    len1 = len(v1Arr)
    len2 = len(v2Arr)
    lenMax = max(len1, len2)
    for x in range(lenMax):
        v1Token = 0
        if x < len1:
            v1Token = int(v1Arr[x])
        v2Token = 0
        if x < len2:
            v2Token = int(v2Arr[x])
        if v1Token < v2Token:
            return -1
        if v1Token > v2Token:
            return 1
    return 0

def prepare_env():
    __prepare_spark_env()
    __prepare_bigdl_env()
<|endoftext|>