<commit_before>from gevent import monkey; monkey.patch_all()

import os
import time
import uuid
import resource
import random
import string
from inspect import getcallargs

import gevent
import msgpack
import rocksdb
from decorator import decorator
from funcserver import RPCServer, RPCClient, BaseHandler

ITERATOR_EXPIRY_CHECK = 5 * 60 # 5 minutes
ITERATOR_EXPIRE = 15 * 60 # 15 minutes
MAX_OPEN_FILES = 500000
ALPHANUM = string.letters + string.digits

def gen_random_seq(length=10):
    return ''.join([random.choice(ALPHANUM) for i in xrange(length)])

@decorator
def ensuretable(fn, self, table, *args, **kwargs):
    if table not in self.tables:
        raise Exception('Table "%s" does not exist' % table)
    return fn(self, self.tables[table], *args, **kwargs)

@decorator
def ensureiter(fn, self, _iter, *args, **kwargs):
    if _iter not in self.iters:
        raise Exception('Iter "%s" does not exist' % _iter)

    _iter = self.iters[_iter]
    _iter.ts_last_activity = time.time()
    return fn(self, _iter, *args, **kwargs)

@decorator
def ensurenewiter(fn, self, *args, **kwargs):
    name = getcallargs(fn, *args, **kwargs)['name'] or gen_random_seq()
    if args: args = list(args); args[0] = name
    else: kwargs['name'] = name

    if name in self.iters:
        raise Exception('iter "%s" exists already!' % name)

    return fn(self, *args, **kwargs)

class Iterator(object):
    NUM_RECORDS = 1000

    def __init__(self, table, prefix=None, type='items', reverse=False):
        '''
        @type: str; can be items/keys/values
        '''
        self.table = table
        self.type = type
        self.prefix = prefix
        self.reverse = reverse
        self.ts_last_activity = time.time()

        iterfn = getattr(self.table.rdb,
            {'keys': 'iterkeys', 'values': 'itervalues'}\
            .get(type, 'iteritems'))

        self._iter = iterfn(prefix=prefix)

        if reverse:
            self._iter = reversed(self._iter)

    def get(self, num=NUM_RECORDS):
        records = []

        for record in self._iter:
            if self.type == 'items':
                key, item = record
                item = self.table.unpackfn(item)
                record = (key, item)
            elif self.type == 'values':
                record = self.table.unpackfn(record)

            records.append(record)
            if len(records) >= num: break

        return records

    def seek(self, key):
        self._iter.seek(key)

    def seek_to_first(self):
        return self._iter.seek_to_first()

    def seek_to_last(self):
        return self._iter.seek_to_last()

class Table(object):
    NAME = 'noname'
    ITER = Iterator

    KEYFN = staticmethod(lambda item: uuid.uuid1().hex)
    PACKFN = staticmethod(msgpack.packb)
    UNPACKFN = staticmethod(msgpack.unpackb)

    def __init__(self, data_dir, db):
        self.data_dir = os.path.join(data_dir, self.NAME)
        self.rdb = self.open()
        self.db = db
        self.iters = {}

        self.keyfn = self.KEYFN
        self.packfn = self.PACKFN
        self.unpackfn = self.UNPACKFN
        self.iter_klass = self.ITER

    def __str__(self):
        return '<Table: %s>' % self.NAME

    def __unicode__(self):
        return str(self)

    @property
    def log(self): return self.db.log

    def open(self):
        opts = self.define_options()
        opts.create_if_missing = True
        return rocksdb.DB(self.data_dir, opts)

    def define_options(self):
        opts = rocksdb.Options()
        return opts

    def put(self, key, item, batch=None,
            keyfn=None, packfn=None):

        packfn = packfn or self.packfn
        keyfn = keyfn or self.keyfn

        db = batch or self.rdb

        if isinstance(item, dict):
            key = key or item.get('_id', None) or keyfn(item)
            item['_id'] = key
        else:
            key = key or keyfn(item)

        value = packfn(item)
        db.put(key, value)

        return key

    def get(self, key, unpackfn=None):
        unpackfn = unpackfn or self.unpackfn

        value = self.rdb.get(key)
        if value is None: return None

        item = unpackfn(value)
        return item

    def delete(self, key, batch=None):
        db = batch or self.rdb
        db.delete(key)

    def put_many(self, data):
        batch = rocksdb.WriteBatch()
        for key, item in data:
            self.put(key, item, batch=batch)
        self.rdb.write(batch)

    def get_many(self, keys):
        data = self.rdb.multi_get(keys)
        for key, value in data.iteritems():
            data[key] = None if value is None else msgpack.unpackb(value)
        return data

    def delete_many(self, keys):
        batch = rocksdb.WriteBatch()
        for key in keys:
            self.delete(key, batch=batch)
        self.rdb.write(batch)

    def delete_all(self):
        batch = rocksdb.WriteBatch()
        _iter = self.rdb.iterkeys()
        _iter.seek_to_first()

        for index, key in enumerate(_iter):
            self.delete(key, batch)

            if index % 2 == 0:
                self.rdb.write(batch)
                time.sleep(0) # allow event loop to gain control

        self.rdb.write(batch)

    # Iteration

    def list_keys(self):
        _iter = self.rdb.iterkeys()
        _iter.seek_to_first()
        return list(_iter)

    def list_values(self):
        _iter = self.rdb.itervalues()
        _iter.seek_to_first()
        return list(self.unpackfn(x) for x in _iter)

    @ensurenewiter
    def iter_keys(self, name=None, prefix=None, reverse=False):
        self.iters[name] = self.iter_klass(self, prefix,
            type='keys', reverse=reverse)
        return name

    @ensurenewiter
    def iter_values(self, name=None, prefix=None, reverse=False):
        self.iters[name] = self.iter_klass(self, prefix,
            type='values', reverse=reverse)
        return name

    @ensurenewiter
    def iter_items(self, name=None, prefix=None, reverse=False):
        self.iters[name] = self.iter_klass(self, prefix,
            type='items', reverse=reverse)
        return name

    def list_iters(self):
        return self.iters.keys()

    def close_iter(self, _iter):
        del self.iters[_iter]

    @ensureiter
    def iter_get(self, _iter, num=Iterator.NUM_RECORDS):
        return _iter.get(num)

    @ensureiter
    def iter_seek(self, _iter, key):
        return _iter.seek(key)

    @ensureiter
    def iter_seek_to_first(self, _iter):
        return _iter.seek_to_first()

    @ensureiter
    def iter_seek_to_last(self, _iter):
        return _iter.seek_to_last()

    # Backup and restore

    def create_backup(self, path):
        backup = rocksdb.BackupEngine(path)
        return backup.create_backup(self.rdb, flush_before_backup=True)

    def stop_backup(self, path):
        backup = rocksdb.BackupEngine(path)
        return backup.stop_backup()

    def delete_backup(self, path, backup_id):
        backup = rocksdb.BackupEngine(path)
        return backup.delete_backup(backup_id)

    def get_backup_info(self, path):
        backup = rocksdb.BackupEngine(path)
        return backup.get_backup_info()

    def restore_backup(self, path, backup_id, db_dir, wal_dir=None):
        backup = rocksdb.BackupEngine(path)
        return backup.restore_backup(backup_id, db_dir,
            wal_dir if wal_dir is not None else db_dir)

    def restore_latest_backup(self, path, db_dir, wal_dir=None):
        backup = rocksdb.BackupEngine(path)
        return backup.restore_latest_backup(db_dir,
            wal_dir if wal_dir is not None else db_dir)

    def purge_old_backups(self, path, num_backups_to_keep):
        backup = rocksdb.BackupEngine(path)
        return backup.purge_old_backups(num_backups_to_keep)

class RocksDBAPI(object):
    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.tables = self.define_tables()

    def define_tables(self):
        return {}

    def list_tables(self):
        return self.tables.keys()

    @ensuretable
    def put(self, table, key, item):
        return table.put(key, item)

    @ensuretable
    def get(self, table, key):
        return table.get(key)

    @ensuretable
    def delete(self, table, key):
        return table.delete(key)

    @ensuretable
    def put_many(self, table, data):
        return table.put_many(data)

    @ensuretable
    def get_many(self, table, keys):
        return table.get_many(keys)

    @ensuretable
    def delete_many(self, table, keys):
        return table.delete_many(keys)

    @ensuretable
    def delete_all(self, table):
        '''
        Deletes all items from the table. Use with caution.
        If the table is very large, this could take a significant
        amount of time.
        '''
        return table.delete_all()

    # Iteration API methods

    @ensuretable
    def iter_keys(self, table, name=None, prefix=None, reverse=False):
        return table.iter_keys(name=name, prefix=prefix, reverse=reverse)

    @ensuretable
    def iter_values(self, table, name=None, prefix=None, reverse=False):
        return table.iter_values(name=name, prefix=prefix, reverse=reverse)

    @ensuretable
    def iter_items(self, table, name=None, prefix=None, reverse=False):
        return table.iter_items(name=name, prefix=prefix, reverse=reverse)

    @ensuretable
    def list_iters(self, table):
        return table.list_iters()

    @ensuretable
    def close_iter(self, table, name):
        return table.close_iter(name)

    @ensuretable
    def iter_get(self, table, name, num=Iterator.NUM_RECORDS):
        return table.iter_get(name, num)

    @ensuretable
    def iter_seek(self, table, name, key):
        return table.iter_seek(name, key)

    @ensuretable
    def iter_seek_to_first(self, table, name):
        return table.iter_seek_to_first(name)

    @ensuretable
    def iter_seek_to_last(self, table, name):
        return table.iter_seek_to_last(name)

    @ensuretable
    def list_keys(self, table):
        '''
        Lists all the keys in the table. This is meant
        to be used only during debugging in development
        and never in production as it loads all the keys
        in table into RAM which might cause memory load
        issues for large tables.
        '''
        return table.list_keys()

    @ensuretable
    def list_values(self, table):
        '''
        Lists all the values in the table. This is meant
        to be used only during debugging in development
        and never in production as it loads all the values
        in table into RAM which might cause memory load
        issues for large tables.
        '''
        return table.list_values()

    # Backup API methods

    @ensuretable
    def create_backup(self, table, path):
        return table.create_backup(path)

    @ensuretable
    def stop_backup(self, table, path):
        return table.stop_backup(path)

    @ensuretable
    def delete_backup(self, table, path, backup_id):
        return table.delete_backup(path, backup_id)

    @ensuretable
    def get_backup_info(self, table, path):
        return table.get_backup_info(path)

    @ensuretable
    def restore_backup(self, table, path, backup_id, db_dir, wal_dir=None):
        return table.restore_backup(path, backup_id, db_dir, wal_dir)

    @ensuretable
    def restore_latest_backup(self, table, path, db_dir, wal_dir=None):
        return table.restore_latest_backup(path, db_dir, wal_dir)

    @ensuretable
    def purge_old_backups(self, table, path, num_backups_to_keep):
        return table.purge_old_backups(path, num_backups_to_keep)

