<commit_before>
import datetime
from django.core.management.color import no_style
from django.db import connection, transaction, models
from django.db.backends.util import truncate_name
from django.db.models.fields import NOT_PROVIDED
from django.dispatch import dispatcher
from django.conf import settings

def alias(attrname):
    """
    Returns a function which calls 'attrname' - for function aliasing.
    We can't just use foo = bar, as this breaks subclassing.
    """
    def func(self, *args, **kwds):
        return getattr(self, attrname)(*args, **kwds)
    return func

class DatabaseOperations(object):

    """
    Generic SQL implementation of the DatabaseOperations.
    Some of this code comes from Django Evolution.
    """

    # We assume the generic DB can handle DDL transactions. MySQL wil change this.
    has_ddl_transactions = True

    def __init__(self):
        self.debug = False
        self.deferred_sql = []
        self.dry_run = False

    def execute(self, sql, params=[]):
        """
        Executes the given SQL statement, with optional parameters.
        If the instance's debug attribute is True, prints out what it executes.
        """
        cursor = connection.cursor()
        if self.debug:
            print "   = %s" % sql, params

        if self.dry_run:
            return []

        cursor.execute(sql, params)
        try:
            return cursor.fetchall()
        except:
            return []

    def add_deferred_sql(self, sql):
        """
        Add a SQL statement to the deferred list, that won't be executed until
        this instance's execute_deferred_sql method is run.
        """
        self.deferred_sql.append(sql)

    def execute_deferred_sql(self):
        """
        Executes all deferred SQL, resetting the deferred_sql list
        """
        for sql in self.deferred_sql:
            self.execute(sql)

        self.deferred_sql = []

    def clear_deferred_sql(self):
        """
        Resets the deferred_sql list to empty.
        """
        self.deferred_sql = []

    def create_table(self, table_name, fields):
        """
        Creates the table 'table_name'. 'fields' is a tuple of fields,
        each repsented by a 2-part tuple of field name and a
        django.db.models.fields.Field object
        """
        qn = connection.ops.quote_name

        # allow fields to be a dictionary
	# removed for now - philosophical reasons (this is almost certainly not what you want)
        #try:
        #    fields = fields.items()
        #except AttributeError:
        #    pass

        columns = [
            self.column_sql(table_name, field_name, field)
            for field_name, field in fields
        ]

        self.execute('CREATE TABLE %s (%s);' % (qn(table_name), ', '.join([col for col in columns if col])))

    add_table = alias('create_table') # Alias for consistency's sake

    def rename_table(self, old_table_name, table_name):
        """
        Renames the table 'old_table_name' to 'table_name'.
        """
        if old_table_name == table_name:
            # No Operation
            return
        qn = connection.ops.quote_name
        params = (qn(old_table_name), qn(table_name))
        self.execute('ALTER TABLE %s RENAME TO %s;' % params)

    def delete_table(self, table_name):
        """
        Deletes the table 'table_name'.
        """
        qn = connection.ops.quote_name
        params = (qn(table_name), )
        self.execute('DROP TABLE %s;' % params)

    drop_table = alias('delete_table')

    def add_column(self, table_name, name, field, keep_default=True):
        """
        Adds the column 'name' to the table 'table_name'.
        Uses the 'field' paramater, a django.db.models.fields.Field instance,
        to generate the necessary sql

        @param table_name: The name of the table to add the column to
        @param name: The name of the column to add
        @param field: The field to use
        """
        qn = connection.ops.quote_name
        sql = self.column_sql(table_name, name, field)
        if sql:
            params = (
                qn(table_name),
                sql,
            )
            sql = 'ALTER TABLE %s ADD COLUMN %s;' % params
            self.execute(sql)

            # Now, drop the default if we need to
            if not keep_default and field.default:
                field.default = NOT_PROVIDED
                self.alter_column(table_name, name, field)

    alter_string_set_type = 'ALTER COLUMN %(column)s TYPE %(type)s'
    alter_string_set_null = 'ALTER COLUMN %(column)s DROP NOT NULL'
    alter_string_drop_null = 'ALTER COLUMN %(column)s SET NOT NULL'
    allows_combined_alters = True

    def alter_column(self, table_name, name, field):
        """
        Alters the given column name so it will match the given field.
        Note that conversion between the two by the database must be possible.

        @param table_name: The name of the table to add the column to
        @param name: The name of the column to alter
        @param field: The new field definition to use
        """

        # hook for the field to do any resolution prior to it's attributes being queried
        if hasattr(field, 'south_init'):
            field.south_init()

        qn = connection.ops.quote_name

        # First, change the type
        params = {
            "column": qn(name),
            "type": field.db_type(),
        }

        # SQLs is a list of (SQL, values) pairs.
        sqls = [(self.alter_string_set_type % params, [])]

        # Next, set any default
        if not field.null and field.has_default():
            default = field.get_default()
            sqls.append(('ALTER COLUMN %s SET DEFAULT %%s ' % (qn(name),), [default]))
        else:
            sqls.append(('ALTER COLUMN %s DROP DEFAULT' % (qn(name),), []))

        # Next, nullity
        params = {
            "column": qn(name),
            "type": field.db_type(),
        }
        if field.null:
            sqls.append((self.alter_string_set_null % params, []))
        else:
            sqls.append((self.alter_string_drop_null % params, []))

        # TODO: Unique

        if self.allows_combined_alters:
            sqls, values = zip(*sqls)
            self.execute(
                "ALTER TABLE %s %s;" % (qn(table_name), ", ".join(sqls)),
                flatten(values),
            )
        else:
            # Databases like e.g. MySQL don't like more than one alter at once.
            for sql, values in sqls:
                self.execute("ALTER TABLE %s %s;" % (qn(table_name), sql), values)

    def column_sql(self, table_name, field_name, field, tablespace=''):
        """
        Creates the SQL snippet for a column. Used by add_column and add_table.
        """
        qn = connection.ops.quote_name

        field.set_attributes_from_name(field_name)

        # hook for the field to do any resolution prior to it's attributes being queried
        if hasattr(field, 'south_init'):
            field.south_init()

        sql = field.db_type()
        if sql:
            field_output = [qn(field.column), sql]
            field_output.append('%sNULL' % (not field.null and 'NOT ' or ''))
            if field.primary_key:
                field_output.append('PRIMARY KEY')
            elif field.unique:
                field_output.append('UNIQUE')

            tablespace = field.db_tablespace or tablespace
            if tablespace and connection.features.supports_tablespaces and field.unique:
                # We must specify the index tablespace inline, because we
                # won't be generating a CREATE INDEX statement for this field.
                field_output.append(connection.ops.tablespace_sql(tablespace, inline=True))

            sql = ' '.join(field_output)
            sqlparams = ()
            # if the field is "NOT NULL" and a default value is provided, create the column with it
            # this allows the addition of a NOT NULL field to a table with existing rows
            if not field.null and field.has_default():
                default = field.get_default()
                if isinstance(default, basestring):
                    default = "'%s'" % default.replace("'", "''")
                elif isinstance(default, datetime.date):
                    default = "'%s'" % default
                sql += " DEFAULT %s"
                sqlparams = (default)

            if field.rel and self.supports_foreign_keys:
                self.add_deferred_sql(
                    self.foreign_key_sql(
                        table_name,
                        field.column,
                        field.rel.to._meta.db_table,
                        field.rel.to._meta.get_field(field.rel.field_name).column
                    )
                )

            if field.db_index and not field.unique:
                self.add_deferred_sql(self.create_index_sql(table_name, [field.column]))

        if hasattr(field, 'post_create_sql'):
            style = no_style()
            for stmt in field.post_create_sql(style, table_name):
                self.add_deferred_sql(stmt)

        if sql:
            return sql % sqlparams
        else:
            return None

    supports_foreign_keys = True

    def foreign_key_sql(self, from_table_name, from_column_name, to_table_name, to_column_name):
        """
        Generates a full SQL statement to add a foreign key constraint
        """
        constraint_name = '%s_refs_%s_%x' % (from_column_name, to_column_name, abs(hash((from_table_name, to_table_name))))
        return 'ALTER TABLE %s ADD CONSTRAINT %s FOREIGN KEY (%s) REFERENCES %s (%s)%s;' % (
            from_table_name,
            truncate_name(constraint_name, connection.ops.max_name_length()),
            from_column_name,
            to_table_name,
            to_column_name,
            connection.ops.deferrable_sql() # Django knows this
        )

    def create_index_name(self, table_name, column_names):
        """
        Generate a unique name for the index
        """
        index_unique_name = ''
        if len(column_names) > 1:
            index_unique_name = '_%x' % abs(hash((table_name, ','.join(column_names))))

        return '%s_%s%s' % (table_name, column_names[0], index_unique_name)

    def create_index_sql(self, table_name, column_names, unique=False, db_tablespace=''):
        """
        Generates a create index statement on 'table_name' for a list of 'column_names'
        """
        if not column_names:
            print "No column names supplied on which to create an index"
            return ''

        if db_tablespace and connection.features.supports_tablespaces:
            tablespace_sql = ' ' + connection.ops.tablespace_sql(db_tablespace)
        else:
            tablespace_sql = ''

        index_name = self.create_index_name(table_name, column_names)
        qn = connection.ops.quote_name
        return 'CREATE %sINDEX %s ON %s (%s)%s;' % (
            unique and 'UNIQUE ' or '',
            index_name,
            table_name,
            ','.join([qn(field) for field in column_names]),
            tablespace_sql
            )

    def create_index(self, table_name, column_names, unique=False, db_tablespace=''):
        """ Executes a create index statement """
        sql = self.create_index_sql(table_name, column_names, unique, db_tablespace)
        self.execute(sql)

    drop_index_string = 'DROP INDEX %(index_name)s'

    def delete_index(self, table_name, column_names, db_tablespace=''):
        """
        Deletes an index created with create_index.
        This is possible using only columns due to the deterministic
        index naming function which relies on column names.
        """
        name = self.create_index_name(table_name, column_names)
        sql = self.drop_index_string % {"index_name": name, "table_name": table_name}
        self.execute(sql)

    drop_index = alias('delete_index')

    def delete_column(self, table_name, name):
        """
        Deletes the column 'column_name' from the table 'table_name'.
        """
        qn = connection.ops.quote_name
        params = (qn(table_name), qn(name))
        self.execute('ALTER TABLE %s DROP COLUMN %s CASCADE;' % params, [])

    drop_column = alias('delete_column')

    def rename_column(self, table_name, old, new):
        """
        Renames the column 'old' from the table 'table_name' to 'new'.
        """
        raise NotImplementedError("rename_column has no generic SQL syntax")

    def start_transaction(self):
        """
        Makes sure the following commands are inside a transaction.
        Must be followed by a (commit|rollback)_transaction call.
        """
        if self.dry_run:
            return
        transaction.commit_unless_managed()
        transaction.enter_transaction_management()
        transaction.managed(True)

    def commit_transaction(self):
        """
        Commits the current transaction.
        Must be preceded by a start_transaction call.
        """
        if self.dry_run:
            return
        transaction.commit()
        transaction.leave_transaction_management()

    def rollback_transaction(self):
        """
        Rolls back the current transaction.
        Must be preceded by a start_transaction call.
        """
        if self.dry_run:
            return
        transaction.rollback()
        transaction.leave_transaction_management()

    def send_create_signal(self, app_label, model_names):
        """
        Sends a post_syncdb signal for the model specified.

        If the model is not found (perhaps it's been deleted?),
        no signal is sent.

        TODO: The behavior of django.contrib.* apps seems flawed in that
        they don't respect created_models.  Rather, they blindly execute
        over all models within the app sending the signal.  This is a
        patch we should push Django to make  For now, this should work.
        """
        app = models.get_app(app_label)
        if not app:
            return

        created_models = []
        for model_name in model_names:
            model = models.get_model(app_label, model_name)
            if model:
                created_models.append(model)

        if created_models:
            # syncdb defaults -- perhaps take these as options?
            verbosity = 1
            interactive = True

            if hasattr(dispatcher, "send"):
                dispatcher.send(signal=models.signals.post_syncdb, sender=app,
                app=app, created_models=created_models,
                verbosity=verbosity, interactive=interactive)
            else:
                models.signals.post_syncdb.send(sender=app,
                app=app, created_models=created_models,
                verbosity=verbosity, interactive=interactive)

    def mock_model(self, model_name, db_table, db_tablespace='',
                    pk_field_name='id', pk_field_type=models.AutoField,
                    pk_field_args=[], pk_field_kwargs={}):
        """
        Generates a MockModel class that provides enough information
        to be used by a foreign key/many-to-many relationship.

        Migrations should prefer to use these rather than actual models
        as models could get deleted over time, but these can remain in
        migration files forever.
        """
        class MockOptions(object):
            def __init__(self):
                self.db_table = db_table
                self.db_tablespace = db_tablespace or settings.DEFAULT_TABLESPACE
                self.object_name = model_name
                self.module_name = model_name.lower()

                if pk_field_type == models.AutoField:
                    pk_field_kwargs['primary_key'] = True

                self.pk = pk_field_type(*pk_field_args, **pk_field_kwargs)
                self.pk.set_attributes_from_name(pk_field_name)
                self.abstract = False

            def get_field_by_name(self, field_name):
                # we only care about the pk field
                return (self.pk, self.model, True, False)

            def get_field(self, name):
                # we only care about the pk field
                return self.pk

        class MockModel(object):
            _meta = None

        # We need to return an actual class object here, not an instance
        MockModel._meta = MockOptions()
        MockModel._meta.model = MockModel
        return MockModel

