<commit_before>from __future__ import division
from bs4 import BeautifulSoup
import general
import math
from scraper import *

class Fleet(Scraper):
    def __init__(self, browser, config):
        super(Fleet, self).__init__(browser, config)

        self.general_client = general.General(browser, config)

    def spy_planet(self, origin_planet, destination_planet, spy_probes_count):

        self.logger.info("Spying planet %s (%s)", destination_planet.name, destination_planet.coordinates)

        result = self.send_fleet(origin_planet, destination_planet.coordinates,
                                 self.missions["spy"], {self.SHIPS_DATA.get('ep'): spy_probes_count})

        return result

    def send_expedition(self, origin_planet, coordinates):
        fleet = {
            # Small expedition fleet
            self.SHIPS_DATA.get('sg'): 1,
            self.SHIPS_DATA.get('lf'): 2,
            self.SHIPS_DATA.get('ep'): 1
        }

        self.logger.info("Sending expedition from planet %s to coordinates %s", origin_planet.name, coordinates)
        self.send_fleet(origin_planet, coordinates, self.missions.get("expedition"), fleet)

    def attack_inactive_planet(self, origin_planet, target_planet):
        fleet = self.get_attack_fleet(origin_planet, target_planet)

        self.logger.info("Attacking planet %s from planet %s", target_planet.planet_name, origin_planet.name)

        result = self.send_fleet(origin_planet, target_planet.coordinates,
                                 self.missions["attack"], fleet)
        return result

    def transport_resources(self, origin_planet, destination_planet, resources):
        planet_resources = self.general_client.get_resources(origin_planet)
        if planet_resources.metal < resources.metal:
            resources.metal = planet_resources.metal
        if planet_resources.crystal < resources.crystal:
            resources.crystal = planet_resources.crystal
        if planet_resources.deuterium < resources.deuterium:
            resources.deuterium = planet_resources.deuterium

        fleet = self.get_tranport_fleet(resources, origin_planet)

        self.send_fleet(origin_planet, destination_planet.coordinates,
                        self.missions["transport"], fleet, resources)

    def send_fleet(self, origin_planet, coordinates, mission, ships, resources=None):
        """
        Missions:
            15 - Expedition,
            7 - Colonization,
            8 - Recycle,
            3 - Transport,
            4 - Transfer,
            6 - Spy,
            5 - Defend planet,
            1 - Attack,
            2 - Alliance attack,
            9 - Destroy star
        """
        if resources is None:
            resources = general.Resources(0, 0)

        if origin_planet.coordinates == coordinates:
            self.logger.error("Origin and destination are the same")
            return FleetResult.WrongParameters

        url = self.url_provider.get_page_url('fleet', origin_planet)
        resp = self.open_url(url)

        soup = BeautifulSoup(resp.read(), "lxml")
        fleet_slots = self.get_fleet_slots_usage(mission, soup)

        if fleet_slots[0] >= fleet_slots[1]:
            self.logger.error('No available slots')
            return FleetResult.NoAvailableSlots

        try:
            self.browser.select_form(name='shipsChosen')
        except mechanize.FormNotFoundError:
            self.logger.error('The planet has no available ships')
            return FleetResult.NoAvailableShips

        # set ships to send
        for ship, amount in ships.iteritems():
            self.logger.info("Adding %d %s to the mission fleet" % (amount, ship.name))
            control_name = "am" + str(ship.id)
            control = self.browser.form.find_control(control_name)
            # If there is no available ships exit
            if not control.readonly:
                self.browser[control_name] = str(amount)
            else:
                self.logger.warning("Not enough %s to send" % ship.name)
                return FleetResult.NoAvailableShips
        self.submit_request()

        # set target planet
        try:
            self.browser.select_form(name='details')
        except mechanize.FormNotFoundError:
            self.logger.error('Error sending ships')
            return FleetResult.NoAvailableShips
        self.browser["galaxy"] = coordinates.split(':')[0]
        self.browser["system"] = coordinates.split(':')[1]
        self.browser["position"] = coordinates.split(':')[2]
        self.submit_request()

        # set mission and resources to send
        try:
            self.browser.select_form(name='sendForm')
        except mechanize.FormNotFoundError:
            self.logger.error('Error sending ships')
            return FleetResult.NoAvailableShips
        self.browser.form.find_control('mission').readonly = False
        self.browser["mission"] = str(mission)
        self.browser["metal"] = str(resources.metal)
        self.browser["crystal"] = str(resources.crystal)
        self.browser["deuterium"] = str(resources.deuterium)
        self.submit_request()
        self.logger.info("Sending %s from planet %s to coordinates %s" % (
            get_ships_list(ships), origin_planet.name, coordinates))
        if not resources.empty():
            self.logger.info("The fleet is transporting %s " % resources)

        return FleetResult.Success

    def get_fleet_slots_usage(self, mission=None, soup=None):
        """
            Get fleet slot usage data.
        """

        if soup is None:
            url = self.url_provider.get_page_url('fleet')
            res = self.open_url(url)
            soup = BeautifulSoup(res.read(), "lxml")

        slots_info = soup.find("div", {"id": "slots"})

        if slots_info is None:
            self.logger.warning("Error Getting fleet slots data")
            return 0, 15

        flefts = slots_info.findAll("div", {"class": "fleft"})

        if mission == self.missions.get("expedition"):
            node = flefts[1].findAll("span", {"class": "tooltip advice"})
        else:
            node = flefts[0].findAll("span", {"class": "tooltip advice"})

        slot_usage = "".join(node[0].findAll(text=True, recursive=False))

        try:
            result = (int(slot_usage.split('/')[0].strip()), int(slot_usage.split('/')[1].strip()))
        except ValueError:
            slot_usage = "".join(node[0].find("span", {"class": "overmark"}).findAll(text=True, recursive=False))
            result = (int(slot_usage.split('/')[0].strip()), int(slot_usage.split('/')[1].strip()))

        return result

    def get_tranport_fleet(self, resources, origin_planet=None):
        """
        Get fleet for transporting resources,
        Will use small cargos if there is enough of them.
        """

        resources_count = resources.total()

        return self.get_cargo_fleet_for_mission(origin_planet, resources_count)

    def get_attack_fleet(self, origin_planet, target_planet):
        """
        Get fleet for attacks to inactive targets.
        Will use small cargos if there is enough of them.
        :param origin_planet: Origin planet
        :param target_planet: Target planet
        :return: Optimized fleet for the mission
        """

        resources = target_planet.resources.total()
        resources_count = resources * target_planet.loot

        return self.get_cargo_fleet_for_mission(origin_planet, resources_count)

    def get_cargo_fleet_for_mission(self, origin_planet, resources_count):
        """
        :param origin_planet: Origin to check for cargos
        :param resources_count: Amount of resources to transport
        :return: Get fleet of cargos for the mission
        """

        small_cargos_count = self.get_ships_count(origin_planet, "sg")
        large_cargos_count = self.get_ships_count(origin_planet, "lg")
        recyclers_count = self.get_ships_count(origin_planet, "209")

        self.logger.info("Computing small cargos / recyclers / large cargos for the mission")
        self.logger.info("Resources : %d" % resources_count)

        self.logger.info("Available small cargos : %d" % small_cargos_count)
        self.logger.info("Available large cargos : %d " % large_cargos_count)
        self.logger.info("Available recyclers : %d " % recyclers_count)

        sel_lg_count, left_count = update_count(large_cargos_count, resources_count, 25000, False)
        sel_recycler_count, left_count = update_count(recyclers_count, left_count, 20000, False)
        sel_sg_count, left_count = update_count(small_cargos_count, left_count, 5000, True)

        self.logger.info("Small cargos : %d " % sel_sg_count)
        self.logger.info("Large cargos : %d " % sel_lg_count)
        self.logger.info("Recyclers : %d " % sel_recycler_count)

        self.logger.info("Leftover resources : %d" % left_count)

        fleet = {}

        if sel_sg_count:
            fleet[ self.SHIPS_DATA.get('sg') ] = sel_sg_count

        if sel_lg_count:
            fleet[ self.SHIPS_DATA.get('lg') ] = sel_lg_count

        if sel_recycler_count:
            fleet[ self.SHIPS_DATA.get('209') ] = sel_recycler_count

        return fleet

    def get_ships_count(self, planet, ship_type):
        """
        :param planet: planet to get the count of ships
        :param ship_type: type of ships to count
        :return:
        """
        ships_count = [ship.amount for ship
                       in planet.ships
                       if ship.item.id == self.SHIPS_DATA.get(ship_type).id]

        return next(iter(ships_count), 0)

def get_ships_list(ships):
    return ", ".join([str(ships.get(ship)) + ' ' + ship.name for ship in ships])

def update_count(sel_count, res_count, sel_size, getCeil):
    if sel_count == 0:
        return 0, res_count

    count = float(res_count) / sel_size
    count = math.ceil(count) if getCeil \
            else math.floor(count)
    count = int(count)

    if count > sel_count:
        count = sel_count

    left = res_count - count * sel_size

    if left < 0:
        left = 0

    return count, left

<commit_msg>Minor update<commit_after>from __future__ import division
from bs4 import BeautifulSoup
import general
import math
from scraper import *