class RocksDBServer(RPCServer):
    NAME = 'RocksDBServer'
    DESC = 'RocksDB Server'

    def __init__(self, *args, **kwargs):
        super(RocksDBServer, self).__init__(*args, **kwargs)

        # make data dir if not already present
        self.data_dir = os.path.abspath(self.args.data_dir)
        if not os.path.exists(self.data_dir):
            os.makedirs(self.data_dir)

    def set_file_limits(self):
        try:
            # ulimit -n unlimited
            resource.setrlimit(resource.RLIMIT_NOFILE,
                (MAX_OPEN_FILES, MAX_OPEN_FILES))
        except ValueError:
            self.log.warning('unable to increase num files limit. run as root?')

    def expire_iters(self):
        while 1:
            ts = time.time()
            for table in self.api.tables.itervalues():
                expired = []

                for iter_name, _iter in table.iters.iteritems():
                    if ts - _iter.ts_last_activity >= ITERATOR_EXPIRE:
                        expired.append(iter_name)

                for iter_name in expired:
                    table.close_iter(iter_name)

            time.sleep(ITERATOR_EXPIRY_CHECK)

    def pre_start(self):
        super(RocksDBServer, self).pre_start()
        self.set_file_limits()

        self.thread_expire_iters = gevent.spawn(self.expire_iters)

    def prepare_api(self):
        return RocksDBAPI(self.args.data_dir)

    def define_args(self, parser):
        parser.add_argument('data_dir', type=str, metavar='data-dir',
            help='Directory path where data is stored')

class RocksDBClient(RPCClient):

    def _iter(self, table, prefix, reverse, fn):
        fn = getattr(self, fn)
        name = fn(table, prefix=prefix, reverse=reverse)

        if prefix is None:
            if reverse:
                self.iter_seek_to_last(table, name)
            else:
                self.iter_seek_to_first(table, name)

        else:
            self.iter_seek(table, name, prefix)

        while 1:
            items = self.iter_get(table, name)
            if not items: break

            for item in items:
                yield item

        self.close_iter(table, name)

    def iterkeys(self, table, prefix=None, reverse=False):
        return self._iter(table, prefix, reverse, 'iter_keys')

    def itervalues(self, table, prefix=None, reverse=False):
        return self._iter(table, prefix, reverse, 'iter_values')

    def iteritems(self, table, prefix=None, reverse=False):
        return self._iter(table, prefix, reverse, 'iter_items')

if __name__ == '__main__':
    RocksDBServer().start()
<commit_msg>fixes #26<commit_after>from gevent import monkey; monkey.patch_all()

import os
import time
import uuid
import resource
import random
import string
from inspect import getcallargs

import gevent
import msgpack
import rocksdb
from decorator import decorator
from funcserver import RPCServer, RPCClient, BaseHandler

ITERATOR_EXPIRY_CHECK = 5 * 60 # 5 minutes
ITERATOR_EXPIRE = 15 * 60 # 15 minutes
MAX_OPEN_FILES = 500000
ALPHANUM = string.letters + string.digits

def gen_random_seq(length=10):
    return ''.join([random.choice(ALPHANUM) for i in xrange(length)])

@decorator
def ensuretable(fn, self, table, *args, **kwargs):
    if table not in self.tables:
        raise Exception('Table "%s" does not exist' % table)
    return fn(self, self.tables[table], *args, **kwargs)

@decorator
def ensureiter(fn, self, _iter, *args, **kwargs):
    if _iter not in self.iters:
        raise Exception('Iter "%s" does not exist' % _iter)

    _iter = self.iters[_iter]
    _iter.ts_last_activity = time.time()
    return fn(self, _iter, *args, **kwargs)

@decorator
def ensurenewiter(fn, self, *args, **kwargs):
    name = getcallargs(fn, *args, **kwargs)['name'] or gen_random_seq()
    if args: args = list(args); args[0] = name
    else: kwargs['name'] = name

    if name in self.iters:
        raise Exception('iter "%s" exists already!' % name)

    return fn(self, *args, **kwargs)

class Iterator(object):
    NUM_RECORDS = 1000

    def __init__(self, table, prefix=None, type='items', reverse=False):
        '''
        @type: str; can be items/keys/values
        '''
        self.table = table
        self.type = type
        self.prefix = prefix
        self.reverse = reverse
        self.ts_last_activity = time.time()

        iterfn = getattr(self.table.rdb,
            {'keys': 'iterkeys', 'values': 'itervalues'}\
            .get(type, 'iteritems'))

        self._iter = iterfn(prefix=prefix)

        if reverse:
            self._iter = reversed(self._iter)

    def get(self, num=NUM_RECORDS):
        records = []

        for record in self._iter:
            if self.type == 'items':
                key, item = record
                item = self.table.unpackfn(item)
                record = (key, item)
            elif self.type == 'values':
                record = self.table.unpackfn(record)

            records.append(record)
            if len(records) >= num: break

        return records

    def seek(self, key):
        self._iter.seek(key)

    def seek_to_first(self):
        return self._iter.seek_to_first()

    def seek_to_last(self):
        return self._iter.seek_to_last()

class Table(object):
    NAME = 'noname'
    ITER = Iterator

    KEYFN = staticmethod(lambda item: uuid.uuid1().hex)
    PACKFN = staticmethod(msgpack.packb)
    UNPACKFN = staticmethod(msgpack.unpackb)

    def __init__(self, data_dir, db):
        self.data_dir = os.path.join(data_dir, self.NAME)
        self.rdb = self.open()
        self.db = db
        self.iters = {}

        self.keyfn = self.KEYFN
        self.packfn = self.PACKFN
        self.unpackfn = self.UNPACKFN
        self.iter_klass = self.ITER

    def __str__(self):
        return '<Table: %s>' % self.NAME

    def __unicode__(self):
        return str(self)

    @property
    def log(self): return self.db.log

    def open(self):
        opts = self.define_options()
        opts.create_if_missing = True
        return rocksdb.DB(self.data_dir, opts)

    def define_options(self):
        opts = rocksdb.Options()
        return opts

    def put(self, key, item, batch=None,
            keyfn=None, packfn=None):

        packfn = packfn or self.packfn
        keyfn = keyfn or self.keyfn

        db = batch or self.rdb

        if isinstance(item, dict):
            key = key or item.get('_id', None) or keyfn(item)
            item['_id'] = key
        else:
            key = key or keyfn(item)

        value = packfn(item)
        db.put(key, value)

        return key

    def get(self, key, unpackfn=None):
        unpackfn = unpackfn or self.unpackfn

        value = self.rdb.get(key)
        if value is None: return None

        item = unpackfn(value)
        return item

    def delete(self, key, batch=None):
        db = batch or self.rdb
        db.delete(key)

    def put_many(self, data):
        batch = rocksdb.WriteBatch()
        for key, item in data:
            self.put(key, item, batch=batch)
        self.rdb.write(batch)

    def get_many(self, keys):
        data = self.rdb.multi_get(keys)
        for key, value in data.iteritems():
            data[key] = None if value is None else msgpack.unpackb(value)
        return data

    def delete_many(self, keys):
        batch = rocksdb.WriteBatch()
        for key in keys:
            self.delete(key, batch=batch)
        self.rdb.write(batch)

    def delete_all(self):
        batch = rocksdb.WriteBatch()
        _iter = self.rdb.iterkeys()
        _iter.seek_to_first()

        for index, key in enumerate(_iter):
            self.delete(key, batch)

            if index % 2 == 0:
                self.rdb.write(batch)
                time.sleep(0) # allow event loop to gain control

        self.rdb.write(batch)

    def count(self):
        _iter = self.rdb.iterkeys()
        _iter.seek_to_first()
        for index, k in enumerate(_iter): pass
        return index + 1

    # Iteration

    def list_keys(self):
        _iter = self.rdb.iterkeys()
        _iter.seek_to_first()
        return list(_iter)

    def list_values(self):
        _iter = self.rdb.itervalues()
        _iter.seek_to_first()
        return list(self.unpackfn(x) for x in _iter)

    @ensurenewiter
    def iter_keys(self, name=None, prefix=None, reverse=False):
        self.iters[name] = self.iter_klass(self, prefix,
            type='keys', reverse=reverse)
        return name

    @ensurenewiter
    def iter_values(self, name=None, prefix=None, reverse=False):
        self.iters[name] = self.iter_klass(self, prefix,
            type='values', reverse=reverse)
        return name

    @ensurenewiter
    def iter_items(self, name=None, prefix=None, reverse=False):
        self.iters[name] = self.iter_klass(self, prefix,
            type='items', reverse=reverse)
        return name

    def list_iters(self):
        return self.iters.keys()

    def close_iter(self, _iter):
        del self.iters[_iter]

    @ensureiter
    def iter_get(self, _iter, num=Iterator.NUM_RECORDS):
        return _iter.get(num)

    @ensureiter
    def iter_seek(self, _iter, key):
        return _iter.seek(key)

    @ensureiter
    def iter_seek_to_first(self, _iter):
        return _iter.seek_to_first()

    @ensureiter
    def iter_seek_to_last(self, _iter):
        return _iter.seek_to_last()

    # Backup and restore

    def create_backup(self, path):
        backup = rocksdb.BackupEngine(path)
        return backup.create_backup(self.rdb, flush_before_backup=True)

    def stop_backup(self, path):
        backup = rocksdb.BackupEngine(path)
        return backup.stop_backup()

    def delete_backup(self, path, backup_id):
        backup = rocksdb.BackupEngine(path)
        return backup.delete_backup(backup_id)

    def get_backup_info(self, path):
        backup = rocksdb.BackupEngine(path)
        return backup.get_backup_info()

    def restore_backup(self, path, backup_id, db_dir, wal_dir=None):
        backup = rocksdb.BackupEngine(path)
        return backup.restore_backup(backup_id, db_dir,
            wal_dir if wal_dir is not None else db_dir)

    def restore_latest_backup(self, path, db_dir, wal_dir=None):
        backup = rocksdb.BackupEngine(path)
        return backup.restore_latest_backup(db_dir,
            wal_dir if wal_dir is not None else db_dir)

    def purge_old_backups(self, path, num_backups_to_keep):
        backup = rocksdb.BackupEngine(path)
        return backup.purge_old_backups(num_backups_to_keep)

