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

import sys
import sqlite3
import argparse

import jcd.common
import jcd.dao

class MinMax(object):

    StationsDayTable = "minmax_stations_day"

    def __init__(self, db, sample_schema, arguments):
        self._db = db
        self._sample_schema = sample_schema
        self._arguments = arguments
        assert self._db is not None
        assert self._sample_schema is not None
        assert self._arguments is not None
        self._create_tables_if_necessary()

    def _create_tables_if_necessary(self):
        if not self._db.has_table(self.StationsDayTable):
            self._create_stations_day_table()

    def _create_stations_day_table(self):
        if self._arguments.verbose:
            print "Creating table", self.StationsDayTable
        try:
            self._db.connection.execute(
                '''
                CREATE TABLE %s (
                date TEXT NOT NULL,
                contract_id INTEGER NOT NULL,
                station_number INTEGER NOT NULL,
                min_bikes INTEGER NOT NULL,
                max_bikes INTEGER NOT NULL,
                min_slots INTEGER NOT NULL,
                max_slots INTEGER NOT NULL,
                num_changes INTEGER NOT NULL,
                PRIMARY KEY (date, contract_id, station_number));
                ''' % self.StationsDayTable)
        except sqlite3.Error as error:
            print "%s: %s" % (type(error).__name__, error)
            raise jcd.common.JcdException(
                "Database error while creating table [%s]" % self.StationsDayTable)

    def _do_stations(self, date):
        try:
            self._db.connection.execute(
                '''
                INSERT OR REPLACE INTO %s
                SELECT ?,
                contract_id,
                station_number,
                MIN(available_bikes),
                MAX(available_bikes),
                MIN(available_bike_stands),
                MAX(available_bike_stands),
                COUNT(timestamp)
                FROM %s.%s
                GROUP BY contract_id, station_number
                ''' % (self.StationsDayTable, self._sample_schema, jcd.dao.ShortSamplesDAO.TableNameArchive),
                (date,))
        except sqlite3.Error as error:
            print "%s: %s" % (type(error).__name__, error)
            raise jcd.common.JcdException(
                "Database error while storing daily min max into table [%s]" % self.StationsDayTable)

    def run(self, date):
        self._do_stations(date)

class Activity(object):

    StationsDayTable = "activity_stations_day"
    ContractsDayTable = "activity_contracts_day"

    def __init__(self, db, sample_schema, arguments):
        self._db = db
        self._sample_schema = sample_schema
        self._arguments = arguments
        assert self._db is not None
        assert self._sample_schema is not None
        assert self._arguments is not None
        self._create_tables_if_necessary()

    def _create_tables_if_necessary(self):
        if not self._db.has_table(self.StationsDayTable):
            self._create_stations_day_table()
        if not self._db.has_table(self.ContractsDayTable):
            self._create_contracts_day_table()

    def _create_stations_day_table(self):
        if self._arguments.verbose:
            print "Creating table", self.StationsDayTable
        try:
            self._db.connection.execute(
                '''
                CREATE TABLE %s (
                date TEXT NOT NULL,
                contract_id INTEGER NOT NULL,
                station_number INTEGER NOT NULL,
                num_changes INTEGER NOT NULL,
                rank_contract INTEGER,
                rank_global INTEGER,
                PRIMARY KEY (date, contract_id, station_number));
                ''' % self.StationsDayTable)
        except sqlite3.Error as error:
            print "%s: %s" % (type(error).__name__, error)
            raise jcd.common.JcdException(
                "Database error while creating table [%s]" % self.StationsDayTable)

    def _create_contracts_day_table(self):
        if self._arguments.verbose:
            print "Creating table", self.ContractsDayTable
        try:
            self._db.connection.execute(
                '''
                CREATE TABLE %s (
                date TEXT NOT NULL,
                contract_id INTEGER NOT NULL,
                num_changes INTEGER NOT NULL,
                rank_global INTEGER,
                PRIMARY KEY (date, contract_id));
                ''' % self.ContractsDayTable)
        except sqlite3.Error as error:
            print "%s: %s" % (type(error).__name__, error)
            raise jcd.common.JcdException(
                "Database error while creating table [%s]" % self.ContractsDayTable)

    def _do_stations(self, date):
        try:
            self._db.connection.execute(
                '''
                INSERT OR REPLACE INTO %s
                    SELECT date(timestamp,'unixepoch') as date,
                        contract_id,
                        station_number,
                        COUNT(timestamp) as num_changes,
                        NULL,
                        NULL
                    FROM %s.%s
                    WHERE date = ?
                    GROUP BY contract_id, station_number
                ''' % (self.StationsDayTable,
                       self._sample_schema,
                       jcd.dao.ShortSamplesDAO.TableNameArchive),
                (date,))
        except sqlite3.Error as error:
            print "%s: %s" % (type(error).__name__, error)
            raise jcd.common.JcdException(
                "Database error while storing daily min max into table [%s]" % self.StationsDayTable)

    def _stations_day_get(self, date):
        try:
            req = self._db.connection.execute(
                '''
                SELECT date,
                    contract_id,
                    station_number,
                    num_changes,
                    rank_contract,
                    rank_global
                FROM %s
                WHERE date = ?
                ORDER BY num_changes DESC
                ''' % (self.StationsDayTable), (date,))
            while True:
                ranks = req.fetchmany(1000)
                if not ranks:
                    break
                for rank in ranks:
                    d_rank = {
                        "date": rank[0],
                        "contract_id": rank[1],
                        "station_number": rank[2],
                        "num_changes": rank[3],
                        "rank_contract": rank[4],
                        "rank_global": rank[5],
                    }
                    yield d_rank
        except sqlite3.Error as error:
            print "%s: %s" % (type(error).__name__, error)
            raise jcd.common.JcdException(
                "Database error while getting daily activity ranks")

    def _stations_day_compute_ranks(self, date):

        contracts = {}
        global_infos = {
            "n_total": None,
            "n_rank": None,
            "n_last": None
        }

        for rank in self._stations_day_get(date):

            ## calculate contract ranks
            # first
            if rank["contract_id"] not in contracts:
                contracts[rank["contract_id"]] = {
                    "c_total": 0,
                    "c_rank": 1,
                    "c_last": rank["num_changes"]
                }
            # quicker lookip
            contract_infos = contracts[rank["contract_id"]]
            # sample done for contract
            contract_infos["c_total"] += 1
            # not the same number of events
            if rank["num_changes"] != contract_infos["c_last"]:
                contract_infos["c_rank"] = contract_infos["c_total"]
                contract_infos["c_last"] = rank["num_changes"]
            # set rank
            rank["rank_contract"] = contract_infos["c_rank"]

            ## calculate global ranks
            # first
            if global_infos["n_last"] is None:
                global_infos["n_total"] = 0
                global_infos["n_rank"] = 1
                global_infos["n_last"] = rank["num_changes"]
            # sample done globally
            global_infos["n_total"] += 1
            # not the same number of events
            if rank["num_changes"] != global_infos["n_last"]:
                global_infos["n_rank"] = global_infos["n_total"]
                global_infos["n_last"] = rank["num_changes"]
            # set rank
            rank["rank_global"] = global_infos["n_rank"]

            ## return value
            yield rank

        # TODO: what to do when there is nothing ?

    def _stations_day_update_ranks(self, date):
        try:
            # update any existing contracts
            req = self._db.connection.executemany(
                '''
                UPDATE %s
                SET rank_global = :rank_global,
                    rank_contract = :rank_contract
                WHERE date = :date AND
                    contract_id = :contract_id AND
                    station_number = :station_number
                ''' % (self.StationsDayTable),
                self._stations_day_compute_ranks(date))
            # return number of inserted records
            return req.rowcount
        except sqlite3.Error as error:
            print "%s: %s" % (type(error).__name__, error)
            raise jcd.common.JcdException("Database error while updating daily station activity rankings")

    def run(self, date):
        self._do_stations(date)
        self._stations_day_update_ranks(date)

class App(object):

    def __init__(self, default_data_path, default_db_filename):
        # construct parser
        self._parser = argparse.ArgumentParser(
            description='Calculate min-max data from jcd stats')
        self._parser.add_argument(
            '--datadir',
            help='choose data folder (default: %s)' % default_data_path,
            default=default_data_path
        )
        self._parser.add_argument(
            '--dbname',
            help='choose db filename (default: %s)' % default_db_filename,
            default=default_db_filename
        )
        self._parser.add_argument(
            '--verbose', '-v',
            action='store_true',
            help='display operationnal informations'
        )
        self._parser.add_argument(
            'date',
            metavar='N',
            type=str,
            nargs='+',
            help='an integer for the accumulator')

    def run(self):
        # parse arguments
        arguments = self._parser.parse_args()
        with jcd.common.SqliteDB(arguments.dbname, arguments.datadir) as db_stats:
            for date in arguments.date:
                if arguments.verbose:
                    print "Processing", date
                # attach samples db
                schema = jcd.dao.ShortSamplesDAO.get_schema_name(date)
                filename = jcd.dao.ShortSamplesDAO.get_db_file_name(schema)
                db_stats.attach_database(filename, schema, arguments.datadir)
                # do processing
                MinMax(db_stats, schema, arguments).run(date)
                Activity(db_stats, schema, arguments).run(date)
                # detach samples db
                db_stats.detach_database(schema)

# main
if __name__ == '__main__':
    try:
        App("~/.jcd_v2", "stats.db").run()
    except KeyboardInterrupt:
        pass
<commit_msg>compute contract daily activity<commit_after>#! /usr/bin/env python

import sys
import sqlite3
import argparse

import jcd.common
import jcd.dao