class Fleet(Scraper):
    def __init__(self, browser, config):
        super(Fleet, self).__init__(browser, config)

        self.general_client = general.General(browser, config)

    def spy_planet(self, origin_planet, destination_planet, spy_probes_count):

        self.logger.info("Spying planet %s (%s)", destination_planet.name, destination_planet.coordinates)

        result = self.send_fleet(origin_planet, destination_planet.coordinates,
                                 self.missions["spy"], {self.SHIPS_DATA.get('ep'): spy_probes_count})

        return result

    def send_expedition(self, origin_planet, coordinates):
        fleet = {
            # Small expedition fleet
            self.SHIPS_DATA.get('sg'): 1,
            self.SHIPS_DATA.get('lf'): 2,
            self.SHIPS_DATA.get('ep'): 1
        }

        self.logger.info("Sending expedition from planet %s to coordinates %s", origin_planet.name, coordinates)
        self.send_fleet(origin_planet, coordinates, self.missions.get("expedition"), fleet)

    def attack_inactive_planet(self, origin_planet, target_planet):
        fleet = self.get_attack_fleet(origin_planet, target_planet)

        self.logger.info("Attacking planet %s from planet %s", target_planet.planet_name, origin_planet.name)

        result = self.send_fleet(origin_planet, target_planet.coordinates,
                                 self.missions["attack"], fleet)
        return result

    def transport_resources(self, origin_planet, destination_planet, resources):
        planet_resources = self.general_client.get_resources(origin_planet)
        if planet_resources.metal < resources.metal:
            resources.metal = planet_resources.metal
        if planet_resources.crystal < resources.crystal:
            resources.crystal = planet_resources.crystal
        if planet_resources.deuterium < resources.deuterium:
            resources.deuterium = planet_resources.deuterium

        fleet = self.get_tranport_fleet(resources, origin_planet)

        self.send_fleet(origin_planet, destination_planet.coordinates,
                        self.missions["transport"], fleet, resources)

    def send_fleet(self, origin_planet, coordinates, mission, ships, resources=None):
        """
        Missions:
            15 - Expedition,
            7 - Colonization,
            8 - Recycle,
            3 - Transport,
            4 - Transfer,
            6 - Spy,
            5 - Defend planet,
            1 - Attack,
            2 - Alliance attack,
            9 - Destroy star
        """
        if resources is None:
            resources = general.Resources(0, 0)

        if origin_planet.coordinates == coordinates:
            self.logger.error("Origin and destination are the same")
            return FleetResult.WrongParameters

        url = self.url_provider.get_page_url('fleet', origin_planet)
        resp = self.open_url(url)

        soup = BeautifulSoup(resp.read(), "lxml")
        fleet_slots = self.get_fleet_slots_usage(mission, soup)

        if fleet_slots[0] >= fleet_slots[1]:
            self.logger.error('No available slots')
            return FleetResult.NoAvailableSlots

        try:
            self.browser.select_form(name='shipsChosen')
        except mechanize.FormNotFoundError:
            self.logger.error('The planet has no available ships')
            return FleetResult.NoAvailableShips

        # set ships to send
        for ship, amount in ships.iteritems():
            self.logger.info("Adding %d %s to the mission fleet" % (amount, ship.name))
            control_name = "am" + str(ship.id)
            control = self.browser.form.find_control(control_name)
            # If there is no available ships exit
            if not control.readonly:
                self.browser[control_name] = str(amount)
            else:
                self.logger.warning("Not enough %s to send" % ship.name)
                return FleetResult.NoAvailableShips
        self.submit_request()

        # set target planet
        try:
            self.browser.select_form(name='details')
        except mechanize.FormNotFoundError:
            self.logger.error('Error sending ships')
            return FleetResult.NoAvailableShips
        self.browser["galaxy"] = coordinates.split(':')[0]
        self.browser["system"] = coordinates.split(':')[1]
        self.browser["position"] = coordinates.split(':')[2]
        self.submit_request()

        # set mission and resources to send
        try:
            self.browser.select_form(name='sendForm')
        except mechanize.FormNotFoundError:
            self.logger.error('Error sending ships')
            return FleetResult.NoAvailableShips
        self.browser.form.find_control('mission').readonly = False
        self.browser["mission"] = str(mission)
        self.browser["metal"] = str(resources.metal)
        self.browser["crystal"] = str(resources.crystal)
        self.browser["deuterium"] = str(resources.deuterium)
        self.submit_request()
        self.logger.info("Sending %s from planet %s to coordinates %s" % (
            get_ships_list(ships), origin_planet.name, coordinates))
        if not resources.empty():
            self.logger.info("The fleet is transporting %s " % resources)

        return FleetResult.Success

    def get_fleet_slots_usage(self, mission=None, soup=None):
        """
            Get fleet slot usage data.
        """

        if soup is None:
            url = self.url_provider.get_page_url('fleet')
            res = self.open_url(url)
            soup = BeautifulSoup(res.read(), "lxml")

        slots_info = soup.find("div", {"id": "slots"})

        if slots_info is None:
            self.logger.warning("Error Getting fleet slots data")
            return 0, 15

        flefts = slots_info.findAll("div", {"class": "fleft"})

        if mission == self.missions.get("expedition"):
            node = flefts[1].findAll("span", {"class": "tooltip advice"})
        else:
            node = flefts[0].findAll("span", {"class": "tooltip advice"})

        slot_usage = "".join(node[0].findAll(text=True, recursive=False))

        try:
            result = (int(slot_usage.split('/')[0].strip()), int(slot_usage.split('/')[1].strip()))
        except ValueError:
            slot_usage = "".join(node[0].find("span", {"class": "overmark"}).findAll(text=True, recursive=False))
            result = (int(slot_usage.split('/')[0].strip()), int(slot_usage.split('/')[1].strip()))

        return result

    def get_tranport_fleet(self, resources, origin_planet=None):
        """
        Get fleet for transporting resources,
        Will use small cargos if there is enough of them.
        """

        resources_count = resources.total()

        return self.get_cargo_fleet_for_mission(origin_planet, resources_count)

    def get_attack_fleet(self, origin_planet, target_planet):
        """
        Get fleet for attacks to inactive targets.
        Will use small cargos if there is enough of them.
        :param origin_planet: Origin planet
        :param target_planet: Target planet
        :return: Optimized fleet for the mission
        """

        resources = target_planet.resources.total()
        resources_count = resources * target_planet.loot

        return self.get_cargo_fleet_for_mission(origin_planet, resources_count)

    def get_cargo_fleet_for_mission(self, origin_planet, resources_count):
        """
        :param origin_planet: Origin to check for cargos
        :param resources_count: Amount of resources to transport
        :return: Get fleet of cargos for the mission
        """

        small_cargos_count = self.get_ships_count(origin_planet, "sg")
        large_cargos_count = self.get_ships_count(origin_planet, "lg")
        recyclers_count = self.get_ships_count(origin_planet, "209")

        self.logger.info("Computing small cargos / recyclers / large cargos for the mission")
        self.logger.info("Resources : %d" % resources_count)

        self.logger.info("Available small cargos : %d" % small_cargos_count)
        self.logger.info("Available large cargos : %d " % large_cargos_count)
        self.logger.info("Available recyclers : %d " % recyclers_count)

        sel_lg_count, left_count = update_count(large_cargos_count, resources_count, 25000, False)
        sel_recycler_count, left_count = update_count(recyclers_count, left_count, 20000, False)
        sel_sg_count, left_count = update_count(small_cargos_count, left_count, 5000, True)

        self.logger.info("Small cargos : %d " % sel_sg_count)
        self.logger.info("Large cargos : %d " % sel_lg_count)
        self.logger.info("Recyclers : %d " % sel_recycler_count)

        self.logger.info("Leftover resources : %d" % left_count)

        fleet = {}

        if sel_sg_count:
            fleet[ self.SHIPS_DATA.get('sg') ] = sel_sg_count

        if sel_lg_count:
            fleet[ self.SHIPS_DATA.get('lg') ] = sel_lg_count

        if sel_recycler_count:
            fleet[ self.SHIPS_DATA.get('209') ] = sel_recycler_count

        if len(fleet) == 0:
            fleet[ self.SHIPS_DATA.get('lg') ] = 0

        return fleet

    def get_ships_count(self, planet, ship_type):
        """
        :param planet: planet to get the count of ships
        :param ship_type: type of ships to count
        :return:
        """
        ships_count = [ship.amount for ship
                       in planet.ships
                       if ship.item.id == self.SHIPS_DATA.get(ship_type).id]

        return next(iter(ships_count), 0)

def get_ships_list(ships):
    return ", ".join([str(ships.get(ship)) + ' ' + ship.name for ship in ships])

def update_count(sel_count, res_count, sel_size, getCeil):
    if sel_count == 0:
        return 0, res_count

    count = float(res_count) / sel_size
    count = math.ceil(count) if getCeil \
            else math.floor(count)
    count = int(count)

    if count > sel_count:
        count = sel_count

    left = res_count - count * sel_size

    if left < 0:
        left = 0

    return count, left

<|endoftext|>


<commit_before>"""
Django settings for openhousens project.

For more information on this file, see
https://docs.djangoproject.com/en/dev/topics/settings/

For the full list of settings and their values, see
https://docs.djangoproject.com/en/dev/ref/settings/
"""

# Build paths inside the project like this: os.path.join(BASE_DIR, ...)
import os
BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# Quick-start development settings - unsuitable for production
# See https://docs.djangoproject.com/en/dev/howto/deployment/checklist/

# SECURITY WARNING: keep the secret key used in production secret!
SECRET_KEY = os.getenv('DJANGO_SECRET_KEY', 'oy8$$i(*9eq73=8)1vq&v+xqpzig0%+b9wb*chpnol@k+pf)tv')

# SECURITY WARNING: don't run with debug turned on in production!
DEBUG = not os.getenv('PRODUCTION', False)

TEMPLATE_DEBUG = DEBUG

ALLOWED_HOSTS = ['*']

# Application definition
try:
    import debug_toolbar
    DEBUG_TOOLBAR = True
except ImportError:
    DEBUG_TOOLBAR = False

INSTALLED_APPS = (
    'django.contrib.auth', # needed by instances
    'django.contrib.contenttypes', # needed by django.contrib.auth
    'django.contrib.staticfiles',
    # @see http://mysociety.github.io/sayit/install/#installing-sayit-as-a-django-app
    'haystack',
    'django_bleach',
    'popolo',
    'instances',
    'speeches',
    'legislature',
)
if DEBUG_TOOLBAR:
    INSTALLED_APPS = INSTALLED_APPS + ('debug_toolbar',)

MIDDLEWARE_CLASSES = (
    'django.middleware.common.CommonMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.middleware.clickjacking.XFrameOptionsMiddleware',
)
if DEBUG_TOOLBAR:
    MIDDLEWARE_CLASSES = MIDDLEWARE_CLASSES + ('debug_toolbar.middleware.DebugToolbarMiddleware',)

ROOT_URLCONF = 'openhousens.urls'

WSGI_APPLICATION = 'openhousens.wsgi.application'

# Database
# https://docs.djangoproject.com/en/dev/ref/settings/#databases

# DATABASES = {
#     'default': {
#         'ENGINE': 'django.db.backends.sqlite3',
#         'NAME': os.path.join(BASE_DIR, 'db.sqlite3'),
#     }
# }

# Internationalization
# https://docs.djangoproject.com/en/dev/topics/i18n/

LANGUAGE_CODE = 'en-us'

TIME_ZONE = 'UTC'

USE_I18N = True

USE_L10N = True

USE_TZ = True

# Static files (CSS, JavaScript, Images)
# https://docs.djangoproject.com/en/dev/howto/static-files/

STATIC_URL = '/static/'

# @see https://devcenter.heroku.com/articles/getting-started-with-django#settings-py
# @see https://devcenter.heroku.com/articles/django-assets

# Parse database configuration from $DATABASE_URL
import dj_database_url
DATABASES = {'default': dj_database_url.config(default='postgres://localhost/openhousens')}

# Honor the 'X-Forwarded-Proto' header for request.is_secure()
SECURE_PROXY_SSL_HEADER = ('HTTP_X_FORWARDED_PROTO', 'https')

# Static asset configuration
STATIC_ROOT = 'staticfiles'
STATICFILES_DIRS = (
    os.path.join(BASE_DIR, 'static'),
)