class RocksDBAPI(object):
    def __init__(self, data_dir):
        self.data_dir = data_dir
        self.tables = self.define_tables()

    def define_tables(self):
        return {}

    def list_tables(self):
        return self.tables.keys()

    @ensuretable
    def put(self, table, key, item):
        return table.put(key, item)

    @ensuretable
    def get(self, table, key):
        return table.get(key)

    @ensuretable
    def delete(self, table, key):
        return table.delete(key)

    @ensuretable
    def put_many(self, table, data):
        return table.put_many(data)

    @ensuretable
    def get_many(self, table, keys):
        return table.get_many(keys)

    @ensuretable
    def delete_many(self, table, keys):
        return table.delete_many(keys)

    @ensuretable
    def delete_all(self, table):
        '''
        Deletes all items from the table. Use with caution.
        If the table is very large, this could take a significant
        amount of time.
        '''
        return table.delete_all()

    @ensuretable
    def count(self, table):
        '''
        Count the number of records (kv pairs)
        '''
        return table.count()

    # Iteration API methods

    @ensuretable
    def iter_keys(self, table, name=None, prefix=None, reverse=False):
        return table.iter_keys(name=name, prefix=prefix, reverse=reverse)

    @ensuretable
    def iter_values(self, table, name=None, prefix=None, reverse=False):
        return table.iter_values(name=name, prefix=prefix, reverse=reverse)

    @ensuretable
    def iter_items(self, table, name=None, prefix=None, reverse=False):
        return table.iter_items(name=name, prefix=prefix, reverse=reverse)

    @ensuretable
    def list_iters(self, table):
        return table.list_iters()

    @ensuretable
    def close_iter(self, table, name):
        return table.close_iter(name)

    @ensuretable
    def iter_get(self, table, name, num=Iterator.NUM_RECORDS):
        return table.iter_get(name, num)

    @ensuretable
    def iter_seek(self, table, name, key):
        return table.iter_seek(name, key)

    @ensuretable
    def iter_seek_to_first(self, table, name):
        return table.iter_seek_to_first(name)

    @ensuretable
    def iter_seek_to_last(self, table, name):
        return table.iter_seek_to_last(name)

    @ensuretable
    def list_keys(self, table):
        '''
        Lists all the keys in the table. This is meant
        to be used only during debugging in development
        and never in production as it loads all the keys
        in table into RAM which might cause memory load
        issues for large tables.
        '''
        return table.list_keys()

    @ensuretable
    def list_values(self, table):
        '''
        Lists all the values in the table. This is meant
        to be used only during debugging in development
        and never in production as it loads all the values
        in table into RAM which might cause memory load
        issues for large tables.
        '''
        return table.list_values()

    # Backup API methods

    @ensuretable
    def create_backup(self, table, path):
        return table.create_backup(path)

    @ensuretable
    def stop_backup(self, table, path):
        return table.stop_backup(path)

    @ensuretable
    def delete_backup(self, table, path, backup_id):
        return table.delete_backup(path, backup_id)

    @ensuretable
    def get_backup_info(self, table, path):
        return table.get_backup_info(path)

    @ensuretable
    def restore_backup(self, table, path, backup_id, db_dir, wal_dir=None):
        return table.restore_backup(path, backup_id, db_dir, wal_dir)

    @ensuretable
    def restore_latest_backup(self, table, path, db_dir, wal_dir=None):
        return table.restore_latest_backup(path, db_dir, wal_dir)

    @ensuretable
    def purge_old_backups(self, table, path, num_backups_to_keep):
        return table.purge_old_backups(path, num_backups_to_keep)

class RocksDBServer(RPCServer):
    NAME = 'RocksDBServer'
    DESC = 'RocksDB Server'

    def __init__(self, *args, **kwargs):
        super(RocksDBServer, self).__init__(*args, **kwargs)

        # make data dir if not already present
        self.data_dir = os.path.abspath(self.args.data_dir)
        if not os.path.exists(self.data_dir):
            os.makedirs(self.data_dir)

    def set_file_limits(self):
        try:
            # ulimit -n unlimited
            resource.setrlimit(resource.RLIMIT_NOFILE,
                (MAX_OPEN_FILES, MAX_OPEN_FILES))
        except ValueError:
            self.log.warning('unable to increase num files limit. run as root?')

    def expire_iters(self):
        while 1:
            ts = time.time()
            for table in self.api.tables.itervalues():
                expired = []

                for iter_name, _iter in table.iters.iteritems():
                    if ts - _iter.ts_last_activity >= ITERATOR_EXPIRE:
                        expired.append(iter_name)

                for iter_name in expired:
                    table.close_iter(iter_name)

            time.sleep(ITERATOR_EXPIRY_CHECK)

    def pre_start(self):
        super(RocksDBServer, self).pre_start()
        self.set_file_limits()

        self.thread_expire_iters = gevent.spawn(self.expire_iters)

    def prepare_api(self):
        return RocksDBAPI(self.args.data_dir)

    def define_args(self, parser):
        parser.add_argument('data_dir', type=str, metavar='data-dir',
            help='Directory path where data is stored')

class RocksDBClient(RPCClient):

    def _iter(self, table, prefix, reverse, fn):
        fn = getattr(self, fn)
        name = fn(table, prefix=prefix, reverse=reverse)

        if prefix is None:
            if reverse:
                self.iter_seek_to_last(table, name)
            else:
                self.iter_seek_to_first(table, name)

        else:
            self.iter_seek(table, name, prefix)

        while 1:
            items = self.iter_get(table, name)
            if not items: break

            for item in items:
                yield item

        self.close_iter(table, name)

    def iterkeys(self, table, prefix=None, reverse=False):
        return self._iter(table, prefix, reverse, 'iter_keys')

    def itervalues(self, table, prefix=None, reverse=False):
        return self._iter(table, prefix, reverse, 'iter_values')

    def iteritems(self, table, prefix=None, reverse=False):
        return self._iter(table, prefix, reverse, 'iter_items')

if __name__ == '__main__':
    RocksDBServer().start()
<|endoftext|>


<commit_before>>>> from sinusoid import SinusoidModel

# there are three ways to set up the model:

# 1. when you create the model

>>> a_model = SinusoidModel(frequencies=[1.2, 2.3],
                        modes=[[1, 0], [0, 1], [1, 1]]
                        )

# note that each mode is a list (or a tuple or a numpy array) with two
# elements, one for each frequency.

# frequencies and modes are immutable

>>> a_model.frequencies
(1.2, 2.3)
>>> a_model.modes
((1, 0), (0, 1), (1, 1))

# The number of components of each mode must match the number of
# frequencies:

>>> bad_model = SinusoidModel(frequencies=[1.2, 2.3],
                              modes=[[1, 0, 0]])
ValueError: Wrong number of modes in mode setter for mode [1, 0, 0]

# 2. After you create the model by adding frequencies and modes

>>> another_model = SinusoidModel()
>>> another_model.add_frequency(1.2, 2.3)
>>> another_model.frequencies
(1.2, 2.3)
>>> another_model.modes
()

# 3. Create a new model using another model as initializer

>>> a_model_extended = SinusoidModel(a_model)

>>> a_model_extended.add_frequency(3.4)

# Adding a frequency when modes already exist adds a weight of zero
# to the new frequency

>>> a_model_extended.modes
((1, 0, 0), (0, 1, 0), (1, 1, 0))

# this behavior can be prevented by setting the appropriate keyword

>>> a_model_extended.add_frequency(3.4, extend_modes=False)

# in this case all of the modes are simply deleted:

>>> a_model_extended.modes

()

# TThere is similar behavior if you set the frequencies directly:
# If there is an unambiguous way to generate additional modes with zero,
# or if the number of frequencies doesn't change, then the frequencies
# are updated and the modes unchanged.
#

>>> a_model_extended = SinusoidModel(a_model)
>>> a_model_extended.frequencies
(1.2, 2.3)

# note that in addition to adding a frequency, I've changed the value
# of another.
>>> a_model_extended.frequencies = [1.1, 2.3, 3.4]

>>> a_model_extended.modes
((1, 0, 0), (0, 1, 0), (1, 1, 0))

# If the number of frequencies set is inconsistent with the number of
# components in each mode then an error is generated.
# For example, if a model has three frequencies and modes defined, and one of
# the frequencies is removed an error is generated.
#

>>> a_model_extended.frequencies = [1.1, 2.3]
ValueError: Number of mode components (3) does not match new number of
#   frequencies. Try first setting modes to None and then change the
#   frequencies.

>>> a_model_extended.modes = None

>>> a_model_extended.frequencies = [1.1, 2.3]

>>> a_model_extended.frequencies
(1.1, 2.3)

# To avoid the error, first set the modes to None then change the frequencies.
# Regardless of how initialization was done, modes are added like this:

>>> a_model.add_mode(2, 0)

>>> a_model.add_mode(2, 0, 0)  # number of entries must match number of frequencies

TypeError: Number of components in mode does not match number of frequencies.

<commit_msg>Fix doctests for add_mode<commit_after>>>> from sinusoid import SinusoidModel

# there are three ways to set up the model:

# 1. when you create the model

>>> a_model = SinusoidModel(frequencies=[1.2, 2.3], modes=[[1, 0], [0, 1], [1, 1]])

# note that each mode is a list (or a tuple or a numpy array) with two
# elements, one for each frequency.

# frequencies and modes are immutable

>>> a_model.frequencies
(1.2, 2.3)
>>> a_model.modes
((1, 0), (0, 1), (1, 1))

# The number of components of each mode must match the number of
# frequencies:

>>> bad_model = SinusoidModel(frequencies=[1.2, 2.3], modes=[[1, 0, 0]])
Traceback (most recent call last):
    ...
ValueError: Wrong number of modes in mode setter for mode [1, 0, 0]

# 2. After you create the model by adding frequencies and modes

>>> another_model = SinusoidModel()
>>> another_model.add_frequency(1.2, 2.3)
>>> another_model.frequencies
(1.2, 2.3)
>>> another_model.modes
()
>>> another_model.add_mode(1, 0)  # add modes one at a time, or...
>>> another_model.modes
((1, 0),)
>>> another_model.add_mode([0, 1], [1, 1])  # several at a time
>>> another_model.modes
((1, 0), (0, 1), (1, 1))

# 3. Create a new model using another model as initializer

>>> a_model_extended = SinusoidModel(a_model)

>>> a_model_extended.add_frequency(3.4)

# Adding a frequency when modes already exist adds a weight of zero
# to the new frequency

>>> a_model_extended.modes
((1, 0, 0), (0, 1, 0), (1, 1, 0))

# this behavior can be prevented by setting the appropriate keyword

>>> a_model_extended.add_frequency(3.4, extend_modes=False)

# in this case all of the modes are simply deleted:

>>> a_model_extended.modes

()

# TThere is similar behavior if you set the frequencies directly:
# If there is an unambiguous way to generate additional modes with zero,
# or if the number of frequencies doesn't change, then the frequencies
# are updated and the modes unchanged.
#

>>> a_model_extended = SinusoidModel(a_model)
>>> a_model_extended.frequencies
(1.2, 2.3)

# note that in addition to adding a frequency, I've changed the value
# of another.
>>> a_model_extended.frequencies = [1.1, 2.3, 3.4]

>>> a_model_extended.modes
((1, 0, 0), (0, 1, 0), (1, 1, 0))

# If the number of frequencies set is inconsistent with the number of
# components in each mode then an error is generated.
# For example, if a model has three frequencies and modes defined, and one of
# the frequencies is removed an error is generated.
#

>>> a_model_extended.frequencies = [1.1, 2.3]
ValueError: Number of mode components (3) does not match new number of
#   frequencies. Try first setting modes to None and then change the
#   frequencies.

>>> a_model_extended.modes = None

>>> a_model_extended.frequencies = [1.1, 2.3]

>>> a_model_extended.frequencies
(1.1, 2.3)

# To avoid the error, first set the modes to None then change the frequencies.
# Regardless of how initialization was done, modes are added like this:

>>> a_model.add_mode(2, 0)

>>> a_model.add_mode(2, 0, 0)  # number of entries must match number of frequencies

TypeError: Number of components in mode does not match number of frequencies.

<|endoftext|>


