<commit_before># G These are the classes that will use SQLalchemy to
# push information to the database

from sqlalchemy import create_engine
from sqlalchemy import MetaData, Column, Table, ForeignKey
from sqlalchemy.orm import sessionmaker, aliased
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.ext.automap import automap_base
from contextlib import contextmanager
import psycopg2
import logging as log
import datetime as TM

class LterTableQuery(object):
    '''
    This class will perform one unit of work and query the
    database so that it will return information about all
    LTER's that exist (probably overkill but why not)
    '''

    def go(self, session, configclass):
        ltertable = configclass
        return session.query(ltertable, ltertable.lterID).all()

class SiteTableQuery(object):
    '''
    This class will perform one unit of work and query the database
    so that it will return information about all the siteID's that
    have been entered into the site table of the database

    Note: As the databse grows this query may return a huge list
    of items and the algorithm that uses this information to
    screen data may take a very long time...may want to watch this
    '''

    def go(self, session, configclass):
        sitetable = configclass
        return session.query(
            sitetable, sitetable.siteID, sitetable.lterID,
            sitetable.lat, sitetable.long, sitetable.descript).all()

class MainTableQuery(object):
    '''
    This class will perform one unit of work and query the database
    so that it will return information about all the siteID's that
    have been entered into the main table of database

    Note: As the databse grows this query may return a huge list
    of items and the algorithm that uses this information to
    screen data may take a very long time...may want to watch this
    '''

    def go(self, session, configclass):
        maintable = configclass
        return session.query(
            maintable, maintable.projID, maintable.siteID,
            maintable.metalink).all()

class TaxaTableQuery(object):
    '''
    This class will perform one unit of work and query the databse
    so that it will return informatoin about the projID's that
    have been entered and those sites that it corresponds too.
    '''

    def go(self, session, configclass):
        taxatable = configclass
        return session.query(
            taxatable, taxatable.taxaID, taxatable.projID,
            taxatable.sppcode, taxatable.kingdom,
            taxatable.phylum, taxatable.order).all()

class RawTableQuery(object):
    '''
    This class will perform one unit of work and query the databse
    so that it will return informatoin about the projID's that
    have been entered and those sites that it corresponds too.
    '''

    def go(self, session, configclass):
        rawtable = configclass
        return session.query(
            rawtable, rawtable.sampleID, rawtable.taxaID,
            rawtable.projID, rawtable.covariate).all()

class UploadToDatabase(object):
    '''
    This class will be used to upload data that was concatenated
    from user information to make tables, compare the data against
    which is already in the database and then proceed to push it
    /append a table base on input data.

    The class takes multiple arguments:
    1) data: A DataFrame object
    2) config: The config objects derived from the configuration file
    i.e. see config.py
    3) table: A table name that will be appended i.e. 'sitetable'
    4) sitelist: a list of unique site abbreviations from user
    5) lter: an lter location supplied to the combobox from user
    '''

    def __init__(self, data, configclass, table, sitelist=None,
                 lter=None, meta=None, taxaprojIDlist=None,
                 rawprojID=None, rawtaxaID=None):
        self.data = data
        self.config = configclass
        self.table = table
        if sitelist is not None:
            self.sitelist = sitelist
        if lter is not None:
            self.lter = lter
        if meta is not None:
            self.meta = meta
        if taxaprojIDlist is not None:
            self.taxaprojID = taxaprojIDlist
        if rawprojID is not None:
            self.rawprojID = rawprojID
        if rawtaxaID is not None:
            self.rawtaxaID = rawtaxaID
        else:
            pass
        self.w = None

    # Creating a method that will take the arguments supplied to
    # the class, perform database queries in a unit of work
    # pattern to retrieve information, then compare
    # the information about to be loaded to a given table
    # and either return and error or pass the check and
    # get prepared to handle the database push
    def check_previous_sites(self):
        # Create session from factory
        session = self.config.Session()

        # Create a list of site abbreviates already
        # present in the database
        previousSitecheck = []

        # Creat a list of metadata urls already present
        # in the database
        previousmetacheck = []

        # Create a list of all lters already present
        # in the database
        allLTERlocation = []

        # Fill list of LTER location
        for row in\
                LterTableQuery().go(session, self.config.ltertable):
            allLTERlocation.append(row.lterID)

        # The table argument supplies the table in the database
        # that information is attempting to be pushed to.
        if self.table == 'sitetable':

            # Update the lists through two queries and
            # and the information retrieved
            for row in\
                    SiteTableQuery().go(session, self.config.sitetable):
                previousSitecheck.append(row.siteID)

            # check the quieried information against that
            # supplied by the user (is this a unique
            # siteID about to be added)
            checkedSite = [
                i for i in previousSitecheck if i in self.sitelist]

        elif self.table == 'maintable':
            # Update the lists through two queries and
            # and the information retrieved
            for row in\
                    MainTableQuery().go(session, self.config.maintable):
                previousSitecheck.append(row.siteID)
                previousmetacheck.append(row.metalink)

            # check the quieried information against that
            # supplied by the user (is this a unique
            # siteID about to be added)
            checkedSite = [
                i for i in set(previousSitecheck) if i in self.sitelist]
            checkedMeta = [
                i for i in set(previousmetacheck) if i in self.meta]

        else:
            pass

        # Closing session
        session.close()

        try:
            # check that the user is from an established lter.
            checkedLTER = [i for i in allLTERlocation if i in self.lter]

        except Exception as e:
            return print(str(e))

        finally:
            boolist = []

            if self.table == 'sitetable':
                # Retrun boolean list.
                # all(TRUE) means that:

                # 1) There are 0 similar site abbreviations in
                # the current database
                boolist.append(len(checkedSite) == 0)
                # 2) The LTER is a previously identified LTER location
                boolist.append(str(self.lter) in checkedLTER)

                return boolist

            elif self.table == 'maintable':
                # Return boolean list.
                # all(TRUE) means that:

                # 1) The metadata url of the records about to
                # be pushed is unique. So if site abbreviations
                # are already present for this then
                # if the metadataurl is unique, the site was
                # used in a different project
                boolist.append(len(checkedMeta) == 0)
                return boolist

    # This needs to be created
    def check_previous_taxa(self):
        # Create session from factory
        session = self.config.Session()

        if self.table == 'taxatable':
            try:
                # Create a list that will have all the projID's that
                # are present in the taxa table
                taxaprojIDcurrent = []

                # Fill in the taxaprojIDcurrent list with queried
                # information (this performs a unit of work)
                for row in\
                        TaxaTableQuery().go(session, self.config.taxatable):
                    taxaprojIDcurrent.append(row.projID)

                taxaprojcheck = [
                    i for i in set(taxaprojIDcurrent) if i in self.taxaprojID]

                if taxaprojcheck == None:
                    boolist = [True]
                    return boolist
                else:
                    pass

                print(list(set(taxaprojcheck)))

                # Creating a boolean list to returnPressed
                boolist = []

                # Appending the boolean list with a our check
                # on the taxa/projID info
                boolist.append(len(list(taxaprojcheck)) == 0)

                return boolist
            except Exception as e:
                print(str(e))
            finally:
                session.close()
        else:
            pass

    def check_previous_rawobs(self):
        session = self.config.Session()

        # Creating a list that will have all the sampleID's
        # that are present in the rawobs table
        projIDcurrent = []
        taxaIDcurrent = []

        try:
            # Query to fill in list created above
            for row in RawTableQuery().go(session, self.config.rawtable):
                projIDcurrent.append(row.projID)
                taxaIDcurrent.append(row.taxaID)

            projIDcheck = [
                i for i in set(projIDcurrent) if i in self.rawprojID]

            taxaIDcheck =[
                i for i in set(taxaIDcurrent) if i in self.rawtaxaID]
        except Exception as e:
            print(str(e))

        finally:
            boolist = []
            boolist.append(len(list(projIDcheck))==0)
            boolist.append(len(list(taxaIDcheck))==0)
            session.close()
            return boolist

    # Method to table the supplied dataframe from the user along
    # with the rest of the arguments and append the databse with
    # it. NOTE: THIS METHOD MUST BE PROCEEDED BY AN IF
    # STATEMENT TO CHECK THE PREVIOUS ENTRIES OF THE DATABASE
    # Otherwise reduant data can be added.
    def push_table_to_postgres(self):
        # Designating the table name to push to based on
        # arguments supplied to the class
        if self.table == 'sitetable':
            try:
                tabletopushto = str(
                    self.config.sitetable.__table__.name)
            except Error as e:
                print(str(e))

        elif self.table == 'maintable':
            try:
                tabletopushto = str(
                    self.config.maintable.__table__.name)
            except Exception as e:
                print(str(e))

        elif self.table == 'taxatable':
            try:
                tabletopushto = str(
                    self.config.taxatable.__table__.name)
            except Exception as e:
                print(str(e))

        elif self.table == 'rawtable':
            try:
                tabletopushto = str(
                    self.config.rawtable.__table__.name)
            except Exception as e:
                print(str(e))
        else:
            pass

        # Try command to push the
        # table created in the previous steps
        # to the database.
        try:
            self.data.to_sql(
                tabletopushto, con=self.config.engine,
                if_exists="append", index=False)
        except Exception as e:
            print(str(e))
<commit_msg>This class was modified with a logger heading.<commit_after># These are the classes that will use SQLalchemy to
# push information to the database

from sqlalchemy import create_engine
from sqlalchemy import MetaData, Column, Table, ForeignKey
from sqlalchemy.orm import sessionmaker, aliased
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.ext.automap import automap_base
from contextlib import contextmanager
import psycopg2
import logging
import datetime as dt

# Setup to log database transactions
date = (str(dt.datetime.now()).split()[0]).replace("-", "_")
logging.basicConfig(
    filename='Logs_DbTransactions/database_log_{}.log'.format(date))
logging.getLogger('sqlalchemy.engine').setLevel(logging.INFO)

class LterTableQuery(object):
    '''
    This class will perform one unit of work and query the
    database so that it will return information about all
    LTER's that exist (probably overkill but why not)
    '''

    def go(self, session, configclass):
        ltertable = configclass
        return session.query(ltertable, ltertable.lterID).all()

class SiteTableQuery(object):
    '''
    This class will perform one unit of work and query the database
    so that it will return information about all the siteID's that
    have been entered into the site table of the database

    Note: As the databse grows this query may return a huge list
    of items and the algorithm that uses this information to
    screen data may take a very long time...may want to watch this
    '''

    def go(self, session, configclass):
        sitetable = configclass
        return session.query(
            sitetable, sitetable.siteID, sitetable.lterID,
            sitetable.lat, sitetable.long, sitetable.descript).all()