# Single-level flattening of lists
def flatten(ls):
    nl = []
    for l in ls:
        nl += l
    return nl

<commit_msg>Oops. Forgot self.<commit_after>
import datetime
from django.core.management.color import no_style
from django.db import connection, transaction, models
from django.db.backends.util import truncate_name
from django.db.models.fields import NOT_PROVIDED
from django.dispatch import dispatcher
from django.conf import settings

def alias(attrname):
    """
    Returns a function which calls 'attrname' - for function aliasing.
    We can't just use foo = bar, as this breaks subclassing.
    """
    def func(self, *args, **kwds):
        return getattr(self, attrname)(*args, **kwds)
    return func

class DatabaseOperations(object):

    """
    Generic SQL implementation of the DatabaseOperations.
    Some of this code comes from Django Evolution.
    """

    # We assume the generic DB can handle DDL transactions. MySQL wil change this.
    has_ddl_transactions = True

    def __init__(self):
        self.debug = False
        self.deferred_sql = []
        self.dry_run = False

    def execute(self, sql, params=[]):
        """
        Executes the given SQL statement, with optional parameters.
        If the instance's debug attribute is True, prints out what it executes.
        """
        cursor = connection.cursor()
        if self.debug:
            print "   = %s" % sql, params

        if self.dry_run:
            return []

        cursor.execute(sql, params)
        try:
            return cursor.fetchall()
        except:
            return []

    def add_deferred_sql(self, sql):
        """
        Add a SQL statement to the deferred list, that won't be executed until
        this instance's execute_deferred_sql method is run.
        """
        self.deferred_sql.append(sql)

    def execute_deferred_sql(self):
        """
        Executes all deferred SQL, resetting the deferred_sql list
        """
        for sql in self.deferred_sql:
            self.execute(sql)

        self.deferred_sql = []

    def clear_deferred_sql(self):
        """
        Resets the deferred_sql list to empty.
        """
        self.deferred_sql = []

    def create_table(self, table_name, fields):
        """
        Creates the table 'table_name'. 'fields' is a tuple of fields,
        each repsented by a 2-part tuple of field name and a
        django.db.models.fields.Field object
        """
        qn = connection.ops.quote_name

        # allow fields to be a dictionary
	# removed for now - philosophical reasons (this is almost certainly not what you want)
        #try:
        #    fields = fields.items()
        #except AttributeError:
        #    pass

        columns = [
            self.column_sql(table_name, field_name, field)
            for field_name, field in fields
        ]

        self.execute('CREATE TABLE %s (%s);' % (qn(table_name), ', '.join([col for col in columns if col])))

    add_table = alias('create_table') # Alias for consistency's sake

    def rename_table(self, old_table_name, table_name):
        """
        Renames the table 'old_table_name' to 'table_name'.
        """
        if old_table_name == table_name:
            # No Operation
            return
        qn = connection.ops.quote_name
        params = (qn(old_table_name), qn(table_name))
        self.execute('ALTER TABLE %s RENAME TO %s;' % params)

    def delete_table(self, table_name):
        """
        Deletes the table 'table_name'.
        """
        qn = connection.ops.quote_name
        params = (qn(table_name), )
        self.execute('DROP TABLE %s;' % params)

    drop_table = alias('delete_table')

    def clear_table(self, table_name):
        """
        Deletes all rows from 'table_name'.
        """
        qn = connection.ops.quote_name
        params = (qn(table_name), )
        self.execute('DELETE FROM %s;' % params)

    add_column_string = 'ALTER TABLE %s ADD COLUMN %s;'

    def add_column(self, table_name, name, field, keep_default=True):
        """
        Adds the column 'name' to the table 'table_name'.
        Uses the 'field' paramater, a django.db.models.fields.Field instance,
        to generate the necessary sql

        @param table_name: The name of the table to add the column to
        @param name: The name of the column to add
        @param field: The field to use
        """
        qn = connection.ops.quote_name
        sql = self.column_sql(table_name, name, field)
        if sql:
            params = (
                qn(table_name),
                sql,
            )
            sql = self.add_column_string % params
            self.execute(sql)

            # Now, drop the default if we need to
            if not keep_default and field.default:
                field.default = NOT_PROVIDED
                self.alter_column(table_name, name, field)

    alter_string_set_type = 'ALTER COLUMN %(column)s TYPE %(type)s'
    alter_string_set_null = 'ALTER COLUMN %(column)s DROP NOT NULL'
    alter_string_drop_null = 'ALTER COLUMN %(column)s SET NOT NULL'
    allows_combined_alters = True

    def alter_column(self, table_name, name, field):
        """
        Alters the given column name so it will match the given field.
        Note that conversion between the two by the database must be possible.

        @param table_name: The name of the table to add the column to
        @param name: The name of the column to alter
        @param field: The new field definition to use
        """

        # hook for the field to do any resolution prior to it's attributes being queried
        if hasattr(field, 'south_init'):
            field.south_init()

        qn = connection.ops.quote_name

        # First, change the type
        params = {
            "column": qn(name),
            "type": field.db_type(),
        }

        # SQLs is a list of (SQL, values) pairs.
        sqls = [(self.alter_string_set_type % params, [])]

        # Next, set any default
        if not field.null and field.has_default():
            default = field.get_default()
            sqls.append(('ALTER COLUMN %s SET DEFAULT %%s ' % (qn(name),), [default]))
        else:
            sqls.append(('ALTER COLUMN %s DROP DEFAULT' % (qn(name),), []))

        # Next, nullity
        params = {
            "column": qn(name),
            "type": field.db_type(),
        }
        if field.null:
            sqls.append((self.alter_string_set_null % params, []))
        else:
            sqls.append((self.alter_string_drop_null % params, []))

        # TODO: Unique

        if self.allows_combined_alters:
            sqls, values = zip(*sqls)
            self.execute(
                "ALTER TABLE %s %s;" % (qn(table_name), ", ".join(sqls)),
                flatten(values),
            )
        else:
            # Databases like e.g. MySQL don't like more than one alter at once.
            for sql, values in sqls:
                self.execute("ALTER TABLE %s %s;" % (qn(table_name), sql), values)

    def column_sql(self, table_name, field_name, field, tablespace=''):
        """
        Creates the SQL snippet for a column. Used by add_column and add_table.
        """
        qn = connection.ops.quote_name

        field.set_attributes_from_name(field_name)

        # hook for the field to do any resolution prior to it's attributes being queried
        if hasattr(field, 'south_init'):
            field.south_init()

        sql = field.db_type()
        if sql:
            field_output = [qn(field.column), sql]
            field_output.append('%sNULL' % (not field.null and 'NOT ' or ''))
            if field.primary_key:
                field_output.append('PRIMARY KEY')
            elif field.unique:
                field_output.append('UNIQUE')

            tablespace = field.db_tablespace or tablespace
            if tablespace and connection.features.supports_tablespaces and field.unique:
                # We must specify the index tablespace inline, because we
                # won't be generating a CREATE INDEX statement for this field.
                field_output.append(connection.ops.tablespace_sql(tablespace, inline=True))

            sql = ' '.join(field_output)
            sqlparams = ()
            # if the field is "NOT NULL" and a default value is provided, create the column with it
            # this allows the addition of a NOT NULL field to a table with existing rows
            if not field.null and field.has_default():
                default = field.get_default()
                if isinstance(default, basestring):
                    default = "'%s'" % default.replace("'", "''")
                elif isinstance(default, datetime.date):
                    default = "'%s'" % default
                sql += " DEFAULT %s"
                sqlparams = (default)

            if field.rel and self.supports_foreign_keys:
                self.add_deferred_sql(
                    self.foreign_key_sql(
                        table_name,
                        field.column,
                        field.rel.to._meta.db_table,
                        field.rel.to._meta.get_field(field.rel.field_name).column
                    )
                )

            if field.db_index and not field.unique:
                self.add_deferred_sql(self.create_index_sql(table_name, [field.column]))

        if hasattr(field, 'post_create_sql'):
            style = no_style()
            for stmt in field.post_create_sql(style, table_name):
                self.add_deferred_sql(stmt)

        if sql:
            return sql % sqlparams
        else:
            return None

    supports_foreign_keys = True

    def foreign_key_sql(self, from_table_name, from_column_name, to_table_name, to_column_name):
        """
        Generates a full SQL statement to add a foreign key constraint
        """
        constraint_name = '%s_refs_%s_%x' % (from_column_name, to_column_name, abs(hash((from_table_name, to_table_name))))
        return 'ALTER TABLE %s ADD CONSTRAINT %s FOREIGN KEY (%s) REFERENCES %s (%s)%s;' % (
            from_table_name,
            truncate_name(constraint_name, connection.ops.max_name_length()),
            from_column_name,
            to_table_name,
            to_column_name,
            connection.ops.deferrable_sql() # Django knows this
        )

    def create_index_name(self, table_name, column_names):
        """
        Generate a unique name for the index
        """
        index_unique_name = ''
        if len(column_names) > 1:
            index_unique_name = '_%x' % abs(hash((table_name, ','.join(column_names))))

        return '%s_%s%s' % (table_name, column_names[0], index_unique_name)

    def create_index_sql(self, table_name, column_names, unique=False, db_tablespace=''):
        """
        Generates a create index statement on 'table_name' for a list of 'column_names'
        """
        if not column_names:
            print "No column names supplied on which to create an index"
            return ''

        if db_tablespace and connection.features.supports_tablespaces:
            tablespace_sql = ' ' + connection.ops.tablespace_sql(db_tablespace)
        else:
            tablespace_sql = ''

        index_name = self.create_index_name(table_name, column_names)
        qn = connection.ops.quote_name
        return 'CREATE %sINDEX %s ON %s (%s)%s;' % (
            unique and 'UNIQUE ' or '',
            index_name,
            table_name,
            ','.join([qn(field) for field in column_names]),
            tablespace_sql
            )

    def create_index(self, table_name, column_names, unique=False, db_tablespace=''):
        """ Executes a create index statement """
        sql = self.create_index_sql(table_name, column_names, unique, db_tablespace)
        self.execute(sql)

    drop_index_string = 'DROP INDEX %(index_name)s'

    def delete_index(self, table_name, column_names, db_tablespace=''):
        """
        Deletes an index created with create_index.
        This is possible using only columns due to the deterministic
        index naming function which relies on column names.
        """
        name = self.create_index_name(table_name, column_names)
        sql = self.drop_index_string % {"index_name": name, "table_name": table_name}
        self.execute(sql)

    drop_index = alias('delete_index')

    def delete_column(self, table_name, name):
        """
        Deletes the column 'column_name' from the table 'table_name'.
        """
        qn = connection.ops.quote_name
        params = (qn(table_name), qn(name))
        self.execute('ALTER TABLE %s DROP COLUMN %s CASCADE;' % params, [])

    drop_column = alias('delete_column')

    def rename_column(self, table_name, old, new):
        """
        Renames the column 'old' from the table 'table_name' to 'new'.
        """
        raise NotImplementedError("rename_column has no generic SQL syntax")

    def start_transaction(self):
        """
        Makes sure the following commands are inside a transaction.
        Must be followed by a (commit|rollback)_transaction call.
        """
        if self.dry_run:
            return
        transaction.commit_unless_managed()
        transaction.enter_transaction_management()
        transaction.managed(True)

    def commit_transaction(self):
        """
        Commits the current transaction.
        Must be preceded by a start_transaction call.
        """
        if self.dry_run:
            return
        transaction.commit()
        transaction.leave_transaction_management()

    def rollback_transaction(self):
        """
        Rolls back the current transaction.
        Must be preceded by a start_transaction call.
        """
        if self.dry_run:
            return
        transaction.rollback()
        transaction.leave_transaction_management()

    def send_create_signal(self, app_label, model_names):
        """
        Sends a post_syncdb signal for the model specified.

        If the model is not found (perhaps it's been deleted?),
        no signal is sent.

        TODO: The behavior of django.contrib.* apps seems flawed in that
        they don't respect created_models.  Rather, they blindly execute
        over all models within the app sending the signal.  This is a
        patch we should push Django to make  For now, this should work.
        """
        app = models.get_app(app_label)
        if not app:
            return

        created_models = []
        for model_name in model_names:
            model = models.get_model(app_label, model_name)
            if model:
                created_models.append(model)

        if created_models:
            # syncdb defaults -- perhaps take these as options?
            verbosity = 1
            interactive = True

            if hasattr(dispatcher, "send"):
                dispatcher.send(signal=models.signals.post_syncdb, sender=app,
                app=app, created_models=created_models,
                verbosity=verbosity, interactive=interactive)
            else:
                models.signals.post_syncdb.send(sender=app,
                app=app, created_models=created_models,
                verbosity=verbosity, interactive=interactive)

    def mock_model(self, model_name, db_table, db_tablespace='',
                    pk_field_name='id', pk_field_type=models.AutoField,
                    pk_field_args=[], pk_field_kwargs={}):
        """
        Generates a MockModel class that provides enough information
        to be used by a foreign key/many-to-many relationship.

        Migrations should prefer to use these rather than actual models
        as models could get deleted over time, but these can remain in
        migration files forever.
        """
        class MockOptions(object):
            def __init__(self):
                self.db_table = db_table
                self.db_tablespace = db_tablespace or settings.DEFAULT_TABLESPACE
                self.object_name = model_name
                self.module_name = model_name.lower()

                if pk_field_type == models.AutoField:
                    pk_field_kwargs['primary_key'] = True

                self.pk = pk_field_type(*pk_field_args, **pk_field_kwargs)
                self.pk.set_attributes_from_name(pk_field_name)
                self.abstract = False

            def get_field_by_name(self, field_name):
                # we only care about the pk field
                return (self.pk, self.model, True, False)

            def get_field(self, name):
                # we only care about the pk field
                return self.pk

        class MockModel(object):
            _meta = None

        # We need to return an actual class object here, not an instance
        MockModel._meta = MockOptions()
        MockModel._meta.model = MockModel
        return MockModel