<commit_before>import scipy as sp
import OpenPNM
ctrl = OpenPNM.Base.Controller()
ctrl.loglevel = 60

class DrainageTest:
    def setup_class(self):
        self.net = OpenPNM.Network.Cubic(shape=[5, 5, 5])
        self.geo = OpenPNM.Geometry.Toray090(network=self.net,
                                             pores=self.net.Ps,
                                             throats=self.net.Ts)
        self.phase = OpenPNM.Phases.Water(network=self.net)
        self.phys = OpenPNM.Physics.Standard(network=self.net,
                                             phase=self.phase,
                                             pores=self.net.Ps,
                                             throats=self.net.Ts)
        self.alg = OpenPNM.Algorithms.Drainage(network=self.net)

    def test_set_inlets_modes(self):
        self.alg.setup(invading_phase=self.phase)

        self.alg.set_inlets(pores=self.net.pores('top'), mode='add')
        assert sp.sum(self.alg['pore.inlets']) == 25

        self.alg.set_inlets(pores=self.net.pores('bottom'), mode='add')
        assert sp.sum(self.alg['pore.inlets']) == 50

        self.alg.set_inlets(pores=self.net.pores('top'), mode='overwrite')
        assert sp.sum(self.alg['pore.inlets']) == 25

        self.alg.set_inlets(pores=self.net.pores('top'), mode='remove')
        assert sp.sum(self.alg['pore.inlets']) == 0

        self.alg.set_inlets(pores=self.net.pores('top'), mode='add')
        self.alg.set_inlets(mode='clear')
        assert sp.sum(self.alg['pore.inlets']) == 0

    def test_set_inlets_conflicting_with_outlets(self):
        self.alg.setup(invading_phase=self.phase)
        self.alg['pore.outlets'][self.net.pores('top')] = True
        flag = False
        try:
            self.alg.set_inlets(pores=self.net.pores('top'), mode='add')
        except:
            flag = True
        assert flag

    def test_set_outlets_conflicting_with_inlets(self):
        self.alg.setup(invading_phase=self.phase)
        self.alg['pore.inlets'][self.net.pores('top')] = True
        flag = False
        try:
            self.alg.set_outlets(pores=self.net.pores('top'))
        except:
            flag = True
        assert flag

    def test_set_outlets_without_trapping(self):
        self.alg.setup(invading_phase=self.phase)
        flag = False
        try:
            self.alg.set_outlets(pores=self.net.pores('top'))
        except:
            flag = True
        assert flag

    def test_set_outlets_modes(self):
        self.alg.setup(invading_phase=self.phase, trapping=True)

        self.alg.set_outlets(pores=self.net.pores('top'), mode='add')
        assert sp.sum(self.alg['pore.outlets']) == 25

        self.alg.set_outlets(pores=self.net.pores('bottom'), mode='add')
        assert sp.sum(self.alg['pore.outlets']) == 50

        self.alg.set_outlets(pores=self.net.pores('top'), mode='overwrite')
        assert sp.sum(self.alg['pore.outlets']) == 25

        self.alg.set_outlets(pores=self.net.pores('top'), mode='remove')
        assert sp.sum(self.alg['pore.outlets']) == 0

        self.alg.set_outlets(pores=self.net.pores('top'), mode='add')
        self.alg.set_outlets(mode='clear')
        assert sp.sum(self.alg['pore.outlets']) == 0

    def test_set_residual_modes(self):
        self.alg.setup(invading_phase=self.phase)

        Ps = sp.random.randint(0, self.net.Np, 10)
        Ts = self.net.find_neighbor_pores(pores=Ps)
        self.alg.set_residual(pores=Ps, throats=Ts, mode='add')
        assert sp.sum(self.alg['pore.residual']) == sp.size(sp.unique(Ps))
        assert sp.sum(self.alg['throat.residual']) == sp.size(sp.unique(Ts))

        Ps = sp.random.randint(0, self.net.Np, 10)
        Ts = self.net.find_neighbor_pores(pores=Ps)
        self.alg.set_residual(pores=Ps, throats=Ts, mode='add')
        assert sp.sum(self.alg['pore.residual']) > sp.size(sp.unique(Ps))
        assert sp.sum(self.alg['throat.residual']) > sp.size(sp.unique(Ts))

        Ps = sp.random.randint(0, self.net.Np, 10)
        Ts = self.net.find_neighbor_pores(pores=Ps)
        self.alg.set_residual(pores=Ps, throats=Ts, mode='overwrite')
        assert sp.sum(self.alg['pore.residual']) == sp.size(sp.unique(Ps))
        assert sp.sum(self.alg['throat.residual']) == sp.size(sp.unique(Ts))

        self.alg.set_residual(pores=Ps, throats=Ts, mode='remove')
        assert sp.sum(self.alg['pore.residual']) == 0

        self.alg.set_residual(pores=Ps, throats=Ts, mode='add')
        self.alg.set_residual(mode='clear')
        assert sp.sum(self.alg['pore.residual']) == 0

    def test_set_boundary_conditions_clear(self):
        self.alg.setup(invading_phase=self.phase, trapping=True)

        self.alg['pore.inlets'] = True
        self.alg['pore.outlets'] = True
        self.alg['pore.residual'] = True
        self.alg['throat.residual'] = True

        self.alg.set_boundary_conditions(mode='clear')
        assert sp.sum(self.alg['pore.inlets']) == 0
        assert sp.sum(self.alg['pore.outlets']) == 0
        assert sp.sum(self.alg['pore.residual']) == 0
        assert sp.sum(self.alg['throat.residual']) == 0

    def test_set_boundary_conditions_bctypes(self):
        self.alg.setup(invading_phase=self.phase, trapping=True)
        Ps = sp.random.randint(0, self.net.Np, 10)

        self.alg.set_boundary_conditions(pores=Ps, bc_type='inlets')
        assert sp.sum(self.alg['pore.inlets']) == sp.size(sp.unique(Ps))
        self.alg['pore.inlets'] = False

        self.alg.set_boundary_conditions(pores=Ps, bc_type='outlets')
        assert sp.sum(self.alg['pore.outlets']) == sp.size(sp.unique(Ps))
        self.alg['pore.outlets'] = False

        self.alg.set_boundary_conditions(pores=Ps, bc_type='residual')
        assert sp.sum(self.alg['pore.residual']) == sp.size(sp.unique(Ps))
        self.alg['pore.residual'] = False

        flag = False
        try:
            self.alg.set_boundary_conditions(pores=Ps, bc_type='bad_type')
        except:
            flag = True
        assert flag

        flag = False
        try:
            self.alg.set_boundary_conditions(bc_type=None, mode='bad_type')
        except:
            flag = True
        assert flag

    def test_run_npts(self):
        self.alg.setup(invading_phase=self.phase)
        Ps = sp.random.randint(0, self.net.Np, 10)
        self.alg.set_boundary_conditions(pores=Ps, bc_type='inlets')
        self.alg.run(npts=20)

    def test_run_inv_pressures(self):
        self.alg.setup(invading_phase=self.phase)
        Ps = sp.random.randint(0, self.net.Np, 10)
        self.alg.set_boundary_conditions(pores=Ps, bc_type='inlets')
        self.alg.run(inv_pressures=range(0, 20000, 1000))
        assert sp.all(self.alg._inv_points == range(0, 20000, 1000))

    def test_run_no_inlets(self):
        self.alg.setup(invading_phase=self.phase)
        flag = False
        try:
            self.alg.run()
        except:
            flag = True
        assert flag

    def test_run_w_trapping_but_no_outlets(self):
        self.alg.setup(invading_phase=self.phase, trapping=True)
        Ps = sp.random.randint(0, self.net.Np, 10)
        self.alg.set_boundary_conditions(pores=Ps, bc_type='inlets')
        flag = False
        try:
            self.alg.run()
        except:
            flag = True
        assert flag

    def test_run_w_trapping(self):
        self.alg.setup(invading_phase=self.phase, trapping=True)
        self.alg.set_boundary_conditions(pores=self.net.pores('top'),
                                         bc_type='inlets')
        self.alg.set_boundary_conditions(pores=self.net.pores('bottom'),
                                         bc_type='outlets')
        self.alg.run()
        data = self.alg.get_drainage_data()
        assert 'capillary_pressure' in data.keys()
        assert 'nonwetting_phase_saturation' in data.keys()

    def test_run_w_residual_pores_and_throats(self):
        self.alg.setup(invading_phase=self.phase)
        self.alg.set_boundary_conditions(pores=self.net.pores('top'),
                                         bc_type='inlets')
        self.alg.set_boundary_conditions(pores=self.net.pores('bottom'),
                                         bc_type='residual')
        self.alg.run()
        data = self.alg.get_drainage_data()
        assert 'capillary_pressure' in data.keys()
        assert 'nonwetting_phase_saturation' in data.keys()
<commit_msg>Finally caught the last line...this should be 100% now.<commit_after>import scipy as sp
import OpenPNM
ctrl = OpenPNM.Base.Controller()
ctrl.loglevel = 60