# @see http://mysociety.github.io/sayit/install/#installing-sayit-as-a-django-app
HAYSTACK_CONNECTIONS = {
    'default': {
        'ENGINE': 'haystack.backends.elasticsearch_backend.ElasticsearchSearchEngine',
        'URL': os.getenv('ELASTICSEARCH_URL', 'http://127.0.0.1:9200/'),
        'INDEX_NAME': 'openhousens',
    },
}
# HAYSTACK_SIGNAL_PROCESSOR = 'haystack.signals.RealtimeSignalProcessor'

BLEACH_ALLOWED_TAGS = [
    'a', 'abbr', 'b', 'i', 'u', 'span', 'sub', 'sup', 'br',
    'p', 'blockquote',
    'ol', 'ul', 'li',
    'table', 'caption', 'tr', 'th', 'td',
]
BLEACH_ALLOWED_ATTRIBUTES = {
    '*': [ 'id', 'title', 'class' ], # class, style
    'a': [ 'href' ],
    'li': [ 'value' ],
}

TEMPLATE_DIRS = (
    os.path.join(BASE_DIR, 'templates'),
)
<commit_msg>remove django-debug-toolbar which is slow, use simple SQL logging to find inefficient pages<commit_after>"""
Django settings for openhousens project.

For more information on this file, see
https://docs.djangoproject.com/en/dev/topics/settings/

For the full list of settings and their values, see
https://docs.djangoproject.com/en/dev/ref/settings/
"""

# Build paths inside the project like this: os.path.join(BASE_DIR, ...)
import os
BASE_DIR = os.path.dirname(os.path.abspath(__file__))

# Quick-start development settings - unsuitable for production
# See https://docs.djangoproject.com/en/dev/howto/deployment/checklist/

# SECURITY WARNING: keep the secret key used in production secret!
SECRET_KEY = os.getenv('DJANGO_SECRET_KEY', 'oy8$$i(*9eq73=8)1vq&v+xqpzig0%+b9wb*chpnol@k+pf)tv')

# SECURITY WARNING: don't run with debug turned on in production!
DEBUG = not os.getenv('PRODUCTION', False)

TEMPLATE_DEBUG = DEBUG

ALLOWED_HOSTS = ['*']

# Application definition

INSTALLED_APPS = (
    'django.contrib.auth', # needed by instances
    'django.contrib.contenttypes', # needed by django.contrib.auth
    'django.contrib.staticfiles',
    # @see http://mysociety.github.io/sayit/install/#installing-sayit-as-a-django-app
    'haystack',
    'django_bleach',
    'popolo',
    'instances',
    'speeches',
    'legislature',
)

MIDDLEWARE_CLASSES = (
    'django.middleware.common.CommonMiddleware',
    'django.middleware.csrf.CsrfViewMiddleware',
    'django.middleware.clickjacking.XFrameOptionsMiddleware',
)

ROOT_URLCONF = 'openhousens.urls'

WSGI_APPLICATION = 'openhousens.wsgi.application'

# Database
# https://docs.djangoproject.com/en/dev/ref/settings/#databases

# DATABASES = {
#     'default': {
#         'ENGINE': 'django.db.backends.sqlite3',
#         'NAME': os.path.join(BASE_DIR, 'db.sqlite3'),
#     }
# }

# Internationalization
# https://docs.djangoproject.com/en/dev/topics/i18n/

LANGUAGE_CODE = 'en-us'

TIME_ZONE = 'UTC'

USE_I18N = True

USE_L10N = True

USE_TZ = True

# Static files (CSS, JavaScript, Images)
# https://docs.djangoproject.com/en/dev/howto/static-files/

STATIC_URL = '/static/'

# @see https://devcenter.heroku.com/articles/getting-started-with-django#settings-py
# @see https://devcenter.heroku.com/articles/django-assets

# Parse database configuration from $DATABASE_URL
import dj_database_url
DATABASES = {'default': dj_database_url.config(default='postgres://localhost/openhousens')}

# Honor the 'X-Forwarded-Proto' header for request.is_secure()
SECURE_PROXY_SSL_HEADER = ('HTTP_X_FORWARDED_PROTO', 'https')

# Static asset configuration
STATIC_ROOT = 'staticfiles'
STATICFILES_DIRS = (
    os.path.join(BASE_DIR, 'static'),
)

# @see http://mysociety.github.io/sayit/install/#installing-sayit-as-a-django-app
HAYSTACK_CONNECTIONS = {
    'default': {
        'ENGINE': 'haystack.backends.elasticsearch_backend.ElasticsearchSearchEngine',
        'URL': os.getenv('ELASTICSEARCH_URL', 'http://127.0.0.1:9200/'),
        'INDEX_NAME': 'openhousens',
    },
}
# HAYSTACK_SIGNAL_PROCESSOR = 'haystack.signals.RealtimeSignalProcessor'

BLEACH_ALLOWED_TAGS = [
    'a', 'abbr', 'b', 'i', 'u', 'span', 'sub', 'sup', 'br',
    'p', 'blockquote',
    'ol', 'ul', 'li',
    'table', 'caption', 'tr', 'th', 'td',
]
BLEACH_ALLOWED_ATTRIBUTES = {
    '*': [ 'id', 'title', 'class' ], # class, style
    'a': [ 'href' ],
    'li': [ 'value' ],
}

TEMPLATE_DIRS = (
    os.path.join(BASE_DIR, 'templates'),
)

if DEBUG:
    LOGGING = {
        'version': 1,
        'handlers': {
            'console': {
                'level': 'DEBUG',
                'class': 'logging.StreamHandler',
            },
        },
        'loggers': {
            'django.db.backends': {
                'level': 'DEBUG',
                'handlers': ['console'],
            },
        }
    }
<|endoftext|>


<commit_before>import logging
import re
import lxml

from django.db import transaction

from document.models import Kamervraag
from document.models import Kamerantwoord
from document.models import Antwoord
from document.models import Vraag
from document.models import CategoryDocument
from document.models import FootNote

from document.models import Document

import website.create
import scraper.documents

logger = logging.getLogger(__name__)

def create_kamervragen(year, max_n=None, skip_if_exists=False):
    logger.info('BEGIN')
    infos = Kamervraag.get_kamervragen_info(year)
    counter = 0
    for info in infos:
        try:
            create_kamervraag(info['document_number'], info['overheidnl_document_id'], skip_if_exists=skip_if_exists)
        except Exception as error:
            logger.error('error for kamervraag id: ' + str(info['overheidnl_document_id']))
            logger.exception(error)
        if max_n and counter >= max_n:
            return
        counter += 1
    logger.info('END')

def create_antwoorden(year, max_n=None, skip_if_exists=False):
    logger.info('BEGIN')
    infos = Kamerantwoord.get_antwoorden_info(year)
    counter = 0
    for info in infos:
        try:
            create_kamerantwoord(info['document_number'], info['overheidnl_document_id'], skip_if_exists=skip_if_exists)
        except Exception as error:
            logger.error('error for kamerantwoord id: ' + str(info['overheidnl_document_id']))
            logger.exception(error)
        if max_n and counter >= max_n:
            break
        counter += 1
    logger.info('END')

@transaction.atomic
def create_kamervraag(document_number, overheidnl_document_id, skip_if_exists=False):
    if skip_if_exists and Kamervraag.objects.filter(document__document_id=document_number).exists():
        return
    document, vraagnummer = create_kamervraag_document(document_number, overheidnl_document_id)
    Kamervraag.objects.filter(document=document).delete()
    kamervraag = Kamervraag.objects.create(
        document=document,
        vraagnummer=vraagnummer,
        receiver=get_receiver_from_title(document.title_full)
    )
    create_vragen_from_kamervraag_html(kamervraag)
    footnotes = create_footnotes(kamervraag.document.content_html)
    for footnote in footnotes:
        FootNote.objects.create(document=document, nr=footnote['nr'], text=footnote['text'], url=footnote['url'])
    return kamervraag

@transaction.atomic
def create_kamerantwoord(document_number, overheidnl_document_id, skip_if_exists=False):
    if skip_if_exists and Kamerantwoord.objects.filter(document__document_id=document_number).exists():
        return
    document, vraagnummer = create_kamervraag_document(document_number, overheidnl_document_id)
    Kamerantwoord.objects.filter(document=document).delete()
    kamerantwoord = Kamerantwoord.objects.create(document=document, vraagnummer=vraagnummer)
    create_antwoorden_from_antwoord_html(kamerantwoord)
    return kamerantwoord

@transaction.atomic
def create_kamervraag_document(document_number, overheidnl_document_id):
    logger.info('BEGIN')
    document_html_url = 'https://zoek.officielebekendmakingen.nl/' + str(overheidnl_document_id) +'.html'
    logger.info(document_html_url)
    document_id, content_html, title = scraper.documents.get_kamervraag_document_id_and_content(document_html_url)
    metadata = scraper.documents.get_metadata(overheidnl_document_id)
    document_id = document_number

    if metadata['date_published']:
        date_published = metadata['date_published']
    elif metadata['date_submitted']:
        date_published = metadata['date_submitted']
    else:
        date_published = None

    if 'submitter' not in metadata:
        metadata['submitter'] = 'undefined'
    if metadata['receiver']:
        metadata['submitter'] = metadata['receiver']

    documents = Document.objects.filter(document_id=document_id).delete()
    if documents:
        logger.warning('document with id: ' + document_id + ' already exists, recreate document.')

    content_html = website.create.update_document_html_links(content_html)

    document = Document.objects.create(
        dossier=None,
        document_id=document_id,
        title_full=title,
        title_short=metadata['title_full'],
        publication_type=metadata['publication_type'],
        publisher=metadata['publisher'],
        date_published=date_published,
        source_url=document_html_url,
        content_html=content_html,
    )
    category_list = website.create.get_categories(text=metadata['category'], category_class=CategoryDocument, sep_char='|')
    document.categories.add(*category_list)

    submitters = metadata['submitter'].split('|')
    for submitter in submitters:
        website.create.create_submitter(document, submitter, date_published)

    logger.info('END')
    return document, metadata['vraagnummer']

def create_title_short(title):
    pattern = 'Vragen\s(.*aan de)\s(.*)over\s.*(\s\(ingezonden.*\)\.)'
    result = re.findall(pattern, title)
    if result:
        title = title.replace(result[0][0], '')
        title = title.replace(result[0][1], '')
        title = title.replace(result[0][2], '')
        title = title.replace('«', '"')
        title = title.replace('»', '"')
    return title.strip()

def get_receiver_from_title(title_full):
    pattern = 'Vragen\s.*aan de\s(.*) over\s.*'
    result = re.findall(pattern, title_full)
    if result:
        return result[0]
    return ''