# Single-level flattening of lists
def flatten(ls):
    nl = []
    for l in ls:
        nl += l
    return nl

<|endoftext|>


<commit_before>"""Utilities using NDG HTTPS Client, including a main module that can be used to
fetch from a URL.
"""
__author__ = "R B Wilkinson"
__date__ = "09/12/11"
__copyright__ = "(C) 2011 Science and Technology Facilities Council"
__license__ = "BSD - see LICENSE file in top-level directory"
__contact__ = "Philip.Kershaw@stfc.ac.uk"
__revision__ = '$Id$'

import cookielib
import httplib
import logging
from optparse import OptionParser
import os
import urllib2
from urllib2 import HTTPHandler, HTTPCookieProcessor

import urlparse

from ndg.httpsclient.urllib2_build_opener import build_opener
from ndg.httpsclient.https import HTTPSContextHandler
from ndg.httpsclient import ssl_context_util

log = logging.getLogger(__name__)

class AccumulatingHTTPCookieProcessor(HTTPCookieProcessor):
    """Cookie processor that adds new cookies (instead of replacing the existing
    ones as HTTPCookieProcessor does)
    """
    def http_request(self, request):
        """Processes cookies for a HTTP request.
        @param request: request to process
        @type request: urllib2.Request
        @return: request
        @rtype: urllib2.Request
        """
        COOKIE_HEADER_NAME = "Cookie"
        tmp_request = urllib2.Request(request.get_full_url(), request.data, {},
                                      request.origin_req_host,
                                      request.unverifiable)
        self.cookiejar.add_cookie_header(tmp_request)
        # Combine existing and new cookies.
        new_cookies = tmp_request.get_header(COOKIE_HEADER_NAME)
        if new_cookies:
            if request.has_header(COOKIE_HEADER_NAME):
                # Merge new cookies with existing ones.
                old_cookies = request.get_header(COOKIE_HEADER_NAME)
                merged_cookies = '; '.join([old_cookies, new_cookies])
                request.add_unredirected_header(COOKIE_HEADER_NAME,
                                                merged_cookies)
            else:
                # No existing cookies so just set new ones.
                request.add_unredirected_header(COOKIE_HEADER_NAME, new_cookies)
        return request

    # Process cookies for HTTPS in the same way.
    https_request = http_request

class URLFetchError(Exception):
    """Error fetching content from URL"""

def fetch_from_url(url, config, data=None):
    """Returns data retrieved from a URL.
    @param url: URL to attempt to open
    @type url: basestring
    @param config: SSL context configuration
    @type config: Configuration
    @return data retrieved from URL or None
    """
    return_code, return_message, response = open_url(url, config, data)
    if return_code and return_code == httplib.OK:
        return_data = response.read()
        response.close()
        return return_data
    else:
        raise URLFetchError(return_message)

def fetch_from_url_to_file(url, config, output_file, data=None):
    """Writes data retrieved from a URL to a file.
    @param url: URL to attempt to open
    @type url: basestring
    @param config: SSL context configuration
    @type config: Configuration
    @param output_file: output file
    @type output_file: basestring
    @return: tuple (
        returned HTTP status code or 0 if an error occurred
        returned message
        boolean indicating whether access was successful)
    """
    return_code, return_message, response = open_url(url, config, data)
    if return_code == httplib.OK:
        return_data = response.read()
        response.close()
        outfile = open(output_file, "w")
        outfile.write(return_data)
        outfile.close()
    return return_code, return_message, return_code == httplib.OK