class MainTableQuery(object):
    '''
    This class will perform one unit of work and query the database
    so that it will return information about all the siteID's that
    have been entered into the main table of database

    Note: As the databse grows this query may return a huge list
    of items and the algorithm that uses this information to
    screen data may take a very long time...may want to watch this
    '''

    def go(self, session, configclass):
        maintable = configclass
        return session.query(
            maintable, maintable.projID, maintable.siteID,
            maintable.metalink).all()

class TaxaTableQuery(object):
    '''
    This class will perform one unit of work and query the databse
    so that it will return informatoin about the projID's that
    have been entered and those sites that it corresponds too.
    '''

    def go(self, session, configclass):
        taxatable = configclass
        return session.query(
            taxatable, taxatable.taxaID, taxatable.projID,
            taxatable.sppcode, taxatable.kingdom,
            taxatable.phylum, taxatable.order).all()

class RawTableQuery(object):
    '''
    This class will perform one unit of work and query the databse
    so that it will return informatoin about the projID's that
    have been entered and those sites that it corresponds too.
    '''

    def go(self, session, configclass):
        rawtable = configclass
        return session.query(
            rawtable, rawtable.sampleID, rawtable.taxaID,
            rawtable.projID, rawtable.covariate).all()

class UploadToDatabase(object):
    '''
    This class will be used to upload data that was concatenated
    from user information to make tables, compare the data against
    which is already in the database and then proceed to push it
    /append a table base on input data.

    The class takes multiple arguments:
    1) data: A DataFrame object
    2) config: The config objects derived from the configuration file
    i.e. see config.py
    3) table: A table name that will be appended i.e. 'sitetable'
    4) sitelist: a list of unique site abbreviations from user
    5) lter: an lter location supplied to the combobox from user
    '''

    def __init__(self, data, configclass, table, sitelist=None,
                 lter=None, meta=None, taxaprojIDlist=None,
                 rawprojID=None, rawtaxaID=None):
        self.data = data
        self.config = configclass
        self.table = table
        if sitelist is not None:
            self.sitelist = sitelist
        if lter is not None:
            self.lter = lter
        if meta is not None:
            self.meta = meta
        if taxaprojIDlist is not None:
            self.taxaprojID = taxaprojIDlist
        if rawprojID is not None:
            self.rawprojID = rawprojID
        if rawtaxaID is not None:
            self.rawtaxaID = rawtaxaID
        else:
            pass
        self.w = None

    # Creating a method that will take the arguments supplied to
    # the class, perform database queries in a unit of work
    # pattern to retrieve information, then compare
    # the information about to be loaded to a given table
    # and either return and error or pass the check and
    # get prepared to handle the database push
    def check_previous_sites(self):
        # Create session from factory
        session = self.config.Session()

        # Create a list of site abbreviates already
        # present in the database
        previousSitecheck = []

        # Creat a list of metadata urls already present
        # in the database
        previousmetacheck = []

        # Create a list of all lters already present
        # in the database
        allLTERlocation = []

        # Fill list of LTER location
        for row in\
                LterTableQuery().go(session, self.config.ltertable):
            allLTERlocation.append(row.lterID)

        # The table argument supplies the table in the database
        # that information is attempting to be pushed to.
        if self.table == 'sitetable':

            # Update the lists through two queries and
            # and the information retrieved
            for row in\
                    SiteTableQuery().go(session, self.config.sitetable):
                previousSitecheck.append(row.siteID)

            # check the quieried information against that
            # supplied by the user (is this a unique
            # siteID about to be added)
            checkedSite = [
                i for i in previousSitecheck if i in self.sitelist]

        elif self.table == 'maintable':
            # Update the lists through two queries and
            # and the information retrieved
            for row in\
                    MainTableQuery().go(session, self.config.maintable):
                previousSitecheck.append(row.siteID)
                previousmetacheck.append(row.metalink)

            # check the quieried information against that
            # supplied by the user (is this a unique
            # siteID about to be added)
            checkedSite = [
                i for i in set(previousSitecheck) if i in self.sitelist]
            checkedMeta = [
                i for i in set(previousmetacheck) if i in self.meta]

        else:
            pass

        # Closing session
        session.close()

        try:
            # check that the user is from an established lter.
            checkedLTER = [i for i in allLTERlocation if i in self.lter]

        except Exception as e:
            return print(str(e))

        finally:
            boolist = []

            if self.table == 'sitetable':
                # Retrun boolean list.
                # all(TRUE) means that:

                # 1) There are 0 similar site abbreviations in
                # the current database
                boolist.append(len(checkedSite) == 0)
                # 2) The LTER is a previously identified LTER location
                boolist.append(str(self.lter) in checkedLTER)

                return boolist

            elif self.table == 'maintable':
                # Return boolean list.
                # all(TRUE) means that:

                # 1) The metadata url of the records about to
                # be pushed is unique. So if site abbreviations
                # are already present for this then
                # if the metadataurl is unique, the site was
                # used in a different project
                boolist.append(len(checkedMeta) == 0)
                return boolist

    # This needs to be created
    def check_previous_taxa(self):
        # Create session from factory
        session = self.config.Session()

        if self.table == 'taxatable':
            try:
                # Create a list that will have all the projID's that
                # are present in the taxa table
                taxaprojIDcurrent = []

                # Fill in the taxaprojIDcurrent list with queried
                # information (this performs a unit of work)
                for row in\
                        TaxaTableQuery().go(session, self.config.taxatable):
                    taxaprojIDcurrent.append(row.projID)

                taxaprojcheck = [
                    i for i in set(taxaprojIDcurrent) if i in self.taxaprojID]

                if taxaprojcheck == None:
                    boolist = [True]
                    return boolist
                else:
                    pass

                print(list(set(taxaprojcheck)))

                # Creating a boolean list to returnPressed
                boolist = []

                # Appending the boolean list with a our check
                # on the taxa/projID info
                boolist.append(len(list(taxaprojcheck)) == 0)

                return boolist
            except Exception as e:
                print(str(e))
            finally:
                session.close()
        else:
            pass

    def check_previous_rawobs(self):
        session = self.config.Session()

        # Creating a list that will have all the sampleID's
        # that are present in the rawobs table
        projIDcurrent = []
        taxaIDcurrent = []

        try:
            # Query to fill in list created above
            for row in RawTableQuery().go(session, self.config.rawtable):
                projIDcurrent.append(row.projID)
                taxaIDcurrent.append(row.taxaID)

            projIDcheck = [
                i for i in set(projIDcurrent) if i in self.rawprojID]

            taxaIDcheck =[
                i for i in set(taxaIDcurrent) if i in self.rawtaxaID]
        except Exception as e:
            print(str(e))

        finally:
            boolist = []
            boolist.append(len(list(projIDcheck))==0)
            boolist.append(len(list(taxaIDcheck))==0)
            session.close()
            return boolist

    # Method to table the supplied dataframe from the user along
    # with the rest of the arguments and append the databse with
    # it. NOTE: THIS METHOD MUST BE PROCEEDED BY AN IF
    # STATEMENT TO CHECK THE PREVIOUS ENTRIES OF THE DATABASE
    # Otherwise reduant data can be added.
    def push_table_to_postgres(self):
        # Designating the table name to push to based on
        # arguments supplied to the class
        if self.table == 'sitetable':
            try:
                tabletopushto = str(
                    self.config.sitetable.__table__.name)
            except Error as e:
                print(str(e))

        elif self.table == 'maintable':
            try:
                tabletopushto = str(
                    self.config.maintable.__table__.name)
            except Exception as e:
                print(str(e))

        elif self.table == 'taxatable':
            try:
                tabletopushto = str(
                    self.config.taxatable.__table__.name)
            except Exception as e:
                print(str(e))

        elif self.table == 'rawtable':
            try:
                tabletopushto = str(
                    self.config.rawtable.__table__.name)
            except Exception as e:
                print(str(e))
        else:
            pass

        # Try command to push the
        # table created in the previous steps
        # to the database.
        try:
            self.data.to_sql(
                tabletopushto, con=self.config.engine,
                if_exists="append", index=False)
        except Exception as e:
            print(str(e))
<|endoftext|>


<commit_before>"""
Rip and replace
"""

import gdata.spreadsheets.client
import gdata.spreadsheets.data
import gdata.docs.data
import gdata.docs.client
import gdata.gauth
import re
import ConfigParser
import sys
import util

config = ConfigParser.ConfigParser()
config.read('mytokens.cfg')

# Get me a client, stat
spr_client = util.get_spreadsheet_client(config)
docs_client = util.get_docs_client(config)

query = gdata.docs.client.DocsQuery(show_collections='true', title='Agronomic Data')
feed = docs_client.GetAllResources(query=query)

for entry in feed:
    feed2 = spr_client.GetWorksheets(entry.id.text.split("/")[-1][14:])
    for entry2 in feed2.entry:
        worksheet = entry2.id.text.split("/")[-1]
        print 'Processing %s WRK: %s Title: %s' % (entry.title.text,
                                                   worksheet, entry2.title.text),
        feed3 = spr_client.get_list_feed(entry.id.text.split("/")[-1][14:], worksheet)
        row = feed3.entry[0]
        data = row.to_dict()
        if data.get('agr7', None) is not None:
            row.set_value('agr7', '[7] Cover crop (rye) biomass at termination (spring) (no significant weeds)')
            spr_client.update(row)
            print ' ... updated'
        else:
            print ' ... not found!'

<commit_msg>cleanups and make more general<commit_after>"""
 Go thru the agronomic data sheets and replace the description label found on
 row 2 with a new value
"""
import gdata.docs.client
import ConfigParser

import util

VARID = 'agr2' # lowercase
NEWVAL = '[2] Soybean final plant population'

config = ConfigParser.ConfigParser()
config.read('mytokens.cfg')

# Get me a client, stat
spr_client = util.get_spreadsheet_client(config)
docs_client = util.get_docs_client(config)

query = gdata.docs.client.DocsQuery(show_collections='true',
                                    title='Agronomic Data')
feed = docs_client.GetAllResources(query=query)

for entry in feed:
    if entry.get_resource_type() != 'spreadsheet':
        continue
    feed2 = spr_client.GetWorksheets(entry.id.text.split("/")[-1][14:])
    for entry2 in feed2.entry:
        worksheet = entry2.id.text.split("/")[-1]
        print 'Processing %s WRK: %s Title: %s' % (entry.title.text,
                                                   worksheet, entry2.title.text),
        feed3 = spr_client.get_list_feed(entry.id.text.split("/")[-1][14:],
                                         worksheet)
        if len(feed3.entry) == 0:
            print 'Skipping as there is no data?!?'
            continue
        row = feed3.entry[0]
        data = row.to_dict()
        if data.get(VARID, None) is not None:
            row.set_value(VARID, NEWVAL)
            spr_client.update(row)
            print ' ... updated'
        else:
            print ' ... not found!'