@transaction.atomic
def link_kamervragen_and_antwoorden():
    logger.info('BEGIN')
    kamerantwoorden = Kamerantwoord.objects.all()
    for kamerantwoord in kamerantwoorden:
        kamervragen = Kamervraag.objects.filter(vraagnummer=kamerantwoord.vraagnummer)
        if kamervragen:
            kamervraag = kamervragen[0]
            kamervraag.kamerantwoord = kamerantwoord
            kamervraag.save()
    logger.info('END')

@transaction.atomic
def create_vragen_from_kamervraag_html(kamervraag):
    logger.info('BEGIN')
    Vraag.objects.filter(kamervraag=kamervraag).delete()
    tree = lxml.html.fromstring(kamervraag.document.content_html)
    elements = tree.xpath('//div[@class="vraag"]')
    counter = 1
    for element in elements:
        vraag_text = ''
        for paragraph in element.iter('p'):
            if paragraph.text is None or paragraph.text == '':
                logger.warning('empty vraag text found for antwoord ' + str(kamervraag.document.document_url) + 'vraat nr: ' + str(counter))
            else:
                vraag_text += paragraph.text_content() + '\n'
        vraag_text = re.sub('\s{2,}', ' ', vraag_text).strip()
        Vraag.objects.create(nr=counter, kamervraag=kamervraag, text=vraag_text)
        counter += 1
    logger.info('END: ' + str(counter) + ' vragen found')

@transaction.atomic
def create_antwoorden_from_antwoord_html(kamerantwoord):
    logger.info('BEGIN')
    Antwoord.objects.filter(kamerantwoord=kamerantwoord).delete()
    tree = lxml.html.fromstring(kamerantwoord.document.content_html)
    elements = tree.xpath('//div[@class="antwoord"]')
    counter = 1
    antwoorden = []
    for element in elements:
        answer_text = ''
        vraag_numbers = element.xpath('h2')[0].text_content().strip()  #
Vraag 1
 or
Vraag 2, 3

        result = re.findall('(\D*)\d', vraag_numbers)
        if result:
            vraag_numbers = vraag_numbers.replace(result[0], '')
        vraag_numbers = vraag_numbers.replace('en', ',').replace(' ', '').split(',')
        for paragraph in element.iter('p'):
            if paragraph.text is None or paragraph.text == '':
                logger.warning('empty vraag text found for antwoord ' + str(kamerantwoord.document.document_url) + 'antwoord nr: ' + str(counter))
            else:
                answer_text += paragraph.text_content() + '\n'
        answer_text = re.sub('\s{2,}', ' ', answer_text).strip()
        counter = 0
        for number in vraag_numbers:
            try:
                first_number = int(vraag_numbers[0])
                number = int(number)
            except ValueError:
                logger.info(vraag_numbers)
                logger.error('could not convert antwoord number to integer: ' + number + ' for document: ' + str(kamerantwoord.document.document_url))
                continue
            if counter > 0:
                answer_text = 'Zie antwoord vraag ' + str(vraag_numbers[0] + '.')
            antwoord = Antwoord.objects.create(nr=number, kamerantwoord=kamerantwoord, text=answer_text, see_answer_nr=first_number)
            antwoorden.append(antwoord)
            counter += 1
    logger.info('END')

def create_footnotes(footnotes_html):
    footnotes = []
    tree = lxml.html.fromstring(footnotes_html)
    elements = tree.xpath('//div[contains(@class, "voet noot")]')
    for element in elements:
        text = element.xpath('p')[0].text_content()
        note_numbers = element.xpath('.//span[@class="nootnum"]')
        if note_numbers:
            nr = note_numbers[0].text_content()
        else:
            nr = None
        links = element.xpath('p/a')
        if links:
            url = links[0].get('href')
        else:
            url = ''
        footnote = {
            'nr': nr,
            'text': text,
            'url': url
        }
        footnotes.append(footnote)
    return footnotes
<commit_msg>fix answer always has a related answer field set to the first answer<commit_after>import logging
import re
import lxml

from django.db import transaction

from document.models import Kamervraag
from document.models import Kamerantwoord
from document.models import Antwoord
from document.models import Vraag
from document.models import CategoryDocument
from document.models import FootNote

from document.models import Document

import website.create
import scraper.documents

logger = logging.getLogger(__name__)

def create_kamervragen(year, max_n=None, skip_if_exists=False):
    logger.info('BEGIN')
    infos = Kamervraag.get_kamervragen_info(year)
    counter = 0
    for info in infos:
        try:
            create_kamervraag(info['document_number'], info['overheidnl_document_id'], skip_if_exists=skip_if_exists)
        except Exception as error:
            logger.error('error for kamervraag id: ' + str(info['overheidnl_document_id']))
            logger.exception(error)
        if max_n and counter >= max_n:
            return
        counter += 1
    logger.info('END')

def create_antwoorden(year, max_n=None, skip_if_exists=False):
    logger.info('BEGIN')
    infos = Kamerantwoord.get_antwoorden_info(year)
    counter = 0
    for info in infos:
        try:
            create_kamerantwoord(info['document_number'], info['overheidnl_document_id'], skip_if_exists=skip_if_exists)
        except Exception as error:
            logger.error('error for kamerantwoord id: ' + str(info['overheidnl_document_id']))
            logger.exception(error)
        if max_n and counter >= max_n:
            break
        counter += 1
    logger.info('END')

@transaction.atomic
def create_kamervraag(document_number, overheidnl_document_id, skip_if_exists=False):
    if skip_if_exists and Kamervraag.objects.filter(document__document_id=document_number).exists():
        return
    document, vraagnummer = create_kamervraag_document(document_number, overheidnl_document_id)
    Kamervraag.objects.filter(document=document).delete()
    kamervraag = Kamervraag.objects.create(
        document=document,
        vraagnummer=vraagnummer,
        receiver=get_receiver_from_title(document.title_full)
    )
    create_vragen_from_kamervraag_html(kamervraag)
    footnotes = create_footnotes(kamervraag.document.content_html)
    for footnote in footnotes:
        FootNote.objects.create(document=document, nr=footnote['nr'], text=footnote['text'], url=footnote['url'])
    return kamervraag

@transaction.atomic
def create_kamerantwoord(document_number, overheidnl_document_id, skip_if_exists=False):
    if skip_if_exists and Kamerantwoord.objects.filter(document__document_id=document_number).exists():
        return
    document, vraagnummer = create_kamervraag_document(document_number, overheidnl_document_id)
    Kamerantwoord.objects.filter(document=document).delete()
    kamerantwoord = Kamerantwoord.objects.create(document=document, vraagnummer=vraagnummer)
    create_antwoorden_from_antwoord_html(kamerantwoord)
    return kamerantwoord

@transaction.atomic
def create_kamervraag_document(document_number, overheidnl_document_id):
    logger.info('BEGIN')
    document_html_url = 'https://zoek.officielebekendmakingen.nl/' + str(overheidnl_document_id) +'.html'
    logger.info(document_html_url)
    document_id, content_html, title = scraper.documents.get_kamervraag_document_id_and_content(document_html_url)
    metadata = scraper.documents.get_metadata(overheidnl_document_id)
    document_id = document_number

    if metadata['date_published']:
        date_published = metadata['date_published']
    elif metadata['date_submitted']:
        date_published = metadata['date_submitted']
    else:
        date_published = None

    if 'submitter' not in metadata:
        metadata['submitter'] = 'undefined'
    if metadata['receiver']:
        metadata['submitter'] = metadata['receiver']

    documents = Document.objects.filter(document_id=document_id).delete()
    if documents:
        logger.warning('document with id: ' + document_id + ' already exists, recreate document.')

    content_html = website.create.update_document_html_links(content_html)

    document = Document.objects.create(
        dossier=None,
        document_id=document_id,
        title_full=title,
        title_short=metadata['title_full'],
        publication_type=metadata['publication_type'],
        publisher=metadata['publisher'],
        date_published=date_published,
        source_url=document_html_url,
        content_html=content_html,
    )
    category_list = website.create.get_categories(text=metadata['category'], category_class=CategoryDocument, sep_char='|')
    document.categories.add(*category_list)

    submitters = metadata['submitter'].split('|')
    for submitter in submitters:
        website.create.create_submitter(document, submitter, date_published)

    logger.info('END')
    return document, metadata['vraagnummer']

def create_title_short(title):
    pattern = 'Vragen\s(.*aan de)\s(.*)over\s.*(\s\(ingezonden.*\)\.)'
    result = re.findall(pattern, title)
    if result:
        title = title.replace(result[0][0], '')
        title = title.replace(result[0][1], '')
        title = title.replace(result[0][2], '')
        title = title.replace('«', '"')
        title = title.replace('»', '"')
    return title.strip()

def get_receiver_from_title(title_full):
    pattern = 'Vragen\s.*aan de\s(.*) over\s.*'
    result = re.findall(pattern, title_full)
    if result:
        return result[0]
    return ''

@transaction.atomic
def link_kamervragen_and_antwoorden():
    logger.info('BEGIN')
    kamerantwoorden = Kamerantwoord.objects.all()
    for kamerantwoord in kamerantwoorden:
        kamervragen = Kamervraag.objects.filter(vraagnummer=kamerantwoord.vraagnummer)
        if kamervragen:
            kamervraag = kamervragen[0]
            kamervraag.kamerantwoord = kamerantwoord
            kamervraag.save()
    logger.info('END')

@transaction.atomic
def create_vragen_from_kamervraag_html(kamervraag):
    logger.info('BEGIN')
    Vraag.objects.filter(kamervraag=kamervraag).delete()
    tree = lxml.html.fromstring(kamervraag.document.content_html)
    elements = tree.xpath('//div[@class="vraag"]')
    counter = 1
    for element in elements:
        vraag_text = ''
        for paragraph in element.iter('p'):
            if paragraph.text is None or paragraph.text == '':
                logger.warning('empty vraag text found for antwoord ' + str(kamervraag.document.document_url) + 'vraat nr: ' + str(counter))
            else:
                vraag_text += paragraph.text_content() + '\n'
        vraag_text = re.sub('\s{2,}', ' ', vraag_text).strip()
        Vraag.objects.create(nr=counter, kamervraag=kamervraag, text=vraag_text)
        counter += 1
    logger.info('END: ' + str(counter) + ' vragen found')

@transaction.atomic
def create_antwoorden_from_antwoord_html(kamerantwoord):
    logger.info('BEGIN')
    Antwoord.objects.filter(kamerantwoord=kamerantwoord).delete()
    tree = lxml.html.fromstring(kamerantwoord.document.content_html)
    elements = tree.xpath('//div[@class="antwoord"]')
    counter = 1
    antwoorden = []
    for element in elements:
        answer_text = ''
        vraag_numbers = element.xpath('h2')[0].text_content().strip()  #
Vraag 1
 or