def fetch_stream_from_url(url, config, data=None):
    """Returns data retrieved from a URL.
    @param url: URL to attempt to open
    @type url: basestring
    @param config: SSL context configuration
    @type config: Configuration
    @return: data retrieved from URL or None
    @rtype: file derived type
    """
    return_code, return_message, response = open_url(url, config, data)
    if return_code and return_code == httplib.OK:
        return response
    else:
        raise URLFetchError(return_message)

def open_url(url, config, data=None):
    """Attempts to open a connection to a specified URL.
    @param url: URL to attempt to open
    @param config: SSL context configuration
    @type config: Configuration
    @param data: HTTP POST data
    @type data: str
    @return: tuple (
        returned HTTP status code or 0 if an error occurred
        returned message or error description
        response object)
    """
    debuglevel = 1 if config.debug else 0

    # Set up handlers for URL opener.
    if config.cookie:
        cj = config.cookie
    else:
        cj = cookielib.CookieJar()
    # Use a cookie processor that accumulates cookies when redirects occur so
    # that an application can redirect for authentication and retain both any
    # cookies for the application and the security system (c.f.,
    # urllib2.HTTPCookieProcessor which replaces cookies).
    cookie_handler = AccumulatingHTTPCookieProcessor(cj)

    handlers = [cookie_handler]

    if config.debug:
        http_handler = HTTPHandler(debuglevel=debuglevel)
        https_handler = HTTPSContextHandler(config.ssl_context,
                                            debuglevel=debuglevel)
        handlers.extend([http_handler, https_handler])

    # Explicitly remove proxy handling if the host is one listed in the value of
    # the no_proxy environment variable because urllib2 does use proxy settings
    # set via http_proxy and https_proxy, but does not take the no_proxy value
    # into account.
    if not _should_use_proxy(url, config.no_proxy):
        handlers.append(urllib2.ProxyHandler({}))
        log.debug("Not using proxy")
    elif config.proxies:
        handlers.append(urllib2.ProxyHandler(config.proxies))
        log.debug("Configuring proxies: %s" % config.proxies)

    opener = build_opener(*handlers, ssl_context=config.ssl_context)

    # Open the URL and check the response.
    return_code = 0
    return_message = ''
    response = None
    try:
        response = opener.open(url, data)
        return_message = response.msg
        return_code = response.code
        if log.isEnabledFor(logging.DEBUG):
            for index, cookie in enumerate(cj):
                log.debug("%s  :  %s", index, cookie)
    except urllib2.HTTPError, exc:
        return_code = exc.code
        return_message = "Error: %s" % exc.msg
        if log.isEnabledFor(logging.DEBUG):
            log.debug("%s %s", exc.code, exc.msg)
    except Exception, exc:
        return_message = "Error: %s" % exc.__str__()
        if log.isEnabledFor(logging.DEBUG):
            import traceback
            log.debug(traceback.format_exc())
    return (return_code, return_message, response)

def _should_use_proxy(url, no_proxy=None):
    """Determines whether a proxy should be used to open a connection to the
    specified URL, based on the value of the no_proxy environment variable.
    @param url: URL
    @type url: basestring or urllib2.Request
    """
    if no_proxy is None:
        no_proxy_effective = os.environ.get('no_proxy', '')
    else:
        no_proxy_effective = no_proxy

    urlObj = urlparse.urlparse(_url_as_string(url))
    for np in [h.strip() for h in no_proxy_effective.split(',')]:
        if urlObj.hostname == np:
            return False

    return True

def _url_as_string(url):
    """Returns the URL string from a URL value that is either a string or
    urllib2.Request..
    @param url: URL
    @type url: basestring or urllib2.Request
    @return: URL string
    @rtype: basestring
    """
    if isinstance(url, urllib2.Request):
        return url.get_full_url()
    elif isinstance(url, basestring):
        return url
    else:
        raise TypeError("Expected type %r or %r" %
                        (basestring, urllib2.Request))

class Configuration(object):
    """Connection configuration.
    """
    def __init__(self, ssl_context, debug=False, proxies=None, no_proxy=None,
                 cookie=None):
        """
        @param ssl_context: SSL context to use with this configuration
        @type ssl_context: OpenSSL.SSL.Context
        @param debug: if True, output debugging information
        @type debug: bool
        @param proxies: proxies to use for
        @type proxies: dict with basestring keys and values
        @param no_proxy: hosts for which a proxy should not be used
        @type no_proxy: basestring
        @param cookie: cookies to set for request
        @type cookie: cookielib.CookieJar
        """
        self.ssl_context = ssl_context
        self.debug = debug
        self.proxies = proxies
        self.no_proxy = no_proxy
        self.cookie = cookie

def main():
    '''Utility to fetch data using HTTP or HTTPS GET from a specified URL.
    '''
    parser = OptionParser(usage="%prog [options] url")
    parser.add_option("-c", "--certificate", dest="cert_file", metavar="FILE",
                      default=os.path.expanduser("~/credentials.pem"),
                      help="Certificate file - defaults to $HOME/credentials.pem")
    parser.add_option("-k", "--private-key", dest="key_file", metavar="FILE",
                      default=None,
                      help="Private key file - defaults to the certificate file")
    parser.add_option("-t", "--ca-certificate-dir", dest="ca_dir",
                      metavar="PATH",
                      default=None,
                      help="Trusted CA certificate file directory")
    parser.add_option("-d", "--debug", action="store_true", dest="debug",
                      default=False,
                      help="Print debug information.")
    parser.add_option("-p", "--post-data-file", dest="data_file",
                      metavar="FILE", default=None,
                      help="POST data file")
    parser.add_option("-f", "--fetch", dest="output_file", metavar="FILE",
                      default=None, help="Output file")
    parser.add_option("-n", "--no-verify-peer", action="store_true",
                      dest="no_verify_peer", default=False,
                      help="Skip verification of peer certificate.")
    (options, args) = parser.parse_args()
    if len(args) != 1:
        parser.error("Incorrect number of arguments")

    url = args[0]

    if options.debug:
        logging.getLogger().setLevel(logging.DEBUG)

    if options.key_file and os.path.exists(options.key_file):
        key_file = options.key_file
    else:
        key_file = None

    if options.cert_file and os.path.exists(options.cert_file):
        cert_file = options.cert_file
    else:
        cert_file = None

    if options.ca_dir and os.path.exists(options.ca_dir):
        ca_dir = options.ca_dir
    else:
        ca_dir = None

    verify_peer = not options.no_verify_peer

    if options.data_file and os.path.exists(options.data_file):
        data_file = open(options.data_file)
        data = data_file.read()
        data_file.close()
    else:
        data = None

    # If a private key file is not specified, the key is assumed to be stored in
    # the certificate file.
    ssl_context = ssl_context_util.make_ssl_context(key_file,
                                                    cert_file,
                                                    None,
                                                    ca_dir,
                                                    verify_peer,
                                                    url)

    config = Configuration(ssl_context, options.debug)
    if options.output_file:
        return_code, return_message = fetch_from_url_to_file(url,
                                                      config,
                                                      options.output_file,
                                                      data)[:2]
        raise SystemExit(return_code, return_message)
    else:
        data = fetch_from_url(url, config)
        print(data)

if __name__=='__main__':
    logging.basicConfig()
    main()
<commit_msg>Extended utils funcs so that they can accept a handlers kw<commit_after>"""Utilities using NDG HTTPS Client, including a main module that can be used to
fetch from a URL.
"""
__author__ = "R B Wilkinson"
__date__ = "09/12/11"
__copyright__ = "(C) 2011 Science and Technology Facilities Council"
__license__ = "BSD - see LICENSE file in top-level directory"
__contact__ = "Philip.Kershaw@stfc.ac.uk"
__revision__ = '$Id$'

import cookielib
import httplib
import logging
from optparse import OptionParser
import os
import urllib2
from urllib2 import HTTPHandler, HTTPCookieProcessor

import urlparse

from ndg.httpsclient.urllib2_build_opener import build_opener
from ndg.httpsclient.https import HTTPSContextHandler
from ndg.httpsclient import ssl_context_util

log = logging.getLogger(__name__)

class AccumulatingHTTPCookieProcessor(HTTPCookieProcessor):
    """Cookie processor that adds new cookies (instead of replacing the existing
    ones as HTTPCookieProcessor does)
    """
    def http_request(self, request):
        """Processes cookies for a HTTP request.
        @param request: request to process
        @type request: urllib2.Request
        @return: request
        @rtype: urllib2.Request
        """
        COOKIE_HEADER_NAME = "Cookie"
        tmp_request = urllib2.Request(request.get_full_url(), request.data, {},
                                      request.origin_req_host,
                                      request.unverifiable)
        self.cookiejar.add_cookie_header(tmp_request)
        # Combine existing and new cookies.
        new_cookies = tmp_request.get_header(COOKIE_HEADER_NAME)
        if new_cookies:
            if request.has_header(COOKIE_HEADER_NAME):
                # Merge new cookies with existing ones.
                old_cookies = request.get_header(COOKIE_HEADER_NAME)
                merged_cookies = '; '.join([old_cookies, new_cookies])
                request.add_unredirected_header(COOKIE_HEADER_NAME,
                                                merged_cookies)
            else:
                # No existing cookies so just set new ones.
                request.add_unredirected_header(COOKIE_HEADER_NAME, new_cookies)
        return request

    # Process cookies for HTTPS in the same way.
    https_request = http_request

class URLFetchError(Exception):
    """Error fetching content from URL"""

def fetch_from_url(url, config, data=None, handlers=None):
    """Returns data retrieved from a URL.
    @param url: URL to attempt to open
    @type url: basestring
    @param config: SSL context configuration
    @type config: Configuration
    @return data retrieved from URL or None
    """
    return_code, return_message, response = open_url(url, config, data=data,
                                                     handlers=handlers)
    if return_code and return_code == httplib.OK:
        return_data = response.read()
        response.close()
        return return_data
    else:
        raise URLFetchError(return_message)

def fetch_from_url_to_file(url, config, output_file, data=None, handlers=None):
    """Writes data retrieved from a URL to a file.
    @param url: URL to attempt to open
    @type url: basestring
    @param config: SSL context configuration
    @type config: Configuration
    @param output_file: output file
    @type output_file: basestring
    @return: tuple (
        returned HTTP status code or 0 if an error occurred
        returned message
        boolean indicating whether access was successful)
    """
    return_code, return_message, response = open_url(url, config, data=data,
                                                     handlers=handlers)
    if return_code == httplib.OK:
        return_data = response.read()
        response.close()
        outfile = open(output_file, "w")
        outfile.write(return_data)
        outfile.close()
    return return_code, return_message, return_code == httplib.OK