<|endoftext|>


<commit_before><commit_msg>Don't discard -add_ast_path linkopt on macOS.<commit_after><|endoftext|>


<commit_before>import pytest
from hungry import eat

def test_catch_one_exception():
    def f(val):
        return int(val[1])

    # should raise IndexError
    with pytest.raises(IndexError):
        f([1])

    # should NOT raise any exceptions
    f([1, 2])

    # decorate 'foo'
    foo = eat(IndexError)(f)

    # this time, the IndexError should be eaten
    foo([1])

    # ...but not ValueError
    foo(['foo'])
<commit_msg>Fix test_catch_one_exception<commit_after>import pytest
from hungry import eat

def test_catch_one_exception():
    def f(val):
        return int(val[1])

    # should raise IndexError
    with pytest.raises(IndexError):
        f([1])

    # should NOT raise any exceptions
    f([1, 2])

    # decorate 'foo'
    foo = eat(IndexError)(f)

    # this time, the IndexError should be eaten
    foo([1])

    # ...but not ValueError
    with pytest.raises(ValueError):
        f([1, 'foo'])
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
JAR manipulation module tests.
"""
import os.path
import pickle
import unittest as test

from cosmos.jar import InvalidJarFile, JarFile
from cosmos.jar_parameters import make_parameter

class CannotOpenJarTestCase(test.TestCase):

    def test_not_fount(self):
        self.assertRaises(IOError, JarFile, 'does/not/exist.jar')

    def test_not_a_jar(self):
        self.assertRaises(InvalidJarFile, JarFile, __file__)

class UseJarTestCase(test.TestCase):

    def setUp(self):
        test_path = os.path.realpath(os.path.dirname(__file__))
        self.prop_jar = JarFile(os.path.join(test_path, 'jar-properties.jar'))
        self.xml_jar = JarFile(os.path.join(test_path, 'jar-xml.jar'))
        self.malformed_jar = JarFile(os.path.join(test_path,
                                                  'jar-malformed-xml.jar'))
        self.nop_jar = JarFile(os.path.join(test_path, 'jar-nop.jar'))

    def tearDown(self):
        self.prop_jar.close()
        self.xml_jar.close()
        self.malformed_jar.close()
        self.nop_jar.close()

    def test_read_manifest(self):
        manifest = self.prop_jar.manifest
        self.assertEquals(manifest['Cosmos-Parameters-Template'],
                                   '/conf/params.properties')
        self.assertTrue('not-a-key' not in manifest.keys())

    def test_read_long_manifest_lines(self):
        self.assertEquals(self.xml_jar.manifest['Cosmos-Parameters-Template'],
                '/conf/long-long-long-long-long-long-long-name-for-having-' +
                'more-than-one-line.xml')

    def test_is_parameterized(self):
        self.assertTrue(self.prop_jar.is_parameterized())
        self.assertTrue(self.xml_jar.is_parameterized())
        self.assertFalse(self.nop_jar.is_parameterized())

    def test_error_on_invalid_parametrization(self):
        self.assertRaises(InvalidJarFile, self.malformed_jar.parameters)

    def test_get_parametrization_from_properties(self):
        params = self.prop_jar.parameters()
        self.assertEquals(len(params), 5)
        self.assertEquals(params[0].name, "foo")
        self.assertEquals(params[0].default_value, None)
        self.assertEquals(params[1].name, "bar")
        self.assertEquals(params[1].default_value, "hola")
        self.assertEquals(params[2].name, "tmp")
        self.assertEquals(params[3].name, "mongo1")
        self.assertEquals(params[4].name, "mongo2")
        self.assertEquals(params[4].default_value, "col_a")

    def test_get_parametrization_from_xml(self):
        params = self.xml_jar.parameters()
        self.assertEquals(len(params), 3)
        self.assertEquals(params[0].name, "foo")
        self.assertEquals(params[1].default_value, "hola")

class JarParametersTestCase(test.TestCase):

    def setUp(self):
        self.foo = make_parameter('foo', 'string')
        self.bar = make_parameter('bar', 'string|value')
        self.tmp = make_parameter('tmp', 'filepath|/tmp')
        self.coll = make_parameter('coll', 'mongocoll')

    def test_pickable(self):
        for param in [self.foo, self.bar, self.tmp, self.coll]:
            serialized_param = pickle.dumps(param)
            deserialized_param = pickle.loads(serialized_param)
            self.assertEquals(param.name, deserialized_param.name)
<commit_msg>Typo fix<commit_after># -*- coding: utf-8 -*-
"""
JAR manipulation module tests.
"""
import os.path
import pickle
import unittest as test

from cosmos.jar import InvalidJarFile, JarFile
from cosmos.jar_parameters import make_parameter

class CannotOpenJarTestCase(test.TestCase):

    def test_not_found(self):
        self.assertRaises(IOError, JarFile, 'does/not/exist.jar')

    def test_not_a_jar(self):
        self.assertRaises(InvalidJarFile, JarFile, __file__)

class UseJarTestCase(test.TestCase):

    def setUp(self):
        test_path = os.path.realpath(os.path.dirname(__file__))
        self.prop_jar = JarFile(os.path.join(test_path, 'jar-properties.jar'))
        self.xml_jar = JarFile(os.path.join(test_path, 'jar-xml.jar'))
        self.malformed_jar = JarFile(os.path.join(test_path,
                                                  'jar-malformed-xml.jar'))
        self.nop_jar = JarFile(os.path.join(test_path, 'jar-nop.jar'))

    def tearDown(self):
        self.prop_jar.close()
        self.xml_jar.close()
        self.malformed_jar.close()
        self.nop_jar.close()

    def test_read_manifest(self):
        manifest = self.prop_jar.manifest
        self.assertEquals(manifest['Cosmos-Parameters-Template'],
                                   '/conf/params.properties')
        self.assertTrue('not-a-key' not in manifest.keys())

    def test_read_long_manifest_lines(self):
        self.assertEquals(self.xml_jar.manifest['Cosmos-Parameters-Template'],
                '/conf/long-long-long-long-long-long-long-name-for-having-' +
                'more-than-one-line.xml')

    def test_is_parameterized(self):
        self.assertTrue(self.prop_jar.is_parameterized())
        self.assertTrue(self.xml_jar.is_parameterized())
        self.assertFalse(self.nop_jar.is_parameterized())

    def test_error_on_invalid_parametrization(self):
        self.assertRaises(InvalidJarFile, self.malformed_jar.parameters)

    def test_get_parametrization_from_properties(self):
        params = self.prop_jar.parameters()
        self.assertEquals(len(params), 5)
        self.assertEquals(params[0].name, "foo")
        self.assertEquals(params[0].default_value, None)
        self.assertEquals(params[1].name, "bar")
        self.assertEquals(params[1].default_value, "hola")
        self.assertEquals(params[2].name, "tmp")
        self.assertEquals(params[3].name, "mongo1")
        self.assertEquals(params[4].name, "mongo2")
        self.assertEquals(params[4].default_value, "col_a")

    def test_get_parametrization_from_xml(self):
        params = self.xml_jar.parameters()
        self.assertEquals(len(params), 3)
        self.assertEquals(params[0].name, "foo")
        self.assertEquals(params[1].default_value, "hola")

class JarParametersTestCase(test.TestCase):

    def setUp(self):
        self.foo = make_parameter('foo', 'string')
        self.bar = make_parameter('bar', 'string|value')
        self.tmp = make_parameter('tmp', 'filepath|/tmp')
        self.coll = make_parameter('coll', 'mongocoll')

    def test_pickable(self):
        for param in [self.foo, self.bar, self.tmp, self.coll]:
            serialized_param = pickle.dumps(param)
            deserialized_param = pickle.loads(serialized_param)
            self.assertEquals(param.name, deserialized_param.name)
<|endoftext|>


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-
"""Misura Language or Mini Language.
Secure minimal Python language subset for conditional evaluation of numerical datasets."""
import numpy as np
from scipy.interpolate import UnivariateSpline

from .env import BaseEnvironment
from .. import csutil

class DataEnvironment(BaseEnvironment):
    # relativi agli oggetti array numpy
    whitelist = BaseEnvironment.whitelist + ['std', 'get', 'set']
    _hdf = False
    """SharedFile instance"""
    prefix = '/'
    """Prefix to be applied to all curves. E.g.: /hsm/sample0/"""
    lslice = slice(None, None)
    """Last applied slicing interval"""
    spline_cache = {}
    """Approximating splines cache"""
    _temperature_path = ''

    @property
    def hdf(self):
        """Trying to anyway access undefined hdf property will raise a exception"""
        if self._hdf is False:
            self.Exit('No data file defined in this environment!')
        return self._hdf

    @hdf.setter
    def hdf(self, f):
        self._hdf = f

    def _cname(self, curve):
        """Retrieve curve name by adding prefix, etc."""
        if not isinstance(curve, str):
            return False
        # curve is Temperature
        if curve in ['T']:
            # but unknown if a local T dataset is available in prefix
            if not self._temperature_path:
                curve = self.prefix + 'T'
                print('Setting temperature path', curve)
                if not self.hdf.has_node(curve):
                    curve = '/kiln/T'
                self._temperature_path = curve
            # Take defined temperature path
            print('Found temperature path', curve, self._temperature_path)
            curve = self._temperature_path
        else:
            curve = self.prefix + curve
        print('_cname returning', curve)
        return curve

    def _c(self, curve0, start_time=0):
        """Returns the array object located at prefix+`curve` on the hdf file.
        The curve is already sliced.
        If `curve` is not a string, returns the unchanged object."""
        curve = self._cname(curve0)
        if curve is False:
            return curve0

        limit_slice = slice(self.hdf.get_time(curve, start_time), None)
        c = self.hdf.col(curve, limit_slice)
        return c
    Curve = _c

    def xy(self, curve, start_time=0):
        """Separately return x and y arrays as tuples"""
        n = self._c(curve, start_time)

        print('xy', type(n))
        return n[:, 0], n[:, 1]

    def Point(self, t=None, curve=False, idx=None):
        """Set output time and temperature at time  `t` of `curve`"""
        if t == None:
            pt = self.AtIndex('T', idx)
            if pt == None:
                return
            t = pt[0]
        else:
            pt = self.AtTime('T', t)
        self.t(pt[0])
        self.T(pt[1])
        if curve is False:
            return
        val = self.At(curve, t)
        self.Value(val)

    def Spline(self, curve0):
        """Returns an UnivariateSpline object from a t,val dataset.
        `curve` is a dataset path, prefix excluded, or a dataset alias."""
        spline = False
        curve = self._cname(curve0)
        if curve:
            spline = self.spline_cache.get(curve, False)
        if spline is not False:
            return spline
        x, y = self.xy(curve)
        if len(x) == 0:
            self.Exit("Zero-length curve requested " + curve)
        spline = UnivariateSpline(x, y)
        return spline

    def _x(self, curve):
        """Retrieve the `x` array constituting the UnivariateSpline of curve"""
        return self._c(curve)[:, 0]
    _t = _x

    def _y(self, curve):
        return self._c(curve)[:, 1]

    def _reset(self):
        BaseEnvironment._reset(self)
        self.spline_cache = {}

    def GetCoolingTimeAndIndex(self):
        ret = self.hdf.max('/kiln/T')
        if not ret:
            print('Cannot find node /kiln/T')
            return False
        idx, t, T = ret
        if idx + 1 >= self.hdf.len('/kiln/T'):
            return False
        return t, idx

    def AtTime(self, curve0, t):
        """t,val of curve at nearest time `t`."""
        print('AtTime0', curve0, t)
        curve = self._cname(curve0)
        if curve is False:
            idx = csutil.find_nearest_val(curve0[:, 0], t)
            return curve0[idx]
        idx = self.hdf.get_time(curve, t)
        r = self.hdf.col_at(curve, idx, raw=True)
        print('AtTime', curve0, curve, t, idx, r)
        return r

    def AtIndex(self, curve0, idx):
        """t,val of curve at index `idx`."""
        curve = self._cname(curve0)
        if curve is False:
            return curve0[idx]
        return self.hdf.col(curve, idx)

    def At(self, curve, t):
        """Value of curve at nearest time `t`."""
        p = self.AtTime(curve, t)
        return p[1]

    def minmax(self, curve0, op=1):
        """Search global minimum/maximum value of curve."""
        curve = self._cname(curve0)
        if curve is False:  # curve is array
            ops = (min, max)
            return ops[op](curve0)
        else:  # curve is an hdf path
            ops = (self.hdf.min, self.hdf.max)
            idx, t, v = ops[op](curve)
            return idx, t, v

    def Max(self, curve):
        """Global maximum value of curve."""
        return self.minmax(curve, 1)

    def Min(self, curve):
        """Global minimum value of curve"""
        return self.minmax(curve, 0)

    def Equals(self, curve0, val):
        """Returns time of curve where its value is val"""
        curve = self._cname(curve0)
        if curve is not False:
            return self.hdf.equals(curve, val)
        # TODO: array was passed

    def Nearest(self, curve0, val):
        curve = self._cname(curve0)
        if curve is not False:
            return self.hdf.nearest(curve, val)
        # TODO: array was passed

    def Drops(self, curve0, val, start_time=0):
        """Returns time where curve value drops below val"""
        curve = self._cname(curve0)
        if curve is not False:
            return self.hdf.drops(curve, val, start_time)
        # TODO: array was passed

    def Raises(self, curve0, val, start_time=0):
        """Returns time where curve value raises above val"""
        curve = self._cname(curve0)
        print('Raises', curve, curve0, val, start_time, self.prefix, self.obj)
        if curve is not False:
            return self.hdf.rises(curve, val, start_time)
        # TODO: array was passed

    def Len(self, curve0):
        """Length of the curve"""
        curve = self._cname(curve0)
        if curve is not False:
            return self.hdf.len(curve)
        return len(curve)

    def TimeDerivative(self, curve, start_time=0):
        """Numerical derivative of curve with respect to its time"""
        x, y = self.xy(curve, start_time=0)
        if len(x) < 2:
            return np.array([])
        r = np.gradient(y) / np.gradient(x)
        print('TimeDerivative', curve, len(r))
        return r

    def Coefficient(self, curve, T0, T1):
        # FIXME
        spline = self.Spline(curve)
        if self.Max('T') < T1:
            self.Exit()
        t0 = self.Nearest('T', T0)
        t1 = self.Nearest('T', T1)
        num = spline(t1) - spline(t0)
        denom = t1 - t0
        return num / denom

    def Ratio(self, numerator, denominator):
        n = self._y(numerator)
        d = self._y(denominator)
        return n / d
<commit_msg>MiLang: Use a summary dataset if available, ESS-1227<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-
"""Misura Language or Mini Language.
Secure minimal Python language subset for conditional evaluation of numerical datasets."""
import numpy as np
from scipy.interpolate import UnivariateSpline

from .env import BaseEnvironment
from .. import csutil

class DataEnvironment(BaseEnvironment):
    # relativi agli oggetti array numpy
    whitelist = BaseEnvironment.whitelist + ['std', 'get', 'set']
    _hdf = False
    """SharedFile instance"""
    prefix = '/'
    """Prefix to be applied to all curves. E.g.: /hsm/sample0/"""
    lslice = slice(None, None)
    """Last applied slicing interval"""
    spline_cache = {}
    """Approximating splines cache"""
    _temperature_path = ''

    @property
    def hdf(self):
        """Trying to anyway access undefined hdf property will raise a exception"""
        if self._hdf is False:
            self.Exit('No data file defined in this environment!')
        return self._hdf

    @hdf.setter
    def hdf(self, f):
        self._hdf = f

    def _cname(self, curve):
        """Retrieve curve name by adding prefix, etc."""
        if not isinstance(curve, str):
            return False
        # curve is Temperature
        if curve in ['T']:
            # but unknown if a local T dataset is available in prefix
            if not self._temperature_path:
                curve = self.prefix + 'T'
                print('Setting temperature path', curve)
                if not self.hdf.has_node(curve):
                    curve = '/kiln/T'
                self._temperature_path = curve
            # Take defined temperature path
            print('Found temperature path', curve, self._temperature_path)
            curve = self._temperature_path
        else:
            curve = self.prefix + curve
        summary = '/summary'+curve
        if self.hdf.has_node(summary):
            curve=summary
        print('_cname returning', curve)
        return curve

    def _c(self, curve0, start_time=0):
        """Returns the array object located at prefix+`curve` on the hdf file.
        The curve is already sliced.
        If `curve` is not a string, returns the unchanged object."""
        curve = self._cname(curve0)
        if curve is False:
            return curve0

        limit_slice = slice(self.hdf.get_time(curve, start_time), None)
        c = self.hdf.col(curve, limit_slice)
        return c
    Curve = _c

    def xy(self, curve, start_time=0):
        """Separately return x and y arrays as tuples"""
        n = self._c(curve, start_time)

        print('xy', type(n))
        return n[:, 0], n[:, 1]

    def Point(self, t=None, curve=False, idx=None):
        """Set output time and temperature at time  `t` of `curve`"""
        if t == None:
            pt = self.AtIndex('T', idx)
            if pt == None:
                return
            t = pt[0]
        else:
            pt = self.AtTime('T', t)
        self.t(pt[0])
        self.T(pt[1])
        if curve is False:
            return
        val = self.At(curve, t)
        self.Value(val)

    def Spline(self, curve0):
        """Returns an UnivariateSpline object from a t,val dataset.
        `curve` is a dataset path, prefix excluded, or a dataset alias."""
        spline = False
        curve = self._cname(curve0)
        if curve:
            spline = self.spline_cache.get(curve, False)
        if spline is not False:
            return spline
        x, y = self.xy(curve)
        if len(x) == 0:
            self.Exit("Zero-length curve requested " + curve)
        spline = UnivariateSpline(x, y)
        return spline

    def _x(self, curve):
        """Retrieve the `x` array constituting the UnivariateSpline of curve"""
        return self._c(curve)[:, 0]
    _t = _x

    def _y(self, curve):
        return self._c(curve)[:, 1]

    def _reset(self):
        BaseEnvironment._reset(self)
        self.spline_cache = {}

    def GetCoolingTimeAndIndex(self):
        ret = self.hdf.max('/kiln/T')
        if not ret:
            print('Cannot find node /kiln/T')
            return False
        idx, t, T = ret
        if idx + 1 >= self.hdf.len('/kiln/T'):
            return False
        return t, idx

    def AtTime(self, curve0, t):
        """t,val of curve at nearest time `t`."""
        print('AtTime0', curve0, t)
        curve = self._cname(curve0)
        if curve is False:
            idx = csutil.find_nearest_val(curve0[:, 0], t)
            return curve0[idx]
        idx = self.hdf.get_time(curve, t)
        r = self.hdf.col_at(curve, idx, raw=True)
        print('AtTime', curve0, curve, t, idx, r)
        return r

    def AtIndex(self, curve0, idx):
        """t,val of curve at index `idx`."""
        curve = self._cname(curve0)
        if curve is False:
            return curve0[idx]
        return self.hdf.col(curve, idx)

    def At(self, curve, t):
        """Value of curve at nearest time `t`."""
        p = self.AtTime(curve, t)
        return p[1]

    def minmax(self, curve0, op=1):
        """Search global minimum/maximum value of curve."""
        curve = self._cname(curve0)
        if curve is False:  # curve is array
            ops = (min, max)
            return ops[op](curve0)
        else:  # curve is an hdf path
            ops = (self.hdf.min, self.hdf.max)
            idx, t, v = ops[op](curve)
            return idx, t, v

    def Max(self, curve):
        """Global maximum value of curve."""
        return self.minmax(curve, 1)

    def Min(self, curve):
        """Global minimum value of curve"""
        return self.minmax(curve, 0)

    def Equals(self, curve0, val):
        """Returns time of curve where its value is val"""
        curve = self._cname(curve0)
        if curve is not False:
            return self.hdf.equals(curve, val)
        # TODO: array was passed

    def Nearest(self, curve0, val):
        curve = self._cname(curve0)
        if curve is not False:
            return self.hdf.nearest(curve, val)
        # TODO: array was passed

    def Drops(self, curve0, val, start_time=0):
        """Returns time where curve value drops below val"""
        curve = self._cname(curve0)
        if curve is not False:
            return self.hdf.drops(curve, val, start_time)
        # TODO: array was passed

    def Raises(self, curve0, val, start_time=0):
        """Returns time where curve value raises above val"""
        curve = self._cname(curve0)
        print('Raises', curve, curve0, val, start_time, self.prefix, self.obj)
        if curve is not False:
            return self.hdf.rises(curve, val, start_time)
        # TODO: array was passed

    def Len(self, curve0):
        """Length of the curve"""
        curve = self._cname(curve0)
        if curve is not False:
            return self.hdf.len(curve)
        return len(curve)

    def TimeDerivative(self, curve, start_time=0):
        """Numerical derivative of curve with respect to its time"""
        x, y = self.xy(curve, start_time=0)
        if len(x) < 2:
            return np.array([])
        r = np.gradient(y) / np.gradient(x)
        print('TimeDerivative', curve, len(r))
        return r

    def Coefficient(self, curve, T0, T1):
        # FIXME
        spline = self.Spline(curve)
        if self.Max('T') < T1:
            self.Exit()
        t0 = self.Nearest('T', T0)
        t1 = self.Nearest('T', T1)
        num = spline(t1) - spline(t0)
        denom = t1 - t0
        return num / denom

    def Ratio(self, numerator, denominator):
        n = self._y(numerator)
        d = self._y(denominator)
        return n / d
<|endoftext|>


<commit_before>import logging
import functools

from datetime import date, timedelta

from django.conf import settings
from django.db import models
from django.template.defaultfilters import date as date_filter
from django.utils.text import slugify
from django.contrib.contenttypes.fields import GenericForeignKey, GenericRelation
from django.contrib.contenttypes.models import ContentType
from django.core.urlresolvers import reverse

from talks.api_ox.api import ApiException, OxfordDateResource, PlacesResource
from talks.api_ox.models import Location, Organisation

logger = logging.getLogger(__name__)

class EventGroupManager(models.Manager):
    def for_events(self, events):
        """Given a QuerySet of `Event`s return all the `EventGroup` related to
        that QuerySet. This function will evaluate the QuerySet `events`.

        Returns a list() of `EventGroup` NOT a QuerySet.
        """
        events = events.select_related('group')
        # NOTE: execs the query
        return set([e.group for e in events])

class EventGroup(models.Model):
    SEMINAR = 'SE'
    CONFERENCE = 'CO'
    EVENT_GROUP_TYPE_CHOICES = (
        (SEMINAR, 'Seminar Series'),
        (CONFERENCE, 'Conference'),
    )

    title = models.CharField(max_length=250)
    slug = models.SlugField()
    description = models.TextField()
    group_type = models.CharField(
        blank=True,
        null=True,
        max_length=2,
        choices=EVENT_GROUP_TYPE_CHOICES
    )

    objects = EventGroupManager()

    def __unicode__(self):
        return self.title

    def get_absolute_url(self):
        return reverse('event-group', args=[self.id])

class SpeakerManager(models.Manager):

    def suggestions(self, query):
        return self.filter(name__icontains=query)

class Speaker(models.Model):
    name = models.CharField(max_length=250)
    slug = models.SlugField()
    bio = models.TextField()
    email_address = models.EmailField(max_length=254)

    objects = SpeakerManager()

    def __unicode__(self):
        return self.name

class Tag(models.Model):

    slug = models.SlugField()
    name = models.CharField(max_length=250, unique=True)
    description = models.TextField()

    def __unicode__(self):
        return self.name

class TagItem(models.Model):

    tag = models.ForeignKey(Tag)
    content_type = models.ForeignKey(ContentType)
    object_id = models.PositiveIntegerField()
    item = GenericForeignKey('content_type', 'object_id')   # atm: Event, EventGroup

class EventManager(models.Manager):

    def todays_events(self):
        today = date.today()
        tomorrow = today + timedelta(days=1)
        return self.filter(start__gte=today, start__lt=tomorrow
                           ).order_by('start')

class Event(models.Model):
    start = models.DateTimeField(null=True, blank=True)
    end = models.DateTimeField(null=True, blank=True)
    title = models.CharField(max_length=250)
    slug = models.SlugField()
    description = models.TextField()
    speakers = models.ManyToManyField(Speaker, null=True, blank=True)
    # TODO audience should it be free text
    # TODO booking information; structure?

    group = models.ForeignKey(EventGroup, null=True, blank=True,
                              related_name='events')
    location = models.ForeignKey(Location, null=True, blank=True)
    # TODO I'm guessing an event can be organised by multiple departments?
    department_organiser = models.ForeignKey(Organisation, null=True, blank=True)

    tags = GenericRelation(TagItem)

    objects = EventManager()

    _cached_resources = {}

    def fetch_resource(self, key, func):
        """Fetch a resource from the API.

        If we have `key` in a cache then pull from there, otherwise call `func`
        """
        if key in self._cached_resources:
            return self._cached_resources[key]
        else:
            try:
                res = func()
                self._cached_resources[key] = res
                return res
            except ApiException:
                logger.warn('Unable to reach API', exc_info=True)
                return None

    @property
    def api_location(self):
        if not self.location:
            return None
        func = functools.partial(PlacesResource.from_identifier,
                                 self.location.identifier)
        return self.fetch_resource(self.location.identifier, func)

    @property
    def api_organisation(self):
        if not self.department_organiser:
            return None
        func = functools.partial(PlacesResource.from_identifier,
                                 self.department_organiser.identifier)
        return self.fetch_resource(self.department_organiser.identifier, func)

    @property
    def oxford_date(self):
        if not self.start:
            return None
        func = functools.partial(OxfordDateResource.from_date, self.start)
        return self.fetch_resource(self.start, func)

    def save(self, *args, **kwargs):
        if not self.id:
            # Newly created object, so set slug
            self.slug = slugify(self.title)
        super(Event, self).save(*args, **kwargs)

    def __unicode__(self):
        return "Event: {title} ({start})".format(title=self.title,
                                          start=self.start.strftime("%Y-%m-%d %H:%M"))

    def get_absolute_url(self):
        return reverse('event', args=[str(self.id)])

    def formatted_date(self):
        if self.start:
            return date_filter(self.start, settings.EVENT_DATETIME_FORMAT)
        else:
            return None

    def formatted_time(self):
        if self.start:
            return date_filter(self.start, settings.EVENT_TIME_FORMAT)
        else:
            return None

    def happening_today(self):
        if self.start:
            return self.start.date() == date.today()
        return False
<commit_msg>Only include if event has a group<commit_after>import logging
import functools

from datetime import date, timedelta

from django.conf import settings
from django.db import models
from django.template.defaultfilters import date as date_filter
from django.utils.text import slugify
from django.contrib.contenttypes.fields import GenericForeignKey, GenericRelation
from django.contrib.contenttypes.models import ContentType
from django.core.urlresolvers import reverse

from talks.api_ox.api import ApiException, OxfordDateResource, PlacesResource
from talks.api_ox.models import Location, Organisation

logger = logging.getLogger(__name__)

class EventGroupManager(models.Manager):
    def for_events(self, events):
        """Given a QuerySet of `Event`s return all the `EventGroup` related to
        that QuerySet. This function will evaluate the QuerySet `events`.

        Returns a list() of `EventGroup` NOT a QuerySet.
        """
        events = events.select_related('group')
        # NOTE: execs the query
        return set([e.group for e in events if e.group])

class EventGroup(models.Model):
    SEMINAR = 'SE'
    CONFERENCE = 'CO'
    EVENT_GROUP_TYPE_CHOICES = (
        (SEMINAR, 'Seminar Series'),
        (CONFERENCE, 'Conference'),
    )

    title = models.CharField(max_length=250)
    slug = models.SlugField()
    description = models.TextField()
    group_type = models.CharField(
        blank=True,
        null=True,
        max_length=2,
        choices=EVENT_GROUP_TYPE_CHOICES
    )

    objects = EventGroupManager()

    def __unicode__(self):
        return self.title

    def get_absolute_url(self):
        return reverse('event-group', args=[self.id])

class SpeakerManager(models.Manager):

    def suggestions(self, query):
        return self.filter(name__icontains=query)

class Speaker(models.Model):
    name = models.CharField(max_length=250)
    slug = models.SlugField()
    bio = models.TextField()
    email_address = models.EmailField(max_length=254)

    objects = SpeakerManager()

    def __unicode__(self):
        return self.name

class Tag(models.Model):

    slug = models.SlugField()
    name = models.CharField(max_length=250, unique=True)
    description = models.TextField()

    def __unicode__(self):
        return self.name

class TagItem(models.Model):

    tag = models.ForeignKey(Tag)
    content_type = models.ForeignKey(ContentType)
    object_id = models.PositiveIntegerField()
    item = GenericForeignKey('content_type', 'object_id')   # atm: Event, EventGroup

class EventManager(models.Manager):

    def todays_events(self):
        today = date.today()
        tomorrow = today + timedelta(days=1)
        return self.filter(start__gte=today, start__lt=tomorrow
                           ).order_by('start')

class Event(models.Model):
    start = models.DateTimeField(null=True, blank=True)
    end = models.DateTimeField(null=True, blank=True)
    title = models.CharField(max_length=250)
    slug = models.SlugField()
    description = models.TextField()
    speakers = models.ManyToManyField(Speaker, null=True, blank=True)
    # TODO audience should it be free text
    # TODO booking information; structure?

    group = models.ForeignKey(EventGroup, null=True, blank=True,
                              related_name='events')
    location = models.ForeignKey(Location, null=True, blank=True)
    # TODO I'm guessing an event can be organised by multiple departments?
    department_organiser = models.ForeignKey(Organisation, null=True, blank=True)

    tags = GenericRelation(TagItem)

    objects = EventManager()

    _cached_resources = {}

    def fetch_resource(self, key, func):
        """Fetch a resource from the API.

        If we have `key` in a cache then pull from there, otherwise call `func`
        """
        if key in self._cached_resources:
            return self._cached_resources[key]
        else:
            try:
                res = func()
                self._cached_resources[key] = res
                return res
            except ApiException:
                logger.warn('Unable to reach API', exc_info=True)
                return None

    @property
    def api_location(self):
        if not self.location:
            return None
        func = functools.partial(PlacesResource.from_identifier,
                                 self.location.identifier)
        return self.fetch_resource(self.location.identifier, func)

    @property
    def api_organisation(self):
        if not self.department_organiser:
            return None
        func = functools.partial(PlacesResource.from_identifier,
                                 self.department_organiser.identifier)
        return self.fetch_resource(self.department_organiser.identifier, func)

    @property
    def oxford_date(self):
        if not self.start:
            return None
        func = functools.partial(OxfordDateResource.from_date, self.start)
        return self.fetch_resource(self.start, func)

    def save(self, *args, **kwargs):
        if not self.id:
            # Newly created object, so set slug
            self.slug = slugify(self.title)
        super(Event, self).save(*args, **kwargs)

    def __unicode__(self):
        return "Event: {title} ({start})".format(title=self.title,
                                          start=self.start.strftime("%Y-%m-%d %H:%M"))

    def get_absolute_url(self):
        return reverse('event', args=[str(self.id)])

    def formatted_date(self):
        if self.start:
            return date_filter(self.start, settings.EVENT_DATETIME_FORMAT)
        else:
            return None

    def formatted_time(self):
        if self.start:
            return date_filter(self.start, settings.EVENT_TIME_FORMAT)
        else:
            return None

    def happening_today(self):
        if self.start:
            return self.start.date() == date.today()
        return False
<|endoftext|>


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-
'''process whole genome bisulfate sequencing FastQ files'''
import argparse, time, urllib2, gzip, shutil
#from pycompss.api.task import task
#from pycompss.api.parameter import *

from fastqreader import *
from FilterReads import *
from bs_index.wg_build import *

class process_wgbs:
    def __init__ (self):
        """
        Initialise the module
        """
        self.ready=None

    def getFastqFiles(self, sra_id, data_dir):
        """
        Function for downloading and extracting the FastQ files from the ENA
        """

        f_index = urllib2.urlopen(
        'http://www.ebi.ac.uk/ena/data/warehouse/filereport?accession=' + str(sra_id) + '&result=read_run&fields=study_accession,run_accession,tax_id,scientific_name,instrument_model,library_layout,fastq_ftp&download=txt')
        data = f_index.read()
        rows = data.split("\n")
        row_count = 0
        files = []
        gzfiles  = []
        for row in rows:
            if row_count == 0:
                row_count += 1
                continue

            row = row.rstrip()
            row = row.split("\t")

            if len(row) < 6:
                continue

            project = row[0]
            srr_id = row[1]
            fastq_files = row[6].split(';')
            row_count += 1

            for fastq_file in fastq_files:
                file_name = fastq_file.split("/")
                print fastq_file
                print data_dir + file_name[-1]
                print file_name[-1]

                req = urllib2.urlopen("ftp://" + fastq_file)
                CHUNK = 16 * 1024

                files.append(data_dir + file_name[-1].replace('.fastq.gz', '.fastq'))
                gzfiles.append(data_dir + file_name[-1])

                with open(data_dir + file_name[-1], 'wb') as fp:
                    while True:
                        chunk = req.read(CHUNK)
                        if not chunk: break
                        fp.write(chunk)

        for gzf in gzfiles:
            with gzip.open(gzf, 'rb') as f_in, open(gzf.replace('.fastq.gz', '.fastq'), 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
            os.remove(gzf)

        return files

    #@task(infile = FILE_IN, outfile = FILE_OUT)
    def FilterFastQReads(self, infile, outfile):
        """
        This is optional, but removes reads that can be problematic for the
        alignment of whole genome datasets.

        If performing RRBS then this step can be skipped
        """
        FilterReads(infile, outfile, True)

    #@constraint(ProcessorCoreCount=8)
    #@task()
    def Builder(self, fasta_file, aligner, aligner_path, ref_path):
        """
        Function to submit the FASTA file for the reference sequence and build the
        required index file used by the aligner.
        """
        builder_exec = os.path.join(aligner_path or aligner_path[aligner],
                                    {BOWTIE   : 'bowtie-build',
                                     BOWTIE2  : 'bowtie2-build',
                                     SOAP     : '2bwt-builder',
                                     RMAP     : '' # do nothing
                                    }[options.aligner])

        build_command = builder_exec + { BOWTIE   : ' -f %(fname)s.fa %(fname)s',
                                         BOWTIE2  : ' -f %(fname)s.fa %(fname)s',
                                         SOAP     : ' %(fname)s.fa'
                                       }[aligner]

        wg_build(fasta_file, build_command, ref_path, aligner)

    def Splitter(self, in_file1, in_file2, tag):
        """
        Function to divide the FastQ files into separte sub files of 1000000
        sequences so that the aligner can get run in parallel.

        Returns: Returns a list of lists of the files that have been generated.
                 Each sub list containing the two paired end files for that
                 subset.
        """

        fqr = fastqreader()
        fqr.openPairedFastQ(file1, file2)
        fqr.createOutputFiles(tag)

        r1 = fqr.next(1)
        r2 = fqr.next(2)

        count_r1 = 0
        count_r2 = 0
        count_r3 = 0
        files_out = []

        while fqr.eof(1) == False and fqr.eof(2) == False:
            r1_id = r1["id"].split(" ")
            r2_id = r2["id"].split(" ")

            if r1_id[0] == r2_id[0]:
                fqr.writeOutput(r1, 1)
                fqr.writeOutput(r2, 2)

                r1 = fqr.next(1)
                r2 = fqr.next(2)

                count_r1 += 1
                count_r2 += 1
                count_r3 += 1
            elif r1_id[0] < r2_id[0]:
                r1 = fqr.next(1)
                count_r1 += 1
            else:
                r2 = fqr.next(2)
                count_r2 += 1

            if count_r3 % 1000000 == 0:
                fqr.incrementOutputFiles()
                f1 = self.fastq1.split("/")
                f1[-1] = f1[-1].replace(".fastq", "." + str(self.output_tag) + "_" + str(self.output_file_count) + ".fastq")
                f1.insert(-1, "tmp")

                f2 = self.fastq2.split("/")
                f2[-1] = f2[-1].replace(".fastq", "." + str(self.output_tag) + "_" + str(self.output_file_count) + ".fastq")
                f2.insert(-1, "tmp")

                files_out.append(["/".join(f1), "/".join(f2)])

        fqr.closePairedFastQ()
        fqr.closeOutputFiles()

        return files_out

    #@constraint(ProcessorCoreCount=8)
    #@task() # No. of CPU, RAM required
    def Aligner(self, input_fastq1, input_fastq2, aligner, aligner_path, genome_fasta):
        """
        Alignment of the paired ends to the reference genome
        Generates bam files for the alignments
        This is performed by running the external program rather than reimplementing
        the code from teh main function to make it easier when it comes to updating
        the changes in BS-Seeker2
        """
        import subprocess

        args = ["python", "bs_seeker2-align.py", "--input_1", input_fastq1, "--input_2", input_fastq2, "--aligner", aligner, "--path", aligner_path_path, "--genome", genome_fasta, "--bt2-p", "4"]
        p = subprocess.Popen(args)
        p.wait()

    #@task() # No. of CPUs, RAM required
    def MethylationCaller(self, bam_file, output_prefix, db_dir):
        """
        Takes the merged and sorted bam file and calls the methylation sites.
        Generates a wig file of the potential sites.
        This is performed by running the external program rather than reimplementing
        the code from teh main function to make it easier when it comes to updating
        the changes in BS-Seeker2
        """
        import subprocess

        args = ["python", "bs_seeker2-call_methylation.py", "--input", str(bam_file), "--output-prefix", str(output_prefix), "--db", db_dir]
        p = subprocess.Popen(args)
        p.wait()

if __name__ == "__main__":
    import sys
    import os
    #from pycompss.api.api import compss_wait_on

    # Set up the command line parameters
    parser = argparse.ArgumentParser(description="Load adjacency list into HDF5 file")
    parser.add_argument("--genome", help="Genome name") #             default="GCA_000001405.22")
    parser.add_argument("--srr_id", help="SRR ID of the dataset") #   default="SRR1536575")
    parser.add_argument("--tmp_dir", help="Temporary data dir")
    parser.add_argument("--data_dir", help="Data directory; location to download SRA FASTQ files and save results")

    # Get the matching parameters from the command line
    args = parser.parse_args()

    srr_id   = args.srr
    genome   = args.genome
    data_dir = args.data_dir
    tmp_dir  = args.tmp_dir

    start = time.time()

    db_dir = ""

    pwgbs = process_wgbs()

    # Optain the paired FastQ files
    in_files = pwgbs.getFastqFiles(srr_id, data_dir)
    in_file1 = in_files[0]
    in_file2 = in_files[1]
    out_file1 = in_file1.replace(".fastq", "_filtered.fastq")
    out_file2 = in_file2.replace(".fastq", "_filtered.fastq")

    # Run the FilterReads.py steps for the individual FastQ files
    x = map(pwgbs.FilterFastQReads, [[in_file1, out_file1], [in_file2, out_file2]])
    x = compss_wait_on(x)

    # Run the bs_seeker2-builder.py steps
    x = pwgbs.Builder(fasta_file, aligner, aligner_path, ref_path)
    x = compss_wait_on(x)

    # Split the paired fastq files
    tmp_fastq = pwgbs.Splitter(in_file1, in_file2, tag)

    bam_sort_files = []
    bam_merge_files = []
    fastq_for_alignment = []
    for bams in tmp_fastq:
        tmp = bams
        tmp.append(aligner, aligner_path, genome_fasta_files)
        fastq_for_alignment.append(tmp)
        bam_root = bams[0] + "_bspe.bam"
        bam_sort_files.append(bams[0] + "_bspe.bam")
        bam_sort_files.append(["-o", bam_root + ".sorted.bam", "-T", bam_root + ".bam_sort", bam_root])

    # Run the bs_seeker2-align.py steps on the split up fastq files
    x = map(pwgbs.Aligner, fastq_for_alignment)
    x = compss_wait_on(x)

    # Sort and merge the aligned bam files
    # Pre-sort the original input bam files
    x = map(pysam.sort, bam_sort_files)
    x = compss_wait_on(x)

    f_bam = in_file1.split("/")
    f_bam[-1] = f_bam[-1].replace(".fastq", ".sorted.bam")
    out_bam_file = "/".join(f_bam)

    pysam.merge(out_bam_file, bam_merge_files)

    pysam.sort("-o", out_bam_file + '.sorted.bam', "-T", out_bam_file + ".bam_sort", out_bam_file)

    pysam.index(out_bam_file)

    # Run the bs_seeker2-call_methylation.py steps
    x = pwgbs.MethylationCaller(out_bam_file, tag, db_dir)
    x = compss_wait_on(x)

    # Tidy up

<commit_msg>Modified the data_dir, plus the script now creates the required directory structure as well<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-
'''process whole genome bisulfate sequencing FastQ files'''
import argparse, time, urllib2, gzip, shutil
#from pycompss.api.task import task
#from pycompss.api.parameter import *

from fastqreader import *
from FilterReads import *
from bs_index.wg_build import *

class process_wgbs:
    def __init__ (self):
        """
        Initialise the module
        """
        self.ready=None

    def getFastqFiles(self, sra_id, data_dir):
        """
        Function for downloading and extracting the FastQ files from the ENA
        """

        f_index = urllib2.urlopen(
        'http://www.ebi.ac.uk/ena/data/warehouse/filereport?accession=' + str(sra_id) + '&result=read_run&fields=study_accession,run_accession,tax_id,scientific_name,instrument_model,library_layout,fastq_ftp&download=txt')
        data = f_index.read()
        rows = data.split("\n")
        row_count = 0
        files = []
        gzfiles  = []
        for row in rows:
            if row_count == 0:
                row_count += 1
                continue

            row = row.rstrip()
            row = row.split("\t")

            if len(row) < 6:
                continue

            project = row[0]
            srr_id = row[1]
            fastq_files = row[6].split(';')
            row_count += 1

            for fastq_file in fastq_files:
                file_name = fastq_file.split("/")
                print fastq_file
                print data_dir + file_name[-1]
                print file_name[-1]

                req = urllib2.urlopen("ftp://" + fastq_file)
                CHUNK = 16 * 1024

                files.append(data_dir + file_name[-1].replace('.fastq.gz', '.fastq'))
                gzfiles.append(data_dir + file_name[-1])

                with open(data_dir + file_name[-1], 'wb') as fp:
                    while True:
                        chunk = req.read(CHUNK)
                        if not chunk: break
                        fp.write(chunk)

        for gzf in gzfiles:
            with gzip.open(gzf, 'rb') as f_in, open(gzf.replace('.fastq.gz', '.fastq'), 'wb') as f_out:
                shutil.copyfileobj(f_in, f_out)
            os.remove(gzf)

        return files

    #@task(infile = FILE_IN, outfile = FILE_OUT)
    def FilterFastQReads(self, infile, outfile):
        """
        This is optional, but removes reads that can be problematic for the
        alignment of whole genome datasets.

        If performing RRBS then this step can be skipped
        """
        FilterReads(infile, outfile, True)

    #@constraint(ProcessorCoreCount=8)
    #@task()
    def Builder(self, fasta_file, aligner, aligner_path, ref_path):
        """
        Function to submit the FASTA file for the reference sequence and build the
        required index file used by the aligner.
        """
        builder_exec = os.path.join(aligner_path or aligner_path[aligner],
                                    {BOWTIE   : 'bowtie-build',
                                     BOWTIE2  : 'bowtie2-build',
                                     SOAP     : '2bwt-builder',
                                     RMAP     : '' # do nothing
                                    }[options.aligner])

        build_command = builder_exec + { BOWTIE   : ' -f %(fname)s.fa %(fname)s',
                                         BOWTIE2  : ' -f %(fname)s.fa %(fname)s',
                                         SOAP     : ' %(fname)s.fa'
                                       }[aligner]

        wg_build(fasta_file, build_command, ref_path, aligner)

    def Splitter(self, in_file1, in_file2, tag):
        """
        Function to divide the FastQ files into separte sub files of 1000000
        sequences so that the aligner can get run in parallel.

        Returns: Returns a list of lists of the files that have been generated.
                 Each sub list containing the two paired end files for that
                 subset.
        """

        fqr = fastqreader()
        fqr.openPairedFastQ(file1, file2)
        fqr.createOutputFiles(tag)

        r1 = fqr.next(1)
        r2 = fqr.next(2)

        count_r1 = 0
        count_r2 = 0
        count_r3 = 0
        files_out = []

        while fqr.eof(1) == False and fqr.eof(2) == False:
            r1_id = r1["id"].split(" ")
            r2_id = r2["id"].split(" ")

            if r1_id[0] == r2_id[0]:
                fqr.writeOutput(r1, 1)
                fqr.writeOutput(r2, 2)

                r1 = fqr.next(1)
                r2 = fqr.next(2)

                count_r1 += 1
                count_r2 += 1
                count_r3 += 1
            elif r1_id[0] < r2_id[0]:
                r1 = fqr.next(1)
                count_r1 += 1
            else:
                r2 = fqr.next(2)
                count_r2 += 1

            if count_r3 % 1000000 == 0:
                fqr.incrementOutputFiles()
                f1 = self.fastq1.split("/")
                f1[-1] = f1[-1].replace(".fastq", "." + str(self.output_tag) + "_" + str(self.output_file_count) + ".fastq")
                f1.insert(-1, "tmp")

                f2 = self.fastq2.split("/")
                f2[-1] = f2[-1].replace(".fastq", "." + str(self.output_tag) + "_" + str(self.output_file_count) + ".fastq")
                f2.insert(-1, "tmp")

                files_out.append(["/".join(f1), "/".join(f2)])

        fqr.closePairedFastQ()
        fqr.closeOutputFiles()

        return files_out

    #@constraint(ProcessorCoreCount=8)
    #@task() # No. of CPU, RAM required
    def Aligner(self, input_fastq1, input_fastq2, aligner, aligner_path, genome_fasta):
        """
        Alignment of the paired ends to the reference genome
        Generates bam files for the alignments
        This is performed by running the external program rather than reimplementing
        the code from teh main function to make it easier when it comes to updating
        the changes in BS-Seeker2
        """
        import subprocess

        args = ["python", "bs_seeker2-align.py", "--input_1", input_fastq1, "--input_2", input_fastq2, "--aligner", aligner, "--path", aligner_path_path, "--genome", genome_fasta, "--bt2-p", "4"]
        p = subprocess.Popen(args)
        p.wait()

    #@task() # No. of CPUs, RAM required
    def MethylationCaller(self, bam_file, output_prefix, db_dir):
        """
        Takes the merged and sorted bam file and calls the methylation sites.
        Generates a wig file of the potential sites.
        This is performed by running the external program rather than reimplementing
        the code from teh main function to make it easier when it comes to updating
        the changes in BS-Seeker2
        """
        import subprocess

        args = ["python", "bs_seeker2-call_methylation.py", "--input", str(bam_file), "--output-prefix", str(output_prefix), "--db", db_dir]
        p = subprocess.Popen(args)
        p.wait()

if __name__ == "__main__":
    import sys
    import os
    #from pycompss.api.api import compss_wait_on

    # Set up the command line parameters
    parser = argparse.ArgumentParser(description="Load adjacency list into HDF5 file")
    parser.add_argument("--genome", help="Genome name") #             default="GCA_000001405.22")
    parser.add_argument("--srr_id", help="SRR ID of the dataset") #   default="SRR1536575")
    parser.add_argument("--tmp_dir", help="Temporary data dir")
    parser.add_argument("--data_dir", help="Data directory; location to download SRA FASTQ files and save results")

    # Get the matching parameters from the command line
    args = parser.parse_args()

    srr_id   = args.srr
    genome   = args.genome
    data_dir = args.data_dir + srr_id
    tmp_dir  = args.tmp_dir

    start = time.time()

    db_dir = ""

    pwgbs = process_wgbs()

    try:
        os.makedirs(data_dir)

        data_dir += "/"
    except:
        pass

    # Optain the paired FastQ files
    in_files = pwgbs.getFastqFiles(srr_id, data_dir)
    in_file1 = in_files[0]
    in_file2 = in_files[1]
    out_file1 = in_file1.replace(".fastq", "_filtered.fastq")
    out_file2 = in_file2.replace(".fastq", "_filtered.fastq")

    # Run the FilterReads.py steps for the individual FastQ files
    x = map(pwgbs.FilterFastQReads, [[in_file1, out_file1], [in_file2, out_file2]])
    x = compss_wait_on(x)

    # Run the bs_seeker2-builder.py steps
    x = pwgbs.Builder(fasta_file, aligner, aligner_path, ref_path)
    x = compss_wait_on(x)

    # Split the paired fastq files
    tmp_fastq = pwgbs.Splitter(in_file1, in_file2, tag)

    bam_sort_files = []
    bam_merge_files = []
    fastq_for_alignment = []
    for bams in tmp_fastq:
        tmp = bams
        tmp.append(aligner, aligner_path, genome_fasta_files)
        fastq_for_alignment.append(tmp)
        bam_root = bams[0] + "_bspe.bam"
        bam_sort_files.append(bams[0] + "_bspe.bam")
        bam_sort_files.append(["-o", bam_root + ".sorted.bam", "-T", bam_root + ".bam_sort", bam_root])

    # Run the bs_seeker2-align.py steps on the split up fastq files
    x = map(pwgbs.Aligner, fastq_for_alignment)
    x = compss_wait_on(x)

    # Sort and merge the aligned bam files
    # Pre-sort the original input bam files
    x = map(pysam.sort, bam_sort_files)
    x = compss_wait_on(x)

    f_bam = in_file1.split("/")
    f_bam[-1] = f_bam[-1].replace(".fastq", ".sorted.bam")
    out_bam_file = "/".join(f_bam)

    pysam.merge(out_bam_file, bam_merge_files)

    pysam.sort("-o", out_bam_file + '.sorted.bam', "-T", out_bam_file + ".bam_sort", out_bam_file)

    pysam.index(out_bam_file)

    # Run the bs_seeker2-call_methylation.py steps
    x = pwgbs.MethylationCaller(out_bam_file, tag, db_dir)
    x = compss_wait_on(x)

    # Tidy up

<|endoftext|>


<commit_before>#!/usr/bin/python
# coding=utf-8
import argparse
import logging
import os.path
import signal
import socket
import sys
import time
import traceback

try:
    # Need simplejson for the object_pairs_hook option in Python 2.6.
    import simplejson as json
except ImportError:
    # Python 2.7+.
    import json

from bson import SON

from mongo_orchestration import __version__
from mongo_orchestration.common import (
    BaseModel,
    DEFAULT_BIND, DEFAULT_PORT, DEFAULT_SERVER, DEFAULT_SOCKET_TIMEOUT,
    PID_FILE, LOG_FILE, LOGGING_FORMAT)
from mongo_orchestration.daemon import Daemon
from mongo_orchestration.servers import Server

# How many times to attempt connecting to mongo-orchestration server.
CONNECT_ATTEMPTS = 5

def read_env():
    """return command-line arguments"""
    parser = argparse.ArgumentParser(description='mongo-orchestration server')
    parser.add_argument('-f', '--config',
                        action='store', default=None, type=str, dest='config')
    parser.add_argument('-e', '--env',
                        action='store', type=str, dest='env', default=None)
    parser.add_argument(action='store', type=str, dest='command',
                        default='start', choices=('start', 'stop', 'restart'))
    parser.add_argument('--no-fork',
                        action='store_true', dest='no_fork', default=False)
    parser.add_argument('-b', '--bind',
                        action='store', dest='bind', type=str,
                        default=DEFAULT_BIND)
    parser.add_argument('-p', '--port',
                        action='store', dest='port', type=int,
                        default=DEFAULT_PORT)
    parser.add_argument('--enable-majority-read-concern', action='store_true',
                        default=False)
    parser.add_argument('-s', '--server',
                        action='store', dest='server', type=str,
                        default=DEFAULT_SERVER, choices=('cherrypy', 'wsgiref'))
    parser.add_argument('--version', action='version',
                        version='Mongo Orchestration v' + __version__)
    parser.add_argument('--socket-timeout-ms', action='store',
                        dest='socket_timeout',
                        type=int, default=DEFAULT_SOCKET_TIMEOUT)
    parser.add_argument('--pidfile', action='store', type=str, dest='pidfile',
                        default=PID_FILE)

    cli_args = parser.parse_args()

    if cli_args.env and not cli_args.config:
        print("Specified release '%s' without a config file" % cli_args.env)
        sys.exit(1)
    if cli_args.command == 'stop' or not cli_args.config:
        return cli_args
    try:
        # read config
        with open(cli_args.config, 'r') as fd:
            config = json.loads(fd.read(), object_pairs_hook=SON)
        if not 'releases' in config:
            print("No releases defined in %s" % cli_args.config)
            sys.exit(1)
        releases = config['releases']
        if cli_args.env is not None and cli_args.env not in releases:
            print("Release '%s' is not defined in %s"
                  % (cli_args.env, cli_args.config))
            sys.exit(1)
        cli_args.releases = releases
        return cli_args
    except (IOError):
        print("config file not found")
        sys.exit(1)
    except (ValueError):
        print("config file is corrupted")
        sys.exit(1)

def setup(releases, default_release):
    """setup storages"""
    from mongo_orchestration import set_releases, cleanup_storage
    set_releases(releases, default_release)
    signal.signal(signal.SIGTERM, cleanup_storage)
    signal.signal(signal.SIGINT, cleanup_storage)

def get_app():
    """return bottle app that includes all sub-apps"""
    from bottle import default_app
    default_app.push()
    for module in ("mongo_orchestration.apps.servers",
                   "mongo_orchestration.apps.replica_sets",
                   "mongo_orchestration.apps.sharded_clusters"):
        __import__(module)
    app = default_app.pop()
    return app

class MyDaemon(Daemon):
    """class uses to run server as daemon"""

    def __init__(self, *args, **kwd):
        super(MyDaemon, self).__init__(*args, **kwd)

    def run(self):
        log = logging.getLogger(__name__)

        from bottle import run
        setup(getattr(self.args, 'releases', {}), self.args.env)
        BaseModel.socket_timeout = self.args.socket_timeout
        if self.args.command in ('start', 'restart'):
            print("Starting Mongo Orchestration on port %d..." % self.args.port)
            try:
                log.debug('Starting HTTP server on host: %s; port: %d',
                          self.args.bind, self.args.port)
                run(get_app(), host=self.args.bind, port=self.args.port,
                    debug=False, reloader=False, quiet=not self.args.no_fork,
                    server=self.args.server)
            except Exception:
                traceback.print_exc(file=sys.stdout)
                log.exception('Could not start a new server.')
                raise

    def set_args(self, args):
        self.args = args

def await_connection(host, port):
    """Wait for the mongo-orchestration server to accept connections."""
    for i in range(CONNECT_ATTEMPTS):
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            try:
                s.connect((host, port))
                return True
            except (IOError, socket.error):
                time.sleep(1)
        finally:
            s.close()
    return False

def main():
    args = read_env()
    Server.enable_majority_read_concern = args.enable_majority_read_concern
    # Silence STDOUT from mongo processes if MO is running as a deamon.
    Server.silence_stdout = not args.no_fork
    # Log both to STDOUT and the log file.
    logging.basicConfig(level=logging.DEBUG, filename=LOG_FILE,
                        format=LOGGING_FORMAT)
    log = logging.getLogger(__name__)

    daemon = MyDaemon(os.path.abspath(args.pidfile), timeout=5,
                      stdout=sys.stdout)
    daemon.set_args(args)
    # Set default bind ip for mongo processes using argument from --bind.
    Server.mongod_default['bind_ip'] = args.bind
    if args.command == 'stop':
        daemon.stop()
    if args.command == 'start' and not args.no_fork:
        print('Preparing to start mongo-orchestration daemon')
        pid = daemon.start()
        print('Daemon process started with pid: %d' % pid)
        if not await_connection(host=args.bind, port=args.port):
            print(
                'Could not connect to daemon running on %s:%d (pid: %d) '
                'within %d attempts.'
                % (args.bind, args.port, pid, CONNECT_ATTEMPTS))
            daemon.stop()
    if args.command == 'start' and args.no_fork:
        log.debug('Starting mongo-orchestration in the foreground')
        daemon.run()
    if args.command == 'restart':
        daemon.restart()

if __name__ == "__main__":
    main()
<commit_msg>Use socket.create_connection to create the connection, supporting IPv6-bound servers and fixes #237. (#258)<commit_after>#!/usr/bin/python
# coding=utf-8
import argparse
import logging
import os.path
import signal
import socket
import sys
import time
import traceback

try:
    # Need simplejson for the object_pairs_hook option in Python 2.6.
    import simplejson as json
except ImportError:
    # Python 2.7+.
    import json

from bson import SON

from mongo_orchestration import __version__
from mongo_orchestration.common import (
    BaseModel,
    DEFAULT_BIND, DEFAULT_PORT, DEFAULT_SERVER, DEFAULT_SOCKET_TIMEOUT,
    PID_FILE, LOG_FILE, LOGGING_FORMAT)
from mongo_orchestration.daemon import Daemon
from mongo_orchestration.servers import Server

# How many times to attempt connecting to mongo-orchestration server.
CONNECT_ATTEMPTS = 5

def read_env():
    """return command-line arguments"""
    parser = argparse.ArgumentParser(description='mongo-orchestration server')
    parser.add_argument('-f', '--config',
                        action='store', default=None, type=str, dest='config')
    parser.add_argument('-e', '--env',
                        action='store', type=str, dest='env', default=None)
    parser.add_argument(action='store', type=str, dest='command',
                        default='start', choices=('start', 'stop', 'restart'))
    parser.add_argument('--no-fork',
                        action='store_true', dest='no_fork', default=False)
    parser.add_argument('-b', '--bind',
                        action='store', dest='bind', type=str,
                        default=DEFAULT_BIND)
    parser.add_argument('-p', '--port',
                        action='store', dest='port', type=int,
                        default=DEFAULT_PORT)
    parser.add_argument('--enable-majority-read-concern', action='store_true',
                        default=False)
    parser.add_argument('-s', '--server',
                        action='store', dest='server', type=str,
                        default=DEFAULT_SERVER, choices=('cherrypy', 'wsgiref'))
    parser.add_argument('--version', action='version',
                        version='Mongo Orchestration v' + __version__)
    parser.add_argument('--socket-timeout-ms', action='store',
                        dest='socket_timeout',
                        type=int, default=DEFAULT_SOCKET_TIMEOUT)
    parser.add_argument('--pidfile', action='store', type=str, dest='pidfile',
                        default=PID_FILE)

    cli_args = parser.parse_args()

    if cli_args.env and not cli_args.config:
        print("Specified release '%s' without a config file" % cli_args.env)
        sys.exit(1)
    if cli_args.command == 'stop' or not cli_args.config:
        return cli_args
    try:
        # read config
        with open(cli_args.config, 'r') as fd:
            config = json.loads(fd.read(), object_pairs_hook=SON)
        if not 'releases' in config:
            print("No releases defined in %s" % cli_args.config)
            sys.exit(1)
        releases = config['releases']
        if cli_args.env is not None and cli_args.env not in releases:
            print("Release '%s' is not defined in %s"
                  % (cli_args.env, cli_args.config))
            sys.exit(1)
        cli_args.releases = releases
        return cli_args
    except (IOError):
        print("config file not found")
        sys.exit(1)
    except (ValueError):
        print("config file is corrupted")
        sys.exit(1)

def setup(releases, default_release):
    """setup storages"""
    from mongo_orchestration import set_releases, cleanup_storage
    set_releases(releases, default_release)
    signal.signal(signal.SIGTERM, cleanup_storage)
    signal.signal(signal.SIGINT, cleanup_storage)

def get_app():
    """return bottle app that includes all sub-apps"""
    from bottle import default_app
    default_app.push()
    for module in ("mongo_orchestration.apps.servers",
                   "mongo_orchestration.apps.replica_sets",
                   "mongo_orchestration.apps.sharded_clusters"):
        __import__(module)
    app = default_app.pop()
    return app

class MyDaemon(Daemon):
    """class uses to run server as daemon"""

    def __init__(self, *args, **kwd):
        super(MyDaemon, self).__init__(*args, **kwd)

    def run(self):
        log = logging.getLogger(__name__)

        from bottle import run
        setup(getattr(self.args, 'releases', {}), self.args.env)
        BaseModel.socket_timeout = self.args.socket_timeout
        if self.args.command in ('start', 'restart'):
            print("Starting Mongo Orchestration on port %d..." % self.args.port)
            try:
                log.debug('Starting HTTP server on host: %s; port: %d',
                          self.args.bind, self.args.port)
                run(get_app(), host=self.args.bind, port=self.args.port,
                    debug=False, reloader=False, quiet=not self.args.no_fork,
                    server=self.args.server)
            except Exception:
                traceback.print_exc(file=sys.stdout)
                log.exception('Could not start a new server.')
                raise

    def set_args(self, args):
        self.args = args

def await_connection(host, port):
    """Wait for the mongo-orchestration server to accept connections."""
    for i in range(CONNECT_ATTEMPTS):
        try:
            conn = socket.create_connection((host, port))
            conn.close()
            return True
        except (IOError, socket.error):
            time.sleep(1)
    return False

def main():
    args = read_env()
    Server.enable_majority_read_concern = args.enable_majority_read_concern
    # Silence STDOUT from mongo processes if MO is running as a deamon.
    Server.silence_stdout = not args.no_fork
    # Log both to STDOUT and the log file.
    logging.basicConfig(level=logging.DEBUG, filename=LOG_FILE,
                        format=LOGGING_FORMAT)
    log = logging.getLogger(__name__)

    daemon = MyDaemon(os.path.abspath(args.pidfile), timeout=5,
                      stdout=sys.stdout)
    daemon.set_args(args)
    # Set default bind ip for mongo processes using argument from --bind.
    Server.mongod_default['bind_ip'] = args.bind
    if args.command == 'stop':
        daemon.stop()
    if args.command == 'start' and not args.no_fork:
        print('Preparing to start mongo-orchestration daemon')
        pid = daemon.start()
        print('Daemon process started with pid: %d' % pid)
        if not await_connection(host=args.bind, port=args.port):
            print(
                'Could not connect to daemon running on %s:%d (pid: %d) '
                'within %d attempts.'
                % (args.bind, args.port, pid, CONNECT_ATTEMPTS))
            daemon.stop()
    if args.command == 'start' and args.no_fork:
        log.debug('Starting mongo-orchestration in the foreground')
        daemon.run()
    if args.command == 'restart':
        daemon.restart()

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