Vraag 2, 3

        result = re.findall('(\D*)\d', vraag_numbers)
        if result:
            vraag_numbers = vraag_numbers.replace(result[0], '')
        vraag_numbers = vraag_numbers.replace('en', ',').replace(' ', '').split(',')
        for paragraph in element.iter('p'):
            if paragraph.text is None or paragraph.text == '':
                logger.warning('empty vraag text found for antwoord ' + str(kamerantwoord.document.document_url) + 'antwoord nr: ' + str(counter))
            else:
                answer_text += paragraph.text_content() + '\n'
        answer_text = re.sub('\s{2,}', ' ', answer_text).strip()
        counter = 0
        see_answer_nr = None
        for number in vraag_numbers:
            try:
                first_number = int(vraag_numbers[0])
                number = int(number)
            except ValueError:
                logger.info(vraag_numbers)
                logger.error('could not convert antwoord number to integer: ' + number + ' for document: ' + str(kamerantwoord.document.document_url))
                continue
            if counter > 0:
                answer_text = 'Zie antwoord vraag ' + str(vraag_numbers[0] + '.')
                see_answer_nr = first_number
            antwoord = Antwoord.objects.create(nr=number, kamerantwoord=kamerantwoord, text=answer_text, see_answer_nr=see_answer_nr)
            antwoorden.append(antwoord)
            counter += 1
    logger.info('END')

def create_footnotes(footnotes_html):
    footnotes = []
    tree = lxml.html.fromstring(footnotes_html)
    elements = tree.xpath('//div[contains(@class, "voet noot")]')
    for element in elements:
        text = element.xpath('p')[0].text_content()
        note_numbers = element.xpath('.//span[@class="nootnum"]')
        if note_numbers:
            nr = note_numbers[0].text_content()
        else:
            nr = None
        links = element.xpath('p/a')
        if links:
            url = links[0].get('href')
        else:
            url = ''
        footnote = {
            'nr': nr,
            'text': text,
            'url': url
        }
        footnotes.append(footnote)
    return footnotes
<|endoftext|>


<commit_before>import numpy as np
import os
import mdtraj as md
from mdtraj.utils.delay_import import import_
import tempfile
from distutils.spawn import find_executable
import simtk.unit as units

from .utils import temporary_directory

PACKMOL_PATH = find_executable("packmol")

HEADER_TEMPLATE = """
# Mixture

tolerance %f
filetype pdb
output %s
add_amber_ter

"""

BOX_TEMPLATE = """
structure %s
  number %d
  inside box 0. 0. 0. %f %f %f
end structure
"""

def standardize_water(mol_traj):
    """Ensure that a water molecule has the correct MDTraj Topology.

    The PDB format doesn't require CONECT records for a water molecule,
    but MDTraj correctly recognize water molecules bonds only if they
    adopt specific residue and atom names. This function standardize
    the names to ensure the Topology is correctly connected.

    Parameters
    ----------
    mol_traj : mdtraj.Trajectory
        A trajectory object describing a single water molecule. If the
        trajectory doesn't describe a water molecule, nothing happens.
        The residue name and atom names are modified to adhere to MDTraj
        standard definition, if this is a water molecule.

    Returns
    -------
    bool
        True if this was a water molecule, False otherwise.

    """
    if mol_traj.topology.n_atoms != 3 or mol_traj.topology.n_residues != 1:
        # This is not a water molecule.
        return False

    # Count oxygen and hydrogens in molecule and save their indices.
    atom_element_ids = {'O': [], 'H': []}
    for atom_index, atom in enumerate(mol_traj.topology.atoms):
        try:
            atom_element_ids[atom.element.symbol].append(atom_index)
        except KeyError:
            # There's an element different than oxygen or hydrogen.
            return False

    # This is water if there are two hydrogens and an oxygen.
    if not (len(atom_element_ids['O']) == 1 and len(atom_element_ids['H']) == 2):
        return False

    # Rename residue and atoms.
    mol_traj.topology.residue(0).name = 'HOH'
    [o_index], [h1_index, h2_index] = atom_element_ids['O'], atom_element_ids['H']
    for index, std_name in zip([o_index, h1_index, h2_index], ['O', 'H1', 'H2']):
        mol_traj.topology.atom(index).name = std_name

    # Update bonds now that water residue is standard.
    mol_traj.topology.create_standard_bonds()
    return True

def pack_box(pdb_filenames_or_trajectories, n_molecules_list, tolerance=2.0, box_size=None):
    """Run packmol to generate a box containing a mixture of molecules.

    Parameters
    ----------
    pdb_filenames_or_trajectories : list({str, Trajectory})
        List of pdb filenames or trajectories for each component of mixture.
        If this is a list of trajectories, the trajectories will be saved to
        as temporary files to be run in packmol. Water molecules must have
        MDTraj-standard residue and atom names as defined in
        mdtraj/formats/pdb/data/pdbNames.xml.
    n_molecules_list : list(int)
        The number of molecules of each mixture component.
    tolerance : float, optional, default=2.0
        The mininum spacing between molecules during packing.  In ANGSTROMS!
    box_size : float, optional
        The size of the box to generate.  In ANGSTROMS.
        Default generates boxes that are very large for increased stability.
        May require extra time for energy minimization and equilibration.

    Returns
    -------
    trj : MDTraj.Trajectory
        Single frame trajectory with mixture box.

    Notes
    -----
    Water molecules must have MDTraj-standard residue and atom names as defined
    in mdtraj/formats/pdb/data/pdbNames.xml, otherwise MDTraj won't be able to
    perceive the bonds and the Topology of the returned Trajectory will be incorrect.
    Be aware that MDTraj uses nanometers internally, but packmol uses angstrom
    units. The present function takes `tolerance` and `box_size` in angstrom
    units, but the output trajectory will have data in nm.
    Also note that OpenMM is pretty picky about the format of unit cell input,
    so use the example in tests/test_packmol.py to ensure that you do the right thing.

    See Also
    --------
    standardize_water
        Standardize residue and atom names of a water molecule.

    """
    assert len(pdb_filenames_or_trajectories) == len(n_molecules_list), "Must input same number of pdb filenames as num molecules"

    if PACKMOL_PATH is None:
        raise(IOError("Packmol not found, cannot run pack_box()"))

    pdb_filenames = []
    trj_i = []

    # We save all the temporary files in a temporary directory
    # that is deleted at the end of the function.
    with temporary_directory() as tmp_dir:

        # We need all molecules as both pdb files (as packmol input)
        # and mdtraj.Trajectory for restoring bonds later.
        for obj in pdb_filenames_or_trajectories:
            try:  # See if MDTraj Trajectory
                tmp_filename = tempfile.mktemp(suffix=".pdb", dir=tmp_dir)
                obj.save_pdb(tmp_filename)
            except AttributeError:  # Not an MDTraj Trajectory, assume filename
                pdb_filenames.append(obj)
                trj_i.append(md.load(obj))
            else:
                pdb_filenames.append(tmp_filename)
                trj_i.append(obj)

        # Approximating volume to initialize box.
        if box_size is None:
            box_size = approximate_volume(pdb_filenames, n_molecules_list)

        # Adjust box_size for periodic box. Packmol does not explicitly
        # support periodic boundary conditions and the suggestion on
        # their docs is to pack in a box 2 angstroms smaller. See
        # http://www.ime.unicamp.br/~martinez/packmol/userguide.shtml#pbc
        packmol_box_size = box_size - 2  # angstroms

        # The path to packmol's output PDB file.
        output_filename = tempfile.mktemp(suffix=".pdb", dir=tmp_dir)

        # Create input file for packmol.
        header = HEADER_TEMPLATE % (tolerance, output_filename)
        for k in range(len(pdb_filenames)):
            filename = pdb_filenames[k]
            n_molecules = n_molecules_list[k]
            header += BOX_TEMPLATE % (filename, n_molecules, packmol_box_size,
                                      packmol_box_size, packmol_box_size)

        print(header)

        packmol_filename = tempfile.mktemp(suffix='.txt', dir=tmp_dir)
        with open(packmol_filename, 'w') as file_handle:
            file_handle.write(header)

        # Run packmol and load output PDB file.
        os.system("%s < %s" % (PACKMOL_PATH, file_handle.name))
        trj = md.load(output_filename)

    assert trj.topology.n_chains == sum(n_molecules_list), "Packmol error: molecules missing from output"

    #Begin hack to introduce bonds for the MISSING CONECT ENTRIES THAT PACKMOL FAILS TO WRITE

    top, bonds = trj.top.to_dataframe()
    bonds_i = [t.top.to_dataframe()[1] for t in trj_i]

    offset = 0
    bonds = []
    for i in range(len(trj_i)):
        n_atoms = trj_i[i].n_atoms
        for j in range(n_molecules_list[i]):
            bonds.extend(bonds_i[i] + offset)
            offset += n_atoms

    bonds = np.array(bonds)
    trj.top = md.Topology.from_dataframe(top, bonds)

    # Set the requested box size.
    trj.unitcell_vectors = np.array([np.eye(3)]) * box_size / 10.

    return trj

def approximate_volume(pdb_filenames, n_molecules_list, box_scaleup_factor=2.0):
    """Approximate the appropriate box size based on the number and types of atoms present.

    Parameters
    ----------
    pdb_filenames : list(str)
        List of pdb filenames for each component of mixture.
    n_molecules_list : list(int)
        The number of molecules of each mixture component.
    box_scaleup_factor : float, optional, default = 2.0
        Factor by which the estimated box size is increased

    Returns
    -------
    box_size : float
        The size of the box to generate.  In ANGSTROMS.

    Notes
    -----
    By default, boxes are very large for increased stability, and therefore may
    require extra time for energy minimization and equilibration.
    """
    volume = 0.0 # in cubic angstroms
    for k, (pdb_file) in enumerate(pdb_filenames):
        molecule_volume = 0.0
        molecule_trj = md.load(pdb_filenames[k])
        for atom in molecule_trj.topology.atoms:
            if atom.element.symbol == 'H':
                molecule_volume += 5.0 # approximated from bondi radius = 1.06 angstroms
            else:
                molecule_volume += 15.0 # approximated from bondi radius of carbon = 1.53 angstroms
        volume += molecule_volume * n_molecules_list[k]
    box_size = volume**(1.0/3.0) * box_scaleup_factor
    return box_size