class MinMax(object):

    StationsDayTable = "minmax_stations_day"

    def __init__(self, db, sample_schema, arguments):
        self._db = db
        self._sample_schema = sample_schema
        self._arguments = arguments
        assert self._db is not None
        assert self._sample_schema is not None
        assert self._arguments is not None
        self._create_tables_if_necessary()

    def _create_tables_if_necessary(self):
        if not self._db.has_table(self.StationsDayTable):
            self._create_stations_day_table()

    def _create_stations_day_table(self):
        if self._arguments.verbose:
            print "Creating table", self.StationsDayTable
        try:
            self._db.connection.execute(
                '''
                CREATE TABLE %s (
                date TEXT NOT NULL,
                contract_id INTEGER NOT NULL,
                station_number INTEGER NOT NULL,
                min_bikes INTEGER NOT NULL,
                max_bikes INTEGER NOT NULL,
                min_slots INTEGER NOT NULL,
                max_slots INTEGER NOT NULL,
                num_changes INTEGER NOT NULL,
                PRIMARY KEY (date, contract_id, station_number));
                ''' % self.StationsDayTable)
        except sqlite3.Error as error:
            print "%s: %s" % (type(error).__name__, error)
            raise jcd.common.JcdException(
                "Database error while creating table [%s]" % self.StationsDayTable)

    def _do_stations(self, date):
        try:
            self._db.connection.execute(
                '''
                INSERT OR REPLACE INTO %s
                SELECT ?,
                contract_id,
                station_number,
                MIN(available_bikes),
                MAX(available_bikes),
                MIN(available_bike_stands),
                MAX(available_bike_stands),
                COUNT(timestamp)
                FROM %s.%s
                GROUP BY contract_id, station_number
                ''' % (self.StationsDayTable, self._sample_schema, jcd.dao.ShortSamplesDAO.TableNameArchive),
                (date,))
        except sqlite3.Error as error:
            print "%s: %s" % (type(error).__name__, error)
            raise jcd.common.JcdException(
                "Database error while storing daily min max into table [%s]" % self.StationsDayTable)

    def run(self, date):
        self._do_stations(date)

class Activity(object):

    StationsDayTable = "activity_stations_day"
    ContractsDayTable = "activity_contracts_day"

    def __init__(self, db, sample_schema, arguments):
        self._db = db
        self._sample_schema = sample_schema
        self._arguments = arguments
        assert self._db is not None
        assert self._sample_schema is not None
        assert self._arguments is not None
        self._create_tables_if_necessary()

    def _create_tables_if_necessary(self):
        if not self._db.has_table(self.StationsDayTable):
            self._create_stations_day_table()
        if not self._db.has_table(self.ContractsDayTable):
            self._create_contracts_day_table()

    def _create_stations_day_table(self):
        if self._arguments.verbose:
            print "Creating table", self.StationsDayTable
        try:
            self._db.connection.execute(
                '''
                CREATE TABLE %s (
                date TEXT NOT NULL,
                contract_id INTEGER NOT NULL,
                station_number INTEGER NOT NULL,
                num_changes INTEGER NOT NULL,
                rank_contract INTEGER,
                rank_global INTEGER,
                PRIMARY KEY (date, contract_id, station_number));
                ''' % self.StationsDayTable)
        except sqlite3.Error as error:
            print "%s: %s" % (type(error).__name__, error)
            raise jcd.common.JcdException(
                "Database error while creating table [%s]" % self.StationsDayTable)

    def _create_contracts_day_table(self):
        if self._arguments.verbose:
            print "Creating table", self.ContractsDayTable
        try:
            self._db.connection.execute(
                '''
                CREATE TABLE %s (
                date TEXT NOT NULL,
                contract_id INTEGER NOT NULL,
                num_changes INTEGER NOT NULL,
                rank_global INTEGER,
                PRIMARY KEY (date, contract_id));
                ''' % self.ContractsDayTable)
        except sqlite3.Error as error:
            print "%s: %s" % (type(error).__name__, error)
            raise jcd.common.JcdException(
                "Database error while creating table [%s]" % self.ContractsDayTable)

    def _do_stations(self, date):
        try:
            self._db.connection.execute(
                '''
                INSERT OR REPLACE INTO %s
                    SELECT date(timestamp,'unixepoch') as date,
                        contract_id,
                        station_number,
                        COUNT(timestamp) as num_changes,
                        NULL,
                        NULL
                    FROM %s.%s
                    WHERE date = ?
                    GROUP BY contract_id, station_number
                ''' % (self.StationsDayTable,
                       self._sample_schema,
                       jcd.dao.ShortSamplesDAO.TableNameArchive),
                (date,))
        except sqlite3.Error as error:
            print "%s: %s" % (type(error).__name__, error)
            raise jcd.common.JcdException(
                "Database error while storing daily min max into table [%s]" % self.StationsDayTable)

    def _do_contracts(self, date):
        try:
            self._db.connection.execute(
                '''
                INSERT OR REPLACE INTO %s
                    SELECT date,
                        contract_id,
                        SUM(num_changes) as num_changes,
                        NULL
                    FROM %s
                    WHERE date = ?
                    GROUP BY contract_id
                ''' % (self.ContractsDayTable,
                       self.StationsDayTable),
                (date,))
        except sqlite3.Error as error:
            print "%s: %s" % (type(error).__name__, error)
            raise jcd.common.JcdException(
                "Database error while storing daily min max into table [%s]" % self.StationsDayTable)

    def _stations_day_get(self, date):
        try:
            req = self._db.connection.execute(
                '''
                SELECT date,
                    contract_id,
                    station_number,
                    num_changes,
                    rank_contract,
                    rank_global
                FROM %s
                WHERE date = ?
                ORDER BY num_changes DESC
                ''' % (self.StationsDayTable), (date,))
            while True:
                ranks = req.fetchmany(1000)
                if not ranks:
                    break
                for rank in ranks:
                    d_rank = {
                        "date": rank[0],
                        "contract_id": rank[1],
                        "station_number": rank[2],
                        "num_changes": rank[3],
                        "rank_contract": rank[4],
                        "rank_global": rank[5],
                    }
                    yield d_rank
        except sqlite3.Error as error:
            print "%s: %s" % (type(error).__name__, error)
            raise jcd.common.JcdException(
                "Database error while getting daily activity ranks")

    def _stations_day_compute_ranks(self, date):

        contracts = {}
        global_infos = {
            "n_total": None,
            "n_rank": None,
            "n_last": None
        }

        for rank in self._stations_day_get(date):

            ## calculate contract ranks
            # first
            if rank["contract_id"] not in contracts:
                contracts[rank["contract_id"]] = {
                    "c_total": 0,
                    "c_rank": 1,
                    "c_last": rank["num_changes"]
                }
            # quicker lookip
            contract_infos = contracts[rank["contract_id"]]
            # sample done for contract
            contract_infos["c_total"] += 1
            # not the same number of events
            if rank["num_changes"] != contract_infos["c_last"]:
                contract_infos["c_rank"] = contract_infos["c_total"]
                contract_infos["c_last"] = rank["num_changes"]
            # set rank
            rank["rank_contract"] = contract_infos["c_rank"]

            ## calculate global ranks
            # first
            if global_infos["n_last"] is None:
                global_infos["n_total"] = 0
                global_infos["n_rank"] = 1
                global_infos["n_last"] = rank["num_changes"]
            # sample done globally
            global_infos["n_total"] += 1
            # not the same number of events
            if rank["num_changes"] != global_infos["n_last"]:
                global_infos["n_rank"] = global_infos["n_total"]
                global_infos["n_last"] = rank["num_changes"]
            # set rank
            rank["rank_global"] = global_infos["n_rank"]

            ## return value
            yield rank

        # TODO: what to do when there is nothing ?

    def _stations_day_update_ranks(self, date):
        try:
            # update any existing contracts
            req = self._db.connection.executemany(
                '''
                UPDATE %s
                SET rank_global = :rank_global,
                    rank_contract = :rank_contract
                WHERE date = :date AND
                    contract_id = :contract_id AND
                    station_number = :station_number
                ''' % (self.StationsDayTable),
                self._stations_day_compute_ranks(date))
            # return number of inserted records
            return req.rowcount
        except sqlite3.Error as error:
            print "%s: %s" % (type(error).__name__, error)
            raise jcd.common.JcdException("Database error while updating daily station activity rankings")

    def run(self, date):
        self._do_stations(date)
        self._do_contracts(date)
        self._stations_day_update_ranks(date)

class App(object):

    def __init__(self, default_data_path, default_db_filename):
        # construct parser
        self._parser = argparse.ArgumentParser(
            description='Calculate min-max data from jcd stats')
        self._parser.add_argument(
            '--datadir',
            help='choose data folder (default: %s)' % default_data_path,
            default=default_data_path
        )
        self._parser.add_argument(
            '--dbname',
            help='choose db filename (default: %s)' % default_db_filename,
            default=default_db_filename
        )
        self._parser.add_argument(
            '--verbose', '-v',
            action='store_true',
            help='display operationnal informations'
        )
        self._parser.add_argument(
            'date',
            metavar='N',
            type=str,
            nargs='+',
            help='an integer for the accumulator')

    def run(self):
        # parse arguments
        arguments = self._parser.parse_args()
        with jcd.common.SqliteDB(arguments.dbname, arguments.datadir) as db_stats:
            for date in arguments.date:
                if arguments.verbose:
                    print "Processing", date
                # attach samples db
                schema = jcd.dao.ShortSamplesDAO.get_schema_name(date)
                filename = jcd.dao.ShortSamplesDAO.get_db_file_name(schema)
                db_stats.attach_database(filename, schema, arguments.datadir)
                # do processing
                MinMax(db_stats, schema, arguments).run(date)
                Activity(db_stats, schema, arguments).run(date)
                # detach samples db
                db_stats.detach_database(schema)