class DrainageTest:
    def setup_class(self):
        self.net = OpenPNM.Network.Cubic(shape=[5, 5, 5])
        self.geo = OpenPNM.Geometry.Toray090(network=self.net,
                                             pores=self.net.Ps,
                                             throats=self.net.Ts)
        self.phase = OpenPNM.Phases.Water(network=self.net)
        self.phys = OpenPNM.Physics.Standard(network=self.net,
                                             phase=self.phase,
                                             pores=self.net.Ps,
                                             throats=self.net.Ts)
        self.alg = OpenPNM.Algorithms.Drainage(network=self.net)

    def test_set_inlets_modes(self):
        self.alg.setup(invading_phase=self.phase)

        self.alg.set_inlets(pores=self.net.pores('top'), mode='add')
        assert sp.sum(self.alg['pore.inlets']) == 25

        self.alg.set_inlets(pores=self.net.pores('bottom'), mode='add')
        assert sp.sum(self.alg['pore.inlets']) == 50

        self.alg.set_inlets(pores=self.net.pores('top'), mode='overwrite')
        assert sp.sum(self.alg['pore.inlets']) == 25

        self.alg.set_inlets(pores=self.net.pores('top'), mode='remove')
        assert sp.sum(self.alg['pore.inlets']) == 0

        self.alg.set_inlets(pores=self.net.pores('top'), mode='add')
        self.alg.set_inlets(mode='clear')
        assert sp.sum(self.alg['pore.inlets']) == 0

    def test_set_inlets_conflicting_with_outlets(self):
        self.alg.setup(invading_phase=self.phase)
        self.alg['pore.outlets'][self.net.pores('top')] = True
        flag = False
        try:
            self.alg.set_inlets(pores=self.net.pores('top'), mode='add')
        except:
            flag = True
        assert flag

    def test_set_outlets_conflicting_with_inlets(self):
        self.alg.setup(invading_phase=self.phase, trapping=True)
        self.alg['pore.inlets'][self.net.pores('top')] = True
        flag = False
        try:
            self.alg.set_outlets(pores=self.net.pores('top'))
        except:
            flag = True
        assert flag

    def test_set_outlets_without_trapping(self):
        self.alg.setup(invading_phase=self.phase)
        flag = False
        try:
            self.alg.set_outlets(pores=self.net.pores('top'))
        except:
            flag = True
        assert flag

    def test_set_outlets_modes(self):
        self.alg.setup(invading_phase=self.phase, trapping=True)

        self.alg.set_outlets(pores=self.net.pores('top'), mode='add')
        assert sp.sum(self.alg['pore.outlets']) == 25

        self.alg.set_outlets(pores=self.net.pores('bottom'), mode='add')
        assert sp.sum(self.alg['pore.outlets']) == 50

        self.alg.set_outlets(pores=self.net.pores('top'), mode='overwrite')
        assert sp.sum(self.alg['pore.outlets']) == 25

        self.alg.set_outlets(pores=self.net.pores('top'), mode='remove')
        assert sp.sum(self.alg['pore.outlets']) == 0

        self.alg.set_outlets(pores=self.net.pores('top'), mode='add')
        self.alg.set_outlets(mode='clear')
        assert sp.sum(self.alg['pore.outlets']) == 0

    def test_set_residual_modes(self):
        self.alg.setup(invading_phase=self.phase)

        Ps = sp.random.randint(0, self.net.Np, 10)
        Ts = self.net.find_neighbor_pores(pores=Ps)
        self.alg.set_residual(pores=Ps, throats=Ts, mode='add')
        assert sp.sum(self.alg['pore.residual']) == sp.size(sp.unique(Ps))
        assert sp.sum(self.alg['throat.residual']) == sp.size(sp.unique(Ts))

        Ps = sp.random.randint(0, self.net.Np, 10)
        Ts = self.net.find_neighbor_pores(pores=Ps)
        self.alg.set_residual(pores=Ps, throats=Ts, mode='add')
        assert sp.sum(self.alg['pore.residual']) > sp.size(sp.unique(Ps))
        assert sp.sum(self.alg['throat.residual']) > sp.size(sp.unique(Ts))

        Ps = sp.random.randint(0, self.net.Np, 10)
        Ts = self.net.find_neighbor_pores(pores=Ps)
        self.alg.set_residual(pores=Ps, throats=Ts, mode='overwrite')
        assert sp.sum(self.alg['pore.residual']) == sp.size(sp.unique(Ps))
        assert sp.sum(self.alg['throat.residual']) == sp.size(sp.unique(Ts))

        self.alg.set_residual(pores=Ps, throats=Ts, mode='remove')
        assert sp.sum(self.alg['pore.residual']) == 0

        self.alg.set_residual(pores=Ps, throats=Ts, mode='add')
        self.alg.set_residual(mode='clear')
        assert sp.sum(self.alg['pore.residual']) == 0

    def test_set_boundary_conditions_clear(self):
        self.alg.setup(invading_phase=self.phase, trapping=True)

        self.alg['pore.inlets'] = True
        self.alg['pore.outlets'] = True
        self.alg['pore.residual'] = True
        self.alg['throat.residual'] = True

        self.alg.set_boundary_conditions(mode='clear')
        assert sp.sum(self.alg['pore.inlets']) == 0
        assert sp.sum(self.alg['pore.outlets']) == 0
        assert sp.sum(self.alg['pore.residual']) == 0
        assert sp.sum(self.alg['throat.residual']) == 0

    def test_set_boundary_conditions_bctypes(self):
        self.alg.setup(invading_phase=self.phase, trapping=True)
        Ps = sp.random.randint(0, self.net.Np, 10)

        self.alg.set_boundary_conditions(pores=Ps, bc_type='inlets')
        assert sp.sum(self.alg['pore.inlets']) == sp.size(sp.unique(Ps))
        self.alg['pore.inlets'] = False

        self.alg.set_boundary_conditions(pores=Ps, bc_type='outlets')
        assert sp.sum(self.alg['pore.outlets']) == sp.size(sp.unique(Ps))
        self.alg['pore.outlets'] = False

        self.alg.set_boundary_conditions(pores=Ps, bc_type='residual')
        assert sp.sum(self.alg['pore.residual']) == sp.size(sp.unique(Ps))
        self.alg['pore.residual'] = False

        flag = False
        try:
            self.alg.set_boundary_conditions(pores=Ps, bc_type='bad_type')
        except:
            flag = True
        assert flag

        flag = False
        try:
            self.alg.set_boundary_conditions(bc_type=None, mode='bad_type')
        except:
            flag = True
        assert flag

    def test_run_npts(self):
        self.alg.setup(invading_phase=self.phase)
        Ps = sp.random.randint(0, self.net.Np, 10)
        self.alg.set_boundary_conditions(pores=Ps, bc_type='inlets')
        self.alg.run(npts=20)

    def test_run_inv_pressures(self):
        self.alg.setup(invading_phase=self.phase)
        Ps = sp.random.randint(0, self.net.Np, 10)
        self.alg.set_boundary_conditions(pores=Ps, bc_type='inlets')
        self.alg.run(inv_pressures=range(0, 20000, 1000))
        assert sp.all(self.alg._inv_points == range(0, 20000, 1000))

    def test_run_no_inlets(self):
        self.alg.setup(invading_phase=self.phase)
        flag = False
        try:
            self.alg.run()
        except:
            flag = True
        assert flag

    def test_run_w_trapping_but_no_outlets(self):
        self.alg.setup(invading_phase=self.phase, trapping=True)
        Ps = sp.random.randint(0, self.net.Np, 10)
        self.alg.set_boundary_conditions(pores=Ps, bc_type='inlets')
        flag = False
        try:
            self.alg.run()
        except:
            flag = True
        assert flag

    def test_run_w_trapping(self):
        self.alg.setup(invading_phase=self.phase, trapping=True)
        self.alg.set_boundary_conditions(pores=self.net.pores('top'),
                                         bc_type='inlets')
        self.alg.set_boundary_conditions(pores=self.net.pores('bottom'),
                                         bc_type='outlets')
        self.alg.run()
        data = self.alg.get_drainage_data()
        assert 'capillary_pressure' in data.keys()
        assert 'nonwetting_phase_saturation' in data.keys()

    def test_run_w_residual_pores_and_throats(self):
        self.alg.setup(invading_phase=self.phase)
        self.alg.set_boundary_conditions(pores=self.net.pores('top'),
                                         bc_type='inlets')
        self.alg.set_boundary_conditions(pores=self.net.pores('bottom'),
                                         bc_type='residual')
        self.alg.run()
        data = self.alg.get_drainage_data()
        assert 'capillary_pressure' in data.keys()
        assert 'nonwetting_phase_saturation' in data.keys()
<|endoftext|>


<commit_before># Copyright (c) 2006-2007 Sippy Software, Inc. All rights reserved.
#
# This file is part of SIPPY, a free RFC3261 SIP stack and B2BUA.
#
# SIPPY is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# For a license to use the SIPPY software under conditions
# other than those described here, or to purchase support for this
# software, please contact Sippy Software, Inc. by e-mail at the
# following addresses: sales@sippysoft.com.
#
# SIPPY is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA.

from signal import signal, SIG_IGN, SIG_DFL
from twisted.internet import reactor
from datetime import datetime
from traceback import print_exc
from sys import stdout

class Signal(object):
    callback = None
    parameters = None
    previous_handler = None

    def __init__(self, signum, callback, *parameters):
        self.callback = callback
        self.parameters = parameters
        self.previous_handler = signal(signum, self.signal_handler)

    def signal_handler(self, signum, *frame):
        try:
            reactor.callFromThread(self.callback, *self.parameters)
        except:
            print datetime.now(), 'Signal: unhandled exception in signal callback'
            print '-' * 70
            print_exc(file = stdout)
            print '-' * 70
            stdout.flush()
        if self.previous_handler not in (SIG_IGN, SIG_DFL):
            try:
                self.previous_handler(signum, *frame)
            except:
                print datetime.now(), 'Signal: unhandled exception in signal chain'
                print '-' * 70
                print_exc(file = stdout)
                print '-' * 70
                stdout.flush()

def log_signal(signum, sip_logger, signal_cb, cb_params):
    sip_logger.write('Dispatching signal %d to handler %s' % (signum, str(signal_cb)))
    return signal_cb(signum, *cb_params)

def LogSignal(sip_logger, signum, signal_cb, *cb_params):
    sip_logger.write('Registering signal %d to handler %s' % (signum, str(signal_cb)))
    return Signal(signum, log_signal, signum, sip_logger, signal_cb, cb_params)

if __name__ == '__main__':
    from signal import SIGHUP
    from os import kill, getpid

    def test(arguments):
        arguments['test'] = True
        reactor.crash()

    arguments = {'test':False}
    Signal(SIGHUP, test, arguments)
    kill(getpid(), SIGHUP)
    reactor.run()
    assert(arguments['test'])
<commit_msg>o Fix log_signal, so it delivers the same parameters as the   regular Signal() function;<commit_after># Copyright (c) 2006-2007 Sippy Software, Inc. All rights reserved.
#
# This file is part of SIPPY, a free RFC3261 SIP stack and B2BUA.
#
# SIPPY is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# For a license to use the SIPPY software under conditions
# other than those described here, or to purchase support for this
# software, please contact Sippy Software, Inc. by e-mail at the
# following addresses: sales@sippysoft.com.
#
# SIPPY is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA.

from signal import signal, SIG_IGN, SIG_DFL
from twisted.internet import reactor
from datetime import datetime
from traceback import print_exc
from sys import stdout

class Signal(object):
    callback = None
    parameters = None
    previous_handler = None

    def __init__(self, signum, callback, *parameters):
        self.callback = callback
        self.parameters = parameters
        self.previous_handler = signal(signum, self.signal_handler)

    def signal_handler(self, signum, *frame):
        try:
            reactor.callFromThread(self.callback, *self.parameters)
        except:
            print datetime.now(), 'Signal: unhandled exception in signal callback'
            print '-' * 70
            print_exc(file = stdout)
            print '-' * 70
            stdout.flush()
        if self.previous_handler not in (SIG_IGN, SIG_DFL):
            try:
                self.previous_handler(signum, *frame)
            except:
                print datetime.now(), 'Signal: unhandled exception in signal chain'
                print '-' * 70
                print_exc(file = stdout)
                print '-' * 70
                stdout.flush()

def log_signal(signum, sip_logger, signal_cb, cb_params):
    sip_logger.write('Dispatching signal %d to handler %s' % (signum, str(signal_cb)))
    return signal_cb(*cb_params)

def LogSignal(sip_logger, signum, signal_cb, *cb_params):
    sip_logger.write('Registering signal %d to handler %s' % (signum, str(signal_cb)))
    return Signal(signum, log_signal, signum, sip_logger, signal_cb, cb_params)

if __name__ == '__main__':
    from signal import SIGHUP, SIGTERM
    from os import kill, getpid

    def test(arguments):
        arguments['test'] = True
        reactor.crash()

    arguments = {'test':False}
    Signal(SIGHUP, test, arguments)
    kill(getpid(), SIGHUP)
    reactor.run()
    assert(arguments['test'])
    from SipLogger import SipLogger
    sip_logger = SipLogger('Signal::selftest')
    arguments = {'test':False}
    LogSignal(sip_logger, SIGTERM, test, arguments)
    kill(getpid(), SIGTERM)
    reactor.run()
    assert(arguments['test'])