def approximate_volume_by_density(smiles_strings, n_molecules_list, density=1.0,
                                  box_scaleup_factor=1.1, box_buffer=2.0):
    """Generate an approximate box size based on the number and molecular weight of molecules present, and a target density for the final solvated mixture. If no density is specified, the target density is assumed to be 1 g/ml.

    Parameters
    ----------
    smiles_strings : list(str)
        List of smiles strings for each component of mixture.
    n_molecules_list : list(int)
        The number of molecules of each mixture component.
    box_scaleup_factor : float, optional, default = 1.1
        Factor by which the estimated box size is increased
    density : float, optional, default 1.0
        Target density for final system in g/ml
    box_buffer : float [ANGSTROMS], optional, default 2.0.
        This quantity is added to the final estimated box size
        (after scale-up). With periodic boundary conditions,
        packmol docs suggests to leave an extra 2 Angstroms
        buffer during packing.

    Returns
    -------
    box_size : float
        The size (edge length) of the box to generate.  In ANGSTROMS.

    Notes
    -----
    By default, boxes are only modestly large. This approach has not been extensively tested for stability but has been used in th Mobley lab for perhaps ~100 different systems without substantial problems.
    """

    oechem = import_("openeye.oechem")

    density = density * units.grams/units.milliliter

    #Load molecules to get molecular weights
    wts = []
    mass = 0.0*units.grams/units.mole * 1./units.AVOGADRO_CONSTANT_NA #For calculating total mass
    for (idx,smi) in enumerate(smiles_strings):
        mol = oechem.OEMol()
        oechem.OEParseSmiles(mol, smi)
        wts.append( oechem.OECalculateMolecularWeight(mol)*units.grams/units.mole )
        mass += n_molecules_list[idx] * wts[idx] * 1./units.AVOGADRO_CONSTANT_NA

    #Estimate volume based on mass and density
    #Density = mass/volume so volume = mass/density (volume units are ml)
    vol = mass/density
    #Convert to box length in angstroms
    edge = vol**(1./3.)

    #Compute final box size
    box_size = edge*box_scaleup_factor/units.angstroms# + box_buffer

    return box_size

def rename_water_atoms( pdb_filename, O_name = 'O', H1_name = 'H1', H2_name = 'H2' ):
    """Rename water atoms in a specified PDB file to have target names. Typically used to ensure a packmol-generated box containing water has water atom names corresponding to what tleap expects for standard water models.

    Parameters
    ----------
    pdb_filename : str
        The target PDB filename to edit
    O_name : str, optional, default 'O'
        Target name to set water oxygen names to
    H1_name : str, optional, default 'H1'
        Target name to set water hydrogen names to, for first hydrogen
    H2_name : str, optional, default 'H2'
        Target name to set water hydrogen names to, for second hydrogen

    Returns
    -------

    Notes
    -------
    Uses ParmEd to makes edits. Identifies waters by reading residues from target PDB file and identifying any residue containing three atoms with names O or O#, H or H#, and H or H# (where # is a digit or sequence of digits) as water molecules.
    """

    parmed = import_("parmed")

    pdb = parmed.load_file( pdb_filename )

    #Find waters and rename
    for residue in pdb.residues:
        if len(residue)==3:
            #Build list of atom types (PDB files don't store these) from names after stripping off digits
            types = []
            for atom in residue.atoms:
                name = atom.name
                while name[-1].isdigit():
                    name = name[:-1]
                types.append(name)
            #See if it's water and, if so, rename
            if 'O' in types and types.count('H')==2:
                hct = 0
                for atom in residue.atoms:
                    if 'O' in atom.name:
                        atom.name = O_name
                    elif 'H' in atom.name:
                        if hct==0:
                            atom.name = H1_name
                        else:
                            atom.name = H2_name
                        hct+=1

    #Write file
    pdb.write_pdb( pdb_filename )

<commit_msg>Fix incompatibility with recent OpenMM versions wherein bond orders are not handled properly for packmol<commit_after>import numpy as np
import os
import mdtraj as md
from mdtraj.utils.delay_import import import_
import tempfile
from distutils.spawn import find_executable
import simtk.unit as units
import copy

from .utils import temporary_directory

PACKMOL_PATH = find_executable("packmol")

HEADER_TEMPLATE = """
# Mixture

tolerance %f
filetype pdb
output %s
add_amber_ter

"""

BOX_TEMPLATE = """
structure %s
  number %d
  inside box 0. 0. 0. %f %f %f
end structure
"""

def standardize_water(mol_traj):
    """Ensure that a water molecule has the correct MDTraj Topology.

    The PDB format doesn't require CONECT records for a water molecule,
    but MDTraj correctly recognize water molecules bonds only if they
    adopt specific residue and atom names. This function standardize
    the names to ensure the Topology is correctly connected.

    Parameters
    ----------
    mol_traj : mdtraj.Trajectory
        A trajectory object describing a single water molecule. If the
        trajectory doesn't describe a water molecule, nothing happens.
        The residue name and atom names are modified to adhere to MDTraj
        standard definition, if this is a water molecule.

    Returns
    -------
    bool
        True if this was a water molecule, False otherwise.

    """
    if mol_traj.topology.n_atoms != 3 or mol_traj.topology.n_residues != 1:
        # This is not a water molecule.
        return False

    # Count oxygen and hydrogens in molecule and save their indices.
    atom_element_ids = {'O': [], 'H': []}
    for atom_index, atom in enumerate(mol_traj.topology.atoms):
        try:
            atom_element_ids[atom.element.symbol].append(atom_index)
        except KeyError:
            # There's an element different than oxygen or hydrogen.
            return False

    # This is water if there are two hydrogens and an oxygen.
    if not (len(atom_element_ids['O']) == 1 and len(atom_element_ids['H']) == 2):
        return False

    # Rename residue and atoms.
    mol_traj.topology.residue(0).name = 'HOH'
    [o_index], [h1_index, h2_index] = atom_element_ids['O'], atom_element_ids['H']
    for index, std_name in zip([o_index, h1_index, h2_index], ['O', 'H1', 'H2']):
        mol_traj.topology.atom(index).name = std_name

    # Update bonds now that water residue is standard.
    mol_traj.topology.create_standard_bonds()
    return True

def pack_box(pdb_filenames_or_trajectories, n_molecules_list, tolerance=2.0, box_size=None):
    """Run packmol to generate a box containing a mixture of molecules.

    Parameters
    ----------
    pdb_filenames_or_trajectories : list({str, Trajectory})
        List of pdb filenames or trajectories for each component of mixture.
        If this is a list of trajectories, the trajectories will be saved to
        as temporary files to be run in packmol. Water molecules must have
        MDTraj-standard residue and atom names as defined in
        mdtraj/formats/pdb/data/pdbNames.xml.
    n_molecules_list : list(int)
        The number of molecules of each mixture component.
    tolerance : float, optional, default=2.0
        The mininum spacing between molecules during packing.  In ANGSTROMS!
    box_size : float, optional
        The size of the box to generate.  In ANGSTROMS.
        Default generates boxes that are very large for increased stability.
        May require extra time for energy minimization and equilibration.

    Returns
    -------
    trj : MDTraj.Trajectory
        Single frame trajectory with mixture box.

    Notes
    -----
    Water molecules must have MDTraj-standard residue and atom names as defined
    in mdtraj/formats/pdb/data/pdbNames.xml, otherwise MDTraj won't be able to
    perceive the bonds and the Topology of the returned Trajectory will be incorrect.
    Be aware that MDTraj uses nanometers internally, but packmol uses angstrom
    units. The present function takes `tolerance` and `box_size` in angstrom
    units, but the output trajectory will have data in nm.
    Also note that OpenMM is pretty picky about the format of unit cell input,
    so use the example in tests/test_packmol.py to ensure that you do the right thing.

    See Also
    --------
    standardize_water
        Standardize residue and atom names of a water molecule.

    """
    assert len(pdb_filenames_or_trajectories) == len(n_molecules_list), "Must input same number of pdb filenames as num molecules"

    if PACKMOL_PATH is None:
        raise(IOError("Packmol not found, cannot run pack_box()"))

    pdb_filenames = []
    trj_i = []

    # We save all the temporary files in a temporary directory
    # that is deleted at the end of the function.
    with temporary_directory() as tmp_dir:

        # We need all molecules as both pdb files (as packmol input)
        # and mdtraj.Trajectory for restoring bonds later.
        for obj in pdb_filenames_or_trajectories:
            try:  # See if MDTraj Trajectory
                tmp_filename = tempfile.mktemp(suffix=".pdb", dir=tmp_dir)
                obj.save_pdb(tmp_filename)
            except AttributeError:  # Not an MDTraj Trajectory, assume filename
                pdb_filenames.append(obj)
                trj_i.append(md.load(obj))
            else:
                pdb_filenames.append(tmp_filename)
                trj_i.append(obj)

        # Approximating volume to initialize box.
        if box_size is None:
            box_size = approximate_volume(pdb_filenames, n_molecules_list)

        # Adjust box_size for periodic box. Packmol does not explicitly
        # support periodic boundary conditions and the suggestion on
        # their docs is to pack in a box 2 angstroms smaller. See
        # http://www.ime.unicamp.br/~martinez/packmol/userguide.shtml#pbc
        packmol_box_size = box_size - 2  # angstroms

        # The path to packmol's output PDB file.
        output_filename = tempfile.mktemp(suffix=".pdb", dir=tmp_dir)

        # Create input file for packmol.
        header = HEADER_TEMPLATE % (tolerance, output_filename)
        for k in range(len(pdb_filenames)):
            filename = pdb_filenames[k]
            n_molecules = n_molecules_list[k]
            header += BOX_TEMPLATE % (filename, n_molecules, packmol_box_size,
                                      packmol_box_size, packmol_box_size)

        print(header)

        packmol_filename = tempfile.mktemp(suffix='.txt', dir=tmp_dir)
        with open(packmol_filename, 'w') as file_handle:
            file_handle.write(header)

        # Run packmol and load output PDB file.
        os.system("%s < %s" % (PACKMOL_PATH, file_handle.name))
        trj = md.load(output_filename)

    assert trj.topology.n_chains == sum(n_molecules_list), "Packmol error: molecules missing from output"

    #Begin hack to introduce bonds for the MISSING CONECT ENTRIES THAT PACKMOL FAILS TO WRITE

    top, bonds = trj.top.to_dataframe()
    bonds_i = [t.top.to_dataframe()[1] for t in trj_i]

    offset = 0
    bonds = []
    for i in range(len(trj_i)):
        n_atoms = trj_i[i].n_atoms
        for j in range(n_molecules_list[i]):
            # Offset atom numbers for the bonds by the current atom number; don't offset bond orders
            thesebonds = copy.deepcopy(bonds_i[i])
            thesebonds[:,0:2] += offset
            # Store
            bonds.extend(thesebonds)
            offset += n_atoms
    bonds = np.array(bonds)
    trj.top = md.Topology.from_dataframe(top, bonds)

    # Set the requested box size.
    trj.unitcell_vectors = np.array([np.eye(3)]) * box_size / 10.

    return trj