# main
if __name__ == '__main__':
    try:
        App("~/.jcd_v2", "stats.db").run()
    except KeyboardInterrupt:
        pass
<|endoftext|>


<commit_before>from copy import copy
from cStringIO import StringIO
import datetime
import inspect
import sys
import traceback

from django.core.management import call_command
from django.db import models

from south import exceptions
from south.db import db
from south.models import MigrationHistory
from south.signals import ran_migration

class Migrator(object):
    def __init__(self, verbosity=0):
        self.verbosity = int(verbosity)

    @staticmethod
    def title(target):
        raise NotImplementedError()

    def print_title(self, target):
        if self.verbosity:
            print self.title(target)

    @staticmethod
    def status(target):
        raise NotImplementedError()

    def print_status(self, migration):
        status = self.status(migration)
        if self.verbosity and status:
            print status

    @staticmethod
    def orm(migration):
        raise NotImplementedError()

    def backwards(self, migration):
        return self._wrap_direction(migration.backwards(), self.orm(migration))

    def direction(self, migration):
        raise NotImplementedError()

    @staticmethod
    def _wrap_direction(direction, orm):
        args = inspect.getargspec(direction)
        if len(args[0]) == 1:
            # Old migration, no ORM should be passed in
            return direction
        return (lambda: direction(orm))

    @staticmethod
    def record(migration):
        raise NotImplementedError()

    def run_migration_error(self, migration, extra_info=''):
        return (' ! Error found during real run of migration! Aborting.\n'
                '\n'
                ' ! Since you have a database that does not support running\n'
                ' ! schema-altering statements in transactions, we have had \n'
                ' ! to leave it in an interim state between migrations.\n'
                '%s\n'
                ' ! The South developers regret this has happened, and would\n'
                ' ! like to gently persuade you to consider a slightly\n'
                ' ! easier-to-deal-with DBMS.\n') % extra_info

    def run_migration(self, migration):
        migration_function = self.direction(migration)
        db.start_transaction()
        try:
            migration_function()
            db.execute_deferred_sql()
        except:
            db.rollback_transaction()
            if not db.has_ddl_transactions:
                print self.run_migration_error(migration)
            raise
        else:
            db.commit_transaction()

    def run(self, migration):
        # Get the correct ORM.
        db.current_orm = self.orm(migration)
        # If the database doesn't support running DDL inside a transaction
        # *cough*MySQL*cough* then do a dry run first.
        if not db.has_ddl_transactions:
            dry_run = DryRunMigrator(migrator=self, ignore_fail=False)
            dry_run.run_migration(migration)
        return self.run_migration(migration)

    def done_migrate(self, migration):
        db.start_transaction()
        try:
            # Record us as having done this
            self.record(migration)
        except:
            db.rollback_transaction()
            raise
        else:
            db.commit_transaction()

    def send_ran_migration(self, migration):
        ran_migration.send(None,
                           app=migration.app_label(),
                           migration=migration,
                           method=self.__class__.__name__.lower())

    def migrate(self, migration):
        """
        Runs the specified migration forwards/backwards, in order.
        """
        app = migration.migrations._migrations
        migration_name = migration.name()
        self.print_status(migration)
        result = self.run(migration)
        self.done_migrate(migration)
        self.send_ran_migration(migration)
        return result

    def migrate_many(self, target, migrations):
        raise NotImplementedError()

class MigratorWrapper(object):
    def __init__(self, migrator, *args, **kwargs):
        self._migrator = copy(migrator)
        attributes = dict([(k, getattr(self, k))
                           for k in self.__class__.__dict__.iterkeys()
                           if not k.startswith('__')])
        self._migrator.__dict__.update(attributes)

    def __getattr__(self, name):
        return getattr(self._migrator, name)

class DryRunMigrator(MigratorWrapper):
    def __init__(self, ignore_fail=True, *args, **kwargs):
        super(DryRunMigrator, self).__init__(*args, **kwargs)
        self._ignore_fail = ignore_fail

    def _run_migration(self, migration):
        if migration.no_dry_run() and self.verbosity:
            print " - Migration '%s' is marked for no-dry-run."
            return
        db.dry_run = True
        db.debug, old_debug = False, db.debug
        pending_creates = db.get_pending_creates()
        db.start_transaction()
        migration_function = self.direction(migration)
        try:
            try:
                migration_function()
            except:
                raise exceptions.FailedDryRun(migration, sys.exc_info())
        finally:
            db.rollback_transactions_dry_run()
            db.debug = old_debug
            db.clear_run_data(pending_creates)
            db.dry_run = False

    def run_migration(self, migration):
        try:
            self._run_migration(migration)
        except exceptions.FailedDryRun:
            if self._ignore_fail:
                return False
            raise

    def done_migrate(self, *args, **kwargs):
        pass

    def send_ran_migration(self, *args, **kwargs):
        pass

class FakeMigrator(MigratorWrapper):
    def run(self, migration):
        if self.verbosity:
            print '   (faked)'

    def send_ran_migration(self, *args, **kwargs):
        pass

class LoadInitialDataMigrator(MigratorWrapper):
    def load_initial_data(self, target):
        if target is None or target != target.migrations[-1]:
            return
        # Load initial data, if we ended up at target
        if self.verbosity:
            print " - Loading initial data for %s." % target.app_label()
        # Override Django's get_apps call temporarily to only load from the
        # current app
        old_get_apps = models.get_apps
        models.get_apps = lambda: [models.get_app(target.app_label())]
        try:
            call_command('loaddata', 'initial_data', verbosity=self.verbosity)
        finally:
            models.get_apps = old_get_apps

    def migrate_many(self, target, migrations):
        migrator = self._migrator
        result = migrator.__class__.migrate_many(migrator, target, migrations)
        if result:
            self.load_initial_data(target)
        return True

class Forwards(Migrator):
    """
    Runs the specified migration forwards, in order.
    """
    torun = 'forwards'

    @staticmethod
    def title(target):
        if target is not None:
            return " - Migrating forwards to %s." % target.name()
        else:
            assert False, "You cannot migrate forwards to zero."

    @staticmethod
    def status(migration):
        return ' > %s' % migration

    @staticmethod
    def orm(migration):
        return migration.orm()

    def forwards(self, migration):
        return self._wrap_direction(migration.forwards(), self.orm(migration))

    direction = forwards

    @staticmethod
    def record(migration):
        # Record us as having done this
        record = MigrationHistory.for_migration(migration)
        record.applied = datetime.datetime.utcnow()
        record.save()

    def format_backwards(self, migration):
        old_debug, old_dry_run = db.debug, db.dry_run
        db.debug = db.dry_run = True
        stdout = sys.stdout
        sys.stdout = StringIO()
        try:
            try:
                self.backwards(migration)()
                return sys.stdout.getvalue()
            except:
                raise
        finally:
            db.debug, db.dry_run = old_debug, old_dry_run
            sys.stdout = stdout

    def run_migration_error(self, migration, extra_info=''):
        extra_info = ('\n'
                      '! You *might* be able to recover with:'
                      '%s'
                      '%s' %
                      (self.format_backwards(migration), extra_info))
        return super(Forwards, self).run_migration_error(migration, extra_info)

    def migrate_many(self, target, migrations):
        try:
            for migration in migrations:
                result = self.migrate(migration)
                if result is False: # The migrations errored, but nicely.
                    return False
        finally:
            # Call any pending post_syncdb signals
            db.send_pending_create_signals()
        return True

class Backwards(Migrator):
    """
    Runs the specified migration backwards, in order.
    """
    torun = 'backwards'

    @staticmethod
    def title(target):
        if target is None:
            return " - Migrating backwards to zero state."
        else:
            return " - Migrating backwards to just after %s." % target.name()

    @staticmethod
    def status(migration):
        return ' < %s' % migration

    @staticmethod
    def orm(migration):
        return migration.prev_orm()

    direction = Migrator.backwards

    @staticmethod
    def record(migration):
        # Record us as having not done this
        record = MigrationHistory.for_migration(migration)
        if record.id is not None:
            record.delete()

    def migrate_many(self, target, migrations):
        for migration in migrations:
            self.migrate(migration)
        return True

<commit_msg>String formatting fix<commit_after>from copy import copy
from cStringIO import StringIO
import datetime
import inspect
import sys
import traceback

from django.core.management import call_command
from django.db import models

from south import exceptions
from south.db import db
from south.models import MigrationHistory
from south.signals import ran_migration