<|endoftext|>


<commit_before># Core site concept heavily inspired by django.contrib.sites

from django.core.context_processors import csrf
from django.core.urlresolvers import reverse
from django.http import HttpResponseRedirect, HttpResponse, HttpResponseNotModified, Http404
from django.shortcuts import render_to_response
from django.template import RequestContext
from django.template.loader import render_to_string
from django.utils.datastructures import SortedDict
from django.utils.http import http_date
from django.views.decorators.cache import never_cache
from django.views.decorators.csrf import csrf_protect
from django.views.static import was_modified_since

try:
    from django.utils.functional import update_wrapper
except ImportError:  # Django>=1.6
    from functools import update_wrapper

from nexus import conf

import mimetypes
import os
import os.path
import posixpath
import stat
import urllib

NEXUS_ROOT = os.path.normpath(os.path.dirname(__file__))

try:
    from django.views.decorators.csrf import ensure_csrf_cookie
except ImportError:  # must be < Django 1.3
    from django.views.decorators.csrf import CsrfViewMiddleware
    from django.middleware.csrf import get_token
    from django.utils.decorators import decorator_from_middleware

    class _EnsureCsrfCookie(CsrfViewMiddleware):
        def _reject(self, request, reason):
            return None

        def process_view(self, request, callback, callback_args, callback_kwargs):
            retval = super(_EnsureCsrfCookie, self).process_view(request, callback, callback_args, callback_kwargs)
            # Forces process_response to send the cookie
            get_token(request)
            return retval

    ensure_csrf_cookie = decorator_from_middleware(_EnsureCsrfCookie)
    ensure_csrf_cookie.__name__ = 'ensure_csrf_cookie'
    ensure_csrf_cookie.__doc__ = """
    Use this decorator to ensure that a view sets a CSRF cookie, whether or not it
    uses the csrf_token template tag, or the CsrfViewMiddleware is used.
    """

class NexusSite(object):
    def __init__(self, name=None, app_name='nexus'):
        self._registry = {}
        self._categories = SortedDict()
        if name is None:
            self.name = 'nexus'
        else:
            self.name = name
        self.app_name = app_name

    def register_category(self, category, label, index=None):
        if index:
            self._categories.insert(index, category, label)
        else:
            self._categories[category] = label

    def register(self, module, namespace=None, category=None):
        module = module(self, category)
        if not namespace:
            namespace = module.get_namespace()
        if namespace:
            module.app_name = module.name = namespace
        self._registry[namespace] = (module, category)
        return module

    def unregister(self, namespace):
        if namespace in self._registry:
            del self._registry[namespace]

    def get_urls(self):
        try:
            from django.conf.urls import patterns, url, include
        except ImportError:  # Django<=1.4
            from django.conf.urls.defaults import patterns, url, include

        base_urls = patterns('',
            url(r'^media/(?P<module>[^/]+)/(?P<path>.+)$', self.media, name='media'),

            url(r'^$', self.as_view(self.dashboard), name='index'),
            url(r'^login/$', self.login, name='login'),
            url(r'^logout/$', self.as_view(self.logout), name='logout'),
        ), self.app_name, self.name

        urlpatterns = patterns('',
            url(r'^', include(base_urls)),
        )
        for namespace, module in self.get_modules():
            urlpatterns += patterns('',
                url(r'^%s/' % namespace, include(module.urls)),
            )

        return urlpatterns

    def urls(self):
        return self.get_urls()

    urls = property(urls)

    def has_permission(self, request, extra_permission=None):
        """
        Returns True if the given HttpRequest has permission to view
        *at least one* page in the admin site.
        """
        permission = request.user.is_active and request.user.is_staff
        if extra_permission:
            permission = permission and request.user.has_perm(extra_permission)
        return permission

    def as_view(self, view, cacheable=False, extra_permission=None):
        """
        Wraps a view in authentication/caching logic

        extra_permission can be used to require an extra permission for this view, such as a module permission
        """
        def inner(request, *args, **kwargs):
            if not self.has_permission(request, extra_permission):
                # show login pane
                return self.login(request)
            return view(request, *args, **kwargs)

        # Mark it as never_cache
        if not cacheable:
            inner = never_cache(inner)

        # We add csrf_protect here so this function can be used as a utility
        # function for any view, without having to repeat 'csrf_protect'.
        if not getattr(view, 'csrf_exempt', False):
            inner = csrf_protect(inner)

        inner = ensure_csrf_cookie(inner)

        return update_wrapper(inner, view)

    def get_context(self, request):
        context = csrf(request)
        context.update({
            'request': request,
            'nexus_site': self,
            'nexus_media_prefix': conf.MEDIA_PREFIX.rstrip('/'),
        })
        return context

    def get_modules(self):
        for k, v in self._registry.iteritems():
            yield k, v[0]

    def get_module(self, module):
        return self._registry[module][0]

    def get_categories(self):
        for k, v in self._categories.iteritems():
            yield k, v

    def get_category_label(self, category):
        return self._categories.get(category, category.title().replace('_', ' '))

    def render_to_string(self, template, context, request, current_app=None):
        if not current_app:
            current_app = self.name
        else:
            current_app = '%s:%s' % (self.name, current_app)

        if request:
            context_instance = RequestContext(request, current_app=current_app)
        else:
            context_instance = None

        context.update(self.get_context(request))

        return render_to_string(template, context,
            context_instance=context_instance
        )

    def render_to_response(self, template, context, request, current_app=None):
        "Shortcut for rendering to response and default context instances"
        if not current_app:
            current_app = self.name
        else:
            current_app = '%s:%s' % (self.name, current_app)

        if request:
            context_instance = RequestContext(request, current_app=current_app)
        else:
            context_instance = None

        context.update(self.get_context(request))

        return render_to_response(template, context,
            context_instance=context_instance
        )

    ## Our views

    def media(self, request, module, path):
        """
        Serve static files below a given point in the directory structure.
        """
        if module == 'nexus':
            document_root = os.path.join(NEXUS_ROOT, 'media')
        else:
            document_root = self.get_module(module).media_root

        path = posixpath.normpath(urllib.unquote(path))
        path = path.lstrip('/')
        newpath = ''
        for part in path.split('/'):
            if not part:
                # Strip empty path components.
                continue
            drive, part = os.path.splitdrive(part)
            head, part = os.path.split(part)
            if part in (os.curdir, os.pardir):
                # Strip '.' and '..' in path.
                continue
            newpath = os.path.join(newpath, part).replace('\\', '/')
        if newpath and path != newpath:
            return HttpResponseRedirect(newpath)
        fullpath = os.path.join(document_root, newpath)
        if os.path.isdir(fullpath):
            raise Http404("Directory indexes are not allowed here.")
        if not os.path.exists(fullpath):
            raise Http404('"%s" does not exist' % fullpath)
        # Respect the If-Modified-Since header.
        statobj = os.stat(fullpath)
        mimetype = mimetypes.guess_type(fullpath)[0] or 'application/octet-stream'
        if not was_modified_since(request.META.get('HTTP_IF_MODIFIED_SINCE'),
                                  statobj[stat.ST_MTIME], statobj[stat.ST_SIZE]):
            return HttpResponseNotModified(content_type=mimetype)
        contents = open(fullpath, 'rb').read()
        response = HttpResponse(contents, content_type=mimetype)
        response["Last-Modified"] = http_date(statobj[stat.ST_MTIME])
        response["Content-Length"] = len(contents)
        return response

    def login(self, request, form_class=None):
        "Login form"
        from django.contrib.auth import login as login_
        from django.contrib.auth.forms import AuthenticationForm

        if form_class is None:
            form_class = AuthenticationForm

        if request.POST:
            form = form_class(request, request.POST)
            if form.is_valid():
                login_(request, form.get_user())
                request.session.save()
                return HttpResponseRedirect(request.POST.get('next') or reverse('nexus:index', current_app=self.name))
            else:
                request.session.set_test_cookie()
        else:
            form = form_class(request)
            request.session.set_test_cookie()

        return self.render_to_response('nexus/login.html', {
            'form': form,
        }, request)
    login = never_cache(login)

    def logout(self, request):
        "Logs out user and redirects them to Nexus home"
        from django.contrib.auth import logout

        logout(request)

        return HttpResponseRedirect(reverse('nexus:index', current_app=self.name))

    def dashboard(self, request):
        "Basic dashboard panel"
        # TODO: these should be ajax
        module_set = []
        for namespace, module in self.get_modules():
            home_url = module.get_home_url(request)

            if hasattr(module, 'render_on_dashboard'):
                # Show by default, unless a permission is required
                if not module.permission or request.user.has_perm(module.permission):
                    module_set.append((module.get_dashboard_title(), module.render_on_dashboard(request), home_url))

        return self.render_to_response('nexus/dashboard.html', {
            'module_set': module_set,
        }, request)

# setup the default site

site = NexusSite()

<commit_msg>ordereddict<commit_after># Core site concept heavily inspired by django.contrib.sites

from collections import OrderedDict

from django.core.context_processors import csrf
from django.core.urlresolvers import reverse
from django.http import HttpResponseRedirect, HttpResponse, HttpResponseNotModified, Http404
from django.shortcuts import render_to_response
from django.template import RequestContext
from django.template.loader import render_to_string
from django.utils.http import http_date
from django.views.decorators.cache import never_cache
from django.views.decorators.csrf import csrf_protect
from django.views.static import was_modified_since

try:
    from django.utils.functional import update_wrapper
except ImportError:  # Django>=1.6
    from functools import update_wrapper

from nexus import conf

import mimetypes
import os
import os.path
import posixpath
import stat
import urllib

NEXUS_ROOT = os.path.normpath(os.path.dirname(__file__))

try:
    from django.views.decorators.csrf import ensure_csrf_cookie
except ImportError:  # must be < Django 1.3
    from django.views.decorators.csrf import CsrfViewMiddleware
    from django.middleware.csrf import get_token
    from django.utils.decorators import decorator_from_middleware

    class _EnsureCsrfCookie(CsrfViewMiddleware):
        def _reject(self, request, reason):
            return None

        def process_view(self, request, callback, callback_args, callback_kwargs):
            retval = super(_EnsureCsrfCookie, self).process_view(request, callback, callback_args, callback_kwargs)
            # Forces process_response to send the cookie
            get_token(request)
            return retval

    ensure_csrf_cookie = decorator_from_middleware(_EnsureCsrfCookie)
    ensure_csrf_cookie.__name__ = 'ensure_csrf_cookie'
    ensure_csrf_cookie.__doc__ = """
    Use this decorator to ensure that a view sets a CSRF cookie, whether or not it
    uses the csrf_token template tag, or the CsrfViewMiddleware is used.
    """