def approximate_volume(pdb_filenames, n_molecules_list, box_scaleup_factor=2.0):
    """Approximate the appropriate box size based on the number and types of atoms present.

    Parameters
    ----------
    pdb_filenames : list(str)
        List of pdb filenames for each component of mixture.
    n_molecules_list : list(int)
        The number of molecules of each mixture component.
    box_scaleup_factor : float, optional, default = 2.0
        Factor by which the estimated box size is increased

    Returns
    -------
    box_size : float
        The size of the box to generate.  In ANGSTROMS.

    Notes
    -----
    By default, boxes are very large for increased stability, and therefore may
    require extra time for energy minimization and equilibration.
    """
    volume = 0.0 # in cubic angstroms
    for k, (pdb_file) in enumerate(pdb_filenames):
        molecule_volume = 0.0
        molecule_trj = md.load(pdb_filenames[k])
        for atom in molecule_trj.topology.atoms:
            if atom.element.symbol == 'H':
                molecule_volume += 5.0 # approximated from bondi radius = 1.06 angstroms
            else:
                molecule_volume += 15.0 # approximated from bondi radius of carbon = 1.53 angstroms
        volume += molecule_volume * n_molecules_list[k]
    box_size = volume**(1.0/3.0) * box_scaleup_factor
    return box_size

def approximate_volume_by_density(smiles_strings, n_molecules_list, density=1.0,
                                  box_scaleup_factor=1.1, box_buffer=2.0):
    """Generate an approximate box size based on the number and molecular weight of molecules present, and a target density for the final solvated mixture. If no density is specified, the target density is assumed to be 1 g/ml.

    Parameters
    ----------
    smiles_strings : list(str)
        List of smiles strings for each component of mixture.
    n_molecules_list : list(int)
        The number of molecules of each mixture component.
    box_scaleup_factor : float, optional, default = 1.1
        Factor by which the estimated box size is increased
    density : float, optional, default 1.0
        Target density for final system in g/ml
    box_buffer : float [ANGSTROMS], optional, default 2.0.
        This quantity is added to the final estimated box size
        (after scale-up). With periodic boundary conditions,
        packmol docs suggests to leave an extra 2 Angstroms
        buffer during packing.

    Returns
    -------
    box_size : float
        The size (edge length) of the box to generate.  In ANGSTROMS.

    Notes
    -----
    By default, boxes are only modestly large. This approach has not been extensively tested for stability but has been used in th Mobley lab for perhaps ~100 different systems without substantial problems.
    """

    oechem = import_("openeye.oechem")

    density = density * units.grams/units.milliliter

    #Load molecules to get molecular weights
    wts = []
    mass = 0.0*units.grams/units.mole * 1./units.AVOGADRO_CONSTANT_NA #For calculating total mass
    for (idx,smi) in enumerate(smiles_strings):
        mol = oechem.OEMol()
        oechem.OEParseSmiles(mol, smi)
        wts.append( oechem.OECalculateMolecularWeight(mol)*units.grams/units.mole )
        mass += n_molecules_list[idx] * wts[idx] * 1./units.AVOGADRO_CONSTANT_NA

    #Estimate volume based on mass and density
    #Density = mass/volume so volume = mass/density (volume units are ml)
    vol = mass/density
    #Convert to box length in angstroms
    edge = vol**(1./3.)

    #Compute final box size
    box_size = edge*box_scaleup_factor/units.angstroms# + box_buffer

    return box_size

def rename_water_atoms( pdb_filename, O_name = 'O', H1_name = 'H1', H2_name = 'H2' ):
    """Rename water atoms in a specified PDB file to have target names. Typically used to ensure a packmol-generated box containing water has water atom names corresponding to what tleap expects for standard water models.

    Parameters
    ----------
    pdb_filename : str
        The target PDB filename to edit
    O_name : str, optional, default 'O'
        Target name to set water oxygen names to
    H1_name : str, optional, default 'H1'
        Target name to set water hydrogen names to, for first hydrogen
    H2_name : str, optional, default 'H2'
        Target name to set water hydrogen names to, for second hydrogen

    Returns
    -------

    Notes
    -------
    Uses ParmEd to makes edits. Identifies waters by reading residues from target PDB file and identifying any residue containing three atoms with names O or O#, H or H#, and H or H# (where # is a digit or sequence of digits) as water molecules.
    """

    parmed = import_("parmed")

    pdb = parmed.load_file( pdb_filename )

    #Find waters and rename
    for residue in pdb.residues:
        if len(residue)==3:
            #Build list of atom types (PDB files don't store these) from names after stripping off digits
            types = []
            for atom in residue.atoms:
                name = atom.name
                while name[-1].isdigit():
                    name = name[:-1]
                types.append(name)
            #See if it's water and, if so, rename
            if 'O' in types and types.count('H')==2:
                hct = 0
                for atom in residue.atoms:
                    if 'O' in atom.name:
                        atom.name = O_name
                    elif 'H' in atom.name:
                        if hct==0:
                            atom.name = H1_name
                        else:
                            atom.name = H2_name
                        hct+=1

    #Write file
    pdb.write_pdb( pdb_filename )
<|endoftext|>


<commit_before># cli commands

from twisted.python import log
from twisted.internet import defer

from opennsa import nsa, error

# this parser should perhaps be somewhere else
def _createSTP(stp_desc):
    network, local_part = stp_desc.rsplit(':',1)
    if '#' in local_part:
        port, label_part = local_part.split('#',1)
        labels = []
        for tvl in label_part.split(';'):
            if not '=' in tvl:
                raise ValueError('Invalid label type-value: %s' % tvl)
            type_, values = tvl.split('=')
            labels.append( nsa.Label( type_, values ) )
    else:
        port = local_part
        labels = None

    return nsa.STP(network, port, labels)

def _createServiceParams(start_time, end_time, src, dst, bandwidth):

    src_stp = _createSTP(src)
    dst_stp = _createSTP(dst)

    return nsa.ServiceParameters(start_time, end_time, src_stp, dst_stp, bandwidth)

@defer.inlineCallbacks
def discover(client, service_url):

    res = yield client.queryNSA(service_url)
    print "-- COMMAND RESULT --"
    print res
    print "--"

@defer.inlineCallbacks
def reserve(client, client_nsa, provider_nsa, src, dst, start_time, end_time, bandwidth, connection_id, global_id):

    service_params = _createServiceParams(start_time, end_time, src, dst, bandwidth)

    log.msg("Connection id: %s  Global id: %s" % (connection_id, global_id))

    try:
        yield client.reserve(client_nsa, provider_nsa, None, global_id, 'Test Connection', connection_id, service_params)
        log.msg("Reservation created at %s" % provider_nsa)
    except error.NSIError, e:
        log.msg('Error reserving %s, %s : %s' % (connection_id, e.__class__.__name__, str(e)))

@defer.inlineCallbacks
def reserveprovision(client, client_nsa, provider_nsa, src, dst, start_time, end_time, bandwidth, connection_id, global_id):

    service_params = _createServiceParams(start_time, end_time, src, dst, bandwidth)

    log.msg("Connection id: %s  Global id: %s" % (connection_id, global_id))

    try:
        yield client.reserve(client_nsa, provider_nsa, None, global_id, 'Test Connection', connection_id, service_params)
        log.msg("Connection reserved at %s" % provider_nsa)
    except error.NSIError, e:
        log.msg('Error reserving %s, %s : %s' % (connection_id, e.__class__.__name__, str(e)))
        defer.returnValue(None)

    try:
        yield client.provision(client_nsa, provider_nsa, None, connection_id)
        log.msg('Connection %s provisioned' % connection_id)
    except error.NSIError, e:
        log.msg('Error provisioning %s, %s : %s' % (connection_id, e.__class__.__name__, str(e)))

@defer.inlineCallbacks
def rprt(client, client_nsa, provider_nsa, src, dst, start_time, end_time, bandwidth, connection_id, global_id):
    # reserve, provision, release,  terminate
    service_params = _createServiceParams(start_time, end_time, src, dst, bandwidth)

    log.msg("Connection id: %s  Global id: %s" % (connection_id, global_id))

    try:
        yield client.reserve(client_nsa, provider_nsa, None, global_id, 'Test Connection', connection_id, service_params)
        log.msg("Connection reserved at %s" % provider_nsa)
    except error.NSIError, e:
        log.msg('Error reserving %s, %s : %s' % (connection_id, e.__class__.__name__, str(e)))
        defer.returnValue(None)

    try:
        yield client.provision(client_nsa, provider_nsa, None, connection_id)
        log.msg('Connection %s provisioned' % connection_id)
    except error.NSIError, e:
        log.msg('Error provisioning %s, %s : %s' % (connection_id, e.__class__.__name__, str(e)))

    try:
        yield client.release(client_nsa, provider_nsa, None, connection_id)
        log.msg('Connection %s released' % connection_id)
    except error.NSIError, e:
        log.msg('Error releasing %s, %s : %s' % (connection_id, e.__class__.__name__, str(e)))

    try:
        yield client.terminate(client_nsa, provider_nsa, None, connection_id)
        log.msg('Connection %s terminated' % connection_id)
    except error.NSIError, e:
        log.msg('Error terminating %s, %s : %s' % (connection_id, e.__class__.__name__, str(e)))

@defer.inlineCallbacks
def provision(client, client_nsa, provider_nsa, connection_id):

    try:
        yield client.provision(client_nsa, provider_nsa, None, connection_id)
        log.msg('Connection %s provisioned' % connection_id)
    except error.NSIError, e:
        log.msg('Error provisioning %s, %s : %s' % (connection_id, e.__class__.__name__, str(e)))

@defer.inlineCallbacks
def release(client, client_nsa, provider_nsa, connection_id):

    try:
        yield client.release(client_nsa, provider_nsa, None, connection_id)
        log.msg('Connection %s released' % connection_id)
    except error.NSIError, e:
        log.msg('Error releasing %s, %s : %s' % (connection_id, e.__class__.__name__, str(e)))

@defer.inlineCallbacks
def terminate(client, client_nsa, provider_nsa, connection_id):

    try:
        yield client.terminate(client_nsa, provider_nsa, None, connection_id)
        log.msg('Connection %s terminated' % connection_id)
    except error.NSIError, e:
        log.msg('Error terminating %s, %s : %s' % (connection_id, e.__class__.__name__, str(e)))

@defer.inlineCallbacks
def querysummary(client, client_nsa, provider_nsa, connection_ids, global_reservation_ids):

    try:
        qc = yield client.query(client_nsa, provider_nsa, None, "Summary", connection_ids, global_reservation_ids)
        log.msg('Query results:')
        for qr in qc:
            log.msg('Connection    : %s' % qr.connectionId)
            if qr.globalReservationId:
                log.msg('  Global ID   : %s' % qr.globalReservationId)
            if qr.description:
                log.msg('  Description : %s' % qr.description)

            log.msg('  Reservation state : %s' % qr.connectionStates.reservationState.state)
            log.msg('  Provision   state : %s' % qr.connectionStates.provisionState.state)
            log.msg('  Activaction state : %s' % qr.connectionStates.activationState.state)
            if qr.children:
                log.msg('  Children    : %s' % qr.children)
    except error.NSIError, e:
        log.msg('Error querying %s, %s : %s' % (connection_ids, e.__class__.__name__, str(e)))