def fetch_stream_from_url(url, config, data=None, handlers=None):
    """Returns data retrieved from a URL.
    @param url: URL to attempt to open
    @type url: basestring
    @param config: SSL context configuration
    @type config: Configuration
    @return: data retrieved from URL or None
    @rtype: file derived type
    """
    return_code, return_message, response = open_url(url, config, data=data,
                                                     handlers=handlers)
    if return_code and return_code == httplib.OK:
        return response
    else:
        raise URLFetchError(return_message)

def open_url(url, config, data=None, handlers=None):
    """Attempts to open a connection to a specified URL.
    @param url: URL to attempt to open
    @param config: SSL context configuration
    @type config: Configuration
    @param data: HTTP POST data
    @type data: str
    @return: tuple (
        returned HTTP status code or 0 if an error occurred
        returned message or error description
        response object)
    """
    debuglevel = 1 if config.debug else 0

    # Set up handlers for URL opener.
    if config.cookie:
        cj = config.cookie
    else:
        cj = cookielib.CookieJar()
    # Use a cookie processor that accumulates cookies when redirects occur so
    # that an application can redirect for authentication and retain both any
    # cookies for the application and the security system (c.f.,
    # urllib2.HTTPCookieProcessor which replaces cookies).
    cookie_handler = AccumulatingHTTPCookieProcessor(cj)

    if not handlers:
        handlers = []

    handlers.append(cookie_handler)

    if config.debug:
        http_handler = HTTPHandler(debuglevel=debuglevel)
        https_handler = HTTPSContextHandler(config.ssl_context,
                                            debuglevel=debuglevel)
        handlers.extend([http_handler, https_handler])

    # Explicitly remove proxy handling if the host is one listed in the value of
    # the no_proxy environment variable because urllib2 does use proxy settings
    # set via http_proxy and https_proxy, but does not take the no_proxy value
    # into account.
    if not _should_use_proxy(url, config.no_proxy):
        handlers.append(urllib2.ProxyHandler({}))
        log.debug("Not using proxy")
    elif config.proxies:
        handlers.append(urllib2.ProxyHandler(config.proxies))
        log.debug("Configuring proxies: %s" % config.proxies)

    opener = build_opener(*handlers, ssl_context=config.ssl_context)

    # Open the URL and check the response.
    return_code = 0
    return_message = ''
    response = None
    try:
        response = opener.open(url, data)
        return_message = response.msg
        return_code = response.code
        if log.isEnabledFor(logging.DEBUG):
            for index, cookie in enumerate(cj):
                log.debug("%s  :  %s", index, cookie)
    except urllib2.HTTPError, exc:
        return_code = exc.code
        return_message = "Error: %s" % exc.msg
        if log.isEnabledFor(logging.DEBUG):
            log.debug("%s %s", exc.code, exc.msg)
    except Exception, exc:
        return_message = "Error: %s" % exc.__str__()
        if log.isEnabledFor(logging.DEBUG):
            import traceback
            log.debug(traceback.format_exc())
    return (return_code, return_message, response)

def _should_use_proxy(url, no_proxy=None):
    """Determines whether a proxy should be used to open a connection to the
    specified URL, based on the value of the no_proxy environment variable.
    @param url: URL
    @type url: basestring or urllib2.Request
    """
    if no_proxy is None:
        no_proxy_effective = os.environ.get('no_proxy', '')
    else:
        no_proxy_effective = no_proxy

    urlObj = urlparse.urlparse(_url_as_string(url))
    for np in [h.strip() for h in no_proxy_effective.split(',')]:
        if urlObj.hostname == np:
            return False

    return True

def _url_as_string(url):
    """Returns the URL string from a URL value that is either a string or
    urllib2.Request..
    @param url: URL
    @type url: basestring or urllib2.Request
    @return: URL string
    @rtype: basestring
    """
    if isinstance(url, urllib2.Request):
        return url.get_full_url()
    elif isinstance(url, basestring):
        return url
    else:
        raise TypeError("Expected type %r or %r" %
                        (basestring, urllib2.Request))

class Configuration(object):
    """Connection configuration.
    """
    def __init__(self, ssl_context, debug=False, proxies=None, no_proxy=None,
                 cookie=None):
        """
        @param ssl_context: SSL context to use with this configuration
        @type ssl_context: OpenSSL.SSL.Context
        @param debug: if True, output debugging information
        @type debug: bool
        @param proxies: proxies to use for
        @type proxies: dict with basestring keys and values
        @param no_proxy: hosts for which a proxy should not be used
        @type no_proxy: basestring
        @param cookie: cookies to set for request
        @type cookie: cookielib.CookieJar
        """
        self.ssl_context = ssl_context
        self.debug = debug
        self.proxies = proxies
        self.no_proxy = no_proxy
        self.cookie = cookie

def main():
    '''Utility to fetch data using HTTP or HTTPS GET from a specified URL.
    '''
    parser = OptionParser(usage="%prog [options] url")
    parser.add_option("-c", "--certificate", dest="cert_file", metavar="FILE",
                      default=os.path.expanduser("~/credentials.pem"),
                      help="Certificate file - defaults to $HOME/credentials.pem")
    parser.add_option("-k", "--private-key", dest="key_file", metavar="FILE",
                      default=None,
                      help="Private key file - defaults to the certificate file")
    parser.add_option("-t", "--ca-certificate-dir", dest="ca_dir",
                      metavar="PATH",
                      default=None,
                      help="Trusted CA certificate file directory")
    parser.add_option("-d", "--debug", action="store_true", dest="debug",
                      default=False,
                      help="Print debug information.")
    parser.add_option("-p", "--post-data-file", dest="data_file",
                      metavar="FILE", default=None,
                      help="POST data file")
    parser.add_option("-f", "--fetch", dest="output_file", metavar="FILE",
                      default=None, help="Output file")
    parser.add_option("-n", "--no-verify-peer", action="store_true",
                      dest="no_verify_peer", default=False,
                      help="Skip verification of peer certificate.")
    (options, args) = parser.parse_args()
    if len(args) != 1:
        parser.error("Incorrect number of arguments")

    url = args[0]

    if options.debug:
        logging.getLogger().setLevel(logging.DEBUG)

    if options.key_file and os.path.exists(options.key_file):
        key_file = options.key_file
    else:
        key_file = None

    if options.cert_file and os.path.exists(options.cert_file):
        cert_file = options.cert_file
    else:
        cert_file = None

    if options.ca_dir and os.path.exists(options.ca_dir):
        ca_dir = options.ca_dir
    else:
        ca_dir = None

    verify_peer = not options.no_verify_peer

    if options.data_file and os.path.exists(options.data_file):
        data_file = open(options.data_file)
        data = data_file.read()
        data_file.close()
    else:
        data = None

    # If a private key file is not specified, the key is assumed to be stored in
    # the certificate file.
    ssl_context = ssl_context_util.make_ssl_context(key_file,
                                                    cert_file,
                                                    None,
                                                    ca_dir,
                                                    verify_peer,
                                                    url)

    config = Configuration(ssl_context, options.debug)
    if options.output_file:
        return_code, return_message = fetch_from_url_to_file(url,
                                                      config,
                                                      options.output_file,
                                                      data)[:2]
        raise SystemExit(return_code, return_message)
    else:
        data = fetch_from_url(url, config)
        print(data)

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


<commit_before>from pymatgen import Molecule
import numpy as np
from pymatgen.io.babel import BabelMolAdaptor
import logging
try:
    import pybel as pb
except:
    logging.warning("Openbabel is needed for molecule models, try 'conda install -c openbabel openbabel' to install it")
    pb = None

class MEGNetMolecule(Molecule):
    def get_all_neighbors(self, cutoff, include_index=True, include_site=False, include_image=True, **kwargs):
        """

        Args:
            cutoff: float, cutoff radius
            include_index: bool, whether to include the site index
            include_site: bool, whether to include site
            include_image: bool, whether to include dummy image

        Returns:
            list of list, neighbors for each site
        """

        dist = self.distance_matrix
        neighbors = []
        for i in dist:
            cond = np.bitwise_and(i <= cutoff, i > 1e-8)
            inds = np.array(np.where(cond)[0], dtype='int')
            d = i[cond]
            images = [(0, 0, 0)] * len(d)
            neighbor = []
            for k, l, m in zip(d, inds, images):
                item = []
                if include_site:
                    item.append(self[l])
                item.append(k)
                if include_index:
                    item += [l]
                if include_image:
                    item += [m]
                neighbor.append(item)
            neighbors.append(neighbor)
        return neighbors

    @classmethod
    def from_pymatgen(cls, mol):
        sites = mol._sites
        return cls.from_sites(sites)

def get_pmg_mol_from_smiles(smiles):
    """
    Get a pymatgen molecule from smiles representation
    Args:
        smiles: (str) smiles representation of molecule
    """
    b_mol = pb.readstring('smi', smiles)
    b_mol.make3D()
    b_mol = b_mol.OBMol
    p_mol = BabelMolAdaptor(b_mol).pymatgen_mol
    return p_mol
<commit_msg>fix MEGNetMolecule<commit_after>from pymatgen import Molecule
import numpy as np
from pymatgen.io.babel import BabelMolAdaptor
import logging
try:
    import pybel as pb
except:
    logging.warning("Openbabel is needed for molecule models, try 'conda install -c openbabel openbabel' to install it")
    pb = None

class MEGNetMolecule(Molecule):
    def get_all_neighbors(self, cutoff, include_index=True, include_image=True, **kwargs):
        """

        Args:
            cutoff: float, cutoff radius
            include_index: bool, whether to include the site index
            include_image: bool, whether to include dummy image

        Returns:
            list of list, neighbors for each site
        """

        dist = self.distance_matrix
        neighbors = []
        for i in dist:
            cond = np.bitwise_and(i <= cutoff, i > 1e-8)
            inds = np.array(np.where(cond)[0], dtype='int')
            d = i[cond]
            images = [(0, 0, 0)] * len(d)
            neighbor = []
            for k, l, m in zip(d, inds, images):
                item = []
                item.append(self[l])
                item.append(k)
                if include_index:
                    item += [l]
                if include_image:
                    item += [m]
                neighbor.append(item)
            neighbors.append(neighbor)
        return neighbors

    @classmethod
    def from_pymatgen(cls, mol):
        sites = mol._sites
        return cls.from_sites(sites)