class NexusSite(object):
    def __init__(self, name=None, app_name='nexus'):
        self._registry = {}
        self._categories = OrderedDict()
        if name is None:
            self.name = 'nexus'
        else:
            self.name = name
        self.app_name = app_name

    def register_category(self, category, label, index=None):
        if index:
            self._categories.insert(index, category, label)
        else:
            self._categories[category] = label

    def register(self, module, namespace=None, category=None):
        module = module(self, category)
        if not namespace:
            namespace = module.get_namespace()
        if namespace:
            module.app_name = module.name = namespace
        self._registry[namespace] = (module, category)
        return module

    def unregister(self, namespace):
        if namespace in self._registry:
            del self._registry[namespace]

    def get_urls(self):
        try:
            from django.conf.urls import patterns, url, include
        except ImportError:  # Django<=1.4
            from django.conf.urls.defaults import patterns, url, include

        base_urls = patterns('',
            url(r'^media/(?P<module>[^/]+)/(?P<path>.+)$', self.media, name='media'),

            url(r'^$', self.as_view(self.dashboard), name='index'),
            url(r'^login/$', self.login, name='login'),
            url(r'^logout/$', self.as_view(self.logout), name='logout'),
        ), self.app_name, self.name

        urlpatterns = patterns('',
            url(r'^', include(base_urls)),
        )
        for namespace, module in self.get_modules():
            urlpatterns += patterns('',
                url(r'^%s/' % namespace, include(module.urls)),
            )

        return urlpatterns

    def urls(self):
        return self.get_urls()

    urls = property(urls)

    def has_permission(self, request, extra_permission=None):
        """
        Returns True if the given HttpRequest has permission to view
        *at least one* page in the admin site.
        """
        permission = request.user.is_active and request.user.is_staff
        if extra_permission:
            permission = permission and request.user.has_perm(extra_permission)
        return permission

    def as_view(self, view, cacheable=False, extra_permission=None):
        """
        Wraps a view in authentication/caching logic

        extra_permission can be used to require an extra permission for this view, such as a module permission
        """
        def inner(request, *args, **kwargs):
            if not self.has_permission(request, extra_permission):
                # show login pane
                return self.login(request)
            return view(request, *args, **kwargs)

        # Mark it as never_cache
        if not cacheable:
            inner = never_cache(inner)

        # We add csrf_protect here so this function can be used as a utility
        # function for any view, without having to repeat 'csrf_protect'.
        if not getattr(view, 'csrf_exempt', False):
            inner = csrf_protect(inner)

        inner = ensure_csrf_cookie(inner)

        return update_wrapper(inner, view)

    def get_context(self, request):
        context = csrf(request)
        context.update({
            'request': request,
            'nexus_site': self,
            'nexus_media_prefix': conf.MEDIA_PREFIX.rstrip('/'),
        })
        return context

    def get_modules(self):
        for k, v in self._registry.iteritems():
            yield k, v[0]

    def get_module(self, module):
        return self._registry[module][0]

    def get_categories(self):
        for k, v in self._categories.iteritems():
            yield k, v

    def get_category_label(self, category):
        return self._categories.get(category, category.title().replace('_', ' '))

    def render_to_string(self, template, context, request, current_app=None):
        if not current_app:
            current_app = self.name
        else:
            current_app = '%s:%s' % (self.name, current_app)

        if request:
            context_instance = RequestContext(request, current_app=current_app)
        else:
            context_instance = None

        context.update(self.get_context(request))

        return render_to_string(template, context,
            context_instance=context_instance
        )

    def render_to_response(self, template, context, request, current_app=None):
        "Shortcut for rendering to response and default context instances"
        if not current_app:
            current_app = self.name
        else:
            current_app = '%s:%s' % (self.name, current_app)

        if request:
            context_instance = RequestContext(request, current_app=current_app)
        else:
            context_instance = None

        context.update(self.get_context(request))

        return render_to_response(template, context,
            context_instance=context_instance
        )

    ## Our views

    def media(self, request, module, path):
        """
        Serve static files below a given point in the directory structure.
        """
        if module == 'nexus':
            document_root = os.path.join(NEXUS_ROOT, 'media')
        else:
            document_root = self.get_module(module).media_root

        path = posixpath.normpath(urllib.unquote(path))
        path = path.lstrip('/')
        newpath = ''
        for part in path.split('/'):
            if not part:
                # Strip empty path components.
                continue
            drive, part = os.path.splitdrive(part)
            head, part = os.path.split(part)
            if part in (os.curdir, os.pardir):
                # Strip '.' and '..' in path.
                continue
            newpath = os.path.join(newpath, part).replace('\\', '/')
        if newpath and path != newpath:
            return HttpResponseRedirect(newpath)
        fullpath = os.path.join(document_root, newpath)
        if os.path.isdir(fullpath):
            raise Http404("Directory indexes are not allowed here.")
        if not os.path.exists(fullpath):
            raise Http404('"%s" does not exist' % fullpath)
        # Respect the If-Modified-Since header.
        statobj = os.stat(fullpath)
        mimetype = mimetypes.guess_type(fullpath)[0] or 'application/octet-stream'
        if not was_modified_since(request.META.get('HTTP_IF_MODIFIED_SINCE'),
                                  statobj[stat.ST_MTIME], statobj[stat.ST_SIZE]):
            return HttpResponseNotModified(content_type=mimetype)
        contents = open(fullpath, 'rb').read()
        response = HttpResponse(contents, content_type=mimetype)
        response["Last-Modified"] = http_date(statobj[stat.ST_MTIME])
        response["Content-Length"] = len(contents)
        return response

    def login(self, request, form_class=None):
        "Login form"
        from django.contrib.auth import login as login_
        from django.contrib.auth.forms import AuthenticationForm

        if form_class is None:
            form_class = AuthenticationForm

        if request.POST:
            form = form_class(request, request.POST)
            if form.is_valid():
                login_(request, form.get_user())
                request.session.save()
                return HttpResponseRedirect(request.POST.get('next') or reverse('nexus:index', current_app=self.name))
            else:
                request.session.set_test_cookie()
        else:
            form = form_class(request)
            request.session.set_test_cookie()

        return self.render_to_response('nexus/login.html', {
            'form': form,
        }, request)
    login = never_cache(login)

    def logout(self, request):
        "Logs out user and redirects them to Nexus home"
        from django.contrib.auth import logout

        logout(request)

        return HttpResponseRedirect(reverse('nexus:index', current_app=self.name))

    def dashboard(self, request):
        "Basic dashboard panel"
        # TODO: these should be ajax
        module_set = []
        for namespace, module in self.get_modules():
            home_url = module.get_home_url(request)

            if hasattr(module, 'render_on_dashboard'):
                # Show by default, unless a permission is required
                if not module.permission or request.user.has_perm(module.permission):
                    module_set.append((module.get_dashboard_title(), module.render_on_dashboard(request), home_url))

        return self.render_to_response('nexus/dashboard.html', {
            'module_set': module_set,
        }, request)

# setup the default site

site = NexusSite()

<|endoftext|>


<commit_before>#!/usr/bin/python
"""Create SITCH feed files."""
import datetime
import gzip
import os
import shutil

from dateutil.parser import parse as dt_parse

import sitchlib

""" Outputs files like state.csv.gz.  These contain CSV data from the
FCC license database.  Use for determining GPS distance from tower.  Also
useful for correlating ARFCN/MCC/MNC for sanity and correct BTS ownership and
licensing representation."""

fcc_fields = ["LICENSE_ID", "SOURCE_SYSTEM", "CALLSIGN", "FACILITY_ID", "FRN",
              "LIC_NAME", "COMMON_NAME", "RADIO_SERVICE_CODE",
              "RADIO_SERVICE_DESC", "ROLLUP_CATEGORY_CODE",
              "ROLLUP_CATEGORY_DESC", "GRANT_DATE", "EXPIRED_DATE",
              "CANCELLATION_DATE", "LAST_ACTION_DATE", "LIC_STATUS_CODE",
              "LIC_STATUS_DESC", "ROLLUP_STATUS_CODE", "ROLLUP_STATUS_DESC",
              "ENTITY_TYPE_CODE", "ENTITY_TYPE_DESC", "ROLLUP_ENTITY_CODE",
              "ROLLUP_ENTITY_DESC", "LIC_ADDRESS", "LIC_CITY", "LIC_STATE",
              "LIC_ZIP_CODE", "CONTACT_COMPANY", "CONTACT_EMAIL",
              "MARKET_CODE", "MARKET_DESC", "CHANNEL_BLOCK", "LOC_TYPE_CODE",
              "LOC_TYPE_DESC", "LOC_CITY", "LOC_COUNTY_CODE",
              "LOC_COUNTY_NAME", "LOC_STATE", "LOC_RADIUS_OP", "LOC_SEQ_ID",
              "LOC_LAT_DEG", "LOC_LAT_MIN", "LOC_LAT_SEC", "LOC_LAT_DIR",
              "LOC_LONG_DEG", "LOC_LONG_MIN", "LOC_LONG_SEC", "LOC_LONG_DIR",
              "HGT_STRUCTURE", "ASR_NUM", "ANTENNA_ID", "ANT_SEQ_ID",
              "ANT_MAKE", "ANT_MODEL", "ANT_TYPE_CODE", "ANT_TYPE_DESC",
              "AZIMUTH", "BEAMWIDTH", "POLARIZATION_CODE", "FREQUENCY_ID",
              "FREQ_SEQ_ID", "FREQ_CLASS_STATION_CODE",
              "FREQ_CLASS_STATION_DESC", "POWER_ERP", "POWER_OUTPUT",
              "FREQUENCY_ASSIGNED", "FREQUENCY_UPPER_BAND", "UNIT_OF_MEASURE",
              "GROUND_ELEVATION", "ARFCN"]

ocid_fields = ["radio", "mcc", "net", "area", "cell", "unit", "lon",
               "lat", "range", "samples", "changeable", "created",
               "updated", "averageSignal", "carrier"]

feed_directory = "/var/production/"  # This is where the finished feed goes

def compress_and_remove_original(infiles):
    """Compress files, remove originals."""
    for uncompressed in list(infiles):
        infile = uncompressed
        outfile = "%s.gz" % infile
        with open(infile, 'rb') as f_in, gzip.open(outfile, 'wb') as f_out:
            shutil.copyfileobj(f_in, f_out)
        print("Written: %s" % outfile)
        os.remove(infile)
        print("Removed: %s" % infile)

def get_now_string():
    """Get ISO string for NOW."""
    now = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    return now

def write_statusfile(file_path, fcc_date, ocid_date):
    """Write status to file."""
    datestring = get_now_string()
    str_1 = "# SITCH Sensor Feed"
    str_2 = "## Processed: %s" % datestring
    str_3 = "Derived from:"
    str_4 = "* The OpenCellID DB http://opencellid.org\n  * Newest record: %s\n  * CGI DB\n  * CC by SA 3.0 https://creativecommons.org/licenses/by-sa/3.0/" % epoch_to_iso8601(ocid_date)  # NOQA
    str_5 = "* The FCC License DB http://data.fcc.gov\n  * Newest record: %s\n  * ARFCN DB" % epoch_to_iso8601(fcc_date)  # NOQA
    str_6 = "* Twilio's API: https://twilio.com\n  * CGI to provider correlation"  # NOQA
    master_str = "\n".join([str_1, str_2, str_3, str_4, str_5, str_6])
    with open(file_path, 'w') as out_file:
        out_file.write(master_str)
    return