@defer.inlineCallbacks
def querydetails(client, client_nsa, provider_nsa, connection_ids, global_reservation_ids):

    try:
        qc = yield client.query(client_nsa, provider_nsa, None, "Details", connection_ids, global_reservation_ids)
        log.msg('Query results:')
        for qr in qc:
            log.msg('Connection: %s' % qr.connectionId)
            log.msg('  States: %s' % qr.connectionStates)
    except error.NSIError, e:
        log.msg('Error querying %s, %s : %s' % (connection_ids, e.__class__.__name__, str(e)))

def path(topology_file, source_stp, dest_stp):

    raise NotImplementedError('Path computation not available for NML yet')
    topo = None

    source_network, source_port = source_stp.split(':',1)
    dest_network,   dest_port   = dest_stp.split(':', 1)

    r_source_stp    = nsa.STP(source_network, source_port)
    r_dest_stp      = nsa.STP(dest_network,   dest_port)

    paths = topo.findPaths(r_source_stp, r_dest_stp)

    for p in sorted(paths, key=lambda p : len(p.network_links)):
        log.msg(str(p))

def topology(topology_file):

    raise NotImplementedError('Topology dump not available for NML yet')
    topo = None

    for nw in topo.networks:
        ns = '%s (%s)' % (nw.name, ','.join( sorted( [ ep.endpoint for ep in nw.endpoints ] ) ) )
        log.msg(ns)

def topologyGraph(topology_file, all_links=False):

    raise NotImplementedError('Topology graph not available for NML yet')
    topo = None

    links = []

    for nw in topo.networks:
        for ep in nw.endpoints:
            if ep.dest_stp:
                nw1 = nw.name.replace('.ets', '').replace('-','_')
                nw2 = ep.dest_stp.network.replace('.ets', '').replace('-', '_')

                l = [ nw1, nw2 ]
                if all_links:
                    if nw1 < nw2: # this prevents us from building double links
                        links.append(l)
                else:
                    l = sorted(l)
                    if not l in links:
                        links.append(l)

    log.msg('graph Network {')
    for l in sorted(links):
        log.msg('  %s -- %s;' % (l[0], l[1]))
    log.msg('}')

<commit_msg>update cli command to do reserve commit in reserve<commit_after># cli commands

from twisted.python import log
from twisted.internet import defer

from opennsa import nsa, error

# this parser should perhaps be somewhere else
def _createSTP(stp_desc):
    network, local_part = stp_desc.rsplit(':',1)
    if '#' in local_part:
        port, label_part = local_part.split('#',1)
        labels = []
        for tvl in label_part.split(';'):
            if not '=' in tvl:
                raise ValueError('Invalid label type-value: %s' % tvl)
            type_, values = tvl.split('=')
            labels.append( nsa.Label( type_, values ) )
    else:
        port = local_part
        labels = None

    return nsa.STP(network, port, labels)

def _createServiceParams(start_time, end_time, src, dst, bandwidth):

    src_stp = _createSTP(src)
    dst_stp = _createSTP(dst)

    return nsa.ServiceParameters(start_time, end_time, src_stp, dst_stp, bandwidth)

@defer.inlineCallbacks
def discover(client, service_url):

    res = yield client.queryNSA(service_url)
    print "-- COMMAND RESULT --"
    print res
    print "--"

@defer.inlineCallbacks
def reserve(client, client_nsa, provider_nsa, src, dst, start_time, end_time, bandwidth, connection_id, global_id):

    service_params = _createServiceParams(start_time, end_time, src, dst, bandwidth)

    log.msg("Connection id: %s  Global id: %s" % (connection_id, global_id))

    try:
        res = yield client.reserve(client_nsa, provider_nsa, None, global_id, 'Test Connection', connection_id, service_params)
        log.msg("Connection held at %s" % provider_nsa)
        yield client.reserveCommit(client_nsa, provider_nsa, connection_id)
        log.msg("Reservation committed at %s" % provider_nsa)

    except error.NSIError, e:
        log.msg('Error reserving %s, %s : %s' % (connection_id, e.__class__.__name__, str(e)))

@defer.inlineCallbacks
def reserveprovision(client, client_nsa, provider_nsa, src, dst, start_time, end_time, bandwidth, connection_id, global_id):

    service_params = _createServiceParams(start_time, end_time, src, dst, bandwidth)

    log.msg("Connection id: %s  Global id: %s" % (connection_id, global_id))

    try:
        yield client.reserve(client_nsa, provider_nsa, None, global_id, 'Test Connection', connection_id, service_params)
        log.msg("Connection reserved at %s" % provider_nsa)
    except error.NSIError, e:
        log.msg('Error reserving %s, %s : %s' % (connection_id, e.__class__.__name__, str(e)))
        defer.returnValue(None)

    try:
        yield client.provision(client_nsa, provider_nsa, None, connection_id)
        log.msg('Connection %s provisioned' % connection_id)
    except error.NSIError, e:
        log.msg('Error provisioning %s, %s : %s' % (connection_id, e.__class__.__name__, str(e)))

@defer.inlineCallbacks
def rprt(client, client_nsa, provider_nsa, src, dst, start_time, end_time, bandwidth, connection_id, global_id):
    # reserve, provision, release,  terminate
    service_params = _createServiceParams(start_time, end_time, src, dst, bandwidth)

    log.msg("Connection id: %s  Global id: %s" % (connection_id, global_id))

    try:
        yield client.reserve(client_nsa, provider_nsa, None, global_id, 'Test Connection', connection_id, service_params)
        log.msg("Connection reserved at %s" % provider_nsa)
    except error.NSIError, e:
        log.msg('Error reserving %s, %s : %s' % (connection_id, e.__class__.__name__, str(e)))
        defer.returnValue(None)

    try:
        yield client.provision(client_nsa, provider_nsa, None, connection_id)
        log.msg('Connection %s provisioned' % connection_id)
    except error.NSIError, e:
        log.msg('Error provisioning %s, %s : %s' % (connection_id, e.__class__.__name__, str(e)))

    try:
        yield client.release(client_nsa, provider_nsa, None, connection_id)
        log.msg('Connection %s released' % connection_id)
    except error.NSIError, e:
        log.msg('Error releasing %s, %s : %s' % (connection_id, e.__class__.__name__, str(e)))

    try:
        yield client.terminate(client_nsa, provider_nsa, None, connection_id)
        log.msg('Connection %s terminated' % connection_id)
    except error.NSIError, e:
        log.msg('Error terminating %s, %s : %s' % (connection_id, e.__class__.__name__, str(e)))

@defer.inlineCallbacks
def provision(client, client_nsa, provider_nsa, connection_id):

    try:
        yield client.provision(client_nsa, provider_nsa, None, connection_id)
        log.msg('Connection %s provisioned' % connection_id)
    except error.NSIError, e:
        log.msg('Error provisioning %s, %s : %s' % (connection_id, e.__class__.__name__, str(e)))

@defer.inlineCallbacks
def release(client, client_nsa, provider_nsa, connection_id):

    try:
        yield client.release(client_nsa, provider_nsa, None, connection_id)
        log.msg('Connection %s released' % connection_id)
    except error.NSIError, e:
        log.msg('Error releasing %s, %s : %s' % (connection_id, e.__class__.__name__, str(e)))

@defer.inlineCallbacks
def terminate(client, client_nsa, provider_nsa, connection_id):

    try:
        yield client.terminate(client_nsa, provider_nsa, None, connection_id)
        log.msg('Connection %s terminated' % connection_id)
    except error.NSIError, e:
        log.msg('Error terminating %s, %s : %s' % (connection_id, e.__class__.__name__, str(e)))

@defer.inlineCallbacks
def querysummary(client, client_nsa, provider_nsa, connection_ids, global_reservation_ids):

    try:
        qc = yield client.query(client_nsa, provider_nsa, None, "Summary", connection_ids, global_reservation_ids)
        log.msg('Query results:')
        for qr in qc:
            log.msg('Connection    : %s' % qr.connectionId)
            if qr.globalReservationId:
                log.msg('  Global ID   : %s' % qr.globalReservationId)
            if qr.description:
                log.msg('  Description : %s' % qr.description)

            log.msg('  Reservation state : %s' % qr.connectionStates.reservationState.state)
            log.msg('  Provision   state : %s' % qr.connectionStates.provisionState.state)
            log.msg('  Activaction state : %s' % qr.connectionStates.activationState.state)
            if qr.children:
                log.msg('  Children    : %s' % qr.children)
    except error.NSIError, e:
        log.msg('Error querying %s, %s : %s' % (connection_ids, e.__class__.__name__, str(e)))

@defer.inlineCallbacks
def querydetails(client, client_nsa, provider_nsa, connection_ids, global_reservation_ids):

    try:
        qc = yield client.query(client_nsa, provider_nsa, None, "Details", connection_ids, global_reservation_ids)
        log.msg('Query results:')
        for qr in qc:
            log.msg('Connection: %s' % qr.connectionId)
            log.msg('  States: %s' % qr.connectionStates)
    except error.NSIError, e:
        log.msg('Error querying %s, %s : %s' % (connection_ids, e.__class__.__name__, str(e)))

def path(topology_file, source_stp, dest_stp):

    raise NotImplementedError('Path computation not available for NML yet')
    topo = None

    source_network, source_port = source_stp.split(':',1)
    dest_network,   dest_port   = dest_stp.split(':', 1)

    r_source_stp    = nsa.STP(source_network, source_port)
    r_dest_stp      = nsa.STP(dest_network,   dest_port)

    paths = topo.findPaths(r_source_stp, r_dest_stp)

    for p in sorted(paths, key=lambda p : len(p.network_links)):
        log.msg(str(p))

def topology(topology_file):

    raise NotImplementedError('Topology dump not available for NML yet')
    topo = None

    for nw in topo.networks:
        ns = '%s (%s)' % (nw.name, ','.join( sorted( [ ep.endpoint for ep in nw.endpoints ] ) ) )
        log.msg(ns)

def topologyGraph(topology_file, all_links=False):

    raise NotImplementedError('Topology graph not available for NML yet')
    topo = None

    links = []

    for nw in topo.networks:
        for ep in nw.endpoints:
            if ep.dest_stp:
                nw1 = nw.name.replace('.ets', '').replace('-','_')
                nw2 = ep.dest_stp.network.replace('.ets', '').replace('-', '_')

                l = [ nw1, nw2 ]
                if all_links:
                    if nw1 < nw2: # this prevents us from building double links
                        links.append(l)
                else:
                    l = sorted(l)
                    if not l in links:
                        links.append(l)

    log.msg('graph Network {')
    for l in sorted(links):
        log.msg('  %s -- %s;' % (l[0], l[1]))
    log.msg('}')

<|endoftext|>