def get_pmg_mol_from_smiles(smiles):
    """
    Get a pymatgen molecule from smiles representation
    Args:
        smiles: (str) smiles representation of molecule
    """
    b_mol = pb.readstring('smi', smiles)
    b_mol.make3D()
    b_mol = b_mol.OBMol
    p_mol = BabelMolAdaptor(b_mol).pymatgen_mol
    return p_mol
<|endoftext|>


<commit_before>from gitmostwanted.app import db

class User(db.Model):
    __tablename__ = 'users'
    __table_args__ = (
        db.UniqueConstraint('email', name='ix_email'),
        db.UniqueConstraint('github_id', name='ix_github_id')
    )

    id = db.Column(db.Integer, primary_key=True)
    email = db.Column(db.String(120))
    username = db.Column(db.String(80))
    github_id = db.Column(db.BigInteger)

class UserAttitude(db.Model):
    __tablename__ = 'users_attitude'

    user_id = db.Column(
        db.Integer,
        db.ForeignKey('users.id', name='fk_users_id', ondelete='CASCADE'),
        primary_key=True
    )
    repo_id = db.Column(
        db.BigInteger,
        db.ForeignKey('repos.id', name='fk_repos_id', ondelete='CASCADE'),
        primary_key=True, index=True
    )
    attitude = db.Column(db.Enum('dislike', 'neutral'), nullable=False)
<commit_msg>The UserAttitude model does not have the like variant, closes #68<commit_after>from gitmostwanted.app import db

class User(db.Model):
    __tablename__ = 'users'
    __table_args__ = (
        db.UniqueConstraint('email', name='ix_email'),
        db.UniqueConstraint('github_id', name='ix_github_id')
    )

    id = db.Column(db.Integer, primary_key=True)
    email = db.Column(db.String(120))
    username = db.Column(db.String(80))
    github_id = db.Column(db.BigInteger)

class UserAttitude(db.Model):
    __tablename__ = 'users_attitude'

    user_id = db.Column(
        db.Integer,
        db.ForeignKey('users.id', name='fk_users_id', ondelete='CASCADE'),
        primary_key=True
    )
    repo_id = db.Column(
        db.BigInteger,
        db.ForeignKey('repos.id', name='fk_repos_id', ondelete='CASCADE'),
        primary_key=True, index=True
    )
    attitude = db.Column(db.Enum('like', 'dislike', 'neutral'), nullable=False)
<|endoftext|>


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

#/* ***** BEGIN LICENSE BLOCK	****
# Version: MPL 1.1/GPL 2.0/LGPL 2.1
#
# The contents of this file are subject to the Mozilla Public License Version
# 1.1 (the "License"); you may not use this file except in compliance with
# the License. You may obtain a copy of the License at
# http://www.mozilla.org/MPL/
#
# Software distributed under the License is distributed on an "AS IS" basis,
# WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
# for the specific language governing rights and limitations under the
# License.
#
# The Original Code is autoBisect.
#
# The Initial Developer of the Original Code is
# Gary Kwong.
# Portions created by the Initial Developer are Copyright (C) 2006-2008
# the Initial Developer. All Rights Reserved.
#
# Contributor(s):
#
# Alternatively, the contents of this file may be used under the terms of
# either the GNU General Public License Version 2 or later (the "GPL"), or
# the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
# in which case the provisions of the GPL or the LGPL are applicable instead
# of those above. If you wish to allow use of your version of this file only
# under the terms of either the GPL or the LGPL, and not to allow others to
# use your version of this file under the terms of the MPL, indicate your
# decision by deleting the provisions above and replace them with the notice
# and other provisions required by the GPL or the LGPL. If you do not delete
# the provisions above, a recipient may use your version of this file under
# the terms of any one of the MPL, the GPL or the LGPL.
#
# * ***** END LICENSE BLOCK	****	/

import os, shutil, subprocess, sys
from optparse import OptionParser

sys.path.append('../jsfunfuzz/')
from fnStartjsfunfuzz import *

def main():
    # Do not support Windows XP because the ~ folder is in "/Documents and Settings/",
    # which contains spaces. This breaks MinGW, which is what MozillaBuild uses.
    # From Windows Vista onwards, the folder is in "/Users/".
    if os.name == 'nt':
        if '5.1' in captureStdout('uname'):
            raise Exception('autoBisect is not supported on Windows XP.')
    verbose = True

    # Parse options and parameters from the command-line.
    filename = sys.argv[-1:][0]
    options = parseOpts()
    (bugOrWfm, compileType, sourceDir, stdoutOutput, resetBool, startRepo, \
     endRepo, archNum, tracingjitBool, methodjitBool, watchExitCode) = options

    os.chdir(sourceDir)

    # Refresh source directory (overwrite all local changes) to default tip if required.
    if resetBool:
        subprocess.call(['hg up -C default'], shell=True)

    # Reverse if a bisect range to find a WFM issue is specified.
    if bugOrWfm == 'wfm':
        tempVal = ''
        tempVal = startRepo
        startRepo = endRepo
        endRepo = tempVal

    # Specify `hg bisect` ranges.
    subprocess.call(['hg bisect -r'], shell=True)
    # If in "bug" mode, this startRepo changeset does not exhibit the issue.
    subprocess.call(['hg bisect -g ' + startRepo], shell=True)
    # If in "bug" mode, this endRepo changeset exhibits the issue.
    stdoutNumOfTests = captureStdout('hg bisect -b ' + endRepo)

    # Find out the number of tests to be executed based on the initial hg bisect output.
    numOfTests = checkNumOfTests(stdoutNumOfTests)

    # For main directory, change into Desktop directory.
    mainDir = os.path.expanduser('~/Desktop/')

    # Change into main directory.
    os.chdir(mainDir)

    for i in xrange(numOfTests):
        autoBisectPath = 'autoBisect-' + compileType + '-' + archNum + '-s' + \
                         startRepo + '-e' + endRepo

        # Create the autoBisect folder.
        try:
            os.makedirs(autoBisectPath)
        except OSError:
            raise Exception('The autoBisect path at "' + autoBisectPath + '" already exists!')

        # Change into autoBisectPath.
        os.chdir(autoBisectPath)

        # Copy the js tree to the autoBisect path.
        # Don't use pymake because older changesets may fail to compile.
        cpJsTreeOrPymakeDir(os.path.expanduser(sourceDir), 'js')
        os.chdir('compilePath')  # Change into compilation directory.

        autoconfRun()

        # Create objdirs within the compilePaths.
        os.mkdir(compileType + '-objdir')
        os.chdir(compileType + '-objdir')

        # Compile the first binary.
        branchType = 'autoBisectBranch'
        if 'jaeger' in sourceDir:
            branchType = 'jm'
        valgrindSupport = False  # Let's disable support for valgrind in the js shell
        threadsafe = False  # Let's disable support for threadsafety in the js shell
        configureJsBinary(archNum, compileType, branchType, valgrindSupport, threadsafe)

        if 'jaeger' in sourceDir:
            branchType = 'autoBisectBranch'  # Reset the branchType

        # Compile and copy the first binary.
        try:
            jsShellName = compileCopy(archNum, compileType, branchType, False)
        except:
            print 'The "good" repository that is currently labelled:', startRepo
            print 'The "bad" repository that is currently labelled:', endRepo
            raise Exception('Compilation failed.')

        # Change back into compilePath.
        os.chdir('../')

        # Test compilePath.
        if verbose:
            print 'DEBUG - This should be the compilePath:'
            print 'DEBUG - %s\n' % os.getcwdu()
            if 'compilePath' not in os.getcwdu():
                raise Exception('We are not in compilePath.')

        os.chdir('../../')  # Change into autoBisectPath directory.
        autoBisectFullPath = os.path.expanduser(os.getcwdu())

        (stdoutStderr, exitCode) = testBinary(jsShellName, filename,
                                              methodjitBool, tracingjitBool)

        # Switch to hg repository directory.
        os.chdir(os.path.expanduser(sourceDir))

        # Label the changeset bad if the exact assert is found (only in debug shells)
        # (Assuming "bad" and not "wfm".)
        # Assertion exit codes: Mac 10.5/10.6 - 133, Linux - 134, WinXP - 3
        if compileType == 'dbg' and stdoutOutput in stdoutStderr and exitCode != 0:
            # Set a random arbitrary value that cannot be a genuine exit code.
            codeToBeObserved = 888888
            if os.name == 'posix':
                if os.uname()[0] == 'Darwin':
                    codeToBeObserved = 133
                elif os.uname()[0] == 'Linux':
                    codeToBeObserved = 134
            elif os.name == 'nt':
                codeToBeObserved = 3

            # Depending on the exit code as per the different platforms, specify
            # if the current changeset is "good" or "bad".
            if exitCode == codeToBeObserved:
                (result, startRepo, endRepo) = bisectLabel(bugOrWfm, 'bad', startRepo, endRepo)

                print 'Now removing autoBisectFullPath, located at:', autoBisectFullPath
                shutil.rmtree(autoBisectFullPath)
                if 'first bad revision' in result:
                    break

        # "Bad" changesets.
        elif (exitCode == 1) or (129 <= exitCode <= 159) or (exitCode == watchExitCode):
            (result, startRepo, endRepo) = bisectLabel(bugOrWfm, 'bad', startRepo, endRepo)

            print 'Now removing autoBisectFullPath, located at:', autoBisectFullPath
            shutil.rmtree(autoBisectFullPath)
            if 'first bad revision' in result:
                break

        # "Good" changesets.
        elif exitCode != watchExitCode:
            if (exitCode == 0) or (3 <= exitCode <= 6):
                (result, startRepo, endRepo) = bisectLabel(bugOrWfm, 'good', startRepo, endRepo)

                print 'Now removing autoBisectFullPath, located at:', autoBisectFullPath
                shutil.rmtree(autoBisectFullPath)
                if 'first bad revision' in result:
                    break

        else:
            raise Exception('Unknown exit code hit:', exitCode)

    # Reset `hg bisect` after finishing everything.
    subprocess.call(['hg bisect -r'], shell=True)