def epoch_to_iso8601(unix_time):
    """Transform epoch time to ISO8601 format."""
    cleaned = float(unix_time)
    return datetime.datetime.utcfromtimestamp(cleaned).isoformat()

def iso8601_to_epoch(iso_time):
    """Transform iso time into a unix timestamp."""
    return int((dt_parse(iso_time) -
                datetime.datetime(1970, 1, 1)).total_seconds())

def main():
    sitchlib.OutfileHandler.ensure_path_exists(feed_directory)
    arfcn_comparator = sitchlib.ArfcnComparator()
    config = sitchlib.ConfigHelper()
    fileout = sitchlib.OutfileHandler(config.base_path,
                                      fcc_fields, ocid_fields)
    twilio_c = sitchlib.TwilioCarriers(config.twilio_sid,
                                       config.twilio_token)
    mcc_mnc_carriers = twilio_c.get_providers_for_country(config.iso_country)
    carrier_enricher = sitchlib.CarrierEnricher(mcc_mnc_carriers)
    fcc_feed_obj = sitchlib.FccCsv(config.fcc_destination_file)
    newest_ocid_record = 0
    newest_fcc_record = 0
    print("Splitting FCC license file into feed files...")
    for row in fcc_feed_obj:
        f_min = row["FREQUENCY_ASSIGNED"]
        f_max = row["FREQUENCY_UPPER_BAND"]
        arfcns = arfcn_comparator.arfcn_from_downlink_range(f_min, f_max)
        net_row = {}
        for column in fcc_fields:
            try:
                net_row[column] = row[column]
            except KeyError:
                pass
        for arfcn in arfcns:
            net_row["ARFCN"] = arfcn
            fileout.write_fcc_record(net_row)
        if iso8601_to_epoch(row["LAST_ACTION_DATE"]) > newest_fcc_record:
            newest_fcc_record = iso8601_to_epoch(row["LAST_ACTION_DATE"])
    print("Compressing FCC feed files")
    compress_and_remove_original(fileout.feed_files)
    fileout = None
    fileout = sitchlib.OutfileHandler(config.base_path,
                                      fcc_fields, ocid_fields)
    ocid_feed_obj = sitchlib.OcidCsv(config.ocid_destination_file)
    print("Splitting OpenCellID feed into MCC files...")
    for row in ocid_feed_obj:
        if row["radio"] != config.target_radio:
            continue
        row["carrier"] = carrier_enricher.get_carrier(row["mcc"], row["net"])
        if int(row["updated"]) > int(newest_ocid_record):
            newest_ocid_record = int(row["updated"])
        fileout.write_ocid_record(row)
    print("Compressing OpenCellID feed files")
    compress_and_remove_original(fileout.feed_files)
    print("Moving to drop directory...")
    staged_files = os.listdir(config.base_path)
    for staged_file in staged_files:
        full_src_file_name = os.path.join(config.base_path, staged_file)
        full_dst_file_name = os.path.join(feed_directory, staged_file)
        if (os.path.isfile(full_src_file_name)):
            shutil.copy(full_src_file_name, full_dst_file_name)
    write_statusfile("/opt/README.md", newest_fcc_record, newest_ocid_record)
    print("ALL DONE!!!")

if __name__ == "__main__":
    main()
<commit_msg>chg: dev: Refactor FCC feed processing. !minor<commit_after>#!/usr/bin/python
"""Create SITCH feed files."""
import datetime
import gzip
import os
import shutil

from dateutil.parser import parse as dt_parse

import sitchlib

""" Outputs files like state.csv.gz.  These contain CSV data from the
FCC license database.  Use for determining GPS distance from tower.  Also
useful for correlating ARFCN/MCC/MNC for sanity and correct BTS ownership and
licensing representation."""

FCC_FIELDS = ["LICENSE_ID", "SOURCE_SYSTEM", "CALLSIGN", "FACILITY_ID", "FRN",
              "LIC_NAME", "COMMON_NAME", "RADIO_SERVICE_CODE",
              "RADIO_SERVICE_DESC", "ROLLUP_CATEGORY_CODE",
              "ROLLUP_CATEGORY_DESC", "GRANT_DATE", "EXPIRED_DATE",
              "CANCELLATION_DATE", "LAST_ACTION_DATE", "LIC_STATUS_CODE",
              "LIC_STATUS_DESC", "ROLLUP_STATUS_CODE", "ROLLUP_STATUS_DESC",
              "ENTITY_TYPE_CODE", "ENTITY_TYPE_DESC", "ROLLUP_ENTITY_CODE",
              "ROLLUP_ENTITY_DESC", "LIC_ADDRESS", "LIC_CITY", "LIC_STATE",
              "LIC_ZIP_CODE", "CONTACT_COMPANY", "CONTACT_EMAIL",
              "MARKET_CODE", "MARKET_DESC", "CHANNEL_BLOCK", "LOC_TYPE_CODE",
              "LOC_TYPE_DESC", "LOC_CITY", "LOC_COUNTY_CODE",
              "LOC_COUNTY_NAME", "LOC_STATE", "LOC_RADIUS_OP", "LOC_SEQ_ID",
              "LOC_LAT_DEG", "LOC_LAT_MIN", "LOC_LAT_SEC", "LOC_LAT_DIR",
              "LOC_LONG_DEG", "LOC_LONG_MIN", "LOC_LONG_SEC", "LOC_LONG_DIR",
              "HGT_STRUCTURE", "ASR_NUM", "ANTENNA_ID", "ANT_SEQ_ID",
              "ANT_MAKE", "ANT_MODEL", "ANT_TYPE_CODE", "ANT_TYPE_DESC",
              "AZIMUTH", "BEAMWIDTH", "POLARIZATION_CODE", "FREQUENCY_ID",
              "FREQ_SEQ_ID", "FREQ_CLASS_STATION_CODE",
              "FREQ_CLASS_STATION_DESC", "POWER_ERP", "POWER_OUTPUT",
              "FREQUENCY_ASSIGNED", "FREQUENCY_UPPER_BAND", "UNIT_OF_MEASURE",
              "GROUND_ELEVATION", "ARFCN"]

OCID_FIELDS = ["radio", "mcc", "net", "area", "cell", "unit", "lon",
               "lat", "range", "samples", "changeable", "created",
               "updated", "averageSignal", "carrier"]

FEED_DIRECTORY = "/var/production/"  # This is where the finished feed goes

def compress_and_remove_original(infiles):
    """Compress files, remove originals."""
    for uncompressed in list(infiles):
        infile = uncompressed
        outfile = "%s.gz" % infile
        with open(infile, 'rb') as f_in, gzip.open(outfile, 'wb') as f_out:
            shutil.copyfileobj(f_in, f_out)
        print("Written: %s" % outfile)
        os.remove(infile)
        print("Removed: %s" % infile)

def get_now_string():
    """Get ISO string for NOW."""
    now = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    return now

def write_statusfile(file_path, fcc_date, ocid_date):
    """Write status to file."""
    datestring = get_now_string()
    body = []
    body.append("# SITCH Sensor Feed\n")
    body.append("## Processed: {}\n".format(datestring))
    body.append("Derived from:\n")
    body.append("* The OpenCellID DB http://opencellid.org")
    body.append("  * Newest record: {}".format(epoch_to_iso8601(ocid_date)))
    body.append("  * CGI DB")
    body.append("  * CC by SA 3.0 https://creativecommons.org/licenses/by-sa/3.0/")  # NOQA
    body.append("* The FCC License DB http://data.fcc.gov")
    body.append("  * Newest record: {}".format(epoch_to_iso8601(fcc_date)))
    body.append("  * ARFCN DB")
    body.append("* Twilio's API: https://twilio.com")
    body.append("  * CGI to provider correlation")
    master_str = "\n".join(body)
    with open(file_path, 'w') as out_file:
        out_file.write(master_str)

def epoch_to_iso8601(unix_time):
    """Transform epoch time to ISO8601 format."""
    cleaned = float(unix_time)
    return datetime.datetime.utcfromtimestamp(cleaned).isoformat()

def iso8601_to_epoch(iso_time):
    """Transform iso time into a unix timestamp."""
    return int((dt_parse(iso_time) -
                datetime.datetime(1970, 1, 1)).total_seconds())

def process_fcc_feed(base_path, fcc_destination_file):
    """Process FCC feed, return newest record timestamp."""
    fcc_feed_obj = sitchlib.FccCsv(fcc_destination_file)
    arfcn_comparator = sitchlib.ArfcnComparator()
    fileout = sitchlib.OutfileHandler(base_path, FCC_FIELDS, OCID_FIELDS)
    newest_fcc_record = 0
    for row in fcc_feed_obj:
        f_min = row["FREQUENCY_ASSIGNED"]
        f_max = row["FREQUENCY_UPPER_BAND"]
        arfcns = arfcn_comparator.arfcn_from_downlink_range(f_min, f_max)
        net_row = {}
        for column in FCC_FIELDS:
            try:
                net_row[column] = row[column]
            except KeyError:
                pass
        for arfcn in arfcns:
            net_row["ARFCN"] = arfcn
            fileout.write_fcc_record(net_row)
        if iso8601_to_epoch(row["LAST_ACTION_DATE"]) > newest_fcc_record:
            newest_fcc_record = iso8601_to_epoch(row["LAST_ACTION_DATE"])
    print("Compressing FCC feed files")
    compress_and_remove_original(fileout.feed_files)
    return newest_fcc_record

def main():
    """Wrap all top-level logic."""
    sitchlib.OutfileHandler.ensure_path_exists(FEED_DIRECTORY)
    config = sitchlib.ConfigHelper()
    twilio_c = sitchlib.TwilioCarriers(config.twilio_sid,
                                       config.twilio_token)
    mcc_mnc_carriers = twilio_c.get_providers_for_country(config.iso_country)
    carrier_enricher = sitchlib.CarrierEnricher(mcc_mnc_carriers)
    newest_ocid_record = 0
    print("Splitting FCC license file into feed files...")
    newest_fcc_record = process_fcc_feed(config.base_path,
                                         config.fcc_destination_file)
    fileout = sitchlib.OutfileHandler(config.base_path,
                                      FCC_FIELDS, OCID_FIELDS)
    ocid_feed_obj = sitchlib.OcidCsv(config.ocid_destination_file)
    print("Splitting OpenCellID feed into MCC files...")
    for row in ocid_feed_obj:
        if row["radio"] != config.target_radio:
            continue
        row["carrier"] = carrier_enricher.get_carrier(row["mcc"], row["net"])
        if int(row["updated"]) > int(newest_ocid_record):
            newest_ocid_record = int(row["updated"])
        fileout.write_ocid_record(row)
    print("Compressing OpenCellID feed files")
    compress_and_remove_original(fileout.feed_files)
    print("Moving to drop directory...")
    staged_files = os.listdir(config.base_path)
    for staged_file in staged_files:
        full_src_file_name = os.path.join(config.base_path, staged_file)
        full_dst_file_name = os.path.join(FEED_DIRECTORY, staged_file)
        if os.path.isfile(full_src_file_name):
            shutil.copy(full_src_file_name, full_dst_file_name)
    write_statusfile("/opt/README.md", newest_fcc_record, newest_ocid_record)
    print("ALL DONE!!!")

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