class Migrator(object):
    def __init__(self, verbosity=0):
        self.verbosity = int(verbosity)

    @staticmethod
    def title(target):
        raise NotImplementedError()

    def print_title(self, target):
        if self.verbosity:
            print self.title(target)

    @staticmethod
    def status(target):
        raise NotImplementedError()

    def print_status(self, migration):
        status = self.status(migration)
        if self.verbosity and status:
            print status

    @staticmethod
    def orm(migration):
        raise NotImplementedError()

    def backwards(self, migration):
        return self._wrap_direction(migration.backwards(), self.orm(migration))

    def direction(self, migration):
        raise NotImplementedError()

    @staticmethod
    def _wrap_direction(direction, orm):
        args = inspect.getargspec(direction)
        if len(args[0]) == 1:
            # Old migration, no ORM should be passed in
            return direction
        return (lambda: direction(orm))

    @staticmethod
    def record(migration):
        raise NotImplementedError()

    def run_migration_error(self, migration, extra_info=''):
        return (' ! Error found during real run of migration! Aborting.\n'
                '\n'
                ' ! Since you have a database that does not support running\n'
                ' ! schema-altering statements in transactions, we have had \n'
                ' ! to leave it in an interim state between migrations.\n'
                '%s\n'
                ' ! The South developers regret this has happened, and would\n'
                ' ! like to gently persuade you to consider a slightly\n'
                ' ! easier-to-deal-with DBMS.\n') % extra_info

    def run_migration(self, migration):
        migration_function = self.direction(migration)
        db.start_transaction()
        try:
            migration_function()
            db.execute_deferred_sql()
        except:
            db.rollback_transaction()
            if not db.has_ddl_transactions:
                print self.run_migration_error(migration)
            raise
        else:
            db.commit_transaction()

    def run(self, migration):
        # Get the correct ORM.
        db.current_orm = self.orm(migration)
        # If the database doesn't support running DDL inside a transaction
        # *cough*MySQL*cough* then do a dry run first.
        if not db.has_ddl_transactions:
            dry_run = DryRunMigrator(migrator=self, ignore_fail=False)
            dry_run.run_migration(migration)
        return self.run_migration(migration)

    def done_migrate(self, migration):
        db.start_transaction()
        try:
            # Record us as having done this
            self.record(migration)
        except:
            db.rollback_transaction()
            raise
        else:
            db.commit_transaction()

    def send_ran_migration(self, migration):
        ran_migration.send(None,
                           app=migration.app_label(),
                           migration=migration,
                           method=self.__class__.__name__.lower())

    def migrate(self, migration):
        """
        Runs the specified migration forwards/backwards, in order.
        """
        app = migration.migrations._migrations
        migration_name = migration.name()
        self.print_status(migration)
        result = self.run(migration)
        self.done_migrate(migration)
        self.send_ran_migration(migration)
        return result

    def migrate_many(self, target, migrations):
        raise NotImplementedError()

class MigratorWrapper(object):
    def __init__(self, migrator, *args, **kwargs):
        self._migrator = copy(migrator)
        attributes = dict([(k, getattr(self, k))
                           for k in self.__class__.__dict__.iterkeys()
                           if not k.startswith('__')])
        self._migrator.__dict__.update(attributes)

    def __getattr__(self, name):
        return getattr(self._migrator, name)

class DryRunMigrator(MigratorWrapper):
    def __init__(self, ignore_fail=True, *args, **kwargs):
        super(DryRunMigrator, self).__init__(*args, **kwargs)
        self._ignore_fail = ignore_fail

    def _run_migration(self, migration):
        if migration.no_dry_run() and self.verbosity:
            print " - Migration '%s' is marked for no-dry-run." % migration
            return
        db.dry_run = True
        db.debug, old_debug = False, db.debug
        pending_creates = db.get_pending_creates()
        db.start_transaction()
        migration_function = self.direction(migration)
        try:
            try:
                migration_function()
            except:
                raise exceptions.FailedDryRun(migration, sys.exc_info())
        finally:
            db.rollback_transactions_dry_run()
            db.debug = old_debug
            db.clear_run_data(pending_creates)
            db.dry_run = False

    def run_migration(self, migration):
        try:
            self._run_migration(migration)
        except exceptions.FailedDryRun:
            if self._ignore_fail:
                return False
            raise

    def done_migrate(self, *args, **kwargs):
        pass

    def send_ran_migration(self, *args, **kwargs):
        pass

class FakeMigrator(MigratorWrapper):
    def run(self, migration):
        if self.verbosity:
            print '   (faked)'

    def send_ran_migration(self, *args, **kwargs):
        pass

class LoadInitialDataMigrator(MigratorWrapper):
    def load_initial_data(self, target):
        if target is None or target != target.migrations[-1]:
            return
        # Load initial data, if we ended up at target
        if self.verbosity:
            print " - Loading initial data for %s." % target.app_label()
        # Override Django's get_apps call temporarily to only load from the
        # current app
        old_get_apps = models.get_apps
        models.get_apps = lambda: [models.get_app(target.app_label())]
        try:
            call_command('loaddata', 'initial_data', verbosity=self.verbosity)
        finally:
            models.get_apps = old_get_apps

    def migrate_many(self, target, migrations):
        migrator = self._migrator
        result = migrator.__class__.migrate_many(migrator, target, migrations)
        if result:
            self.load_initial_data(target)
        return True

class Forwards(Migrator):
    """
    Runs the specified migration forwards, in order.
    """
    torun = 'forwards'

    @staticmethod
    def title(target):
        if target is not None:
            return " - Migrating forwards to %s." % target.name()
        else:
            assert False, "You cannot migrate forwards to zero."

    @staticmethod
    def status(migration):
        return ' > %s' % migration

    @staticmethod
    def orm(migration):
        return migration.orm()

    def forwards(self, migration):
        return self._wrap_direction(migration.forwards(), self.orm(migration))

    direction = forwards

    @staticmethod
    def record(migration):
        # Record us as having done this
        record = MigrationHistory.for_migration(migration)
        record.applied = datetime.datetime.utcnow()
        record.save()

    def format_backwards(self, migration):
        old_debug, old_dry_run = db.debug, db.dry_run
        db.debug = db.dry_run = True
        stdout = sys.stdout
        sys.stdout = StringIO()
        try:
            try:
                self.backwards(migration)()
                return sys.stdout.getvalue()
            except:
                raise
        finally:
            db.debug, db.dry_run = old_debug, old_dry_run
            sys.stdout = stdout

    def run_migration_error(self, migration, extra_info=''):
        extra_info = ('\n'
                      '! You *might* be able to recover with:'
                      '%s'
                      '%s' %
                      (self.format_backwards(migration), extra_info))
        return super(Forwards, self).run_migration_error(migration, extra_info)

    def migrate_many(self, target, migrations):
        try:
            for migration in migrations:
                result = self.migrate(migration)
                if result is False: # The migrations errored, but nicely.
                    return False
        finally:
            # Call any pending post_syncdb signals
            db.send_pending_create_signals()
        return True

class Backwards(Migrator):
    """
    Runs the specified migration backwards, in order.
    """
    torun = 'backwards'

    @staticmethod
    def title(target):
        if target is None:
            return " - Migrating backwards to zero state."
        else:
            return " - Migrating backwards to just after %s." % target.name()

    @staticmethod
    def status(migration):
        return ' < %s' % migration

    @staticmethod
    def orm(migration):
        return migration.prev_orm()

    direction = Migrator.backwards

    @staticmethod
    def record(migration):
        # Record us as having not done this
        record = MigrationHistory.for_migration(migration)
        if record.id is not None:
            record.delete()

    def migrate_many(self, target, migrations):
        for migration in migrations:
            self.migrate(migration)
        return True

<|endoftext|>


<commit_before>import math, re

# uses m as base, each list[-1] is the number to multiply to convert to meters
# metric heights
mHeights = [['mm','millimetre','millimeter', 0.001],
			['cm','centimetre','centimeter', 0.01],
			['dm','decimetre','decimeter', 0.1],
			['m','metre','meter', 1]]
# imperial heights
iHeights = [['\'','ft','feet','foot', 0.3048],
			['\"','in','inch', 0.0254],
			['yr','yard', 0.9144]]

# uses kg as base, each list[-1] is the number to multiply to convert to kg
# metric masses
mMasses = [['mg','milligram', 0.000001],
			['cg','centigram', 0.00001],
			['dg','decigram', 0.0001],
			['g','gram', 0.001],
			['kg','kilogram', 1]]
# imperial masses
iMasses = [['oz','ou','ounce', 0.0283495],
			['lb','pound', 0.453592],
			['st','stone', 6.35029]]

# combined heights and masses
heights = mHeights + iHeights
masses = mMasses + iMasses

# lists of units for height and mass
metricUnits = [i for s in (mHeights + mMasses) for i in s if type(i) == str]
imperialUnits = [i for s in (iHeights + iMasses) for i in s if type(i) == str]
heightUnits = [i for s in heights for i in s if type(i) == str]
massUnits = [i for s in masses for i in s if type(i) == str]

def declare():
	return {'bmi': 'privmsg'}

def callback(self):
	user = self.user.split('!')[0].lower()
	message = self.message.split(self.command, 1)[1]

	co = count(message)
	bmi = co[0]
	mass = co[1]
	height = co[2]
	met = co[3]
	imp = co[4]

	if len(message) > 0:
		query = message.split()[0].lower()
		if query == 'set':
			bmi = mass / (height ** 2)
			if (bmi > 30 or bmi < 15) and not self.isop:
				return self.msg(self.channel, 'Please ask a bot operator to set your BMI for you.')

			try:
				self.locker.bmi[user] = bmi
			except:
				self.locker.bmi = {user: bmi}

			return self.msg(self.channel, 'Your BMI has been set to %s, which is \002\003%s\017.' % (format(bmi, '.2f'), classifyBmi(bmi)))

		if height and mass:
			bmi = mass / (height ** 2)
			return self.msg(self.channel, 'Your BMI is %s, you are \002\003%s\017.' % (format(bmi, '.2f'), classifyBmi(bmi)))

		if height and bmi:
			mass = bmi * (height ** 2)
			if 'imperial' in message or (imp > met and 'metric' not in message):
				return self.msg(self.channel, 'Your mass is %slbs.' % format(mass / 0.453592, '.2f'))
			return self.msg(self.channel, 'Your mass is %skg.' % format(mass, '.2f'))

		if mass and bmi:
			height = math.sqrt(mass / bmi)
			if 'imperial' in message or (imp > met and 'metric' not in message):
				iHeight = mToFtIn(height)
				return self.msg(self.channel, 'Your height is %d\'%d".' % (iHeight[0], iHeight[1]))
			return self.msg(self.channel, 'Your height is %sm.' % format(height, '.2f'))
	else:
		query = user

	try:
		if query in self.locker.bmi:
			bmi = self.locker.bmi[query]
			if query == user:
				return self.msg(self.channel, 'Your BMI is %s, which is \002\003%s\017.' % (format(bmi,'.2f'), classifyBmi(bmi)))
			else:
				return self.msg(self.channel, '%s\'s BMI is %s, which is \002\003%s\017.' % (message.split()[0], format(bmi,'.2f'), classifyBmi(bmi)))
		else:
			return self.msg(self.channel, 'This user has not set a BMI yet.')
	except:
		self.locker.bmi = dict()