def parseOpts():
    usage = 'Usage: %prog [options] filename'
    parser = OptionParser(usage)
    # See http://docs.python.org/library/optparse.html#optparse.OptionParser.disable_interspersed_args
    parser.disable_interspersed_args()

    # autoBisect details
    parser.add_option('-b', '--bugOrWfm',
                      dest='bugOrWfm',
                      type='choice',
                      choices=['bug', 'wfm'],
                      default='bug',
                      help='Bisect to find a bug or WFM issue. ' + \
                           'Only accepts "bug" or "wfm". ' + \
                           'Defaults to "bug"')
    parser.add_option('-c', '--compileType',
                      dest='compileType',
                      type='choice',
                      choices=['dbg', 'opt'],
                      default='dbg',
                      help='js shell compile type. Defaults to "dbg"')
    parser.add_option('-d', '--dir',
                      dest='dir',
                      default=os.path.expanduser('~/tracemonkey/'),
                      help='Source code directory. Defaults to "~/tracemonkey/"')
    parser.add_option('-o', '--output',
                      dest='output',
                      help='Stdout or stderr output to be observed. ' + \
                           'For assertions, set to "ssertion fail"')
    parser.add_option('-r', '--resetToTipFirstBool',
                      dest='resetBool',
                      action='store_true',
                      default=False,
                      help='First reset to default tip overwriting all local changes. ' + \
                           'Equivalent to first executing `hg update -C default`. ' + \
                           'Defaults to "False"')

    # Define the start and end repositories.
    parser.add_option('-s', '--start',
                      dest='startRepo',
                      help='Start repository (earlier)')
    parser.add_option('-e', '--end',
                      dest='endRepo',
                      default='tip',
                      help='End repository (later). Defaults to "tip"')

    # Define the architecture to be tested.
    parser.add_option('-a', '--architecture',
                      dest='archi',
                      type='choice',
                      choices=['32', '64'],
                      help='Test architecture. Only accepts "32" or "64"')

    # Define parameters to be passed to the binary.
    parser.add_option('-j', '--tracingjit',
                      dest='tracingjitBool',
                      action='store_true',
                      default=False,
                      help='Enable -j, tracing JIT when autoBisecting. Defaults to "False"')
    parser.add_option('-m', '--methodjit',
                      dest='methodjitBool',
                      action='store_true',
                      default=False,
                      help='Enable -m, method JIT when autoBisecting. Defaults to "False"')

    # Special case in which a specific exit code needs to be observed.
    parser.add_option('-w', '--watchExitCode',
                      dest='watchExitCode',
                      type='choice',
                      choices=['3', '4', '5', '6'],
                      help='Look out for a specific exit code in the range [3,6]')

    (options, args) = parser.parse_args()
    if len(args) != 1:
        parser.error('There is a wrong number of arguments.')
    if options.startRepo == None:
        parser.error('Please specify an earlier start repository for the bisect range.')
    return options.bugOrWfm, options.compileType, options.dir, options.output, \
            options.resetBool, options.startRepo, options.endRepo, options.archi, \
            options.tracingjitBool, options.methodjitBool, options.watchExitCode

def checkNumOfTests(str):
    # Sample bisect range message:
    # "Testing changeset 41831:4f4c01fb42c3 (2 changesets remaining, ~1 tests)"
    # This function looks for the number just after the "~".
    testNum = 0
    i = 0
    while i < len(str):
        if (str[i] == '~'):
            # This works no matter it is a one-digit or two-digit number.
            testNum = int(str[i+1] + str[i+2])
            # Sometimes estimation is not entirely accurate, one more test
            # round may be needed.
            # It will be checked to stop when the First bad changeset is found.
            testNum = testNum + 1;
            break
        i = i + 1
    if testNum == 0:
        raise Exception('The number of tests to be executed should not be 0.')
    return testNum

# Run the testcase on the compiled js binary.
def testBinary(shell, file, methodjitBool, tracingjitBool):
    methodJit = '-m' if methodjitBool else ''
    tracingJit = '-j' if tracingjitBool else ''
    testBinaryCmd = './' + shell, methodJit, tracingJit, file
    print 'The testing command is:', testBinaryCmd

    # Capture stdout and stderr into the same string.
    p = subprocess.Popen([testBinaryCmd], stderr=subprocess.STDOUT,
                            stdout=subprocess.PIPE, shell=True)
    output = p.communicate()[0]
    retCode = p.returncode
    print 'The exit code is:', retCode
    print 'The first output is:', output

    # Switch to interactive input mode similar to `cat testcase.js | ./js -j -i`.
    if retCode == 0:
        print 'Switching to interactive input mode in case passing as a CLI ' + \
                'argument does not reproduce the issue..'
        testBinaryCmd2 = subprocess.Popen(['cat', file], stdout=PIPE)
        testBinaryCmd3 = subprocess.Popen(['./' + shell, methodJit, tracingJit, '-i'],
            stdin=testBinaryCmd2.STDOUT, stdout=subprocess.PIPE)
        output2 = testBinaryCmd3.communicate()[0]
        retCode = testBinaryCmd3.returncode
        print 'The exit code is:', retCode
        print 'The second output is:', output2
    return output, retCode

# This function labels a changeset as "good" or "bad" depending on parameters.
def bisectLabel(bugOrWfm, gdBad, startRepo, endRepo):
    bisectLabelTuple = ()
    if bugOrWfm == 'bug':
        if gdBad == 'bad':
            bisectLabelTuple = ('BAD', '-b')
        elif gdBad == 'good':
            bisectLabelTuple = ('GOOD', '-g')
    elif bugOrWfm == 'wfm':
        if gdBad == 'bad':
            bisectLabelTuple = ('GOOD', '-g')
        elif gdBad == 'good':
            bisectLabelTuple = ('BAD', '-b')

    print bisectLabelTuple[0], 'changeset: hg bisect', bisectLabelTuple[1]
    outputResult = captureStdout('hg bisect ' + bisectLabelTuple[1])
    print outputResult

    print 'autoBisect is now currently in hg revision:',
    currRev = captureStdout(['hg identify -n'])

    # Update the startRepo/endRepo values.
    if bugOrWfm == 'bug':
        if gdBad == 'bad':
            endRepo = currRev
        elif gdBad == 'good':
            startRepo = currRev
    elif bugOrWfm == 'wfm':
        if gdBad == 'bad':
            startRepo = currRev
        elif gdBad == 'good':
            endRepo = currRev
    return outputResult, startRepo, endRepo

if __name__ == '__main__':
    main()
<commit_msg>Continue to try to make autoBisect.py at least partially work.<commit_after>#!/usr/bin/env python

#/* ***** BEGIN LICENSE BLOCK	****
# Version: MPL 1.1/GPL 2.0/LGPL 2.1
#
# The contents of this file are subject to the Mozilla Public License Version
# 1.1 (the "License"); you may not use this file except in compliance with
# the License. You may obtain a copy of the License at
# http://www.mozilla.org/MPL/
#
# Software distributed under the License is distributed on an "AS IS" basis,
# WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
# for the specific language governing rights and limitations under the
# License.
#
# The Original Code is autoBisect.
#
# The Initial Developer of the Original Code is
# Gary Kwong.
# Portions created by the Initial Developer are Copyright (C) 2006-2008
# the Initial Developer. All Rights Reserved.
#
# Contributor(s):
#
# Alternatively, the contents of this file may be used under the terms of
# either the GNU General Public License Version 2 or later (the "GPL"), or
# the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
# in which case the provisions of the GPL or the LGPL are applicable instead
# of those above. If you wish to allow use of your version of this file only
# under the terms of either the GPL or the LGPL, and not to allow others to
# use your version of this file under the terms of the MPL, indicate your
# decision by deleting the provisions above and replace them with the notice
# and other provisions required by the GPL or the LGPL. If you do not delete
# the provisions above, a recipient may use your version of this file under
# the terms of any one of the MPL, the GPL or the LGPL.
#
# * ***** END LICENSE BLOCK	****	/

import os, shutil, subprocess, sys
from optparse import OptionParser

sys.path.append('../jsfunfuzz/')
from fnStartjsfunfuzz import *

