<commit_before>import errno
import fcntl
import msgpack
import os
import select
import shlex
from subprocess import Popen, PIPE
import sys
import tempfile
import traceback

from . import __version__

from .helpers import Error, IntegrityError
from .repository import Repository

BUFSIZE = 10 * 1024 * 1024

class ConnectionClosed(Error):
    """Connection closed by remote host"""

class PathNotAllowed(Error):
    """Repository path not allowed"""

class InvalidRPCMethod(Error):
    """RPC method is not valid"""

class RepositoryServer:  # pragma: no cover
    rpc_methods = (
        '__len__',
        'check',
        'commit',
        'delete',
        'destroy',
        'get',
        'list',
        'negotiate',
        'open',
        'put',
        'repair',
        'rollback',
        'save_key',
        'load_key',
    )

    def __init__(self, restrict_to_paths):
        self.repository = None
        self.restrict_to_paths = restrict_to_paths

    def serve(self):
        stdin_fd = sys.stdin.fileno()
        stdout_fd = sys.stdout.fileno()
        # Make stdin non-blocking
        fl = fcntl.fcntl(stdin_fd, fcntl.F_GETFL)
        fcntl.fcntl(stdin_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
        # Make stdout blocking
        fl = fcntl.fcntl(stdout_fd, fcntl.F_GETFL)
        fcntl.fcntl(stdout_fd, fcntl.F_SETFL, fl & ~os.O_NONBLOCK)
        unpacker = msgpack.Unpacker(use_list=False)
        while True:
            r, w, es = select.select([stdin_fd], [], [], 10)
            if r:
                data = os.read(stdin_fd, BUFSIZE)
                if not data:
                    return
                unpacker.feed(data)
                for unpacked in unpacker:
                    if not (isinstance(unpacked, tuple) and len(unpacked) == 4):
                        raise Exception("Unexpected RPC data format.")
                    type, msgid, method, args = unpacked
                    method = method.decode('ascii')
                    try:
                        if method not in self.rpc_methods:
                            raise InvalidRPCMethod(method)
                        try:
                            f = getattr(self, method)
                        except AttributeError:
                            f = getattr(self.repository, method)
                        res = f(*args)
                    except BaseException as e:
                        exc = "Remote Traceback by Borg %s%s%s" % (__version__, os.linesep, traceback.format_exc())
                        os.write(stdout_fd, msgpack.packb((1, msgid, e.__class__.__name__, exc)))
                    else:
                        os.write(stdout_fd, msgpack.packb((1, msgid, None, res)))
            if es:
                return

    def negotiate(self, versions):
        return 1

    def open(self, path, create=False):
        path = os.fsdecode(path)
        if path.startswith('/~'):
            path = path[1:]
        path = os.path.realpath(os.path.expanduser(path))
        if self.restrict_to_paths:
            for restrict_to_path in self.restrict_to_paths:
                if path.startswith(os.path.realpath(restrict_to_path)):
                    break
            else:
                raise PathNotAllowed(path)
        self.repository = Repository(path, create)
        return self.repository.id

class RemoteRepository:
    extra_test_args = []
    remote_path = 'borg'
    # default umask, overriden by --umask, defaults to read/write only for owner
    umask = 0o077

    class RPCError(Exception):
        def __init__(self, name):
            self.name = name

    def __init__(self, location, create=False):
        self.location = location
        self.preload_ids = []
        self.msgid = 0
        self.to_send = b''
        self.cache = {}
        self.ignore_responses = set()
        self.responses = {}
        self.unpacker = msgpack.Unpacker(use_list=False)
        self.p = None
        # XXX: ideally, the testsuite would subclass Repository and
        # override ssh_cmd() instead of this crude hack, although
        # __testsuite__ is not a valid domain name so this is pretty
        # safe.
        if location.host == '__testsuite__':
            args = [sys.executable, '-m', 'borg.archiver', 'serve' ] + self.extra_test_args
        else:  # pragma: no cover
            args = self.ssh_cmd()
        self.p = Popen(args, bufsize=0, stdin=PIPE, stdout=PIPE)
        self.stdin_fd = self.p.stdin.fileno()
        self.stdout_fd = self.p.stdout.fileno()
        fcntl.fcntl(self.stdin_fd, fcntl.F_SETFL, fcntl.fcntl(self.stdin_fd, fcntl.F_GETFL) | os.O_NONBLOCK)
        fcntl.fcntl(self.stdout_fd, fcntl.F_SETFL, fcntl.fcntl(self.stdout_fd, fcntl.F_GETFL) | os.O_NONBLOCK)
        self.r_fds = [self.stdout_fd]
        self.x_fds = [self.stdin_fd, self.stdout_fd]

        try:
            version = self.call('negotiate', 1)
        except ConnectionClosed:
            raise Exception('Server immediately closed connection - is Borg installed and working on the server?')
        if version != 1:
            raise Exception('Server insisted on using unsupported protocol version %d' % version)
        self.id = self.call('open', location.path, create)

    def __del__(self):
        self.close()

    def __repr__(self):
        return '<%s %s>' % (self.__class__.__name__, self.location.canonical_path())

    def umask_flag(self):
        return ['--umask', '%03o' % self.umask]

    def ssh_cmd(self, location):
        args = shlex.split(os.environ.get('BORG_RSH', 'ssh'))
        if location.port:
            args += ['-p', str(location.port)]
        if location.user:
            args.append('%s@%s' % (location.user, location.host))
        else:
            args.append('%s' % location.host)
        # use local umask also for the remote process
        args += [self.remote_path, 'serve'] + self.umask_flag()
        return args

    def call(self, cmd, *args, **kw):
        for resp in self.call_many(cmd, [args], **kw):
            return resp

    def call_many(self, cmd, calls, wait=True, is_preloaded=False):
        if not calls:
            return

        def fetch_from_cache(args):
            msgid = self.cache[args].pop(0)
            if not self.cache[args]:
                del self.cache[args]
            return msgid

        calls = list(calls)
        waiting_for = []
        w_fds = [self.stdin_fd]
        while wait or calls:
            while waiting_for:
                try:
                    error, res = self.responses.pop(waiting_for[0])
                    waiting_for.pop(0)
                    if error:
                        if error == b'DoesNotExist':
                            raise Repository.DoesNotExist(self.location.orig)
                        elif error == b'AlreadyExists':
                            raise Repository.AlreadyExists(self.location.orig)
                        elif error == b'CheckNeeded':
                            raise Repository.CheckNeeded(self.location.orig)
                        elif error == b'IntegrityError':
                            raise IntegrityError(res)
                        elif error == b'PathNotAllowed':
                            raise PathNotAllowed(*res)
                        elif error == b'ObjectNotFound':
                            raise Repository.ObjectNotFound(res[0], self.location.orig)
                        elif error == b'InvalidRPCMethod':
                            raise InvalidRPCMethod(*res)
                        else:
                            raise self.RPCError(res.decode('utf-8'))
                    else:
                        yield res
                        if not waiting_for and not calls:
                            return
                except KeyError:
                    break
            r, w, x = select.select(self.r_fds, w_fds, self.x_fds, 1)
            if x:
                raise Exception('FD exception occurred')
            if r:
                data = os.read(self.stdout_fd, BUFSIZE)
                if not data:
                    raise ConnectionClosed()
                self.unpacker.feed(data)
                for unpacked in self.unpacker:
                    if not (isinstance(unpacked, tuple) and len(unpacked) == 4):
                        raise Exception("Unexpected RPC data format.")
                    type, msgid, error, res = unpacked
                    if msgid in self.ignore_responses:
                        self.ignore_responses.remove(msgid)
                    else:
                        self.responses[msgid] = error, res
            if w:
                while not self.to_send and (calls or self.preload_ids) and len(waiting_for) < 100:
                    if calls:
                        if is_preloaded:
                            if calls[0] in self.cache:
                                waiting_for.append(fetch_from_cache(calls.pop(0)))
                        else:
                            args = calls.pop(0)
                            if cmd == 'get' and args in self.cache:
                                waiting_for.append(fetch_from_cache(args))
                            else:
                                self.msgid += 1
                                waiting_for.append(self.msgid)
                                self.to_send = msgpack.packb((1, self.msgid, cmd, args))
                    if not self.to_send and self.preload_ids:
                        args = (self.preload_ids.pop(0),)
                        self.msgid += 1
                        self.cache.setdefault(args, []).append(self.msgid)
                        self.to_send = msgpack.packb((1, self.msgid, cmd, args))

                if self.to_send:
                    try:
                        self.to_send = self.to_send[os.write(self.stdin_fd, self.to_send):]
                    except OSError as e:
                        # io.write might raise EAGAIN even though select indicates
                        # that the fd should be writable
                        if e.errno != errno.EAGAIN:
                            raise
                if not self.to_send and not (calls or self.preload_ids):
                    w_fds = []
        self.ignore_responses |= set(waiting_for)

    def check(self, repair=False):
        return self.call('check', repair)

    def commit(self, *args):
        return self.call('commit')

    def rollback(self, *args):
        return self.call('rollback')

    def destroy(self):
        return self.call('destroy')

    def __len__(self):
        return self.call('__len__')

    def list(self, limit=None, marker=None):
        return self.call('list', limit, marker)

    def get(self, id_):
        for resp in self.get_many([id_]):
            return resp

    def get_many(self, ids, is_preloaded=False):
        for resp in self.call_many('get', [(id_,) for id_ in ids], is_preloaded=is_preloaded):
            yield resp

    def put(self, id_, data, wait=True):
        return self.call('put', id_, data, wait=wait)

    def delete(self, id_, wait=True):
        return self.call('delete', id_, wait=wait)

    def save_key(self, keydata):
        return self.call('save_key', keydata)

    def load_key(self):
        return self.call('load_key')

    def close(self):
        if self.p:
            self.p.stdin.close()
            self.p.stdout.close()
            self.p.wait()
            self.p = None

    def preload(self, ids):
        self.preload_ids += ids

class RepositoryCache:
    """A caching Repository wrapper

    Caches Repository GET operations using a local temporary Repository.
    """
    def __init__(self, repository):
        self.repository = repository
        tmppath = tempfile.mkdtemp(prefix='borg-tmp')
        self.caching_repo = Repository(tmppath, create=True, exclusive=True)

    def __del__(self):
        self.caching_repo.destroy()

    def get(self, key):
        return next(self.get_many([key]))

    def get_many(self, keys):
        unknown_keys = [key for key in keys if key not in self.caching_repo]
        repository_iterator = zip(unknown_keys, self.repository.get_many(unknown_keys))
        for key in keys:
            try:
                yield self.caching_repo.get(key)
            except Repository.ObjectNotFound:
                for key_, data in repository_iterator:
                    if key_ == key:
                        self.caching_repo.put(key, data)
                        yield data
                        break
        # Consume any pending requests
        for _ in repository_iterator:
            pass

def cache_if_remote(repository):
    if isinstance(repository, RemoteRepository):
        return RepositoryCache(repository)
    return repository
<commit_msg>ssh_cmd: fix wrong caller, fixes #255<commit_after>import errno
import fcntl
import msgpack
import os
import select
import shlex
from subprocess import Popen, PIPE
import sys
import tempfile
import traceback

from . import __version__

from .helpers import Error, IntegrityError
from .repository import Repository

BUFSIZE = 10 * 1024 * 1024

class ConnectionClosed(Error):
    """Connection closed by remote host"""

class PathNotAllowed(Error):
    """Repository path not allowed"""

class InvalidRPCMethod(Error):
    """RPC method is not valid"""

class RepositoryServer:  # pragma: no cover
    rpc_methods = (
        '__len__',
        'check',
        'commit',
        'delete',
        'destroy',
        'get',
        'list',
        'negotiate',
        'open',
        'put',
        'repair',
        'rollback',
        'save_key',
        'load_key',
    )

    def __init__(self, restrict_to_paths):
        self.repository = None
        self.restrict_to_paths = restrict_to_paths

    def serve(self):
        stdin_fd = sys.stdin.fileno()
        stdout_fd = sys.stdout.fileno()
        # Make stdin non-blocking
        fl = fcntl.fcntl(stdin_fd, fcntl.F_GETFL)
        fcntl.fcntl(stdin_fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
        # Make stdout blocking
        fl = fcntl.fcntl(stdout_fd, fcntl.F_GETFL)
        fcntl.fcntl(stdout_fd, fcntl.F_SETFL, fl & ~os.O_NONBLOCK)
        unpacker = msgpack.Unpacker(use_list=False)
        while True:
            r, w, es = select.select([stdin_fd], [], [], 10)
            if r:
                data = os.read(stdin_fd, BUFSIZE)
                if not data:
                    return
                unpacker.feed(data)
                for unpacked in unpacker:
                    if not (isinstance(unpacked, tuple) and len(unpacked) == 4):
                        raise Exception("Unexpected RPC data format.")
                    type, msgid, method, args = unpacked
                    method = method.decode('ascii')
                    try:
                        if method not in self.rpc_methods:
                            raise InvalidRPCMethod(method)
                        try:
                            f = getattr(self, method)
                        except AttributeError:
                            f = getattr(self.repository, method)
                        res = f(*args)
                    except BaseException as e:
                        exc = "Remote Traceback by Borg %s%s%s" % (__version__, os.linesep, traceback.format_exc())
                        os.write(stdout_fd, msgpack.packb((1, msgid, e.__class__.__name__, exc)))
                    else:
                        os.write(stdout_fd, msgpack.packb((1, msgid, None, res)))
            if es:
                return

    def negotiate(self, versions):
        return 1

    def open(self, path, create=False):
        path = os.fsdecode(path)
        if path.startswith('/~'):
            path = path[1:]
        path = os.path.realpath(os.path.expanduser(path))
        if self.restrict_to_paths:
            for restrict_to_path in self.restrict_to_paths:
                if path.startswith(os.path.realpath(restrict_to_path)):
                    break
            else:
                raise PathNotAllowed(path)
        self.repository = Repository(path, create)
        return self.repository.id

class RemoteRepository:
    extra_test_args = []
    remote_path = 'borg'
    # default umask, overriden by --umask, defaults to read/write only for owner
    umask = 0o077

    class RPCError(Exception):
        def __init__(self, name):
            self.name = name

    def __init__(self, location, create=False):
        self.location = location
        self.preload_ids = []
        self.msgid = 0
        self.to_send = b''
        self.cache = {}
        self.ignore_responses = set()
        self.responses = {}
        self.unpacker = msgpack.Unpacker(use_list=False)
        self.p = None
        # XXX: ideally, the testsuite would subclass Repository and
        # override ssh_cmd() instead of this crude hack, although
        # __testsuite__ is not a valid domain name so this is pretty
        # safe.
        if location.host == '__testsuite__':
            args = [sys.executable, '-m', 'borg.archiver', 'serve' ] + self.extra_test_args
        else:  # pragma: no cover
            args = self.ssh_cmd(location)
        self.p = Popen(args, bufsize=0, stdin=PIPE, stdout=PIPE)
        self.stdin_fd = self.p.stdin.fileno()
        self.stdout_fd = self.p.stdout.fileno()
        fcntl.fcntl(self.stdin_fd, fcntl.F_SETFL, fcntl.fcntl(self.stdin_fd, fcntl.F_GETFL) | os.O_NONBLOCK)
        fcntl.fcntl(self.stdout_fd, fcntl.F_SETFL, fcntl.fcntl(self.stdout_fd, fcntl.F_GETFL) | os.O_NONBLOCK)
        self.r_fds = [self.stdout_fd]
        self.x_fds = [self.stdin_fd, self.stdout_fd]

        try:
            version = self.call('negotiate', 1)
        except ConnectionClosed:
            raise Exception('Server immediately closed connection - is Borg installed and working on the server?')
        if version != 1:
            raise Exception('Server insisted on using unsupported protocol version %d' % version)
        self.id = self.call('open', location.path, create)

    def __del__(self):
        self.close()

    def __repr__(self):
        return '<%s %s>' % (self.__class__.__name__, self.location.canonical_path())

    def umask_flag(self):
        return ['--umask', '%03o' % self.umask]

    def ssh_cmd(self, location):
        args = shlex.split(os.environ.get('BORG_RSH', 'ssh'))
        if location.port:
            args += ['-p', str(location.port)]
        if location.user:
            args.append('%s@%s' % (location.user, location.host))
        else:
            args.append('%s' % location.host)
        # use local umask also for the remote process
        args += [self.remote_path, 'serve'] + self.umask_flag()
        return args

    def call(self, cmd, *args, **kw):
        for resp in self.call_many(cmd, [args], **kw):
            return resp

    def call_many(self, cmd, calls, wait=True, is_preloaded=False):
        if not calls:
            return

        def fetch_from_cache(args):
            msgid = self.cache[args].pop(0)
            if not self.cache[args]:
                del self.cache[args]
            return msgid

        calls = list(calls)
        waiting_for = []
        w_fds = [self.stdin_fd]
        while wait or calls:
            while waiting_for:
                try:
                    error, res = self.responses.pop(waiting_for[0])
                    waiting_for.pop(0)
                    if error:
                        if error == b'DoesNotExist':
                            raise Repository.DoesNotExist(self.location.orig)
                        elif error == b'AlreadyExists':
                            raise Repository.AlreadyExists(self.location.orig)
                        elif error == b'CheckNeeded':
                            raise Repository.CheckNeeded(self.location.orig)
                        elif error == b'IntegrityError':
                            raise IntegrityError(res)
                        elif error == b'PathNotAllowed':
                            raise PathNotAllowed(*res)
                        elif error == b'ObjectNotFound':
                            raise Repository.ObjectNotFound(res[0], self.location.orig)
                        elif error == b'InvalidRPCMethod':
                            raise InvalidRPCMethod(*res)
                        else:
                            raise self.RPCError(res.decode('utf-8'))
                    else:
                        yield res
                        if not waiting_for and not calls:
                            return
                except KeyError:
                    break
            r, w, x = select.select(self.r_fds, w_fds, self.x_fds, 1)
            if x:
                raise Exception('FD exception occurred')
            if r:
                data = os.read(self.stdout_fd, BUFSIZE)
                if not data:
                    raise ConnectionClosed()
                self.unpacker.feed(data)
                for unpacked in self.unpacker:
                    if not (isinstance(unpacked, tuple) and len(unpacked) == 4):
                        raise Exception("Unexpected RPC data format.")
                    type, msgid, error, res = unpacked
                    if msgid in self.ignore_responses:
                        self.ignore_responses.remove(msgid)
                    else:
                        self.responses[msgid] = error, res
            if w:
                while not self.to_send and (calls or self.preload_ids) and len(waiting_for) < 100:
                    if calls:
                        if is_preloaded:
                            if calls[0] in self.cache:
                                waiting_for.append(fetch_from_cache(calls.pop(0)))
                        else:
                            args = calls.pop(0)
                            if cmd == 'get' and args in self.cache:
                                waiting_for.append(fetch_from_cache(args))
                            else:
                                self.msgid += 1
                                waiting_for.append(self.msgid)
                                self.to_send = msgpack.packb((1, self.msgid, cmd, args))
                    if not self.to_send and self.preload_ids:
                        args = (self.preload_ids.pop(0),)
                        self.msgid += 1
                        self.cache.setdefault(args, []).append(self.msgid)
                        self.to_send = msgpack.packb((1, self.msgid, cmd, args))

                if self.to_send:
                    try:
                        self.to_send = self.to_send[os.write(self.stdin_fd, self.to_send):]
                    except OSError as e:
                        # io.write might raise EAGAIN even though select indicates
                        # that the fd should be writable
                        if e.errno != errno.EAGAIN:
                            raise
                if not self.to_send and not (calls or self.preload_ids):
                    w_fds = []
        self.ignore_responses |= set(waiting_for)

    def check(self, repair=False):
        return self.call('check', repair)

    def commit(self, *args):
        return self.call('commit')

    def rollback(self, *args):
        return self.call('rollback')

    def destroy(self):
        return self.call('destroy')

    def __len__(self):
        return self.call('__len__')

    def list(self, limit=None, marker=None):
        return self.call('list', limit, marker)

    def get(self, id_):
        for resp in self.get_many([id_]):
            return resp

    def get_many(self, ids, is_preloaded=False):
        for resp in self.call_many('get', [(id_,) for id_ in ids], is_preloaded=is_preloaded):
            yield resp

    def put(self, id_, data, wait=True):
        return self.call('put', id_, data, wait=wait)

    def delete(self, id_, wait=True):
        return self.call('delete', id_, wait=wait)

    def save_key(self, keydata):
        return self.call('save_key', keydata)

    def load_key(self):
        return self.call('load_key')

    def close(self):
        if self.p:
            self.p.stdin.close()
            self.p.stdout.close()
            self.p.wait()
            self.p = None

    def preload(self, ids):
        self.preload_ids += ids

class RepositoryCache:
    """A caching Repository wrapper

    Caches Repository GET operations using a local temporary Repository.
    """
    def __init__(self, repository):
        self.repository = repository
        tmppath = tempfile.mkdtemp(prefix='borg-tmp')
        self.caching_repo = Repository(tmppath, create=True, exclusive=True)

    def __del__(self):
        self.caching_repo.destroy()

    def get(self, key):
        return next(self.get_many([key]))

    def get_many(self, keys):
        unknown_keys = [key for key in keys if key not in self.caching_repo]
        repository_iterator = zip(unknown_keys, self.repository.get_many(unknown_keys))
        for key in keys:
            try:
                yield self.caching_repo.get(key)
            except Repository.ObjectNotFound:
                for key_, data in repository_iterator:
                    if key_ == key:
                        self.caching_repo.put(key, data)
                        yield data
                        break
        # Consume any pending requests
        for _ in repository_iterator:
            pass

def cache_if_remote(repository):
    if isinstance(repository, RemoteRepository):
        return RepositoryCache(repository)
    return repository
<|endoftext|>


<commit_before>####
#
#   get followers, classify them as bots, aggregate.
#
###

from __future__ import print_function
import json
import tweepy
import requests

auth = tweepy.OAuthHandler(consumer_key', 'consumer_secret')
auth.set_access_token()

twitter_api = tweepy.API(auth, parser=tweepy.parsers.JSONParser())

screen_name = '@clayadavis'

def get_botOrNot(screen_name):
    user_timeline = twitter_api.user_timeline(screen_name, count=200)

    if user_timeline:
        user_data = user_timeline[0]['user']
    else:
        user_data = twitter_api.get_user(screen_name)
    if not screen_name.startswith('@'):
        screen_name = '@' + screen_name
    search = twitter_api.search(screen_name, count=100)
    mentions = search['statuses']
    tweets = user_timeline + mentions
    post_body = {'content': tweets,
             'meta': {
                'user_id': user_data['id_str'],
                'screen_name': screen_name,
                },
            }

    bon_url = 'http://truthy.indiana.edu/botornot/api/1/check_account'
    bon_response = requests.post(bon_url, data=json.dumps(post_body))
    bon_response.status_code
    return(bon_response.json())

<commit_msg>(oops) Forgot to commit the deleted file<commit_after><|endoftext|>


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

"""
Calculates distances between neighboring residues of a ligand molecule and produces a set of
unambiguous distance restraints for HADDOCK to keep it in place during semi-flexible refinement.

Produces, at most, one restraint per ligand atom.

Creates a file containing the restraints, named after the original input PDB file and the ligand
residue name, with '.tbl' extension.
"""

from __future__ import print_function
import argparse
import os
import sys

try:
    import numpy as np
    from Bio.PDB import PDBParser
    from Bio.PDB import NeighborSearch
    from Bio.PDB.Polypeptide import is_aa
except ImportError as e:
    print('[!] Could not import module \'biopython\': {0}'.format(e), file=sys.stderr)
    sys.exit(1)

cline_parser = argparse.ArgumentParser(description=__doc__)
cline_parser.add_argument('pdbf', help='PDB file')
cline_parser.add_argument('-l', '--ligand', help='Ligand residue name', required=True)
cline_parser.add_argument('-p', '--pml', action='store_true', help='Write Pymol file with restraints')
args = cline_parser.parse_args()

if not os.path.isfile(args.pdbf):
    print('[!!] File not found: {0}'.format(args.pdbf))
    sys.exit(1)
else:
    root_name, _ = os.path.splitext(os.path.basename(args.pdbf))

# Read in structure
pdb_parser = PDBParser(QUIET=1)
structure = pdb_parser.get_structure(root_name, args.pdbf)

# Remove hydrogens
atom_lst = list(structure.get_atoms())
for atom in atom_lst:
    if atom.element == 'H':
        res = atom.parent
        res.detach_child(atom.name)

ligand = None
for residue in structure.get_residues():
    if residue.resname.strip() == args.ligand.strip():
        ligand = residue
        break

if not ligand:
    print('[!!] Ligand residue \'{0}\' not found in structure'.format(args.ligand), file=sys.stderr)
    sys.exit(1)

# Calculate center of mass of the ligand
ligand_com = map(lambda x: sum(x)/len(x), zip(*[at.coord for at in ligand]))
ligand_com = np.asarray(ligand_com, dtype=np.float32)

# Calculate neighbors considering only aminoacid/nucleotide atoms (excl. waters, other ligands, etc)
sel_atoms = [at for at in structure.get_atoms() if at.parent.id[0] == ' ']
ns = NeighborSearch(sel_atoms)
neighbors = ns.search(ligand_com, 10.0, level='R')  # 10A radius, return residues

# Calculate residue closer to each ligand atom and the respective distance
ligand_atoms = ligand.child_list
min_dist_list, _seen = [], set()

for l_at in ligand_atoms:
    distances = []
    for residue in neighbors:
        for r_at in residue:
            distances.append((r_at, l_at, r_at - l_at))

    distances.sort(key=lambda x: x[-1])
    min_dist = distances[0]
    # One restraint per residue to keep the number of restraints small
    if min_dist[0].parent not in _seen:
        min_dist_list.append(min_dist)
        _seen.add(min_dist[0].parent)

# Output
# TBL: use chain ID as segid ID
with open('{}_{}.tbl'.format(root_name, args.ligand), 'w') as tbl_handle:
    _str = 'assi (segi {:4s} and resi {:4d} and name {:6s})\n' + \
           '     (segi {:4s} and resi {:4d} and name {:6s}) {:6.3f} 1.0 1.0'

    print('! Restraints to fix \'{}\' in its initial position'.format(args.ligand), file=tbl_handle)
    for dist in min_dist_list:
        r_at, l_at, d = dist
        print(_str.format(r_at.parent.parent.id, r_at.parent.id[1], '"' + r_at.name + '"',
                          l_at.parent.parent.id, l_at.parent.id[1], '"' + l_at.name + '"',
                          d), file=tbl_handle)

# PML (if requested)
if args.pml:
    with open('{}_{}.pml'.format(root_name, args.ligand), 'w') as pml_handle:
        _str = 'distance rst_{}, c. {} and i. {} and n. "{}", c. {} and i. {} and n. "{}"'
        for irest, dist in enumerate(min_dist_list, start=1):
            r_at, l_at, _ = dist
            print(_str.format(irest,
                              r_at.parent.parent.id, r_at.parent.id[1], r_at.name,
                              l_at.parent.parent.id, l_at.parent.id[1], l_at.name), file=pml_handle)

<commit_msg>fix the issue with map object<commit_after>#!/usr/bin/env python

"""
Calculates distances between neighboring residues of a ligand molecule and produces a set of
unambiguous distance restraints for HADDOCK to keep it in place during semi-flexible refinement.

Produces, at most, one restraint per ligand atom.

Creates a file containing the restraints, named after the original input PDB file and the ligand
residue name, with '.tbl' extension.
"""

from __future__ import print_function
import argparse
import os
import sys

try:
    import numpy as np
    from Bio.PDB import PDBParser
    from Bio.PDB import NeighborSearch
    from Bio.PDB.Polypeptide import is_aa
except ImportError as e:
    print('[!] Could not import module \'biopython\': {0}'.format(e), file=sys.stderr)
    sys.exit(1)

cline_parser = argparse.ArgumentParser(description=__doc__)
cline_parser.add_argument('pdbf', help='PDB file')
cline_parser.add_argument('-l', '--ligand', help='Ligand residue name', required=True)
cline_parser.add_argument('-p', '--pml', action='store_true', help='Write Pymol file with restraints')
args = cline_parser.parse_args()

if not os.path.isfile(args.pdbf):
    print('[!!] File not found: {0}'.format(args.pdbf))
    sys.exit(1)
else:
    root_name, _ = os.path.splitext(os.path.basename(args.pdbf))

# Read in structure
pdb_parser = PDBParser(QUIET=1)
structure = pdb_parser.get_structure(root_name, args.pdbf)

# Remove hydrogens
atom_lst = list(structure.get_atoms())
for atom in atom_lst:
    if atom.element == 'H':
        res = atom.parent
        res.detach_child(atom.name)

ligand = None
for residue in structure.get_residues():
    if residue.resname.strip() == args.ligand.strip():
        ligand = residue
        break

if not ligand:
    print('[!!] Ligand residue \'{0}\' not found in structure'.format(args.ligand), file=sys.stderr)
    sys.exit(1)

# Calculate center of mass of the ligand
ligand_com = list(map(lambda x: sum(x)/len(x), zip(*[at.coord for at in ligand])))
ligand_com = np.asarray(ligand_com, dtype=np.float32)

# Calculate neighbors considering only aminoacid/nucleotide atoms (excl. waters, other ligands, etc)
sel_atoms = [at for at in structure.get_atoms() if at.parent.id[0] == ' ']
ns = NeighborSearch(sel_atoms)
neighbors = ns.search(ligand_com, 10.0, level='R')  # 10A radius, return residues

# Calculate residue closer to each ligand atom and the respective distance
ligand_atoms = ligand.child_list
min_dist_list, _seen = [], set()

for l_at in ligand_atoms:
    distances = []
    for residue in neighbors:
        for r_at in residue:
            distances.append((r_at, l_at, r_at - l_at))

    distances.sort(key=lambda x: x[-1])
    min_dist = distances[0]
    # One restraint per residue to keep the number of restraints small
    if min_dist[0].parent not in _seen:
        min_dist_list.append(min_dist)
        _seen.add(min_dist[0].parent)

# Output
# TBL: use chain ID as segid ID
with open('{}_{}.tbl'.format(root_name, args.ligand), 'w') as tbl_handle:
    _str = 'assi (segi {:4s} and resi {:4d} and name {:6s})\n' + \
           '     (segi {:4s} and resi {:4d} and name {:6s}) {:6.3f} 1.0 1.0'

    print('! Restraints to fix \'{}\' in its initial position'.format(args.ligand), file=tbl_handle)
    for dist in min_dist_list:
        r_at, l_at, d = dist
        print(_str.format(r_at.parent.parent.id, r_at.parent.id[1], '"' + r_at.name + '"',
                          l_at.parent.parent.id, l_at.parent.id[1], '"' + l_at.name + '"',
                          d), file=tbl_handle)

# PML (if requested)
if args.pml:
    with open('{}_{}.pml'.format(root_name, args.ligand), 'w') as pml_handle:
        _str = 'distance rst_{}, c. {} and i. {} and n. "{}", c. {} and i. {} and n. "{}"'
        for irest, dist in enumerate(min_dist_list, start=1):
            r_at, l_at, _ = dist
            print(_str.format(irest,
                              r_at.parent.parent.id, r_at.parent.id[1], r_at.name,
                              l_at.parent.parent.id, l_at.parent.id[1], l_at.name), file=pml_handle)

<|endoftext|>


<commit_before>#script to download library and dll needed for stiuniverse
#the purpose is to remove binary from git repository

import zipfile
import os
import urllib2
#hard coded url
url = 'stigame_extern_x86.zip'

def download(url , output):

    req = urllib2.Request(url)

    resp = urllib2.urlopen(req)

    fp = open(output, 'wb')

    buf_size = 16384

    length = resp.info().getheaders('Content-Length')[0]
    length = int(length)
    current = 0
    print "Starting download"
    chunk = resp.read(buf_size)
    last_prc = None
    i = 0
    while len(chunk) == buf_size:
        current = current + len(chunk)
        prc = (current*100) / length

        if not prc == last_prc:
            prefix = ""
            nb = i / 5
            for x in range(nb):
                prefix = prefix + "="
            i = i + 1

            print prefix + "> " + str(prc) + "%"
            last_prc = prc

        fp.write(chunk)

        chunk = resp.read(buf_size)

    fp.write(chunk)

    fp.close()

    print "Download complete !"

def extract(zipname, dst = None):

    zp = zipfile.ZipFile(zipname, 'r')

    names = zp.namelist()

    for n in names:
        if not n.endswith('/'):
            cdir = os.path.dirname(n)
            if not dst == None:
                cdir = os.path.join(dst, cdir)
            if not os.path.exists(cdir):
                os.makedirs(cdir)
            print "Extracting : " + n
            if dst == None:
                zp.extract(n)
            else:
                zp.extract(n, dst)

    zp.close()

if __name__ == '__main__':
    download(url, 'stigame_extern_x86.zip')
    extract('stigame_extern_x86.zip')

<commit_msg> extern url fixed 	modified:   retrieve_extern.py<commit_after>#script to download library and dll needed for stiuniverse
#the purpose is to remove binary from git repository

import zipfile
import os
import urllib2
#hard coded url
url = 'https://docs.google.com/uc?authuser=0&id=0B_P5byse05MmS01LUi1HcEdsZnc&export=download'

def download(url , output):

    req = urllib2.Request(url)

    resp = urllib2.urlopen(req)

    fp = open(output, 'wb')

    buf_size = 16384

    length = resp.info().getheaders('Content-Length')[0]
    length = int(length)
    current = 0
    print "Starting download"
    chunk = resp.read(buf_size)
    last_prc = None
    i = 0
    while len(chunk) == buf_size:
        current = current + len(chunk)
        prc = (current*100) / length

        if not prc == last_prc:
            prefix = ""
            nb = i / 5
            for x in range(nb):
                prefix = prefix + "="
            i = i + 1

            print prefix + "> " + str(prc) + "%"
            last_prc = prc

        fp.write(chunk)

        chunk = resp.read(buf_size)

    fp.write(chunk)

    fp.close()

    print "Download complete !"

def extract(zipname, dst = None):

    zp = zipfile.ZipFile(zipname, 'r')

    names = zp.namelist()

    for n in names:
        if not n.endswith('/'):
            cdir = os.path.dirname(n)
            if not dst == None:
                cdir = os.path.join(dst, cdir)
            if not os.path.exists(cdir):
                os.makedirs(cdir)
            print "Extracting : " + n
            if dst == None:
                zp.extract(n)
            else:
                zp.extract(n, dst)

    zp.close()

if __name__ == '__main__':
    download(url, 'stigame_extern_x86.zip')
    extract('stigame_extern_x86.zip')

<|endoftext|>


<commit_before>import networkx as nx
from copy import deepcopy
import numpy as np
from petsc4py import PETSc

class GeometricProperties:
    def __init__(self, n_x, n_y, size_x, size_y):
        self.n_x = n_x
        self.n_y = n_y
        self.size_x = size_x
        self.size_y = size_y

class PhysicalProperties:
    def __init__(self, k, rho, c_p):
        self.k = k
        self.rho = rho
        self.c_p = c_p

class ConstantInitialCondition:
    def __init__(self, T):
        self.T = T

class GhostNodeBoundaryCondition:
    def __init__(self, T_E, T_W, T_N, T_S):
        self.T_E = T_E
        self.T_W = T_W
        self.T_N = T_N
        self.T_S = T_S

class TimestepProperties:
    def __init__(self, delta_t, final_time):
        self.delta_t = delta_t
        self.final_time = final_time

def build_graph(geometric_properties, physical_properties, initial_condition, boundary_condition):
    '''
    Builds the problem's non-oriented graph G, where each node represents a control volume, and each
    edge represent that two control volumes are neighbors.
    '''
    # Aliases
    size_x = geometric_properties.size_x
    n_x = geometric_properties.n_x
    size_y = geometric_properties.size_y
    n_y = geometric_properties.n_y

    # Computed variables
    dx = size_x / n_x
    dy = size_y / n_y

    # Constants data among the grid
    G = {}
    G['dx'] = dx
    G['dy'] = dy
    G['n_x'] = n_x
    G['n_y'] = n_y
    G['k'] = physical_properties.k
    G['rho'] = physical_properties.rho
    G['c_p'] = physical_properties.c_p
    G['T'] = np.ones(shape=(n_x, n_y)) * initial_condition.T

    return G

def build_matrix_structure(graph):
    return None

def solve(
    geometric_properties,
    physical_properties,
    initial_condition,
    boundary_condition,
    timestep_properties,

    # Output options
    on_timestep_callback=None,

    # Solver options
    use_multigrid=True,
    use_newton=False,
):
    options = PETSc.Options()
    options.clear()

    if use_multigrid:
        options.setValue('pc_type', 'gamg')
        options.setValue('pc_gamg_reuse_interpolation', 'True')
    else:
        options.setValue('pc_type', 'ilu')

    options.setValue('ksp_type', 'gmres')

    if use_newton:
        options.setValue('snes_type', 'newtonls')
    else:
        options.setValue('snes_type', 'ksponly')

    options.setValue('ksp_atol', 1e-7)

    if on_timestep_callback is None:
        on_timestep_callback = lambda *args, **kwargs: None

    graph = build_graph(
        geometric_properties,
        physical_properties,
        initial_condition,
        boundary_condition,
    )
    old_graph = deepcopy(graph)

    current_time = 0.0
    on_timestep_callback(current_time, graph['T'][:])

    n_x = graph['n_x']
    n_y = graph['n_y']
    dmda = PETSc.DMDA().create([n_x, n_y], dof=1, stencil_width=1, stencil_type='star')

    dt = timestep_properties.delta_t
    dx = graph['dx']
    dy = graph['dy']
    k = graph['k']
    rho = graph['rho']
    c_p = graph['c_p']

    # Build the matrix coeficients
    A_P = rho * dx * dy / dt + \
        2 * (k * dy) / (c_p * dx) + \
        2 * (k * dx) / (c_p * dy)
    A_E = (k * dy) / (c_p * dx)
    A_W = (k * dy) / (c_p * dx)
    A_N = (k * dx) / (c_p * dy)
    A_S = (k * dx) / (c_p * dy)

    snes = PETSc.SNES().create()
    snes.setFromOptions()

    r = PETSc.Vec().createSeq(n_x * n_y)  # residual vector
    x = PETSc.Vec().createSeq(n_x * n_y)  # solution vector
    b = PETSc.Vec().createSeq(n_x * n_y)  # right-hand side

    # solution = initial_condition
    while current_time < timestep_properties.final_time:
        # Aliases
        B_T = (rho * dx * dy * old_graph['T']) / dt

        # Build the equation system
        def residual_function(snes, X, f):
            x = X.getArray(readonly=True).reshape((n_x, n_y))
            eqs = np.zeros(shape=(n_x, n_y))
            eqs[1:-1, 1:-1] = \
                  A_P * x[1:-1, 1:-1] \
                - A_W * x[1:-1, :-2] \
                - A_E * x[1:-1, 2:] \
                - A_N * x[:-2, 1:-1] \
                - A_S * x[2:, 1:-1] \
                - B_T[1:-1, 1:-1]

            # Boundary conditions
            # Left
            eqs[1:-1, :1] = \
                  A_P * x[1:-1, :1] \
                - A_W * boundary_condition.T_W(current_time) \
                - A_E * x[1:-1, 1:2] \
                - A_N * x[:-2, :1] \
                - A_S * x[2:, :1] \
                - B_T[1:-1, :1]
            # Right
            eqs[1:-1, -1:] = \
                  A_P * x[1:-1, -1:] \
                - A_W * x[1:-1, -2:-1] \
                - A_E * boundary_condition.T_E(current_time) \
                - A_N * x[:-2, -1:] \
                - A_S * x[2:, -1:] \
                - B_T[1:-1, -1:]
            # Bottom
            eqs[-1:, 1:-1] = \
                  A_P * x[-1:, 1:-1] \
                - A_W * x[-1:, :-2] \
                - A_E * x[-1:, 2:] \
                - A_N * x[-2:-1, 1:-1] \
                - A_S * boundary_condition.T_S(current_time) \
                - B_T[-1:, 1:-1]
            # Top
            eqs[:1, 1:-1] = \
                  A_P * x[:1, 1:-1] \
                - A_W * x[:1, :-2] \
                - A_E * x[:1, 2:] \
                - A_N * boundary_condition.T_N(current_time) \
                - A_S * x[1:2, 1:-1] \
                - B_T[:1, 1:-1]

            # Top-Left
            eqs[:1, :1] = \
                  A_P * x[:1, :1] \
                - A_W * boundary_condition.T_W(current_time) \
                - A_E * x[:1, 1:2] \
                - A_N * boundary_condition.T_N(current_time) \
                - A_S * x[1:2, :1] \
                - B_T[:1, :1]
            # Top-Right
            eqs[:1, -1:] = \
                  A_P * x[:1, -1:] \
                - A_W * x[:1, -2:-1] \
                - A_E * boundary_condition.T_E(current_time) \
                - A_N * boundary_condition.T_N(current_time) \
                - A_S * x[1:2, -1:] \
                - B_T[:1, -1:]
            # Bottom-Left
            eqs[-1:, :1] = \
                  A_P * x[-1:, :1] \
                - A_W * boundary_condition.T_W(current_time) \
                - A_E * x[-1:, 1:2] \
                - A_N * x[-2:-1, :1] \
                - A_S * boundary_condition.T_S(current_time) \
                - B_T[-1:, :1]
            # Bottom-Right
            eqs[-1:, -1:] = \
                  A_P * x[-1:, -1:] \
                - A_W * x[-1:, -2:-1] \
                - A_E * boundary_condition.T_E(current_time) \
                - A_N * x[-2:-1, -1:] \
                - A_S * boundary_condition.T_S(current_time) \
                - B_T[-1:, -1:]

            f[:] = eqs.reshape(n_x * n_y)

        # Solve Ax = B
        snes.setFunction(residual_function, r)

        snes.setDM(dmda)

        initial_guess = np.array(old_graph['T'])
        x.setArray(initial_guess)
        b.set(0)
        snes.solve(b, x)
        solution = x[:].reshape(n_x, n_y)
        solution.flags['WRITEABLE'] = False

        # Retrieve the solution
        old_graph = deepcopy(graph)
        graph['T'][:] = solution

        # Advance in time
        current_time += timestep_properties.delta_t

        on_timestep_callback(current_time, solution)

    return graph
<commit_msg>Fix boundary conditions<commit_after>from copy import deepcopy
import numpy as np
from petsc4py import PETSc

class GeometricProperties:
    def __init__(self, n_x, n_y, size_x, size_y):
        self.n_x = n_x
        self.n_y = n_y
        self.size_x = size_x
        self.size_y = size_y

class PhysicalProperties:
    def __init__(self, k, rho, c_p):
        self.k = k
        self.rho = rho
        self.c_p = c_p

class ConstantInitialCondition:
    def __init__(self, T):
        self.T = T

class GhostNodeBoundaryCondition:
    def __init__(self, T_E, T_W, T_N, T_S):
        self.T_E = T_E
        self.T_W = T_W
        self.T_N = T_N
        self.T_S = T_S

class TimestepProperties:
    def __init__(self, delta_t, final_time):
        self.delta_t = delta_t
        self.final_time = final_time

def build_grid(geometric_properties, physical_properties, initial_condition, boundary_condition):
    # Aliases
    size_x = geometric_properties.size_x
    n_x = geometric_properties.n_x
    size_y = geometric_properties.size_y
    n_y = geometric_properties.n_y

    # Computed variables
    dx = size_x / n_x
    dy = size_y / n_y

    # Data among the grid
    G = {}
    G['dx'] = dx
    G['dy'] = dy
    G['n_x'] = n_x
    G['n_y'] = n_y
    G['k'] = physical_properties.k
    G['rho'] = physical_properties.rho
    G['c_p'] = physical_properties.c_p
    G['T'] = np.ones(shape=(n_x, n_y)) * initial_condition.T

    return G

def solve(
    geometric_properties,
    physical_properties,
    initial_condition,
    boundary_condition,
    timestep_properties,

    # Output options
    on_timestep_callback=None,

    # Solver options
    use_multigrid=True,
    use_newton=False,
):
    options = PETSc.Options()
    options.clear()

    if use_multigrid:
        options.setValue('pc_type', 'gamg')
        options.setValue('pc_gamg_reuse_interpolation', 'True')
    else:
        options.setValue('pc_type', 'ilu')

    options.setValue('ksp_type', 'gmres')

    if use_newton:
        options.setValue('snes_type', 'newtonls')
    else:
        options.setValue('snes_type', 'ksponly')

    options.setValue('ksp_atol', 1e-7)

    if on_timestep_callback is None:
        on_timestep_callback = lambda *args, **kwargs: None

    graph = build_grid(
        geometric_properties,
        physical_properties,
        initial_condition,
        boundary_condition,
    )
    old_graph = deepcopy(graph)

    current_time = 0.0
    on_timestep_callback(current_time, graph['T'][:])

    n_x = graph['n_x']
    n_y = graph['n_y']
    dmda = PETSc.DMDA().create([n_x, n_y], dof=1, stencil_width=1, stencil_type='star')

    dt = timestep_properties.delta_t
    dx = graph['dx']
    dy = graph['dy']
    k = graph['k']
    rho = graph['rho']
    c_p = graph['c_p']

    # Build the matrix coeficients
    A_P = rho * dx * dy / dt + \
        2 * (k * dy) / (c_p * dx) + \
        2 * (k * dx) / (c_p * dy)
    A_E = (k * dy) / (c_p * dx)
    A_W = (k * dy) / (c_p * dx)
    A_N = (k * dx) / (c_p * dy)
    A_S = (k * dx) / (c_p * dy)

    snes = PETSc.SNES().create()
    snes.setFromOptions()

    r = PETSc.Vec().createSeq(n_x * n_y)  # residual vector
    x = PETSc.Vec().createSeq(n_x * n_y)  # solution vector
    b = PETSc.Vec().createSeq(n_x * n_y)  # right-hand side

    # solution = initial_condition
    while current_time < timestep_properties.final_time:
        # Aliases
        B_T = (rho * dx * dy * old_graph['T']) / dt

        # Build the equation system
        def residual_function(snes, X, f):
            x = X.getArray(readonly=True).reshape((n_x, n_y))
            eqs = np.zeros(shape=(n_x, n_y))
            eqs[1:-1, 1:-1] = \
                  A_P * x[1:-1, 1:-1] \
                - A_W * x[1:-1, :-2] \
                - A_E * x[1:-1, 2:] \
                - A_N * x[:-2, 1:-1] \
                - A_S * x[2:, 1:-1] \
                - B_T[1:-1, 1:-1]

            # Boundary conditions
            # Left
            eqs[1:-1, :1] = \
                  A_P * x[1:-1, :1] \
                - boundary_condition.T_W(current_time) \
                - A_E * x[1:-1, 1:2] \
                - A_N * x[:-2, :1] \
                - A_S * x[2:, :1] \
                - B_T[1:-1, :1]
            # Right
            eqs[1:-1, -1:] = \
                  A_P * x[1:-1, -1:] \
                - A_W * x[1:-1, -2:-1] \
                - boundary_condition.T_E(current_time) \
                - A_N * x[:-2, -1:] \
                - A_S * x[2:, -1:] \
                - B_T[1:-1, -1:]
            # Bottom
            eqs[-1:, 1:-1] = \
                  A_P * x[-1:, 1:-1] \
                - A_W * x[-1:, :-2] \
                - A_E * x[-1:, 2:] \
                - A_N * x[-2:-1, 1:-1] \
                - boundary_condition.T_S(current_time) \
                - B_T[-1:, 1:-1]
            # Top
            eqs[:1, 1:-1] = \
                  A_P * x[:1, 1:-1] \
                - A_W * x[:1, :-2] \
                - A_E * x[:1, 2:] \
                - boundary_condition.T_N(current_time) \
                - A_S * x[1:2, 1:-1] \
                - B_T[:1, 1:-1]

            # Top-Left
            eqs[:1, :1] = \
                  A_P * x[:1, :1] \
                - boundary_condition.T_W(current_time) \
                - A_E * x[:1, 1:2] \
                - boundary_condition.T_N(current_time) \
                - A_S * x[1:2, :1] \
                - B_T[:1, :1]
            # Top-Right
            eqs[:1, -1:] = \
                  A_P * x[:1, -1:] \
                - A_W * x[:1, -2:-1] \
                - boundary_condition.T_E(current_time) \
                - boundary_condition.T_N(current_time) \
                - A_S * x[1:2, -1:] \
                - B_T[:1, -1:]
            # Bottom-Left
            eqs[-1:, :1] = \
                  A_P * x[-1:, :1] \
                - boundary_condition.T_W(current_time) \
                - A_E * x[-1:, 1:2] \
                - A_N * x[-2:-1, :1] \
                - boundary_condition.T_S(current_time) \
                - B_T[-1:, :1]
            # Bottom-Right
            eqs[-1:, -1:] = \
                  A_P * x[-1:, -1:] \
                - A_W * x[-1:, -2:-1] \
                - boundary_condition.T_E(current_time) \
                - A_N * x[-2:-1, -1:] \
                - boundary_condition.T_S(current_time) \
                - B_T[-1:, -1:]

            f[:] = eqs.reshape(n_x * n_y)

        # Solve Ax = B
        snes.setFunction(residual_function, r)

        snes.setDM(dmda)

        initial_guess = np.array(old_graph['T'])
        x.setArray(initial_guess)
        b.set(0)
        snes.solve(b, x)
        solution = x[:].reshape(n_x, n_y)
        solution.flags['WRITEABLE'] = False

        # Retrieve the solution
        old_graph = deepcopy(graph)
        graph['T'][:] = solution

        # Advance in time
        current_time += timestep_properties.delta_t

        on_timestep_callback(current_time, solution)

    return graph
<|endoftext|>


<commit_before>import numpy as np
import sympy
from sympy.printing.codeprinter import Assignment

from pyne import nucname
from pyne import data
from pyne import rxname

t, G = sympy.symbols('t G')
phi = sympy.MatrixSymbol('phi', G, 1)

decay_rxs = ['bminus', 'bplus', 'ec', 'alpha', 'it', 'sf', 'bminus_n']
xs_rxs = ['gamma', 'z_2n', 'z_3n', 'alpha', 'fission', 'proton', 'gamma_1', 'z_2n_1']

def child_decays(nuc):
    childname = nucname.name(nuc)
    expr = 0
    for rx in decay_rxs:
        try:
            parent = rxname.parent(nuc, rx, b'decay')
        except RuntimeError:
            continue
        if data.branch_ratio(parent, nuc) < 1e-16:
            continue
        parname = nucname.name(parent)
        par0, gamma, lambda_par = sympy.symbols('{0}_0 gamma_{0}_{1}_{2} lambda_{0}'.format(parname, childname, rx))
        expr += gamma * sympy.exp(lambda_par * t) * par0
    return expr

def child_xss(nuc):
    expr = 0
    for rx in xs_rxs:
        try:
            parent = rxname.parent(nuc, rx)
        except RuntimeError:
            continue
        parname = nucname.name(parent)
        par0 = sympy.symbols('{0}_0'.format(parname))
        sigma_rx_par = sympy.MatrixSymbol('sigma_{0}_{1}'.format(rx, parname), 1, G)
        expr += sympy.exp((sigma_rx_par*phi)[0] * t) * par0
    return expr

def gennuc(nuc):
    name = nucname.name(nuc)
    nuc0, nuc1 = sympy.symbols('{0}_0 {0}_1'.format(name))
    lambda_nuc = sympy.symbols('lambda_{0}'.format(name))
    sigma_a_nuc = sympy.MatrixSymbol('sigma_a_{0}'.format(name), 1, G)
    rhs = sympy.exp(-((sigma_a_nuc*phi)[0] + lambda_nuc)*t) * nuc0
    rhs += child_decays(nuc)
    rhs += child_xss(nuc)
    eq = Assignment(nuc1, rhs)
    return eq

if __name__ == '__main__':
    data.atomic_mass('U235')
    nucs = set(data.atomic_mass_map.keys())
    for nuc in data.atomic_mass_map:
        nucm = nuc + 1
        if nucname.anum(nuc) == 0 or data.decay_const(nucm) < 1e-16 or data.decay_const(nuc) == data.decay_const(nucm):
            continue
        nucs.add(nucm)
    nucs = [nuc for nuc in nucs if nucname.anum(nuc) > 0 and
                                not np.isnan(data.decay_const(nuc)) and
                                nuc < 200000000]
    nucs.sort()

    system = list(map(gennuc, nucs))

    with open('system.txt', 'w') as f:
        for eq in system:
            f.write(str(eq) + '\n')
<commit_msg>Use CodeBlock<commit_after>import numpy as np
import sympy
from sympy.printing.codeprinter import Assignment, CodeBlock

from pyne import nucname
from pyne import data
from pyne import rxname

t, G = sympy.symbols('t G')
phi = sympy.MatrixSymbol('phi', G, 1)

decay_rxs = ['bminus', 'bplus', 'ec', 'alpha', 'it', 'sf', 'bminus_n']
xs_rxs = ['gamma', 'z_2n', 'z_3n', 'alpha', 'fission', 'proton', 'gamma_1', 'z_2n_1']

def child_decays(nuc):
    childname = nucname.name(nuc)
    expr = 0
    for rx in decay_rxs:
        try:
            parent = rxname.parent(nuc, rx, b'decay')
        except RuntimeError:
            continue
        if data.branch_ratio(parent, nuc) < 1e-16:
            continue
        parname = nucname.name(parent)
        par0, gamma, lambda_par = sympy.symbols('{0}_0 gamma_{0}_{1}_{2} lambda_{0}'.format(parname, childname, rx))
        expr += gamma * sympy.exp(lambda_par * t) * par0
    return expr

def child_xss(nuc):
    expr = 0
    for rx in xs_rxs:
        try:
            parent = rxname.parent(nuc, rx)
        except RuntimeError:
            continue
        parname = nucname.name(parent)
        par0 = sympy.symbols('{0}_0'.format(parname))
        sigma_rx_par = sympy.MatrixSymbol('sigma_{0}_{1}'.format(rx, parname), 1, G)
        expr += sympy.exp((sigma_rx_par*phi)[0] * t) * par0
    return expr

def gennuc(nuc):
    name = nucname.name(nuc)
    nuc0, nuc1 = sympy.symbols('{0}_0 {0}_1'.format(name))
    lambda_nuc = sympy.symbols('lambda_{0}'.format(name))
    sigma_a_nuc = sympy.MatrixSymbol('sigma_a_{0}'.format(name), 1, G)
    rhs = sympy.exp(-((sigma_a_nuc*phi)[0] + lambda_nuc)*t) * nuc0
    rhs += child_decays(nuc)
    rhs += child_xss(nuc)
    eq = Assignment(nuc1, rhs)
    return eq

if __name__ == '__main__':
    data.atomic_mass('U235')
    nucs = set(data.atomic_mass_map.keys())
    for nuc in data.atomic_mass_map:
        nucm = nuc + 1
        if nucname.anum(nuc) == 0 or data.decay_const(nucm) < 1e-16 or data.decay_const(nuc) == data.decay_const(nucm):
            continue
        nucs.add(nucm)
    nucs = [nuc for nuc in nucs if nucname.anum(nuc) > 0 and
                                not np.isnan(data.decay_const(nuc)) and
                                nuc < 200000000]
    nucs.sort()

    system = CodeBlock(*map(gennuc, nucs))

    with open('system.txt', 'w') as f:
        for eq in system.args:
            f.write(sympy.srepr(eq) + '\n')

    with open('system-C.txt', 'w') as f:
        f.write(sympy.ccode(system))

    system_cse = system.cse()

    with open('system-cse.txt', 'w') as f:
        for eq in system_cse.args:
            f.write(str(eq) + '\n')

    with open('system-cse-C.txt', 'w') as f:
        f.write(sympy.ccode(system_cse))
<|endoftext|>


<commit_before># Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""This script is used to synthesize generated parts of this library."""

import synthtool as s
import synthtool.gcp as gcp
import synthtool.languages.java as java

gapic = gcp.GAPICGenerator()
common_templates = gcp.CommonTemplates()

versions = ['v1beta1']
service = 'webrisk'

for version in versions:
    library = gapic.java_library(
        service=service,
        version=version,
        config_path=f'/google/cloud/webrisk/artman_webrisk.yaml',
        artman_output_name='')

    s.copy(library / f'gapic-google-cloud-{service}-{version}/src', 'src')
    s.copy(library / f'grpc-google-cloud-{service}-{version}/src', f'../../google-api-grpc/grpc-google-cloud-{service}-{version}/src')
    s.copy(library / f'proto-google-cloud-{service}-{version}/src', f'../../google-api-grpc/proto-google-cloud-{service}-{version}/src')

    java.format_code('./src')
    java.format_code(f'../../google-api-grpc/grpc-google-cloud-{service}-{version}/src')
    java.format_code(f'../../google-api-grpc/proto-google-cloud-{service}-{version}/src')<commit_msg>Change .yaml file name for webrisk client. (#4898)<commit_after># Copyright 2018 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""This script is used to synthesize generated parts of this library."""

import synthtool as s
import synthtool.gcp as gcp
import synthtool.languages.java as java

gapic = gcp.GAPICGenerator()
common_templates = gcp.CommonTemplates()

versions = ['v1beta1']
service = 'webrisk'

for version in versions:
    library = gapic.java_library(
        service=service,
        version=version,
        config_path=f'/google/cloud/webrisk/artman_webrisk_{version}.yaml',
        artman_output_name='')

    s.copy(library / f'gapic-google-cloud-{service}-{version}/src', 'src')
    s.copy(library / f'grpc-google-cloud-{service}-{version}/src', f'../../google-api-grpc/grpc-google-cloud-{service}-{version}/src')
    s.copy(library / f'proto-google-cloud-{service}-{version}/src', f'../../google-api-grpc/proto-google-cloud-{service}-{version}/src')

    java.format_code('./src')
    java.format_code(f'../../google-api-grpc/grpc-google-cloud-{service}-{version}/src')
    java.format_code(f'../../google-api-grpc/proto-google-cloud-{service}-{version}/src')<|endoftext|>


<commit_before>from tremendous.bindings import lib
from tremendous.bindings import ffi

def apply_format(color, body):
    s = lib.apply_format(color, body)
    return ffi.string(s)
<commit_msg>add 256 color functions<commit_after>from tremendous.bindings import lib
from tremendous.bindings import ffi

def apply_format(color, body):
    s = lib.apply_format(color, body)
    return ffi.string(s)

def apply_256(rgb_t, body):
    s = lib.foreground_256(rgb_t, body)
    return ffi.string(s)

def apply_256_bg(rgb_t, body):
    s = lib.background_256(rgb_t, body)
    return ffi.string(s)

def apply_256_hl(rgb_t, body):
    s = lib.highlight_256(rgb_t, body)
    return ffi.string(s)
<|endoftext|>


<commit_before>from rflint import RfLint, ERROR, WARNING
from .parser import RobotFile
from common import SuiteRule, TestRule, KeywordRule, GeneralRule

<commit_msg>use relative import for all imports<commit_after>from .rflint import RfLint, ERROR, WARNING
from .parser import RobotFile
from .common import SuiteRule, TestRule, KeywordRule, GeneralRule

<|endoftext|>


<commit_before>from __future__ import print_function, division, absolute_import #, unicode_literals # not casa compatible
from builtins import bytes, dict, object, range, map, input#, str # not casa compatible
from future.utils import itervalues, viewitems, iteritems, listvalues, listitems
from io import open

import distributed
from dask import delayed, compute

from rfpipe import state, source, search, util

import logging
logger = logging.getLogger(__name__)
vys_timeout_default = 10

def pipeline_scan_distributed(st, segments=None, host='cbe-node-01',
                              cfile=None, vys_timeout=vys_timeout_default):
    """ Given rfpipe state and dask distributed client, run search pipline """

    saved = []
    if not isinstance(segments, list):
        segments = range(st.nsegment)

    for segment in segments:
        saved.append(pipeline_seg(st, segment, host=host, cfile=cfile,
                     vys_timeout=vys_timeout))

    return saved

def pipeline_vystest(wait, host='cbe-node-01', preffile=None,
                     cfile=None, **kwargs):
    """ Start one segment vysmaw jobs reading a segment each after time wait
    Uses example realfast scan configuration from files.
    """

    st = state.state_vystest(wait, preffile=preffile,
                             **kwargs)

    saved = pipeline_scan_distributed(st, host=host, cfile=cfile)

    return saved

def pipeline_seg(st, segment, host='cbe-node-01', cfile=None,
                 vys_timeout=vys_timeout_default):
    """ Build DAG from delayed objects and execute at end with preferred
        scheduler
    """

    logger.info('Building dask for observation {0}.'.format(st.fileroot))
    cl = distributed.Client('{0}:{1}'.format(host, '8786'))

    # plan fft
    wisdom = cl.submit(search.set_wisdom, st.npixx, st.npixy, pure=True)

    data = cl.submit(source.read_segment, st, segment, timeout=vys_timeout,
                     cfile=cfile, pure=True)
    data_prep = cl.submit(source.data_prep, st, data, pure=True)

    # **TODO: need to add condition on data_prep being nonzero
    saved = []
    for dmind in range(len(st.dmarr)):
        delay = cl.submit(util.calc_delay, st.freq, st.freq.max(),
                          st.dmarr[dmind], st.inttime, pure=True)
        data_dm = cl.submit(search.dedisperse, data_prep, delay, pure=True)

        for dtind in range(len(st.dtarr)):
            data_dmdt = cl.submit(search.resample, data_dm, st.dtarr[dtind],
                                  pure=True)
            canddatalist = cl.submit(search.search_thresh, st, data_dmdt,
                                     segment, dmind, dtind, wisdom=wisdom)

            candidates = cl.submit(search.calc_features, canddatalist,
                                   pure=True)
            saved.append(cl.submit(search.save_cands, st, candidates,
                                   canddatalist))

    return saved

def pipeline_seg_delayed(st, segment, cl=None, cfile=None,
                         vys_timeout=vys_timeout_default):
    """ Build DAG from delayed objects and execute at end with preferred
        scheduler
    """

    # ** how to use distributed arguments?
    # (workers=workers, allow_other_workers=allow_other_workers)

    logger.info('Building dask for observation {0}.'.format(st.fileroot))

    # plan fft
    wisdom = delayed(search.set_wisdom, pure=True)(st.npixx, st.npixy)

    data = delayed(source.read_segment, pure=True)(st, segment,
                                                   timeout=vys_timeout,
                                                   cfile=cfile)
    data_prep = delayed(source.data_prep, pure=True)(st, data)

    # **TODO: need to add condition on data_prep being nonzero
    saved = []
    for dmind in range(len(st.dmarr)):
        delay = delayed(util.calc_delay, pure=True)(st.freq, st.freq.max(),
                                                    st.dmarr[dmind],
                                                    st.inttime)
        data_dm = delayed(search.dedisperse, pure=True)(data_prep, delay)

        for dtind in range(len(st.dtarr)):
            data_dmdt = delayed(search.resample, pure=True)(data_dm,
                                                            st.dtarr[dtind])
            canddatalist = delayed(search.search_thresh,
                                   pure=True)(st, data_dmdt, segment, dmind,
                                              dtind, wisdom=wisdom)

            candidates = delayed(search.calc_features, pure=True)(canddatalist)
            saved.append(delayed(search.save_cands, pure=True)(st, candidates,
                                                               canddatalist))

    if cl:
        # if using distributed client, return collection with futures
        return cl.persist(saved)
    else:
        # otherwise return the delayed objects
        return compute(*saved, num_workers=st.prefs.nthread)
<commit_msg>going back to delayed pipeline_seg<commit_after>from __future__ import print_function, division, absolute_import #, unicode_literals # not casa compatible
from builtins import bytes, dict, object, range, map, input#, str # not casa compatible
from future.utils import itervalues, viewitems, iteritems, listvalues, listitems
from io import open

import distributed
from dask import delayed, compute

from rfpipe import source, search, util

import logging
logger = logging.getLogger(__name__)
vys_timeout_default = 10

def pipeline_scan_distributed(st, segments=None, host='cbe-node-01',
                              cfile=None, vys_timeout=vys_timeout_default):
    """ Given rfpipe state and dask distributed client, run search pipline """

    cl = distributed.Client('{0}:{1}'.format(host, '8786'))

    saved = []
    if not isinstance(segments, list):
        segments = range(st.nsegment)

    for segment in segments:
        saved.append(pipeline_seg(st, segment, cl=cl, cfile=cfile,
                     vys_timeout=vys_timeout))

    return saved

def pipeline_seg(st, segment, cl=None, cfile=None,
                 vys_timeout=vys_timeout_default):
    """ Build DAG from delayed objects and execute at end with preferred
        scheduler. Can use distributed client or compute locally.
    """

    # ** how to use distributed arguments?
    # (workers=workers, allow_other_workers=allow_other_workers)

    logger.info('Building dask for observation {0}.'.format(st.fileroot))

    # plan fft
    wisdom = delayed(search.set_wisdom, pure=True)(st.npixx, st.npixy)

    data = delayed(source.read_segment, pure=True)(st, segment,
                                                   timeout=vys_timeout,
                                                   cfile=cfile)
    data_prep = delayed(source.data_prep, pure=True)(st, data)

    # **TODO: need to add condition on data_prep being nonzero
    saved = []
    for dmind in range(len(st.dmarr)):
        delay = delayed(util.calc_delay, pure=True)(st.freq, st.freq.max(),
                                                    st.dmarr[dmind],
                                                    st.inttime)
        data_dm = delayed(search.dedisperse, pure=True)(data_prep, delay)

        for dtind in range(len(st.dtarr)):
            data_dmdt = delayed(search.resample, pure=True)(data_dm,
                                                            st.dtarr[dtind])
            canddatalist = delayed(search.search_thresh,
                                   pure=True)(st, data_dmdt, segment, dmind,
                                              dtind, wisdom=wisdom)

            candidates = delayed(search.calc_features, pure=True)(canddatalist)
            saved.append(delayed(search.save_cands, pure=True)(st, candidates,
                                                               canddatalist))

    if cl:
        # if using distributed client, return collection with futures
        return cl.persist(saved)
    else:
        # otherwise return the delayed objects
        return compute(*saved, num_workers=st.prefs.nthread)
<|endoftext|>


<commit_before>from __future__ import division
import os
import time

import functools

from flask import Flask
from flask import request
from flask import Response
from flask import redirect

import rigorwebapp.plugin
import rigorwebapp.backend
from rigorwebapp.utils import debug_detail, debug_error

def get_app(rigor_config):
	"""Builds and returns a Flask app object which gets its configuration from the given config file.
	rigor_config is an instance of rigor.config.RigorDefaultConfiguration(path_to_config_file)
	"""

	#================================================================================
	# FLASK HELPERS

	def check_auth(username, password):
		"""This function is called to check if a username /
		password combination is valid.
		"""
		return username == rigor_config.get('webapp', 'http_auth_user') and password == rigor_config.get('webapp', 'http_auth_password')

	def authenticate():
		"""Sends a 401 response that enables basic auth"""
		return Response(
			'Could not verify your access level for that URL.\n'
			'You have to login with proper credentials', 401,
			{'WWW-Authenticate': 'Basic realm="Login Required"'}
		)

	def use_basic_auth(fn):
		@functools.wraps(fn)
		def decorated(*args, **kwargs):
			if rigor_config.get('webapp', 'use_http_auth'):
				auth = request.authorization
				if not auth or not check_auth(auth.username, auth.password):
					return authenticate()
			return fn(*args, **kwargs)
		return decorated

	def simulate_slow(fn):
		@functools.wraps(fn)
		def decorated(*args, **kwargs):
			latency = int(rigor_config.get('webapp', 'fake_latency'))
			if latency > 0:
				debug_detail('adding {} seconds of fake delay'.format(latency))
				time.sleep(latency)
			return fn(*args, **kwargs)
		return decorated

	#================================================================================
	# FLASK SETUP

	share_folder = None
	share_folder_locations = ['../../share', '../../../../share', '../share']
	for path in share_folder_locations:
		path = os.path.abspath(os.path.join(os.path.dirname(__file__), path))
		if os.path.exists(path):
			share_folder = os.path.abspath(path)
			break
	if share_folder is None:
		debug_error('ERROR: Share folder not found')
		raise IOError('Share folder not found')
	else:
		debug_detail('found share folder at {}'.format(share_folder))

	app = Flask(
		__name__,
		static_folder=share_folder + '/static',
		template_folder=share_folder + '/templates',
	)
	app.config['SECRET_KEY'] = rigor_config.get('webapp', 'flask_secret_key')

	#================================================================================
	# PLUGINS

	debug_detail('loading backend')
	backend = rigorwebapp.backend.Backend(rigor_config)

	plugin_instances = rigorwebapp.plugin.load_plugin_instances(rigor_config, backend)
	debug_detail('loaded plugin classes:')
	for plugin_instance in plugin_instances:
		debug_detail('    {}'.format(plugin_instance.__class__.__name__))

	#================================================================================
	# ROUTING

	# let plugins add routes
	debug_detail('adding plugin routes')
	for plugin_instance in plugin_instances:
		plugin_instance.add_routes(app, backend, plugin_instances)

	@app.route('/')
	def index_redirect():
		initial_db = rigor_config.get('webapp', 'initial_db')
		initial_db = backend._encode_name_for_db(initial_db)
		return redirect('/db/{}/perceptsearch'.format(initial_db))

	app.config['PROPAGATE_EXCEPTIONS'] = True
	debug_detail('ready')

	# set up logging for sqlalchemy
	from logging import getLogger
	from logging import StreamHandler
	getLogger('sqlalchemy').addHandler(StreamHandler())

	return app
<commit_msg>fixed basic http auth for flask deployment<commit_after>from __future__ import division
import os
import time

import functools

from flask import Flask
from flask import request
from flask import Response
from flask import redirect

import rigorwebapp.plugin
import rigorwebapp.backend
from rigorwebapp.utils import debug_detail, debug_error

def get_app(rigor_config):
	"""Builds and returns a Flask app object which gets its configuration from the given config file.
	rigor_config is an instance of rigor.config.RigorDefaultConfiguration(path_to_config_file)
	"""

	#================================================================================
	# FLASK HELPERS

	def check_auth(username, password):
		"""This function is called to check if a username /
		password combination is valid.
		"""
		return username == rigor_config.get('webapp', 'http_auth_user') and password == rigor_config.get('webapp', 'http_auth_password')

	def authenticate():
		"""Sends a 401 response that enables basic auth"""
		return Response(
			'Could not verify your access level for that URL.\n'
			'You have to login with proper credentials', 401,
			{'WWW-Authenticate': 'Basic realm="Login Required"'}
		)

	def simulate_slow(fn):
		@functools.wraps(fn)
		def decorated(*args, **kwargs):
			latency = int(rigor_config.get('webapp', 'fake_latency'))
			if latency > 0:
				debug_detail('adding {} seconds of fake delay'.format(latency))
				time.sleep(latency)
			return fn(*args, **kwargs)
		return decorated

	#================================================================================
	# FLASK SETUP

	share_folder = None
	share_folder_locations = ['../../share', '../../../../share', '../share']
	for path in share_folder_locations:
		path = os.path.abspath(os.path.join(os.path.dirname(__file__), path))
		if os.path.exists(path):
			share_folder = os.path.abspath(path)
			break
	if share_folder is None:
		debug_error('ERROR: Share folder not found')
		raise IOError('Share folder not found')
	else:
		debug_detail('found share folder at {}'.format(share_folder))

	app = Flask(
		__name__,
		static_folder=share_folder + '/static',
		template_folder=share_folder + '/templates',
	)
	app.config['SECRET_KEY'] = rigor_config.get('webapp', 'flask_secret_key')

	#================================================================================
	# PLUGINS

	debug_detail('loading backend')
	backend = rigorwebapp.backend.Backend(rigor_config)

	plugin_instances = rigorwebapp.plugin.load_plugin_instances(rigor_config, backend)
	debug_detail('loaded plugin classes:')
	for plugin_instance in plugin_instances:
		debug_detail('    {}'.format(plugin_instance.__class__.__name__))

	#================================================================================
	# ROUTING

	# let plugins add routes
	debug_detail('adding plugin routes')
	for plugin_instance in plugin_instances:
		plugin_instance.add_routes(app, backend, plugin_instances)

	@app.before_request
	# see http://stackoverflow.com/a/13451233/856925 for more thoughts on this
	def require_auth_everywhere():
		if rigor_config.getboolean('webapp', 'use_http_auth'):
			auth = request.authorization
			if not auth or not check_auth(auth.username, auth.password):
				return authenticate()

	@app.route('/')
	def index_redirect():
		initial_db = rigor_config.get('webapp', 'initial_db')
		initial_db = backend._encode_name_for_db(initial_db)
		return redirect('/db/{}/perceptsearch'.format(initial_db))

	app.config['PROPAGATE_EXCEPTIONS'] = True
	debug_detail('ready')

	# set up logging for sqlalchemy
	from logging import getLogger
	from logging import StreamHandler
	getLogger('sqlalchemy').addHandler(StreamHandler())

	return app
<|endoftext|>


<commit_before>
import arrow
import superdesk

from flask import json
from bson import ObjectId
from bson.errors import InvalidId
from eve.utils import str_to_date
from eve.io.mongo import MongoJSONEncoder
from eve_elastic import ElasticJSONSerializer

class SuperdeskJSONEncoder(MongoJSONEncoder, ElasticJSONSerializer):
    """Custom JSON encoder for elastic that can handle `bson.ObjectId`s."""

    pass

def try_cast(v):
    """Cast string value to date or ObjectId if possible.

    :param v: string value
    """
    try:
        str_to_date(v)  # try if it matches format
        return arrow.get(v).datetime  # return timezone aware time
    except ValueError:
        try:
            return ObjectId(v)
        except InvalidId:
            return v

def cast_item(o):
    with superdesk.app.app_context():
        if isinstance(o, (int, float, bool)):
            return o
        elif isinstance(o, str):
            return try_cast(o)
        elif isinstance(o, list):
            for i, v in enumerate(o):
                o[i] = cast_item(v)
            return o
        elif isinstance(o, dict):
            for k, v in o.items():
                o[k] = cast_item(v)
            return o
        else:
            return o

def loads(s):
    o = json.loads(s)

    if isinstance(o, list):
        for i, v in enumerate(o):
            o[i] = cast_item(v)
        return o
    elif isinstance(o, dict):
        for k, v in o.items():
            o[k] = cast_item(v)
        return o
    else:
        return cast_item(o)

def dumps(o):
    with superdesk.app.app_context():
        return MongoJSONEncoder().encode(o)
<commit_msg>(fix):Modified the json loads function to handle arrow parsererror.<commit_after>
import arrow
import superdesk

from arrow.parser import ParserError
from flask import json
from bson import ObjectId
from bson.errors import InvalidId
from eve.utils import str_to_date
from eve.io.mongo import MongoJSONEncoder
from eve_elastic import ElasticJSONSerializer

class SuperdeskJSONEncoder(MongoJSONEncoder, ElasticJSONSerializer):
    """Custom JSON encoder for elastic that can handle `bson.ObjectId`s."""

    pass

def try_cast(v):
    """Cast string value to date or ObjectId if possible.

    :param v: string value
    """
    try:
        str_to_date(v)  # try if it matches format
        return arrow.get(v).datetime  # return timezone aware time
    except (ValueError, ParserError):
        try:
            return ObjectId(v)
        except InvalidId:
            return v

def cast_item(o):
    with superdesk.app.app_context():
        if isinstance(o, (int, float, bool)):
            return o
        elif isinstance(o, str):
            return try_cast(o)
        elif isinstance(o, list):
            for i, v in enumerate(o):
                o[i] = cast_item(v)
            return o
        elif isinstance(o, dict):
            for k, v in o.items():
                o[k] = cast_item(v)
            return o
        else:
            return o

def loads(s):
    o = json.loads(s)

    if isinstance(o, list):
        for i, v in enumerate(o):
            o[i] = cast_item(v)
        return o
    elif isinstance(o, dict):
        for k, v in o.items():
            o[k] = cast_item(v)
        return o
    else:
        return cast_item(o)

def dumps(o):
    with superdesk.app.app_context():
        return MongoJSONEncoder().encode(o)
<|endoftext|>