def classifyBmi(bmi):
	if bmi < 18.5:
		return '08underweight'
	if bmi < 25.0:
		return '09normal'
	if bmi < 30.0:
		return '07FAT'
	else:
		return '04FAT AS FUCK'

def mToFtIn(meters):
	feet = int(meters // 0.3048)
	inches = int(meters / 0.3048 % 1 * 12)
	return [feet, inches]

def count(message):
	# counters for accumulated mass/height/bmi values
	mass = 0.0
	height = 0.0
	bmi = 0.0

	# counters for number of metric/imperial units
	metric = 0
	imperial = 0

	parameters = parseMessage(message)

	for parameter in parameters:
		p1 = ''.join(parameter[1:]).rstrip('es')

		# check if unit is imperial or metric (used for returning mass/height)
		if p1 in metricUnits:
			metric += 1
		elif p1 in imperialUnits:
			imperial += 1

		# check if unit is a height/mass unit or is 'bmi'
		if p1 in [u2.rstrip('es') for u2 in [u1 for u1 in heightUnits]]:
			for unit in heights:
				if p1 in [u.rstrip('es') for u in unit[:-1]]:
					height += parameter[0] * unit[-1]
					break
		if p1 in [u2.rstrip('es') for u2 in [u1 for u1 in massUnits]]:
			for unit in masses:
				if p1 in [u.rstrip('es') for u in unit[:-1]]:
					mass += parameter[0] * unit[-1]
					break
		elif p1 == 'bmi':
			bmi += parameter[0]

	return [bmi, mass, height, metric, imperial]

def split(text, separators):
	for separator in separators:
		text = text.replace(separator, ' %s ' % separator)

	return [i.strip() for i in text.split(' ')]

def parseMessage(message):
	parameters = []

	reFloat = re.compile('(\d+)?[.]?\d+')
	reString = re.compile('[^\.\d]+')

	messageSplit = split(message.lower(), heightUnits + massUnits + ['bmi'])

	for item in messageSplit:
		if item:
			floatSearch = reFloat.search(item)
			stringSearch = reString.search(item)

			# handles cases where no space between value and unit
			if floatSearch and stringSearch:
				parameters.append([float(reFloat.search(item).group()), reString.search(item).group()])
			elif floatSearch:
				parameters.append([float(reFloat.search(item).group())])
			# handles cases where there is a space between the value and unit
			elif stringSearch and len(parameters) > 0:
				parameters[-1].append(reString.search(item).group())

	# handles cases where inches is not included (ie: 5'6 for 5'6")
	for i, item in enumerate(parameters):
		if i > 0 and parameters[i - 1][1] == '\'' and len(item) < 2:
			item.append('\"')

	return parameters

class api:
	def msg(self, channel, text):
		return "[%s] %s" % (channel, text)
class empty:
	pass

'''
# interactive testing:
api = api()
setattr(api, 'type', 'privmsg')
setattr(api, 'command', 'bmi')
setattr(api, 'channel', "#test")
setattr(api, 'locker', empty)
setattr(api, 'user', 'joe!username@hostmask')
while(True):
	_input = raw_input('Enter message here: ')
	setattr(api, 'message', _input)
	print callback(api)
'''

if __name__ == "__main__":
	api = api()
	setattr(api, 'isop', True)
	setattr(api, 'type', 'privmsg')
	setattr(api, 'command', 'bmi')
	setattr(api, 'channel', "#test")
	setattr(api, 'locker', empty)

	setattr(api, 'user', 'joe!username@hostmask')
	setattr(api, 'message', '^bmi 5\'6\" 130pounds')
	if 'Your BMI is 20.98, you are' not in callback(api):
		exit(1)

	setattr(api, 'message', '^bmi set 5\' 6\" 130lbs')
	if 'Your BMI has been set to 20.98, which is' not in callback(api):
		exit(2)

	setattr(api, 'message', '^bmi')
	if 'Your BMI is 20.98, which is' not in callback(api):
		exit(3)

	setattr(api, 'message', '^bmi Joe')
	if 'Your BMI is 20.98, which is' not in callback(api) and 'Joe' not in callback(api):
		exit(4)

	setattr(api, 'user', 'blow!username@hostmask')
	setattr(api, 'message', '^bmi joe')
	if 'joe\'s BMI is 20.98, which is' not in callback(api):
		exit(5)

	setattr(api, 'message', '^bmi set 5\'6\" 280lbs')
	if 'Your BMI has been set to ' not in callback(api):
		exit(6)

	setattr(api, 'isop', False)
	setattr(api, 'message', '^bmi set 5\'6\" 280lbs')
	if 'Please ask a bot operator to set your BMI for you.' not in callback(api):
		exit(7)

	setattr(api, 'message', '^bmi 5\'6\" 20.98bmi')
	if 'Your mass is 129.99lbs.' not in callback(api):
		exit(8)

	setattr(api, 'message', '^bmi 130lbs 20.98bmi')
	if 'Your height is 5\'6".' not in callback(api):
		exit(9)
<commit_msg>Fixed error messages, allow ops to set other's BMIs<commit_after>import math, re

# uses m as base, each list[-1] is the number to multiply to convert to meters
# metric heights
mHeights = [['mm','millimetre','millimeter', 0.001],
			['cm','centimetre','centimeter', 0.01],
			['dm','decimetre','decimeter', 0.1],
			['m','metre','meter', 1]]
# imperial heights
iHeights = [['\'','ft','feet','foot', 0.3048],
			['\"','in','inch', 0.0254],
			['yr','yard', 0.9144]]

# uses kg as base, each list[-1] is the number to multiply to convert to kg
# metric masses
mMasses = [['mg','milligram', 0.000001],
			['cg','centigram', 0.00001],
			['dg','decigram', 0.0001],
			['g','gram', 0.001],
			['kg','kilogram', 1]]
# imperial masses
iMasses = [['oz','ou','ounce', 0.0283495],
			['lb','pound', 0.453592],
			['st','stone', 6.35029]]

# combined heights and masses
heights = mHeights + iHeights
masses = mMasses + iMasses

# lists of units for height and mass
metricUnits = [i for s in (mHeights + mMasses) for i in s if type(i) == str]
imperialUnits = [i for s in (iHeights + iMasses) for i in s if type(i) == str]
heightUnits = [i for s in heights for i in s if type(i) == str]
massUnits = [i for s in masses for i in s if type(i) == str]

def declare():
	return {'bmi': 'privmsg'}

def callback(self):
	user = self.user.split('!')[0].lower()
	message = self.message.split(self.command, 1)[1]

	co = count(message)
	bmi = co[0]
	mass = co[1]
	height = co[2]
	met = co[3]
	imp = co[4]

	if len(message) > 0:
		query = message.split()[0].lower()
		if query == 'set':
			bmi = mass / (height ** 2)
			if (bmi > 30 or bmi < 15):
				if self.isop and type(message.split()[1].lower()) == str:
					user = message.split()[1].lower()
				return self.msg(self.channel, 'Please ask a bot operator to set your BMI for you.')

			try:
				self.locker.bmi[user] = bmi
			except:
				self.locker.bmi = {user: bmi}

			if user == self.user.split('!')[0].lower():
				return self.msg(self.channel, 'Your BMI has been set to %s, which is \002\003%s\017.' % (userformat(bmi, '.2f'), classifyBmi(bmi)))
			else:
				return self.msg(self.channel, '%s\'s BMI has been set to %s, which is \002\003%s\017.' % (message.split()[1], userformat(bmi, '.2f'), classifyBmi(bmi)))

		if height and mass:
			bmi = mass / (height ** 2)
			return self.msg(self.channel, 'Your BMI is %s, you are \002\003%s\017.' % (format(bmi, '.2f'), classifyBmi(bmi)))

		if height and bmi:
			mass = bmi * (height ** 2)
			if 'imperial' in message or (imp > met and 'metric' not in message):
				return self.msg(self.channel, 'Your mass is %slbs.' % format(mass / 0.453592, '.2f'))
			return self.msg(self.channel, 'Your mass is %skg.' % format(mass, '.2f'))

		if mass and bmi:
			height = math.sqrt(mass / bmi)
			if 'imperial' in message or (imp > met and 'metric' not in message):
				iHeight = mToFtIn(height)
				return self.msg(self.channel, 'Your height is %d\'%d".' % (iHeight[0], iHeight[1]))
			return self.msg(self.channel, 'Your height is %sm.' % format(height, '.2f'))
	else:
		query = user

	try:
		if query in self.locker.bmi:
			bmi = self.locker.bmi[query]
			if query == user:
				return self.msg(self.channel, 'Your BMI is %s, which is \002\003%s\017.' % (format(bmi,'.2f'), classifyBmi(bmi)))
			else:
				return self.msg(self.channel, '%s\'s BMI is %s, which is \002\003%s\017.' % (message.split()[0], format(bmi,'.2f'), classifyBmi(bmi)))
		else:
			if len(message.split()) < 2:
				return self.msg(self.channel, 'This user has not set a BMI yet.')
			else:
				return self.msg(self.channel, 'Invalid input.')
	except:
		self.locker.bmi = dict()
		return self.msg(self.channel, 'No user\'s BMI has been recorded yet.')

def classifyBmi(bmi):
	if bmi < 18.5:
		return '08underweight'
	if bmi < 25.0:
		return '09normal'
	if bmi < 30.0:
		return '07FAT'
	else:
		return '04FAT AS FUCK'

def mToFtIn(meters):
	feet = int(meters // 0.3048)
	inches = int(meters / 0.3048 % 1 * 12)
	return [feet, inches]

def count(message):
	# counters for accumulated mass/height/bmi values
	mass = 0.0
	height = 0.0
	bmi = 0.0

	# counters for number of metric/imperial units
	metric = 0
	imperial = 0

	parameters = parseMessage(message)

	for parameter in parameters:
		p1 = ''.join(parameter[1:]).rstrip('es')

		# check if unit is imperial or metric (used for returning mass/height)
		if p1 in metricUnits:
			metric += 1
		elif p1 in imperialUnits:
			imperial += 1

		# check if unit is a height/mass unit or is 'bmi'
		if p1 in [u2.rstrip('es') for u2 in [u1 for u1 in heightUnits]]:
			for unit in heights:
				if p1 in [u.rstrip('es') for u in unit[:-1]]:
					height += parameter[0] * unit[-1]
					break
		if p1 in [u2.rstrip('es') for u2 in [u1 for u1 in massUnits]]:
			for unit in masses:
				if p1 in [u.rstrip('es') for u in unit[:-1]]:
					mass += parameter[0] * unit[-1]
					break
		elif p1 == 'bmi':
			bmi += parameter[0]

	return [bmi, mass, height, metric, imperial]

def split(text, separators):
	for separator in separators:
		text = text.replace(separator, ' %s ' % separator)

	return [i.strip() for i in text.split(' ')]

def parseMessage(message):
	parameters = []

	reFloat = re.compile('(\d+)?[.]?\d+')
	reString = re.compile('[^\.\d]+')

	messageSplit = split(message.lower(), heightUnits + massUnits + ['bmi'])

	for item in messageSplit:
		if item:
			floatSearch = reFloat.search(item)
			stringSearch = reString.search(item)

			# handles cases where no space between value and unit
			if floatSearch and stringSearch:
				parameters.append([float(reFloat.search(item).group()), reString.search(item).group()])
			elif floatSearch:
				parameters.append([float(reFloat.search(item).group())])
			# handles cases where there is a space between the value and unit
			elif stringSearch and len(parameters) > 0:
				parameters[-1].append(reString.search(item).group())

	# handles cases where inches is not included (ie: 5'6 for 5'6")
	for i, item in enumerate(parameters):
		if i > 0 and parameters[i - 1][1] == '\'' and len(item) < 2:
			item.append('\"')

	return parameters

class api:
	def msg(self, channel, text):
		return "[%s] %s" % (channel, text)
class empty:
	pass

'''
# interactive testing:
api = api()
setattr(api, 'type', 'privmsg')
setattr(api, 'command', 'bmi')
setattr(api, 'channel', "#test")
setattr(api, 'locker', empty)
setattr(api, 'user', 'joe!username@hostmask')
while(True):
	_input = raw_input('Enter message here: ')
	setattr(api, 'message', _input)
	print callback(api)
'''

if __name__ == "__main__":
	api = api()
	setattr(api, 'isop', True)
	setattr(api, 'type', 'privmsg')
	setattr(api, 'command', 'bmi')
	setattr(api, 'channel', "#test")
	setattr(api, 'locker', empty)

	setattr(api, 'user', 'joe!username@hostmask')
	setattr(api, 'message', '^bmi 5\'6\" 130pounds')
	if 'Your BMI is 20.98, you are' not in callback(api):
		exit(1)

	setattr(api, 'message', '^bmi set 5\' 6\" 130lbs')
	if 'Your BMI has been set to 20.98, which is' not in callback(api):
		exit(2)

	setattr(api, 'message', '^bmi')
	if 'Your BMI is 20.98, which is' not in callback(api):
		exit(3)

	setattr(api, 'message', '^bmi Joe')
	if 'Your BMI is 20.98, which is' not in callback(api) and 'Joe' not in callback(api):
		exit(4)

	setattr(api, 'user', 'blow!username@hostmask')
	setattr(api, 'message', '^bmi joe')
	if 'joe\'s BMI is 20.98, which is' not in callback(api):
		exit(5)

	setattr(api, 'message', '^bmi set 5\'6\" 280lbs')
	if 'Your BMI has been set to ' not in callback(api):
		exit(6)

	setattr(api, 'isop', False)
	setattr(api, 'message', '^bmi set 5\'6\" 280lbs')
	if 'Please ask a bot operator to set your BMI for you.' not in callback(api):
		exit(7)

	setattr(api, 'isop', True)
	setattr(api, 'message', '^bmi set Frank 5\'6\" 280lbs')
	if 'Frank\'s BMI has been set to' not in callback(api):
		exit(8)

	setattr(api, 'message', '^bmi 5\'6\" 20.98bmi')
	if 'Your mass is 129.99lbs.' not in callback(api):
		exit(9)

	setattr(api, 'message', '^bmi 130lbs 20.98bmi')
	if 'Your height is 5\'6".' not in callback(api):
		exit(10)
<|endoftext|>


<commit_before>import csv
import re
from asyncio import gather, get_event_loop
from functools import partial
from io import BytesIO, TextIOWrapper
from itertools import islice
from os import link, mkfifo, path
from ruamel import yaml
from socket import socket, AF_UNIX, SOCK_STREAM, SOCK_NONBLOCK
from threading import RLock
from zipfile import ZipFile, BadZipFile

from jd4._compare import compare_stream
from jd4.cgroup import wait_cgroup
from jd4.compile import build
from jd4.error import FormatError
from jd4.pool import get_sandbox, put_sandbox
from jd4.status import STATUS_ACCEPTED, STATUS_WRONG_ANSWER, \
                       STATUS_TIME_LIMIT_EXCEEDED, STATUS_MEMORY_LIMIT_EXCEEDED, \
                       STATUS_RUNTIME_ERROR, STATUS_SYSTEM_ERROR
from jd4.util import read_pipe, parse_memory_bytes, parse_time_ns

CHUNK_SIZE = 32768
MAX_STDERR_SIZE = 8192
DEFAULT_TIME_NS = 1000000000
DEFAULT_MEMORY_BYTES = 268435456
PROCESS_LIMIT = 64

class CaseBase:
    def __init__(self, time_limit_ns, memory_limit_bytes, process_limit, score):
        self.time_limit_ns = time_limit_ns
        self.memory_limit_bytes = memory_limit_bytes
        self.process_limit = process_limit
        self.score = score

    async def judge(self, package):
        loop = get_event_loop()
        sandbox, = await get_sandbox(1)
        try:
            executable = await package.install(sandbox)
            stdin_file = path.join(sandbox.in_dir, 'stdin')
            mkfifo(stdin_file)
            stdout_file = path.join(sandbox.in_dir, 'stdout')
            mkfifo(stdout_file)
            stderr_file = path.join(sandbox.in_dir, 'stderr')
            mkfifo(stderr_file)
            with socket(AF_UNIX, SOCK_STREAM | SOCK_NONBLOCK) as cgroup_sock:
                cgroup_sock.bind(path.join(sandbox.in_dir, 'cgroup'))
                cgroup_sock.listen()
                execute_task = loop.create_task(executable.execute(
                    sandbox,
                    stdin_file='/in/stdin',
                    stdout_file='/in/stdout',
                    stderr_file='/in/stderr',
                    cgroup_file='/in/cgroup'))
                others_task = gather(
                    loop.run_in_executor(None, self.do_input, stdin_file),
                    loop.run_in_executor(None, self.do_output, stdout_file),
                    read_pipe(stderr_file, MAX_STDERR_SIZE),
                    wait_cgroup(cgroup_sock,
                                execute_task,
                                self.time_limit_ns,
                                self.time_limit_ns,
                                self.memory_limit_bytes,
                                self.process_limit))
                execute_status = await execute_task
                _, correct, stderr, (time_usage_ns, memory_usage_bytes) = \
                    await others_task
            if memory_usage_bytes >= self.memory_limit_bytes:
                status = STATUS_MEMORY_LIMIT_EXCEEDED
                score = 0
            elif time_usage_ns >= self.time_limit_ns:
                status = STATUS_TIME_LIMIT_EXCEEDED
                score = 0
            elif execute_status:
                status = STATUS_RUNTIME_ERROR
                score = 0
            elif not correct:
                status = STATUS_WRONG_ANSWER
                score = 0
            else:
                status = STATUS_ACCEPTED
                score = self.score
            return status, score, time_usage_ns, memory_usage_bytes, stderr
        finally:
            put_sandbox(sandbox)

def dos2unix(src, dst):
    while True:
        buf = src.read(CHUNK_SIZE)
        if not buf:
            break
        buf = buf.replace(b'\r', b'')
        dst.write(buf)

class DefaultCase(CaseBase):
    def __init__(self, open_input, open_output, time_ns, memory_bytes, score):
        super().__init__(time_ns, memory_bytes, PROCESS_LIMIT, score)
        self.open_input = open_input
        self.open_output = open_output

    def do_input(self, input_file):
        try:
            with open(input_file, 'wb') as dst, self.open_input() as src:
                dos2unix(src, dst)
        except BrokenPipeError:
            pass

    def do_output(self, output_file):
        with open(output_file, 'rb') as out, self.open_output() as ans:
            return compare_stream(ans, out)

class CustomJudgeCase:
    def __init__(self, open_input, time_ns, memory_bytes, open_judge, judge_lang):
        self.open_input = open_input
        self.time_ns = time_ns
        self.memory_bytes = memory_bytes
        self.open_judge = open_judge
        self.judge_lang = judge_lang

    async def judge(self, user_package):
        loop = get_event_loop()
        judge_package, message, _, _ = await build(
            self.judge_lang,
            await loop.run_in_executor(None, lambda: self.open_judge().read()))
        if not judge_package:
            return STATUS_SYSTEM_ERROR, 0, 0, 0, message
        user_sandbox, judge_sandbox = await get_sandbox(2)
        try:
            async def prepare_user_sandbox():
                await user_sandbox.reset()
                return await user_package.install(user_sandbox)

            async def prepare_judge_sandbox():
                await judge_sandbox.reset()
                return await judge_package.install(judge_sandbox)

            user_executable, judge_executable = \
                await gather(prepare_user_sandbox(), prepare_judge_sandbox())
            user_stdin_file = path.join(user_sandbox.in_dir, 'stdin')
            mkfifo(user_stdin_file)
            user_stdout_file = path.join(user_sandbox.in_dir, 'stdout')
            mkfifo(user_stdout_file)
            judge_stdin_file = path.join(judge_sandbox.in_dir, 'stdin')
            link(user_stdout_file, judge_stdin_file)
            user_stderr_file = path.join(user_sandbox.in_dir, 'stderr')
            mkfifo(user_stderr_file)
            judge_stdout_file = path.join(judge_sandbox.in_dir, 'stdout')
            mkfifo(judge_stdout_file)
            judge_stderr_file = path.join(judge_sandbox.in_dir, 'stderr')
            mkfifo(judge_stderr_file)
            judge_extra_file = path.join(judge_sandbox.in_dir, 'extra')
            mkfifo(judge_extra_file)
            with socket(AF_UNIX, SOCK_STREAM | SOCK_NONBLOCK) as user_cgroup_sock, \
                 socket(AF_UNIX, SOCK_STREAM | SOCK_NONBLOCK) as judge_cgroup_sock:
                user_cgroup_sock.bind(path.join(user_sandbox.in_dir, 'cgroup'))
                judge_cgroup_sock.bind(path.join(judge_sandbox.in_dir, 'cgroup'))
                user_cgroup_sock.listen()
                judge_cgroup_sock.listen()
                user_execute_task = loop.create_task(user_executable.execute(
                    user_sandbox,
                    stdin_file='/in/stdin',
                    stdout_file='/in/stdout',
                    stderr_file='/in/stderr',
                    cgroup_file='/in/cgroup'))
                judge_execute_task = loop.create_task(judge_executable.execute(
                    judge_sandbox,
                    stdin_file='/in/stdin',
                    stdout_file='/in/stdout',
                    stderr_file='/in/stderr',
                    extra_file='/in/extra',
                    cgroup_file='/in/cgroup'))
                others_task = gather(
                    loop.run_in_executor(None, self.do_input, user_stdin_file),
                    loop.run_in_executor(None, self.do_input, judge_extra_file),
                    read_pipe(user_stderr_file, MAX_STDERR_SIZE),
                    read_pipe(judge_stdout_file, MAX_STDERR_SIZE),
                    read_pipe(judge_stderr_file, MAX_STDERR_SIZE),
                    wait_cgroup(user_cgroup_sock,
                                user_execute_task,
                                self.time_ns,
                                self.time_ns,
                                self.memory_bytes,
                                PROCESS_LIMIT),
                    wait_cgroup(judge_cgroup_sock,
                                judge_execute_task,
                                DEFAULT_TIME_NS,
                                self.time_ns + DEFAULT_TIME_NS,
                                DEFAULT_MEMORY_BYTES,
                                PROCESS_LIMIT))
                user_execute_status, judge_execute_status = await gather(
                    user_execute_task, judge_execute_task)
                _, _, user_stderr, judge_stdout, judge_stderr, \
                (user_time_usage_ns, user_memory_usage_bytes), \
                (judge_time_usage_ns, judge_memory_usage_bytes) = \
                    await others_task
            if (judge_execute_status or
                judge_memory_usage_bytes >= DEFAULT_MEMORY_BYTES or
                judge_time_usage_ns >= DEFAULT_TIME_NS):
                status = STATUS_SYSTEM_ERROR
                score = 0
            elif user_memory_usage_bytes >= self.memory_bytes:
                status = STATUS_MEMORY_LIMIT_EXCEEDED
                score = 0
            elif user_time_usage_ns >= self.time_ns:
                status = STATUS_TIME_LIMIT_EXCEEDED
                score = 0
            elif user_execute_status:
                status = STATUS_RUNTIME_ERROR
                score = 0
            else:
                try:
                    status, score = map(int, judge_stdout.split())
                except SystemError:
                    status = STATUS_SYSTEM_ERROR
                    score = 0
            return status, score, user_time_usage_ns, user_memory_usage_bytes, user_stderr
        finally:
            put_sandbox(user_sandbox, judge_sandbox)

    def do_input(self, input_file):
        try:
            with open(input_file, 'wb') as dst, self.open_input() as src:
                dos2unix(src, dst)
        except BrokenPipeError:
            pass

class APlusBCase(CaseBase):
    def __init__(self, a, b, time_limit_ns, memory_limit_bytes, score):
        super().__init__(time_limit_ns, memory_limit_bytes, PROCESS_LIMIT, score)
        self.a = a
        self.b = b

    def do_input(self, input_file):
        try:
            with open(input_file, 'w') as file:
                file.write('{} {}\n'.format(self.a, self.b))
        except BrokenPipeError:
            pass

    def do_output(self, output_file):
        with open(output_file, 'rb') as file:
            return compare_stream(BytesIO(str(self.a + self.b).encode()), file)

def read_legacy_cases(config, open):
    num_cases = int(config.readline())
    for line in islice(csv.reader(config, delimiter='|'), num_cases):
        input, output, time_str, score_str = line[:4]
        try:
            memory_bytes = int(float(line[4]) * 1024)
        except (IndexError, ValueError):
            memory_bytes = DEFAULT_MEMORY_BYTES
        yield DefaultCase(partial(open, path.join('input', input)),
                          partial(open, path.join('output', output)),
                          int(float(time_str) * 1000000000),
                          memory_bytes,
                          int(score_str))

def read_yaml_cases(config, open):
    for case in yaml.safe_load(config)['cases']:
        if 'judge' not in case:
            yield DefaultCase(partial(open, case['input']),
                              partial(open, case['output']),
                              parse_time_ns(case['time']),
                              parse_memory_bytes(case['memory']),
                              int(case['score']))
        else:
            yield CustomJudgeCase(partial(open, case['input']),
                                  parse_time_ns(case['time']),
                                  parse_memory_bytes(case['memory']),
                                  partial(open, case['judge']),
                                  path.splitext(case['judge'])[1][1:])

def read_cases(file):
    zip_file = ZipFile(file)
    canonical_dict = dict((name.lower(), name)
                          for name in zip_file.namelist())

    def open(name):
        try:
            return zip_file.open(canonical_dict[name.lower()])
        except KeyError:
            raise FileNotFoundError(name) from None
    try:
        config = TextIOWrapper(open('config.ini'),
                               encoding='utf-8', errors='replace')
        return read_legacy_cases(config, open)
    except FileNotFoundError:
        pass
    try:
        config = open('config.yaml')
        return read_yaml_cases(config, open)
    except FileNotFoundError:
        pass
    raise FormatError('config file not found')
<commit_msg>remove unused vars (#81)<commit_after>import csv
from asyncio import gather, get_event_loop
from functools import partial
from io import BytesIO, TextIOWrapper
from itertools import islice
from os import link, mkfifo, path
from ruamel import yaml
from socket import socket, AF_UNIX, SOCK_STREAM, SOCK_NONBLOCK
from zipfile import ZipFile

from jd4._compare import compare_stream
from jd4.cgroup import wait_cgroup
from jd4.compile import build
from jd4.error import FormatError
from jd4.pool import get_sandbox, put_sandbox
from jd4.status import STATUS_ACCEPTED, STATUS_WRONG_ANSWER, \
                       STATUS_TIME_LIMIT_EXCEEDED, STATUS_MEMORY_LIMIT_EXCEEDED, \
                       STATUS_RUNTIME_ERROR, STATUS_SYSTEM_ERROR
from jd4.util import read_pipe, parse_memory_bytes, parse_time_ns

CHUNK_SIZE = 32768
MAX_STDERR_SIZE = 8192
DEFAULT_TIME_NS = 1000000000
DEFAULT_MEMORY_BYTES = 268435456
PROCESS_LIMIT = 64

class CaseBase:
    def __init__(self, time_limit_ns, memory_limit_bytes, process_limit, score):
        self.time_limit_ns = time_limit_ns
        self.memory_limit_bytes = memory_limit_bytes
        self.process_limit = process_limit
        self.score = score

    async def judge(self, package):
        loop = get_event_loop()
        sandbox, = await get_sandbox(1)
        try:
            executable = await package.install(sandbox)
            stdin_file = path.join(sandbox.in_dir, 'stdin')
            mkfifo(stdin_file)
            stdout_file = path.join(sandbox.in_dir, 'stdout')
            mkfifo(stdout_file)
            stderr_file = path.join(sandbox.in_dir, 'stderr')
            mkfifo(stderr_file)
            with socket(AF_UNIX, SOCK_STREAM | SOCK_NONBLOCK) as cgroup_sock:
                cgroup_sock.bind(path.join(sandbox.in_dir, 'cgroup'))
                cgroup_sock.listen()
                execute_task = loop.create_task(executable.execute(
                    sandbox,
                    stdin_file='/in/stdin',
                    stdout_file='/in/stdout',
                    stderr_file='/in/stderr',
                    cgroup_file='/in/cgroup'))
                others_task = gather(
                    loop.run_in_executor(None, self.do_input, stdin_file),
                    loop.run_in_executor(None, self.do_output, stdout_file),
                    read_pipe(stderr_file, MAX_STDERR_SIZE),
                    wait_cgroup(cgroup_sock,
                                execute_task,
                                self.time_limit_ns,
                                self.time_limit_ns,
                                self.memory_limit_bytes,
                                self.process_limit))
                execute_status = await execute_task
                _, correct, stderr, (time_usage_ns, memory_usage_bytes) = \
                    await others_task
            if memory_usage_bytes >= self.memory_limit_bytes:
                status = STATUS_MEMORY_LIMIT_EXCEEDED
                score = 0
            elif time_usage_ns >= self.time_limit_ns:
                status = STATUS_TIME_LIMIT_EXCEEDED
                score = 0
            elif execute_status:
                status = STATUS_RUNTIME_ERROR
                score = 0
            elif not correct:
                status = STATUS_WRONG_ANSWER
                score = 0
            else:
                status = STATUS_ACCEPTED
                score = self.score
            return status, score, time_usage_ns, memory_usage_bytes, stderr
        finally:
            put_sandbox(sandbox)

def dos2unix(src, dst):
    while True:
        buf = src.read(CHUNK_SIZE)
        if not buf:
            break
        buf = buf.replace(b'\r', b'')
        dst.write(buf)

class DefaultCase(CaseBase):
    def __init__(self, open_input, open_output, time_ns, memory_bytes, score):
        super().__init__(time_ns, memory_bytes, PROCESS_LIMIT, score)
        self.open_input = open_input
        self.open_output = open_output

    def do_input(self, input_file):
        try:
            with open(input_file, 'wb') as dst, self.open_input() as src:
                dos2unix(src, dst)
        except BrokenPipeError:
            pass

    def do_output(self, output_file):
        with open(output_file, 'rb') as out, self.open_output() as ans:
            return compare_stream(ans, out)

class CustomJudgeCase:
    def __init__(self, open_input, time_ns, memory_bytes, open_judge, judge_lang):
        self.open_input = open_input
        self.time_ns = time_ns
        self.memory_bytes = memory_bytes
        self.open_judge = open_judge
        self.judge_lang = judge_lang

    async def judge(self, user_package):
        loop = get_event_loop()
        judge_package, message, _, _ = await build(
            self.judge_lang,
            await loop.run_in_executor(None, lambda: self.open_judge().read()))
        if not judge_package:
            return STATUS_SYSTEM_ERROR, 0, 0, 0, message
        user_sandbox, judge_sandbox = await get_sandbox(2)
        try:
            async def prepare_user_sandbox():
                await user_sandbox.reset()
                return await user_package.install(user_sandbox)

            async def prepare_judge_sandbox():
                await judge_sandbox.reset()
                return await judge_package.install(judge_sandbox)

            user_executable, judge_executable = \
                await gather(prepare_user_sandbox(), prepare_judge_sandbox())
            user_stdin_file = path.join(user_sandbox.in_dir, 'stdin')
            mkfifo(user_stdin_file)
            user_stdout_file = path.join(user_sandbox.in_dir, 'stdout')
            mkfifo(user_stdout_file)
            judge_stdin_file = path.join(judge_sandbox.in_dir, 'stdin')
            link(user_stdout_file, judge_stdin_file)
            user_stderr_file = path.join(user_sandbox.in_dir, 'stderr')
            mkfifo(user_stderr_file)
            judge_stdout_file = path.join(judge_sandbox.in_dir, 'stdout')
            mkfifo(judge_stdout_file)
            judge_stderr_file = path.join(judge_sandbox.in_dir, 'stderr')
            mkfifo(judge_stderr_file)
            judge_extra_file = path.join(judge_sandbox.in_dir, 'extra')
            mkfifo(judge_extra_file)
            with socket(AF_UNIX, SOCK_STREAM | SOCK_NONBLOCK) as user_cgroup_sock, \
                 socket(AF_UNIX, SOCK_STREAM | SOCK_NONBLOCK) as judge_cgroup_sock:
                user_cgroup_sock.bind(path.join(user_sandbox.in_dir, 'cgroup'))
                judge_cgroup_sock.bind(path.join(judge_sandbox.in_dir, 'cgroup'))
                user_cgroup_sock.listen()
                judge_cgroup_sock.listen()
                user_execute_task = loop.create_task(user_executable.execute(
                    user_sandbox,
                    stdin_file='/in/stdin',
                    stdout_file='/in/stdout',
                    stderr_file='/in/stderr',
                    cgroup_file='/in/cgroup'))
                judge_execute_task = loop.create_task(judge_executable.execute(
                    judge_sandbox,
                    stdin_file='/in/stdin',
                    stdout_file='/in/stdout',
                    stderr_file='/in/stderr',
                    extra_file='/in/extra',
                    cgroup_file='/in/cgroup'))
                others_task = gather(
                    loop.run_in_executor(None, self.do_input, user_stdin_file),
                    loop.run_in_executor(None, self.do_input, judge_extra_file),
                    read_pipe(user_stderr_file, MAX_STDERR_SIZE),
                    read_pipe(judge_stdout_file, MAX_STDERR_SIZE),
                    read_pipe(judge_stderr_file, MAX_STDERR_SIZE),
                    wait_cgroup(user_cgroup_sock,
                                user_execute_task,
                                self.time_ns,
                                self.time_ns,
                                self.memory_bytes,
                                PROCESS_LIMIT),
                    wait_cgroup(judge_cgroup_sock,
                                judge_execute_task,
                                DEFAULT_TIME_NS,
                                self.time_ns + DEFAULT_TIME_NS,
                                DEFAULT_MEMORY_BYTES,
                                PROCESS_LIMIT))
                user_execute_status, judge_execute_status = await gather(
                    user_execute_task, judge_execute_task)
                _, _, user_stderr, judge_stdout, judge_stderr, \
                (user_time_usage_ns, user_memory_usage_bytes), \
                (judge_time_usage_ns, judge_memory_usage_bytes) = \
                    await others_task
            if (judge_execute_status or
                judge_memory_usage_bytes >= DEFAULT_MEMORY_BYTES or
                judge_time_usage_ns >= DEFAULT_TIME_NS):
                status = STATUS_SYSTEM_ERROR
                score = 0
            elif user_memory_usage_bytes >= self.memory_bytes:
                status = STATUS_MEMORY_LIMIT_EXCEEDED
                score = 0
            elif user_time_usage_ns >= self.time_ns:
                status = STATUS_TIME_LIMIT_EXCEEDED
                score = 0
            elif user_execute_status:
                status = STATUS_RUNTIME_ERROR
                score = 0
            else:
                try:
                    status, score = map(int, judge_stdout.split())
                except SystemError:
                    status = STATUS_SYSTEM_ERROR
                    score = 0
            return status, score, user_time_usage_ns, user_memory_usage_bytes, user_stderr
        finally:
            put_sandbox(user_sandbox, judge_sandbox)

    def do_input(self, input_file):
        try:
            with open(input_file, 'wb') as dst, self.open_input() as src:
                dos2unix(src, dst)
        except BrokenPipeError:
            pass

class APlusBCase(CaseBase):
    def __init__(self, a, b, time_limit_ns, memory_limit_bytes, score):
        super().__init__(time_limit_ns, memory_limit_bytes, PROCESS_LIMIT, score)
        self.a = a
        self.b = b

    def do_input(self, input_file):
        try:
            with open(input_file, 'w') as file:
                file.write('{} {}\n'.format(self.a, self.b))
        except BrokenPipeError:
            pass

    def do_output(self, output_file):
        with open(output_file, 'rb') as file:
            return compare_stream(BytesIO(str(self.a + self.b).encode()), file)

def read_legacy_cases(config, open):
    num_cases = int(config.readline())
    for line in islice(csv.reader(config, delimiter='|'), num_cases):
        input, output, time_str, score_str = line[:4]
        try:
            memory_bytes = int(float(line[4]) * 1024)
        except (IndexError, ValueError):
            memory_bytes = DEFAULT_MEMORY_BYTES
        yield DefaultCase(partial(open, path.join('input', input)),
                          partial(open, path.join('output', output)),
                          int(float(time_str) * 1000000000),
                          memory_bytes,
                          int(score_str))

def read_yaml_cases(config, open):
    for case in yaml.safe_load(config)['cases']:
        if 'judge' not in case:
            yield DefaultCase(partial(open, case['input']),
                              partial(open, case['output']),
                              parse_time_ns(case['time']),
                              parse_memory_bytes(case['memory']),
                              int(case['score']))
        else:
            yield CustomJudgeCase(partial(open, case['input']),
                                  parse_time_ns(case['time']),
                                  parse_memory_bytes(case['memory']),
                                  partial(open, case['judge']),
                                  path.splitext(case['judge'])[1][1:])

def read_cases(file):
    zip_file = ZipFile(file)
    canonical_dict = dict((name.lower(), name)
                          for name in zip_file.namelist())

    def open(name):
        try:
            return zip_file.open(canonical_dict[name.lower()])
        except KeyError:
            raise FileNotFoundError(name) from None
    try:
        config = TextIOWrapper(open('config.ini'),
                               encoding='utf-8', errors='replace')
        return read_legacy_cases(config, open)
    except FileNotFoundError:
        pass
    try:
        config = open('config.yaml')
        return read_yaml_cases(config, open)
    except FileNotFoundError:
        pass
    raise FormatError('config file not found')
<|endoftext|>