def main():
    # Do not support Windows XP because the ~ folder is in "/Documents and Settings/",
    # which contains spaces. This breaks MinGW, which is what MozillaBuild uses.
    # From Windows Vista onwards, the folder is in "/Users/".
    if os.name == 'nt':
        if '5.1' in captureStdout('uname'):
            raise Exception('autoBisect is not supported on Windows XP.')
    verbose = True

    # Parse options and parameters from the command-line.
    filename = sys.argv[-1:][0]
    options = parseOpts()
    (bugOrWfm, compileType, sourceDir, stdoutOutput, resetBool, startRepo, \
     endRepo, archNum, tracingjitBool, methodjitBool, watchExitCode) = options

    os.chdir(sourceDir)

    # Refresh source directory (overwrite all local changes) to default tip if required.
    if resetBool:
        subprocess.call(['hg up -C default'], shell=True)

    # Reverse if a bisect range to find a WFM issue is specified.
    if bugOrWfm == 'wfm':
        tempVal = ''
        tempVal = startRepo
        startRepo = endRepo
        endRepo = tempVal

    # Specify `hg bisect` ranges.
    subprocess.call(['hg bisect -r'], shell=True)
    # If in "bug" mode, this startRepo changeset does not exhibit the issue.
    subprocess.call(['hg bisect -g ' + startRepo], shell=True)
    # If in "bug" mode, this endRepo changeset exhibits the issue.
    stdoutNumOfTests = captureStdout('hg bisect -b ' + endRepo)

    # Find out the number of tests to be executed based on the initial hg bisect output.
    numOfTests = checkNumOfTests(stdoutNumOfTests)

    # For main directory, change into Desktop directory.
    mainDir = os.path.expanduser('~/Desktop/')

    # Change into main directory.
    os.chdir(mainDir)

    for i in xrange(numOfTests):
        autoBisectPath = 'autoBisect-' + compileType + '-' + archNum + '-s' + \
                         startRepo + '-e' + endRepo

        # Create the autoBisect folder.
        try:
            os.makedirs(autoBisectPath)
        except OSError:
            raise Exception('The autoBisect path at "' + autoBisectPath + '" already exists!')

        # Change into autoBisectPath.
        os.chdir(autoBisectPath)

        # Copy the js tree to the autoBisect path.
        # Don't use pymake because older changesets may fail to compile.
        cpJsTreeOrPymakeDir(os.path.expanduser(sourceDir), 'js')
        os.chdir('compilePath')  # Change into compilation directory.

        autoconfRun()

        # Create objdirs within the compilePaths.
        os.mkdir(compileType + '-objdir')
        os.chdir(compileType + '-objdir')

        # Compile the first binary.
        branchType = 'autoBisectBranch'
        if 'jaeger' in sourceDir:
            branchType = 'jm'
        valgrindSupport = False  # Let's disable support for valgrind in the js shell
        threadsafe = False  # Let's disable support for threadsafety in the js shell
        configureJsBinary(archNum, compileType, branchType, valgrindSupport, threadsafe)

        if 'jaeger' in sourceDir:
            branchType = 'autoBisectBranch'  # Reset the branchType

        # Compile and copy the first binary.
        try:
            jsShellName = compileCopy(archNum, compileType, branchType, False)
        except:
            print 'The "good" repository that is currently labelled:', startRepo
            print 'The "bad" repository that is currently labelled:', endRepo
            raise Exception('Compilation failed.')

        # Change back into compilePath.
        os.chdir('../')

        # Test compilePath.
        if verbose:
            print 'DEBUG - This should be the compilePath:'
            print 'DEBUG - %s\n' % os.getcwdu()
            if 'compilePath' not in os.getcwdu():
                raise Exception('We are not in compilePath.')

        os.chdir('../../')  # Change into autoBisectPath directory.
        autoBisectFullPath = os.path.expanduser(os.getcwdu())

        (stdoutStderr, exitCode) = testBinary(jsShellName, filename,
                                              methodjitBool, tracingjitBool)

        # Switch to hg repository directory.
        os.chdir(os.path.expanduser(sourceDir))

        # Label the changeset bad if the exact assert is found (only in debug shells)
        # (Assuming "bad" and not "wfm".)
        # Assertion exit codes: Mac 10.5/10.6 - 133, Linux - 134, WinXP - 3
        if compileType == 'dbg' and stdoutOutput in stdoutStderr and exitCode != 0:
            # Set a random arbitrary value that cannot be a genuine exit code.
            codeToBeObserved = 888888
            if os.name == 'posix':
                if os.uname()[0] == 'Darwin':
                    codeToBeObserved = 133
                elif os.uname()[0] == 'Linux':
                    codeToBeObserved = 134
            elif os.name == 'nt':
                codeToBeObserved = 3

            # Depending on the exit code as per the different platforms, specify
            # if the current changeset is "good" or "bad".
            if exitCode == codeToBeObserved:
                (result, startRepo, endRepo) = bisectLabel(bugOrWfm, 'bad', startRepo, endRepo)

                print 'Now removing autoBisectFullPath, located at:', autoBisectFullPath
                shutil.rmtree(autoBisectFullPath)
                if 'first bad revision' in result:
                    break

        # "Bad" changesets.
        elif (exitCode == 1) or (129 <= exitCode <= 159) or (exitCode == watchExitCode):
            (result, startRepo, endRepo) = bisectLabel(bugOrWfm, 'bad', startRepo, endRepo)

            print 'Now removing autoBisectFullPath, located at:', autoBisectFullPath
            shutil.rmtree(autoBisectFullPath)
            if 'first bad revision' in result:
                break

        # "Good" changesets.
        elif exitCode != watchExitCode:
            if (exitCode == 0) or (3 <= exitCode <= 6):
                (result, startRepo, endRepo) = bisectLabel(bugOrWfm, 'good', startRepo, endRepo)

                print 'Now removing autoBisectFullPath, located at:', autoBisectFullPath
                shutil.rmtree(autoBisectFullPath)
                if 'first bad revision' in result:
                    break

        else:
            raise Exception('Unknown exit code hit:', exitCode)

    # Reset `hg bisect` after finishing everything.
    subprocess.call(['hg bisect -r'], shell=True)

def parseOpts():
    usage = 'Usage: %prog [options] filename'
    parser = OptionParser(usage)
    # See http://docs.python.org/library/optparse.html#optparse.OptionParser.disable_interspersed_args
    parser.disable_interspersed_args()

    # autoBisect details
    parser.add_option('-b', '--bugOrWfm',
                      dest='bugOrWfm',
                      type='choice',
                      choices=['bug', 'wfm'],
                      default='bug',
                      help='Bisect to find a bug or WFM issue. ' + \
                           'Only accepts "bug" or "wfm". ' + \
                           'Defaults to "bug"')
    parser.add_option('-c', '--compileType',
                      dest='compileType',
                      type='choice',
                      choices=['dbg', 'opt'],
                      default='dbg',
                      help='js shell compile type. Defaults to "dbg"')
    parser.add_option('-d', '--dir',
                      dest='dir',
                      default=os.path.expanduser('~/tracemonkey/'),
                      help='Source code directory. Defaults to "~/tracemonkey/"')
    parser.add_option('-o', '--output',
                      dest='output',
                      help='Stdout or stderr output to be observed. ' + \
                           'For assertions, set to "ssertion fail"')
    parser.add_option('-r', '--resetToTipFirstBool',
                      dest='resetBool',
                      action='store_true',
                      default=False,
                      help='First reset to default tip overwriting all local changes. ' + \
                           'Equivalent to first executing `hg update -C default`. ' + \
                           'Defaults to "False"')

    # Define the start and end repositories.
    parser.add_option('-s', '--start',
                      dest='startRepo',
                      help='Start repository (earlier)')
    parser.add_option('-e', '--end',
                      dest='endRepo',
                      default='tip',
                      help='End repository (later). Defaults to "tip"')

    # Define the architecture to be tested.
    parser.add_option('-a', '--architecture',
                      dest='archi',
                      type='choice',
                      choices=['32', '64'],
                      help='Test architecture. Only accepts "32" or "64"')

    # Define parameters to be passed to the binary.
    parser.add_option('-j', '--tracingjit',
                      dest='tracingjitBool',
                      action='store_true',
                      default=False,
                      help='Enable -j, tracing JIT when autoBisecting. Defaults to "False"')
    parser.add_option('-m', '--methodjit',
                      dest='methodjitBool',
                      action='store_true',
                      default=False,
                      help='Enable -m, method JIT when autoBisecting. Defaults to "False"')

    # Special case in which a specific exit code needs to be observed.
    parser.add_option('-w', '--watchExitCode',
                      dest='watchExitCode',
                      type='choice',
                      choices=['3', '4', '5', '6'],
                      help='Look out for a specific exit code in the range [3,6]')

    (options, args) = parser.parse_args()
    if len(args) != 1:
        parser.error('There is a wrong number of arguments.')
    if options.startRepo == None:
        parser.error('Please specify an earlier start repository for the bisect range.')
    return options.bugOrWfm, options.compileType, options.dir, options.output, \
            options.resetBool, options.startRepo, options.endRepo, options.archi, \
            options.tracingjitBool, options.methodjitBool, options.watchExitCode

def checkNumOfTests(str):
    # Sample bisect range message:
    # "Testing changeset 41831:4f4c01fb42c3 (2 changesets remaining, ~1 tests)"
    # This function looks for the number just after the "~".
    testNum = 0
    i = 0
    while i < len(str):
        if (str[i] == '~'):
            # This works no matter it is a one-digit or two-digit number.
            testNum = int(str[i+1] + str[i+2])
            # Sometimes estimation is not entirely accurate, one more test
            # round may be needed.
            # It will be checked to stop when the First bad changeset is found.
            testNum = testNum + 1;
            break
        i = i + 1
    if testNum == 0:
        raise Exception('The number of tests to be executed should not be 0.')
    return testNum

# Run the testcase on the compiled js binary.
def testBinary(shell, file, methodjitBool, tracingjitBool):
    methodJit = ' -m' if methodjitBool else ''
    tracingJit = ' -j' if tracingjitBool else ''
    testBinaryCmd = './' + shell + methodJit + tracingJit + ' ' + file
    print 'The testing command is:', testBinaryCmd

    # Capture stdout and stderr into the same string.
    p = subprocess.Popen([testBinaryCmd], stderr=subprocess.STDOUT,
                            stdout=subprocess.PIPE, shell=True)
    output = p.communicate()[0]
    retCode = p.returncode
    print 'The exit code is:', retCode
    print 'The first output is:', output

    # Switch to interactive input mode similar to `cat testcase.js | ./js -j -i`.
    if retCode == 0:
        print 'Switching to interactive input mode in case passing as a CLI ' + \
                'argument does not reproduce the issue..'
        testBinaryCmd2 = subprocess.Popen(['cat', file], stdout=PIPE)
        testBinaryCmd3 = subprocess.Popen(['./' + shell, methodJit, tracingJit, '-i'],
            stdin=testBinaryCmd2.STDOUT, stdout=subprocess.PIPE)
        output2 = testBinaryCmd3.communicate()[0]
        retCode = testBinaryCmd3.returncode
        print 'The exit code is:', retCode
        print 'The second output is:', output2
    return output, retCode

# This function labels a changeset as "good" or "bad" depending on parameters.
def bisectLabel(bugOrWfm, gdBad, startRepo, endRepo):
    bisectLabelTuple = ()
    if bugOrWfm == 'bug':
        if gdBad == 'bad':
            bisectLabelTuple = ('BAD', '-b')
        elif gdBad == 'good':
            bisectLabelTuple = ('GOOD', '-g')
    elif bugOrWfm == 'wfm':
        if gdBad == 'bad':
            bisectLabelTuple = ('GOOD', '-g')
        elif gdBad == 'good':
            bisectLabelTuple = ('BAD', '-b')

    print bisectLabelTuple[0], 'changeset: hg bisect', bisectLabelTuple[1]
    outputResult = captureStdout('hg bisect ' + bisectLabelTuple[1])
    print outputResult

    print 'autoBisect is now currently in hg revision:',
    currRev = captureStdout(['hg identify -n'])

    # Update the startRepo/endRepo values.
    if bugOrWfm == 'bug':
        if gdBad == 'bad':
            endRepo = currRev
        elif gdBad == 'good':
            startRepo = currRev
    elif bugOrWfm == 'wfm':
        if gdBad == 'bad':
            startRepo = currRev
        elif gdBad == 'good':
            endRepo = currRev
    return outputResult, startRepo, endRepo

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