<commit_before>## Copyright 2002 Andrew Loewenstern, All Rights Reserved

from const import reactor
import time

from ktable import KTable, K
from knode import KNode as Node

from hash import newID

from actions import FindNode, GetValue
from twisted.web import xmlrpc
from twisted.internet.defer import Deferred
from twisted.python import threadable
from twisted.internet.app import Application
from twisted.web import server
threadable.init()

from bsddb3 import db ## find this at http://pybsddb.sf.net/
from bsddb3._db import DBNotFoundError

# don't ping unless it's been at least this many seconds since we've heard from a peer
MAX_PING_INTERVAL = 60 * 15 # fifteen minutes

# this is the main class!
class Khashmir(xmlrpc.XMLRPC):
    __slots__ = ['listener', 'node', 'table', 'store', 'app']
    def __init__(self, host, port):
	self.node = Node(newID(), host, port)
	self.table = KTable(self.node)
	self.app = Application("xmlrpc")
	self.app.listenTCP(port, server.Site(self))
	self.store = db.DB()
	self.store.open(None, None, db.DB_BTREE)

    def render(self, request):
	"""
	    Override the built in render so we can have access to the request object!
	    note, crequest is probably only valid on the initial call (not after deferred!)
	"""
	self.crequest = request
	return xmlrpc.XMLRPC.render(self, request)

    #######
    #######  LOCAL INTERFACE    - use these methods!
    def addContact(self, host, port):
	"""
	 ping this node and add the contact info to the table on pong!
	"""
	n =Node(" "*20, host, port)  # note, we
	self.sendPing(n)

    ## this call is async!
    def findNode(self, id, callback, errback=None):
	""" returns the contact info for node, or the k closest nodes, from the global table """
	# get K nodes out of local table/cache, or the node we want
	nodes = self.table.findNodes(id)
	d = Deferred()
	d.addCallbacks(callback, errback)
	if len(nodes) == 1 and nodes[0].id == id :
	    d.callback(nodes)
	else:
	    # create our search state
	    state = FindNode(self, id, d.callback)
	    reactor.callFromThread(state.goWithNodes, nodes)

    ## also async
    def valueForKey(self, key, callback):
	""" returns the values found for key in global table """
	nodes = self.table.findNodes(key)
	# create our search state
	state = GetValue(self, key, callback)
	reactor.callFromThread(state.goWithNodes, nodes)

    ## async, but in the current implementation there is no guarantee a store does anything so there is no callback right now
    def storeValueForKey(self, key, value):
	""" stores the value for key in the global table, returns immediately, no status
	    in this implementation, peers respond but don't indicate status to storing values
	    values are stored in peers on a first-come first-served basis
	    this will probably change so more than one value can be stored under a key
	"""
	def _storeValueForKey(nodes, key=key, value=value, response= self._storedValueHandler, default= lambda t: "didn't respond"):
	    for node in nodes:
		if node.id != self.node.id:
		    df = node.storeValue(key, value, self.node.senderDict())
		    df.addCallbacks(response, default)
	# this call is asynch
	self.findNode(key, _storeValueForKey)

    def insertNode(self, n):
	"""
	insert a node in our local table, pinging oldest contact in bucket, if necessary

	If all you have is a host/port, then use addContact, which calls this method after
	receiving the PONG from the remote node.  The reason for the seperation is we can't insert
	a node into the table without it's peer-ID.  That means of course the node passed into this
	method needs to be a properly formed Node object with a valid ID.
	"""
	old = self.table.insertNode(n)
	if old and (time.time() - old.lastSeen) > MAX_PING_INTERVAL and old.id != self.node.id:
	    # the bucket is full, check to see if old node is still around and if so, replace it

	    ## these are the callbacks used when we ping the oldest node in a bucket
	    def _staleNodeHandler(oldnode=old, newnode = n):
		""" called if the pinged node never responds """
		self.table.replaceStaleNode(old, newnode)

	    def _notStaleNodeHandler(sender, old=old):
		""" called when we get a ping from the remote node """
		if sender['id'] == old.id:
		    self.table.insertNode(old)

	    df = old.ping()
	    df.addCallbacks(_notStaleNodeHandler, self._staleNodeHandler)

    def sendPing(self, node):
	"""
	    ping a node
	"""
	df = node.ping(self.node.senderDict())
	## these are the callbacks we use when we issue a PING
	def _pongHandler(sender, id=node.id, host=node.host, port=node.port, table=self.table):
	    if id != 20 * ' ' and id != sender['id']:
		# whoah, got response from different peer than we were expecting
		pass
	    else:
		#print "Got PONG from %s at %s:%s" % (`msg['id']`, t.target.host, t.target.port)
		n = Node(sender['id'], host, port)
		table.insertNode(n)
	    return
	def _defaultPong(err):
	    # this should probably increment a failed message counter and dump the node if it gets over a threshold
	    return

	df.addCallbacks(_pongHandler,_defaultPong)

    def findCloseNodes(self):
	"""
	    This does a findNode on the ID one away from our own.
	    This will allow us to populate our table with nodes on our network closest to our own.
	    This is called as soon as we start up with an empty table
	"""
	id = self.node.id[:-1] + chr((ord(self.node.id[-1]) + 1) % 256)
	def callback(nodes):
	    pass
	self.findNode(id, callback)

    def refreshTable(self):
	"""

	"""
	def callback(nodes):
	    pass

	for bucket in self.table.buckets:
	    if time.time() - bucket.lastAccessed >= 60 * 60:
		id = randRange(bucket.min, bucket.max)
		self.findNode(id, callback)

    #####
    ##### INCOMING MESSAGE HANDLERS

    def xmlrpc_ping(self, sender):
	"""
	    takes sender dict = {'id', <id>, 'port', port} optional keys = 'ip'
	    returns sender dict
	"""
	ip = self.crequest.getClientIP()
	n = Node(sender['id'], ip, sender['port'])
	self.insertNode(n)
	return self.node.senderDict()

    def xmlrpc_find_node(self, target, sender):
	nodes = self.table.findNodes(target)
	nodes = map(lambda node: node.senderDict(), nodes)
	ip = self.crequest.getClientIP()
	n = Node(sender['id'], ip, sender['port'])
	self.insertNode(n)
	return nodes, self.node.senderDict()

    def xmlrpc_store_value(self, key, value, sender):
	if not self.store.has_key(key):
	    self.store.put(key, value)
	ip = self.crequest.getClientIP()
	n = Node(sender['id'], ip, sender['port'])
	self.insertNode(n)
	return self.node.senderDict()

    def xmlrpc_find_value(self, key, sender):
    	ip = self.crequest.getClientIP()
	n = Node(sender['id'], ip, sender['port'])
	self.insertNode(n)
	if self.store.has_key(key):
	    return {'values' : self.store[key]}, self.node.senderDict()
	else:
	    nodes = self.table.findNodes(key)
	    nodes = map(lambda node: node.senderDict(), nodes)
	    return {'nodes' : nodes}, self.node.senderDict()

    ###
    ### message response callbacks
    # called when we get a response to store value
    def _storedValueHandler(self, sender):
	pass

#------ testing

def test_build_net(quiet=0, peers=256, pause=1):
    from whrandom import randrange
    import thread
    port = 2001
    l = []

    if not quiet:
	print "Building %s peer table." % peers

    for i in xrange(peers):
	a = Khashmir('localhost', port + i)
	l.append(a)

    thread.start_new_thread(l[0].app.run, ())
    time.sleep(1)
    for peer in l[1:]:
	peer.app.run()
	#time.sleep(.25)

    print "adding contacts...."

    for peer in l[1:]:
	n = l[randrange(0, len(l))].node
	peer.addContact(n.host, n.port)
	n = l[randrange(0, len(l))].node
	peer.addContact(n.host, n.port)
	n = l[randrange(0, len(l))].node
	peer.addContact(n.host, n.port)
	if pause:
	    time.sleep(.30)

    time.sleep(5)
    print "finding close nodes...."

    for peer in l:
	peer.findCloseNodes()
	if pause:
	    time.sleep(1)
#    for peer in l:
#	peer.refreshTable()
    return l

def test_find_nodes(l, quiet=0):
    import threading, sys
    from whrandom import randrange
    flag = threading.Event()

    n = len(l)

    a = l[randrange(0,n)]
    b = l[randrange(0,n)]

    def callback(nodes, l=l, flag=flag):
	if (len(nodes) >0) and (nodes[0].id == b.node.id):
	    print "test_find_nodes	PASSED"
	else:
	    print "test_find_nodes	FAILED"
	flag.set()
    a.findNode(b.node.id, callback)
    flag.wait()

def test_find_value(l, quiet=0):
    from whrandom import randrange
    from sha import sha
    import time, threading, sys

    fa = threading.Event()
    fb = threading.Event()
    fc = threading.Event()

    n = len(l)
    a = l[randrange(0,n)]
    b = l[randrange(0,n)]
    c = l[randrange(0,n)]
    d = l[randrange(0,n)]

    key = sha(`randrange(0,100000)`).digest()
    value = sha(`randrange(0,100000)`).digest()
    if not quiet:
	print "inserting value...",
	sys.stdout.flush()
    a.storeValueForKey(key, value)
    time.sleep(3)
    print "finding..."

    def mc(flag, value=value):
	def callback(values, f=flag, val=value):
	    try:
		if(len(values) == 0):
		    print "find                FAILED"
		else:
		    if values != val:
			print "find                FAILED"
		    else:
			print "find                FOUND"
	    finally:
		f.set()
	return callback
    b.valueForKey(key, mc(fa))
    fa.wait()
    c.valueForKey(key, mc(fb))
    fb.wait()
    d.valueForKey(key, mc(fc))
    fc.wait()

if __name__ == "__main__":
    l = test_build_net()
    time.sleep(3)
    print "finding nodes..."
    test_find_nodes(l)
    test_find_nodes(l)
    test_find_nodes(l)
    print "inserting and fetching values..."
    test_find_value(l)
    test_find_value(l)
    test_find_value(l)
    test_find_value(l)
    test_find_value(l)
    test_find_value(l)
<commit_msg>now we gracefully deal with multiple values per key<commit_after>## Copyright 2002 Andrew Loewenstern, All Rights Reserved

from const import reactor
import time
from pickle import loads, dumps
from sha import sha

from ktable import KTable, K
from knode import KNode as Node

from hash import newID

from actions import FindNode, GetValue
from twisted.web import xmlrpc
from twisted.internet.defer import Deferred
from twisted.python import threadable
from twisted.internet.app import Application
from twisted.web import server
threadable.init()

from bsddb3 import db ## find this at http://pybsddb.sf.net/
from bsddb3._db import DBNotFoundError

# don't ping unless it's been at least this many seconds since we've heard from a peer
MAX_PING_INTERVAL = 60 * 15 # fifteen minutes

# this is the main class!
class Khashmir(xmlrpc.XMLRPC):
    __slots__ = ['listener', 'node', 'table', 'store', 'itime', 'kw', 'app']
    def __init__(self, host, port):
	self.node = Node(newID(), host, port)
	self.table = KTable(self.node)
	self.app = Application("xmlrpc")
	self.app.listenTCP(port, server.Site(self))

	## these databases may be more suited to on-disk rather than in-memory
	# h((key, value)) -> (key, value, time) mappings
	self.store = db.DB()
	self.store.open(None, None, db.DB_BTREE)

	# <insert time> -> h((key, value))
	self.itime = db.DB()
	self.itime.set_flags(db.DB_DUP)
	self.itime.open(None, None, db.DB_BTREE)

	# key -> h((key, value))
	self.kw = db.DB()
	self.kw.set_flags(db.DB_DUP)
	self.kw.open(None, None, db.DB_BTREE)

    def render(self, request):
	"""
	    Override the built in render so we can have access to the request object!
	    note, crequest is probably only valid on the initial call (not after deferred!)
	"""
	self.crequest = request
	return xmlrpc.XMLRPC.render(self, request)

    #######
    #######  LOCAL INTERFACE    - use these methods!
    def addContact(self, host, port):
	"""
	 ping this node and add the contact info to the table on pong!
	"""
	n =Node(" "*20, host, port)  # note, we
	self.sendPing(n)

    ## this call is async!
    def findNode(self, id, callback, errback=None):
	""" returns the contact info for node, or the k closest nodes, from the global table """
	# get K nodes out of local table/cache, or the node we want
	nodes = self.table.findNodes(id)
	d = Deferred()
	d.addCallbacks(callback, errback)
	if len(nodes) == 1 and nodes[0].id == id :
	    d.callback(nodes)
	else:
	    # create our search state
	    state = FindNode(self, id, d.callback)
	    reactor.callFromThread(state.goWithNodes, nodes)

    ## also async
    def valueForKey(self, key, callback):
	""" returns the values found for key in global table """
	nodes = self.table.findNodes(key)
	# create our search state
	state = GetValue(self, key, callback)
	reactor.callFromThread(state.goWithNodes, nodes)

    ## async, but in the current implementation there is no guarantee a store does anything so there is no callback right now
    def storeValueForKey(self, key, value):
	""" stores the value for key in the global table, returns immediately, no status
	    in this implementation, peers respond but don't indicate status to storing values
	    values are stored in peers on a first-come first-served basis
	    this will probably change so more than one value can be stored under a key
	"""
	def _storeValueForKey(nodes, key=key, value=value, response= self._storedValueHandler, default= lambda t: "didn't respond"):
	    for node in nodes:
		if node.id != self.node.id:
		    df = node.storeValue(key, value, self.node.senderDict())
		    df.addCallbacks(response, default)
	# this call is asynch
	self.findNode(key, _storeValueForKey)

    def insertNode(self, n):
	"""
	insert a node in our local table, pinging oldest contact in bucket, if necessary

	If all you have is a host/port, then use addContact, which calls this method after
	receiving the PONG from the remote node.  The reason for the seperation is we can't insert
	a node into the table without it's peer-ID.  That means of course the node passed into this
	method needs to be a properly formed Node object with a valid ID.
	"""
	old = self.table.insertNode(n)
	if old and (time.time() - old.lastSeen) > MAX_PING_INTERVAL and old.id != self.node.id:
	    # the bucket is full, check to see if old node is still around and if so, replace it

	    ## these are the callbacks used when we ping the oldest node in a bucket
	    def _staleNodeHandler(oldnode=old, newnode = n):
		""" called if the pinged node never responds """
		self.table.replaceStaleNode(old, newnode)

	    def _notStaleNodeHandler(sender, old=old):
		""" called when we get a ping from the remote node """
		if sender['id'] == old.id:
		    self.table.insertNode(old)

	    df = old.ping(self.node.senderDict())
	    df.addCallbacks(_notStaleNodeHandler, self._staleNodeHandler)

    def sendPing(self, node):
	"""
	    ping a node
	"""
	df = node.ping(self.node.senderDict())
	## these are the callbacks we use when we issue a PING
	def _pongHandler(sender, id=node.id, host=node.host, port=node.port, table=self.table):
	    if id != 20 * ' ' and id != sender['id']:
		# whoah, got response from different peer than we were expecting
		pass
	    else:
		#print "Got PONG from %s at %s:%s" % (`msg['id']`, t.target.host, t.target.port)
		n = Node(sender['id'], host, port)
		table.insertNode(n)
	    return
	def _defaultPong(err):
	    # this should probably increment a failed message counter and dump the node if it gets over a threshold
	    return

	df.addCallbacks(_pongHandler,_defaultPong)

    def findCloseNodes(self):
	"""
	    This does a findNode on the ID one away from our own.
	    This will allow us to populate our table with nodes on our network closest to our own.
	    This is called as soon as we start up with an empty table
	"""
	id = self.node.id[:-1] + chr((ord(self.node.id[-1]) + 1) % 256)
	def callback(nodes):
	    pass
	self.findNode(id, callback)

    def refreshTable(self):
	"""

	"""
	def callback(nodes):
	    pass

	for bucket in self.table.buckets:
	    if time.time() - bucket.lastAccessed >= 60 * 60:
		id = randRange(bucket.min, bucket.max)
		self.findNode(id, callback)

    #####
    ##### INCOMING MESSAGE HANDLERS

    def xmlrpc_ping(self, sender):
	"""
	    takes sender dict = {'id', <id>, 'port', port} optional keys = 'ip'
	    returns sender dict
	"""
	ip = self.crequest.getClientIP()
	n = Node(sender['id'], ip, sender['port'])
	self.insertNode(n)
	return self.node.senderDict()

    def xmlrpc_find_node(self, target, sender):
	nodes = self.table.findNodes(target)
	nodes = map(lambda node: node.senderDict(), nodes)
	ip = self.crequest.getClientIP()
	n = Node(sender['id'], ip, sender['port'])
	self.insertNode(n)
	return nodes, self.node.senderDict()

    def xmlrpc_store_value(self, key, value, sender):
	h1 = sha(key+value).digest()
	t = `time.time()`
	if not self.store.has_key(h1):
	    v = dumps((key, value, t))
	    self.store.put(h1, v)
	    self.itime.put(t, h1)
	    self.kw.put(key, h1)
	else:
	    # update last insert time
	    tup = loads(self.store[h1])
	    self.store[h1] = dumps((tup[0], tup[1], t))
	    self.itime.put(t, h1)

	ip = self.crequest.getClientIP()
	n = Node(sender['id'], ip, sender['port'])
	self.insertNode(n)
	return self.node.senderDict()

    def xmlrpc_find_value(self, key, sender):
    	ip = self.crequest.getClientIP()
	n = Node(sender['id'], ip, sender['port'])
	self.insertNode(n)
	if self.kw.has_key(key):
	    c = self.kw.cursor()
	    tup = c.set_range(key)
	    l = []
	    while(tup):
		h1 = tup[1]
		v = loads(self.store[h1])[1]
		l.append(v)
		tup = c.next()
	    return {'values' : l}, self.node.senderDict()
	else:
	    nodes = self.table.findNodes(key)
	    nodes = map(lambda node: node.senderDict(), nodes)
	    return {'nodes' : nodes}, self.node.senderDict()

    ###
    ### message response callbacks
    # called when we get a response to store value
    def _storedValueHandler(self, sender):
	pass

#------ testing

def test_build_net(quiet=0, peers=64, pause=1):
    from whrandom import randrange
    import thread
    port = 2001
    l = []

    if not quiet:
	print "Building %s peer table." % peers

    for i in xrange(peers):
	a = Khashmir('localhost', port + i)
	l.append(a)

    thread.start_new_thread(l[0].app.run, ())
    time.sleep(1)
    for peer in l[1:]:
	peer.app.run()
	#time.sleep(.25)

    print "adding contacts...."

    for peer in l[1:]:
	n = l[randrange(0, len(l))].node
	peer.addContact(n.host, n.port)
	n = l[randrange(0, len(l))].node
	peer.addContact(n.host, n.port)
	n = l[randrange(0, len(l))].node
	peer.addContact(n.host, n.port)
	if pause:
	    time.sleep(.30)

    time.sleep(1)
    print "finding close nodes...."

    for peer in l:
	peer.findCloseNodes()
	if pause:
	    time.sleep(.5)
    if pause:
	    time.sleep(10)
#    for peer in l:
#	peer.refreshTable()
    return l

def test_find_nodes(l, quiet=0):
    import threading, sys
    from whrandom import randrange
    flag = threading.Event()

    n = len(l)

    a = l[randrange(0,n)]
    b = l[randrange(0,n)]

    def callback(nodes, flag=flag):
	if (len(nodes) >0) and (nodes[0].id == b.node.id):
	    print "test_find_nodes	PASSED"
	else:
	    print "test_find_nodes	FAILED"
	flag.set()
    a.findNode(b.node.id, callback)
    flag.wait()

def test_find_value(l, quiet=0):
    from whrandom import randrange
    from sha import sha
    import time, threading, sys

    fa = threading.Event()
    fb = threading.Event()
    fc = threading.Event()

    n = len(l)
    a = l[randrange(0,n)]
    b = l[randrange(0,n)]
    c = l[randrange(0,n)]
    d = l[randrange(0,n)]

    key = sha(`randrange(0,100000)`).digest()
    value = sha(`randrange(0,100000)`).digest()
    if not quiet:
	print "inserting value..."
	sys.stdout.flush()
    a.storeValueForKey(key, value)
    time.sleep(3)
    print "finding..."
    sys.stdout.flush()

    class cb:
	def __init__(self, flag, value=value):
	    self.flag = flag
	    self.val = value
	    self.found = 0
	def callback(self, values):
	    try:
		if(len(values) == 0):
		    if not self.found:
			print "find                FAILED"
		    else:
			print "find                FOUND"
		    sys.stdout.flush()

		else:
		    if self.val in values:
			self.found = 1
	    finally:
		self.flag.set()

    b.valueForKey(key, cb(fa).callback)
    fa.wait()
    c.valueForKey(key, cb(fb).callback)
    fb.wait()
    d.valueForKey(key, cb(fc).callback)
    fc.wait()

def test_one(port):
    import thread
    k = Khashmir('localhost', port)
    thread.start_new_thread(k.app.run, ())
    return k

if __name__ == "__main__":
    l = test_build_net()
    time.sleep(3)
    print "finding nodes..."
    for i in range(10):
	test_find_nodes(l)
    print "inserting and fetching values..."
    for i in range(10):
	test_find_value(l)
<|endoftext|>


<commit_before>'''
Wrapper for YAJL C library version 2.x.
'''

from ctypes import Structure, c_uint, c_ubyte, c_int, c_long, c_double, \
                   c_void_p, c_char_p, CFUNCTYPE, POINTER, byref, string_at, cast , \
                   cdll, util, c_char
from decimal import Decimal

from ijson import common, backends
from ijson.compat import b2s

yajl = backends.find_yajl(2)

yajl.yajl_alloc.restype = POINTER(c_char)
yajl.yajl_get_error.restype = POINTER(c_char)

C_EMPTY = CFUNCTYPE(c_int, c_void_p)
C_INT = CFUNCTYPE(c_int, c_void_p, c_int)
C_LONG = CFUNCTYPE(c_int, c_void_p, c_long)
C_DOUBLE = CFUNCTYPE(c_int, c_void_p, c_double)
C_STR = CFUNCTYPE(c_int, c_void_p, POINTER(c_ubyte), c_uint)

def number(value):
    '''
    Helper function casting a string that represents any Javascript number
    into appropriate Python value: either int or Decimal.
    '''
    try:
        return int(value)
    except ValueError:
        return Decimal(value)

_callback_data = [
    # Mapping of JSON parser events to callback C types and value converters.
    # Used to define the Callbacks structure and actual callback functions
    # inside the parse function.
    ('null', C_EMPTY, lambda: None),
    ('boolean', C_INT, lambda v: bool(v)),
    # "integer" and "double" aren't actually yielded by yajl since "number"
    # takes precedence if defined
    ('integer', C_LONG, lambda v, l: int(string_at(v, l))),
    ('double', C_DOUBLE, lambda v, l: float(string_at(v, l))),
    ('number', C_STR, lambda v, l: number(b2s(string_at(v, l)))),
    ('string', C_STR, lambda v, l: string_at(v, l).decode('utf-8')),
    ('start_map', C_EMPTY, lambda: None),
    ('map_key', C_STR, lambda v, l: b2s(string_at(v, l))),
    ('end_map', C_EMPTY, lambda: None),
    ('start_array', C_EMPTY, lambda: None),
    ('end_array', C_EMPTY, lambda: None),
]

class Callbacks(Structure):
    _fields_ = [(name, type) for name, type, func in _callback_data]

YAJL_OK = 0
YAJL_CANCELLED = 1
YAJL_INSUFFICIENT_DATA = 2
YAJL_ERROR = 3

YAJL_ALLOW_COMMENTS = 1

def basic_parse(f, allow_comments=False, buf_size=64 * 1024):
    '''
    Iterator yielding unprefixed events.

    Parameters:

    - f: a readable file-like object with JSON input
    - allow_comments: tells parser to allow comments in JSON input
    - buf_size: a size of an input buffer
    '''
    events = []

    def callback(event, func_type, func):
        def c_callback(context, *args):
            events.append((event, func(*args)))
            return 1
        return func_type(c_callback)

    callbacks = Callbacks(*[callback(*data) for data in _callback_data])
    handle = yajl.yajl_alloc(byref(callbacks), None, None)
    if allow_comments:
        yajl.yajl_config(handle, YAJL_ALLOW_COMMENTS, 1)
    try:
        while True:
            buffer = f.read(buf_size)
            if buffer:
                result = yajl.yajl_parse(handle, buffer, len(buffer))
            else:
                result = yajl.yajl_complete_parse(handle)
            if result == YAJL_ERROR:
                perror = yajl.yajl_get_error(handle, 1, buffer, len(buffer))
                error = cast(perror, c_char_p).value
                yajl.yajl_free_error(handle, perror)
                raise common.JSONError(error)
            if not buffer and not events:
                if result == YAJL_INSUFFICIENT_DATA:
                    raise common.IncompleteJSONError()
                break

            for event in events:
                yield event
            events = []
    finally:
        yajl.yajl_free(handle)

def parse(file, **kwargs):
    '''
    Backend-specific wrapper for ijson.common.parse.
    '''
    return common.parse(basic_parse(file, **kwargs))

def items(file, prefix):
    '''
    Backend-specific wrapper for ijson.common.items.
    '''
    return common.items(parse(file), prefix)
<commit_msg>Added multiple value support to yajl2 backend<commit_after>'''
Wrapper for YAJL C library version 2.x.
'''

from ctypes import Structure, c_uint, c_ubyte, c_int, c_long, c_double, \
                   c_void_p, c_char_p, CFUNCTYPE, POINTER, byref, string_at, cast , \
                   cdll, util, c_char
from decimal import Decimal

from ijson import common, backends
from ijson.compat import b2s

yajl = backends.find_yajl(2)

yajl.yajl_alloc.restype = POINTER(c_char)
yajl.yajl_get_error.restype = POINTER(c_char)

C_EMPTY = CFUNCTYPE(c_int, c_void_p)
C_INT = CFUNCTYPE(c_int, c_void_p, c_int)
C_LONG = CFUNCTYPE(c_int, c_void_p, c_long)
C_DOUBLE = CFUNCTYPE(c_int, c_void_p, c_double)
C_STR = CFUNCTYPE(c_int, c_void_p, POINTER(c_ubyte), c_uint)

def number(value):
    '''
    Helper function casting a string that represents any Javascript number
    into appropriate Python value: either int or Decimal.
    '''
    try:
        return int(value)
    except ValueError:
        return Decimal(value)

_callback_data = [
    # Mapping of JSON parser events to callback C types and value converters.
    # Used to define the Callbacks structure and actual callback functions
    # inside the parse function.
    ('null', C_EMPTY, lambda: None),
    ('boolean', C_INT, lambda v: bool(v)),
    # "integer" and "double" aren't actually yielded by yajl since "number"
    # takes precedence if defined
    ('integer', C_LONG, lambda v, l: int(string_at(v, l))),
    ('double', C_DOUBLE, lambda v, l: float(string_at(v, l))),
    ('number', C_STR, lambda v, l: number(b2s(string_at(v, l)))),
    ('string', C_STR, lambda v, l: string_at(v, l).decode('utf-8')),
    ('start_map', C_EMPTY, lambda: None),
    ('map_key', C_STR, lambda v, l: b2s(string_at(v, l))),
    ('end_map', C_EMPTY, lambda: None),
    ('start_array', C_EMPTY, lambda: None),
    ('end_array', C_EMPTY, lambda: None),
]

class Callbacks(Structure):
    _fields_ = [(name, type) for name, type, func in _callback_data]

YAJL_OK = 0
YAJL_CANCELLED = 1
YAJL_INSUFFICIENT_DATA = 2
YAJL_ERROR = 3

# constants defined in yajl_parse.h
YAJL_ALLOW_COMMENTS = 1
YAJL_MULTIPLE_VALUES = 8

def basic_parse(f, allow_comments=False, buf_size=64 * 1024,
                multiple_values=False):
    '''
    Iterator yielding unprefixed events.

    Parameters:

    - f: a readable file-like object with JSON input
    - allow_comments: tells parser to allow comments in JSON input
    - buf_size: a size of an input buffer
    - multiple_values: allows the parser to parse multiple JSON objects
    '''
    events = []

    def callback(event, func_type, func):
        def c_callback(context, *args):
            events.append((event, func(*args)))
            return 1
        return func_type(c_callback)

    callbacks = Callbacks(*[callback(*data) for data in _callback_data])
    handle = yajl.yajl_alloc(byref(callbacks), None, None)
    if allow_comments:
        yajl.yajl_config(handle, YAJL_ALLOW_COMMENTS, 1)
    if multiple_values:
        yajl.yajl_config(handle, YAJL_MULTIPLE_VALUES, 1)
    try:
        while True:
            buffer = f.read(buf_size)
            if buffer:
                result = yajl.yajl_parse(handle, buffer, len(buffer))
            else:
                result = yajl.yajl_complete_parse(handle)
            if result == YAJL_ERROR:
                perror = yajl.yajl_get_error(handle, 1, buffer, len(buffer))
                error = cast(perror, c_char_p).value
                yajl.yajl_free_error(handle, perror)
                raise common.JSONError(error)
            if not buffer and not events:
                if result == YAJL_INSUFFICIENT_DATA:
                    raise common.IncompleteJSONError()
                break

            for event in events:
                yield event
            events = []
    finally:
        yajl.yajl_free(handle)

def parse(file, **kwargs):
    '''
    Backend-specific wrapper for ijson.common.parse.
    '''
    return common.parse(basic_parse(file, **kwargs))

def items(file, prefix):
    '''
    Backend-specific wrapper for ijson.common.items.
    '''
    return common.items(parse(file), prefix)
<|endoftext|>


<commit_before>from lasagne import *

from craynn import Expression

from ..subnetworks import achain

__all__ = [
  'Net', 'net',
  'get_input_layer'
]

def get_input_layer(shape_or_layer, index=None):
  if hasattr(shape_or_layer, '__iter__') and all([ type(s) is int for s in shape_or_layer ]):
    name = 'input' if index is None else 'input%d' % index
    return layers.InputLayer(shape=shape_or_layer, name=name)
  else:
    return shape_or_layer

class Net(Expression):
  def __init__(self, factory, inputs):
    if not hasattr(inputs, '__iter__'):
      input_layer = get_input_layer(inputs)
      outputs = factory(input_layer)
      input_layers = [input_layer]
    else:
      input_layers = [
        get_input_layer(input, i)
        for i, input in enumerate(inputs)
      ]
      outputs = factory(input_layers)

    if not hasattr(outputs, '__iter__'):
      outputs = [outputs]

    super(Net, self).__init__(input_layers, outputs)

net = lambda inputs: lambda *factory: Net(achain(*factory), inputs)<commit_msg>now, fixed<commit_after>from lasagne import *

from craynn import Expression

from ..subnetworks import achain

__all__ = [
  'Net', 'net',
  'get_input_layer'
]

def get_input_layer(shape_or_layer, index=None):
  if hasattr(shape_or_layer, '__iter__') and all([ type(s) is int for s in shape_or_layer ]):
    name = 'input' if index is None else 'input%d' % index
    return layers.InputLayer(shape=shape_or_layer, name=name)
  else:
    return shape_or_layer

class Net(Expression):
  def __init__(self, factory, inputs):
    ### either single layer instance or one shape
    if not hasattr(inputs, '__iter__') or all([ type(s) is int for s in inputs ]):
      input_layer = get_input_layer(inputs)
      outputs = factory(input_layer)
      input_layers = [input_layer]
    else:
      input_layers = [
        get_input_layer(input, i)
        for i, input in enumerate(inputs)
      ]
      outputs = factory(input_layers)

    if not hasattr(outputs, '__iter__'):
      outputs = [outputs]

    super(Net, self).__init__(input_layers, outputs)

net = lambda inputs: lambda *factory: Net(achain(*factory), inputs)<|endoftext|>


<commit_before>#!/usr/bin/env python
# Copyright 2016 The Fontbakery Authors
#
# 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.
#
from os.path import basename
from argparse import ArgumentParser
from fontTools.ttLib import TTFont
from fontTools.ttLib.tables._c_m_a_p import CmapSubtable

description = """

fontbakery-fix-cmap.py
~~~~~~~~~~~~~~~~~~~~~~

Manipulate a collection of fonts' cmap tables.
"""

def convert_cmap_subtables_to_v4(font):
  cmap = font['cmap']
  outtables = []
  for table in cmap.tables:
    newtable = CmapSubtable.newSubtable(4)
    newtable.platformID = table.platformID
    newtable.platEncID = table.platEncID
    newtable.language = table.language
    newtable.cmap = table.cmap
    outtables.append(newtable)
  font['cmap'].tables = outtables

def remove_cmap_subtable(font, plat_id, enc_id):
  for table in font['cmap'].tables:
    if table.platformID == plat_id and table.platEncID == enc_id:
      font['cmap'].tables.remove(table)

def main():
  parser = ArgumentParser(description=description)
  parser.add_argument('fonts', nargs='+')
  parser.add_argument('--format-4-subtables', '-f4', default=False,
                      action='store_true',
                      help="Convert cmap subtables to format 4")
  parser.add_argument('-drop-mac-subtable', '-dm', default=False,
                      action='store_true',
                      help='Drop Mac cmap subtables')
  args = parser.parse_args()

  for path in args.fonts:
    font = TTFont(path)
    font_filename = basename(path)
    fixit = False

    if args.format_4_subtables:
      print('%s: Converting Cmap subtables to format 4' % font_filename)
      convert_cmap_subtables_to_v4(font)
      fixit = True

    if args.drop_mac_subtable:
      if font['cmap'].getcmap(1, 0):
        print('%s: Dropping Cmap Mac subtable' % font_filename)
        remove_cmap_subtable(font, 1, 0)
        fixit = True

    if fixit:
      print('Saving %s to %s.fix' % (font_filename, path))
      font.save(path + '.fix')

if __name__ == '__main__':
  main()
<commit_msg>use double dash on --drop-mac-subtable long attribute<commit_after>#!/usr/bin/env python
# Copyright 2016 The Fontbakery Authors
#
# 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.
#
from os.path import basename
from argparse import ArgumentParser
from fontTools.ttLib import TTFont
from fontTools.ttLib.tables._c_m_a_p import CmapSubtable

description = """

fontbakery-fix-cmap.py
~~~~~~~~~~~~~~~~~~~~~~

Manipulate a collection of fonts' cmap tables.
"""

def convert_cmap_subtables_to_v4(font):
  cmap = font['cmap']
  outtables = []
  for table in cmap.tables:
    newtable = CmapSubtable.newSubtable(4)
    newtable.platformID = table.platformID
    newtable.platEncID = table.platEncID
    newtable.language = table.language
    newtable.cmap = table.cmap
    outtables.append(newtable)
  font['cmap'].tables = outtables

def remove_cmap_subtable(font, plat_id, enc_id):
  for table in font['cmap'].tables:
    if table.platformID == plat_id and table.platEncID == enc_id:
      font['cmap'].tables.remove(table)

def main():
  parser = ArgumentParser(description=description)
  parser.add_argument('fonts', nargs='+')
  parser.add_argument('--format-4-subtables', '-f4', default=False,
                      action='store_true',
                      help="Convert cmap subtables to format 4")
  parser.add_argument('--drop-mac-subtable', '-dm', default=False,
                      action='store_true',
                      help='Drop Mac cmap subtables')
  args = parser.parse_args()

  for path in args.fonts:
    font = TTFont(path)
    font_filename = basename(path)
    fixit = False

    if args.format_4_subtables:
      print('%s: Converting Cmap subtables to format 4' % font_filename)
      convert_cmap_subtables_to_v4(font)
      fixit = True

    if args.drop_mac_subtable:
      if font['cmap'].getcmap(1, 0):
        print('%s: Dropping Cmap Mac subtable' % font_filename)
        remove_cmap_subtable(font, 1, 0)
        fixit = True

    if fixit:
      print('Saving %s to %s.fix' % (font_filename, path))
      font.save(path + '.fix')

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


<commit_before>import logging
import os
from tempfile import mkdtemp
from typing import Callable, Dict, List, Optional, Tuple, Type

from imhotep.repositories import Repository
from imhotep.tools import Tool

from .repositories import AuthenticatedRepository, Repository

log = logging.getLogger(__name__)

class RepoManager:
    """
    Manages creation and deletion of `Repository` objects.
    """

    to_cleanup: Dict[str, str] = {}

    def __init__(
        self,
        authenticated: bool = False,
        cache_directory: None = None,
        tools: Optional[List[Tool]] = None,
        executor: Optional[Callable] = None,
        shallow_clone: bool = False,
        domain: Optional[str] = None,
    ) -> None:
        self.should_cleanup = cache_directory is None
        self.authenticated = authenticated
        self.cache_directory = cache_directory
        self.tools = tools or []
        self.executor = executor
        self.shallow = shallow_clone
        self.domain = domain

    def get_repo_class(self) -> Type[Repository]:
        if self.authenticated:
            return AuthenticatedRepository
        return Repository

    def clone_dir(self, repo_name: str) -> str:
        dired_repo_name = repo_name.replace("/", "__")
        if not self.cache_directory:
            dirname = mkdtemp(suffix=dired_repo_name)
        else:
            dirname = os.path.abspath(f"{self.cache_directory}/{dired_repo_name}")
        return dirname

    def fetch(self, dirname, remote_name, ref):
        log.debug("Fetching %s %s", remote_name, ref)
        self.executor(f"cd {dirname} && git fetch --depth=1 {remote_name} {ref}")

    def pull(self, dirname):
        log.debug("Pulling all %s", dirname)
        self.executor("cd %s && git pull --all" % dirname)

    def add_remote(self, dirname, name, url):
        log.debug("Adding remote %s url: %s", name, url)
        self.executor(f"cd {dirname} && git remote add {name} {url}")

    def set_up_clone(self, repo_name: str, remote_repo: None) -> Tuple[str, Repository]:
        """Sets up the working directory and returns a tuple of
        (dirname, repo )"""
        dirname = self.clone_dir(repo_name)
        self.to_cleanup[repo_name] = dirname
        klass = self.get_repo_class()
        repo = klass(
            repo_name,
            dirname,
            self.tools,
            self.executor,
            shallow=self.shallow_clone,
            domain=self.domain,
        )
        return (dirname, repo)

    def clone_repo(self, repo_name: str, remote_repo: None, ref: str) -> Repository:
        """Clones the given repo and returns the Repository object."""
        self.shallow_clone = False
        dirname, repo = self.set_up_clone(repo_name, remote_repo)

        if self.executor is None:
            log.error("Executor does not exist.")
            raise RuntimeError
        if os.path.isdir("%s/.git" % dirname):
            log.debug("Updating %s to %s", repo.download_location, dirname)
            self.executor("cd %s && git switch master" % dirname)
            self.pull(dirname)
        else:
            log.debug("Cloning %s to %s", repo.download_location, dirname)
            self.executor(f"git clone {repo.download_location} {dirname}")

        if remote_repo is not None:
            log.debug("Pulling remote branch from %s", remote_repo.url)
            self.add_remote(dirname, remote_repo.name, remote_repo.url)
            self.pull(dirname)
        return repo

    def cleanup(self) -> None:
        if self.should_cleanup:
            for repo_dir in self.to_cleanup.values():
                log.debug("Cleaning up %s", repo_dir)
                if self.executor is None:
                    log.error("Executor does not exist.")
                    raise RuntimeError
                self.executor("rm -rf %s" % repo_dir)

class ShallowRepoManager(RepoManager):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def clone_repo(self, repo_name, remote_repo, ref):
        self.shallow_clone = True
        dirname, repo = self.set_up_clone(repo_name, remote_repo)
        remote_name = "origin"
        log.debug("Shallow cloning.")
        download_location = repo.download_location
        log.debug("Creating stub git repo at %s" % (dirname))
        self.executor(f"mkdir -p {dirname}")
        self.executor(f"cd {dirname} && git init")
        log.debug("Adding origin repo %s " % (download_location))
        self.add_remote(dirname, "origin", download_location)

        if remote_repo:
            self.add_remote(dirname, remote_repo.name, remote_repo.url)
            remote_name = remote_repo.name
        self.fetch(dirname, "origin", "HEAD")
        self.fetch(dirname, remote_name, ref)
        return repo
<commit_msg>docstring<commit_after>import logging
import os
from tempfile import mkdtemp
from typing import Callable, Dict, List, Optional, Tuple, Type

from imhotep.repositories import Repository
from imhotep.tools import Tool

from .repositories import AuthenticatedRepository, Repository

log = logging.getLogger(__name__)

class RepoManager:
    """
    Manages creation and deletion of `Repository` objects.
    """

    to_cleanup: Dict[str, str] = {}

    def __init__(
        self,
        authenticated: bool = False,
        cache_directory: None = None,
        tools: Optional[List[Tool]] = None,
        executor: Optional[Callable] = None,
        shallow_clone: bool = False,
        domain: Optional[str] = None,
    ) -> None:
        self.should_cleanup = cache_directory is None
        self.authenticated = authenticated
        self.cache_directory = cache_directory
        self.tools = tools or []
        self.executor = executor
        self.shallow = shallow_clone
        self.domain = domain

    def get_repo_class(self) -> Type[Repository]:
        if self.authenticated:
            return AuthenticatedRepository
        return Repository

    def clone_dir(self, repo_name: str) -> str:
        """Gets a full path of the target directory to clone into."""
        dired_repo_name = repo_name.replace("/", "__")
        if not self.cache_directory:
            dirname = mkdtemp(suffix=dired_repo_name)
        else:
            dirname = os.path.abspath(f"{self.cache_directory}/{dired_repo_name}")
        return dirname

    def fetch(self, dirname, remote_name, ref):
        log.debug("Fetching %s %s", remote_name, ref)
        self.executor(f"cd {dirname} && git fetch --depth=1 {remote_name} {ref}")

    def pull(self, dirname):
        log.debug("Pulling all %s", dirname)
        self.executor("cd %s && git pull --all" % dirname)

    def add_remote(self, dirname, name, url):
        log.debug("Adding remote %s url: %s", name, url)
        self.executor(f"cd {dirname} && git remote add {name} {url}")

    def set_up_clone(self, repo_name: str, remote_repo: None) -> Tuple[str, Repository]:
        """Sets up the working directory and returns a tuple of (dirname, repo)."""
        dirname = self.clone_dir(repo_name)
        self.to_cleanup[repo_name] = dirname
        klass = self.get_repo_class()
        repo = klass(
            repo_name,
            dirname,
            self.tools,
            self.executor,
            shallow=self.shallow_clone,
            domain=self.domain,
        )
        return (dirname, repo)

    def clone_repo(self, repo_name: str, remote_repo: None, ref: str) -> Repository:
        """Clones the given repo and returns the Repository object."""
        self.shallow_clone = False
        dirname, repo = self.set_up_clone(repo_name, remote_repo)

        if self.executor is None:
            log.error("Executor does not exist.")
            raise RuntimeError
        if os.path.isdir("%s/.git" % dirname):
            log.debug("Updating %s to %s", repo.download_location, dirname)
            self.executor("cd %s && git switch master" % dirname)
            self.pull(dirname)
        else:
            log.debug("Cloning %s to %s", repo.download_location, dirname)
            self.executor(f"git clone {repo.download_location} {dirname}")

        if remote_repo is not None:
            log.debug("Pulling remote branch from %s", remote_repo.url)
            self.add_remote(dirname, remote_repo.name, remote_repo.url)
            self.pull(dirname)
        return repo

    def cleanup(self) -> None:
        if self.should_cleanup:
            for repo_dir in self.to_cleanup.values():
                log.debug("Cleaning up %s", repo_dir)
                if self.executor is None:
                    log.error("Executor does not exist.")
                    raise RuntimeError
                self.executor("rm -rf %s" % repo_dir)

class ShallowRepoManager(RepoManager):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def clone_repo(self, repo_name, remote_repo, ref):
        self.shallow_clone = True
        dirname, repo = self.set_up_clone(repo_name, remote_repo)
        remote_name = "origin"
        log.debug("Shallow cloning.")
        download_location = repo.download_location
        log.debug("Creating stub git repo at %s" % (dirname))
        self.executor(f"mkdir -p {dirname}")
        self.executor(f"cd {dirname} && git init")
        log.debug("Adding origin repo %s " % (download_location))
        self.add_remote(dirname, "origin", download_location)

        if remote_repo:
            self.add_remote(dirname, remote_repo.name, remote_repo.url)
            remote_name = remote_repo.name
        self.fetch(dirname, "origin", "HEAD")
        self.fetch(dirname, remote_name, ref)
        return repo
<|endoftext|>


<commit_before># coding=utf-8
from __future__ import unicode_literals

import os.path
import warnings

__version__ = (0, 5, 0)

def _get_git_revision(path):
    revision_file = os.path.join(path, 'refs', 'heads', 'master')
    if not os.path.exists(revision_file):
        return None
    fh = open(revision_file, 'r')
    try:
        return fh.read()
    finally:
        fh.close()

def get_revision():
    """
    :returns: Revision number of this branch/checkout, if available. None if
        no revision number can be determined.
    """
    package_dir = os.path.dirname(__file__)
    checkout_dir = os.path.normpath(os.path.join(package_dir, '..'))
    path = os.path.join(checkout_dir, '.git')
    if os.path.exists(path):
        return _get_git_revision(path)
    return None

__build__ = get_revision()

def lazy_object(location):
    def inner(*args, **kwargs):
        parts = location.rsplit('.', 1)
        warnings.warn('`djangoratings.%s` is deprecated. Please use `%s` instead.' % (parts[1], location), DeprecationWarning)
        try:
            imp = __import__(parts[0], globals(), locals(), [parts[1]], -1)
        except:
            imp = __import__(parts[0], globals(), locals(), [parts[1]])
        func = getattr(imp, parts[1])
        if callable(func):
            return func(*args, **kwargs)
        return func
    return inner

RatingField = lazy_object('xratings.fields.RatingField')
AnonymousRatingField = lazy_object('xratings.fields.AnonymousRatingField')
Rating = lazy_object('xratings.fields.Rating')
<commit_msg>Update __init__.py<commit_after># coding=utf-8
from __future__ import unicode_literals

import os.path
import warnings

__version__ = (0, 5, 0)

def _get_git_revision(path):
    revision_file = os.path.join(path, 'refs', 'heads', 'master')
    if not os.path.exists(revision_file):
        return None
    fh = open(revision_file, 'r')
    try:
        return fh.read()
    finally:
        fh.close()

def get_revision():
    """
    :returns: Revision number of this branch/checkout, if available. None if
        no revision number can be determined.
    """
    package_dir = os.path.dirname(__file__)
    checkout_dir = os.path.normpath(os.path.join(package_dir, '..'))
    path = os.path.join(checkout_dir, '.git')
    if os.path.exists(path):
        return _get_git_revision(path)
    return None

__build__ = get_revision()

def lazy_object(location):
    def inner(*args, **kwargs):
        parts = location.rsplit('.', 1)
        warnings.warn('`xratings.%s` is deprecated. Please use `%s` instead.' % (parts[1], location), DeprecationWarning)
        try:
            imp = __import__(parts[0], globals(), locals(), [parts[1]], -1)
        except:
            imp = __import__(parts[0], globals(), locals(), [parts[1]])
        func = getattr(imp, parts[1])
        if callable(func):
            return func(*args, **kwargs)
        return func
    return inner

RatingField = lazy_object('xratings.fields.RatingField')
AnonymousRatingField = lazy_object('xratings.fields.AnonymousRatingField')
Rating = lazy_object('xratings.fields.Rating')
<|endoftext|>


<commit_before>import argparse
import os
import logging
from joblib import Parallel, delayed

import simulators
from bin.coadd_analysis_script import main as anaylsis_main
from bin.coadd_chi2_db import main as db_main
from simulators.fake_simulator import main as fake_generator
from simulators.iam_script import main as iam_script_main

def make_fake_parameter_file(info):
    name = os.path.join(simulators.paths["parameters"], "{}_params.dat".format(info["star"].upper()))

    with open(name, "w") as f:
        for key, value in info.items():
            f.write("{0} \t= {1}\n".format(key, value))

def _parser():
    """Take care of all the argparse stuff.

    :returns: the args
    """
    parser = argparse.ArgumentParser(description='Best host modelling.')
    parser.add_argument("star", help='Star name.', type=str)
    parser.add_argument("obsnum", help='Star observation number.', type=int)
    parser.add_argument("teff", help='Temperature of Star.', type=int)
    parser.add_argument("logg", help='Logg of Star.', type=float)
    parser.add_argument("feh", help='Feh of Star.', type=float)
    parser.add_argument("teff2", help='Temperature of companion.', type=int)
    parser.add_argument("logg2", help='Logg of  companion.', type=float)
    parser.add_argument("feh2", help='Feh of companion.', type=float)
    parser.add_argument('rv', help='rv radial velocity of companion', type=float, default=0)
    parser.add_argument('gamma', help='Gamma radial velocity', type=float, default=0)
    parser.add_argument('-i', "--independent", help='Independent rv of companion', action="store_true")
    parser.add_argument('-s', '--suffix', type=str, default="",
                        help='Extra name identifier.')
    parser.add_argument('-n', '--noise',
                        help='SNR value int', type=int, default=None)
    parser.add_argument('-r', '--replace',
                        help='Replace old fake spectra.', action="store_true")
    parser.add_argument('-a', '--area_scale',
                        help='Disable area_scaling.', action="store_false")
    parser.add_argument('-j', '--n_jobs',
                        help='Number of parallel jobs.', type=int, default=4)
    parser.add_argument('-f', '--fudge',
                        help='Fudge factor.', default=None)
    return parser.parse_args()

def main(star, obsnum, teff, logg, feh, teff2, logg2, feh2, gamma=0, rv=0,
         noise=False, suffix="", replace=False, independent=False, fudge=None, area_scale=True, n_jobs=4):
    chips = range(1, 5)

    starinfo = {"star": star, "temp": teff, "logg": logg, "fe_h": feh, "comp_temp": teff2}
    make_fake_parameter_file(starinfo)

    params1 = "{}, {}, {}".format(teff, logg, feh)
    params2 = "{}, {}, {}".format(teff2, logg2, feh2)

    fake_generator(star=star, sim_num=obsnum, params1=params1, params2=params2, rv=rv, gamma=gamma, noise=noise,
                   replace=replace, noplots=True, mode="iam", independent=independent, fudge=fudge,
                   area_scale=area_scale)

    # iam_script
    # for chip in chips:
    #     iam_script_main(star=star, obsnum=obsnum, chip=chip, suffix=suffix, area_scale=area_scale)
    Parallel(n_jobs=n_jobs)(
        delayed(iam_script_main)(star=star, obsnum=obsnum, chip=chip, suffix=suffix, area_scale=area_scale, betasigma=True)
        for chip in chips)

    # Generate db
    db_main(star=star, obsnum=obsnum, suffix=suffix, move=True, replace=True)

    # Selected Analysis
    # anaylsis_main(star=star, obsnum=obsnum, suffix=suffix, mode="smallest_chi2")
    # anaylsis_main(star=star, obsnum=obsnum, suffix=suffix, mode="compare_spectra")
    try:
        anaylsis_main(star=star, obsnum=obsnum, suffix=suffix, mode="chi2_parabola")
    except:
        pass
    try:
        anaylsis_main(star=star, obsnum=obsnum, suffix=suffix, mode="all")
    except:
        pass
    # anaylsis_main(star=star, obsnum=obsnum, suffix=suffix, mode="contrast")

    print("Noise level =", noise)

if __name__ == "__main__":
    args = vars(_parser())
    opts = {k: args[k] for k in args}

    main(**opts)
    print("iam fake analysis")
    print("Original opts", opts)
<commit_msg>Add normalization to fake simulation recovery<commit_after>import argparse
import os
import logging
from joblib import Parallel, delayed

import simulators
from bin.coadd_analysis_script import main as anaylsis_main
from bin.coadd_chi2_db import main as db_main
from simulators.fake_simulator import main as fake_generator
from simulators.iam_script import main as iam_script_main

def make_fake_parameter_file(info):
    name = os.path.join(simulators.paths["parameters"], "{}_params.dat".format(info["star"].upper()))

    with open(name, "w") as f:
        for key, value in info.items():
            f.write("{0} \t= {1}\n".format(key, value))

def _parser():
    """Take care of all the argparse stuff.

    :returns: the args
    """
    parser = argparse.ArgumentParser(description='Best host modelling.')
    parser.add_argument("star", help='Star name.', type=str)
    parser.add_argument("obsnum", help='Star observation number.', type=int)
    parser.add_argument("teff", help='Temperature of Star.', type=int)
    parser.add_argument("logg", help='Logg of Star.', type=float)
    parser.add_argument("feh", help='Feh of Star.', type=float)
    parser.add_argument("teff2", help='Temperature of companion.', type=int)
    parser.add_argument("logg2", help='Logg of  companion.', type=float)
    parser.add_argument("feh2", help='Feh of companion.', type=float)
    parser.add_argument('rv', help='rv radial velocity of companion', type=float, default=0)
    parser.add_argument('gamma', help='Gamma radial velocity', type=float, default=0)
    parser.add_argument('-i', "--independent", help='Independent rv of companion', action="store_true")
    parser.add_argument('-s', '--suffix', type=str, default="",
                        help='Extra name identifier.')
    parser.add_argument('-n', '--noise',
                        help='SNR value int', type=int, default=None)
    parser.add_argument('-r', '--replace',
                        help='Replace old fake spectra.', action="store_true")
    parser.add_argument('-a', '--area_scale',
                        help='Disable area_scaling.', action="store_false")
    parser.add_argument('-j', '--n_jobs',
                        help='Number of parallel jobs.', type=int, default=4)
    parser.add_argument('-f', '--fudge',
                        help='Fudge factor.', default=None)
    parser.add_argument("--renormalize", help="renormalize before chi2", action="store_true")
    parser.add_argument("-m", "--norm_method", help="Re-normalization method flux to models. Default=scalar",
                        choices=["scalar", "linear"], default="scalar")
    return parser.parse_args()

def main(star, obsnum, teff, logg, feh, teff2, logg2, feh2, gamma=0, rv=0,
         noise=False, suffix="", replace=False, independent=False,
         fudge=None, area_scale=True, n_jobs=4,
         renormalize=False, norm_method="scalar"):
    chips = range(1, 5)

    starinfo = {"star": star, "temp": teff, "logg": logg, "fe_h": feh, "comp_temp": teff2}
    make_fake_parameter_file(starinfo)

    params1 = "{}, {}, {}".format(teff, logg, feh)
    params2 = "{}, {}, {}".format(teff2, logg2, feh2)

    fake_generator(star=star, sim_num=obsnum, params1=params1, params2=params2, rv=rv, gamma=gamma, noise=noise,
                   replace=replace, noplots=True, mode="iam", independent=independent, fudge=fudge,
                   area_scale=area_scale)

    # iam_script
    # for chip in chips:
    #     iam_script_main(star=star, obsnum=obsnum, chip=chip, suffix=suffix, area_scale=area_scale)
    Parallel(n_jobs=n_jobs)(
        delayed(iam_script_main)(star=star, obsnum=obsnum, chip=chip, suffix=suffix,
                                 area_scale=area_scale, betasigma=True,
                                 renormalize=renormalize, norm_method=norm_method)
        for chip in chips)

    # Generate db
    db_main(star=star, obsnum=obsnum, suffix=suffix, move=True, replace=True)

    # Selected Analysis
    # anaylsis_main(star=star, obsnum=obsnum, suffix=suffix, mode="smallest_chi2")
    # anaylsis_main(star=star, obsnum=obsnum, suffix=suffix, mode="compare_spectra")
    try:
        anaylsis_main(star=star, obsnum=obsnum, suffix=suffix, mode="chi2_parabola")
    except:
        pass
    try:
        anaylsis_main(star=star, obsnum=obsnum, suffix=suffix, mode="all")
    except:
        pass
    # anaylsis_main(star=star, obsnum=obsnum, suffix=suffix, mode="contrast")

    print("Noise level =", noise)

if __name__ == "__main__":
    args = vars(_parser())
    opts = {k: args[k] for k in args}

    main(**opts)
    print("iam fake analysis")
    print("Original opts", opts)
<|endoftext|>


<commit_before>import abc

class RODataSet(object):
    """
    An :py:mod:`abc` that represents a read-only set of
    :py:class:`~.importtools.importables.Importable` instances.

    """

    @abc.abstractmethod
    def __iter__(self):
        """
        Iterate over all the :py:class:`~.importtools.importables.Importable`
        instances in this set.

        """

    @abc.abstractmethod
    def __contains__(self, importable):
        """
        Check if the dataset contains an equal
        :py:class:`~.importtools.importables.Importable` instance.

        See :py:meth:`~.importtools.importables.Importable.__cmp__` method for
        more details about the equality test between two
        :py:class:`~.importtools.importables.Importable` instances.

        """

    @abc.abstractmethod
    def get(self, importable, default=None):
        """
        Return an equal :py:class:`~.importtools.importables.Importable` from
        the dataset or the default value if no such instances are found.

        """

class DataSet(RODataSet):
    """
    An :py:mod:`abc` that represents a mutable set of
    :py:class:`~.importtools.importables.Importable` instances.

    """

    @abc.abstractmethod
    def add(self, importable):
        """
        Add the :py:class:`~.importtools.importables.Importable` in the dataset
        or replaces an existing and equal one.

        """

    @abc.abstractmethod
    def pop(self, importable):
        """
        Remove and return an equal
        :py:class:`~.importtools.importables.Importable` from the dataset.

        """

class MemoryDataSet(dict, DataSet):
    """
    A simple in-memory :py:class:`dict`-based :py:class:`DataSet`
    implementation.

    Since this particular :py:class:`DataSet` implementation only uses the
    :py:meth:`__hash__` and :py:meth:`__cmp__` methods of an
    :py:class:`~.importtools.importables.Importable` :py:mod:`abc` we can
    exemplify this class functionality using :py:class:`tuple` objects:

    >>> from importtools import MemoryDataSet
    >>> item1 = (1, )
    >>> item2 = (1, )
    >>> item3 = (2, )

    At first, a newly created :py:class:`MemoryDataSet` instance has no
    elements:

    >>> mds = MemoryDataSet()
    >>> list(mds)
    []

    After creation, it can be populated and the elements in the dataset can be
    retrieved using other equal elements.

    >>> mds = MemoryDataSet()
    >>> mds.add(item1)
    >>> mds.get(item1) == item1
    True
    >>> mds.get(item1) is item1
    True
    >>> mds.get(item2) == item2
    True
    >>> mds.get(item2) is item1
    True

    Trying to get an inexistent item should return the default value or
    :py:class:`None`:

    >>> mds = MemoryDataSet()
    >>> mds.add(item1)
    >>> mds.get(item3) is None
    True
    >>> mds.get(item3, 'default')
    'default'

    Initial data can be passed when constructing instances:

    >>> mds = MemoryDataSet((item1, item2, item3))

    """
    def __init__(self, data_loader=None):
        data_iter = tuple()
        if data_loader is not None:
            data_iter = iter(data_loader)
        mapping = ((i, i) for i in data_iter)
        super(MemoryDataSet, self).__init__(mapping)

    def add(self, importable):
        self[importable] = importable

class DiffDataSet(MemoryDataSet):
    """
    A :py:class:`DataSet` implementation that remembers all the changes,
    additions and removals done to it.

    Using a :py:class:`DiffDataSet` as the destination of the import algorithm
    allows optimal persistence of the changes by grouping them in a way suited
    for batch processing.

    In order to exemplify this class functionality we start with a
    instance that we prepopulate:

    An :py:class:`~.importtools.importables.Importable` mock is needed to
    populate the datasets:

    >>> from importtools import (
    ...     Importable, MockImportable, MemoryDataSet, DiffDataSet
    ... )
    >>> dds = DiffDataSet([
    ...     MockImportable(0),
    ...     MockImportable(1),
    ...     MockImportable(2),
    ...     MockImportable(3)
    ... ])

    >>> sorted(dds, key=lambda x: x.args)
    [<(0,) IMPORTED>, <(1,) IMPORTED>, <(2,) IMPORTED>, <(3,) IMPORTED>]

    Fetching existing elements from ``dds`` should work correctly:

    >>> dds.get(MockImportable(1)) == MockImportable(1)
    True

    New :py:class:`~.importtools.importables.Importable` can still added in
    ``dds``:

    >>> dds.add(MockImportable(4))
    >>> dds.add(MockImportable(5))
    >>> sorted(dds, key=lambda x: x.args)
    ... # doctest: +NORMALIZE_WHITESPACE
    [<(0,) IMPORTED>, <(1,) IMPORTED>, <(2,) IMPORTED>,
     <(3,) IMPORTED>, <(4,) IMPORTED>, <(5,) IMPORTED>]

    Check if the new elements were added:

    >>> dds.get(MockImportable(4))
    <(4,) IMPORTED>
    >>> dds.get(MockImportable(5))
    <(5,) IMPORTED>

    A lookup with an inexistent element should return the default value:

    >>> dds.get(MockImportable(100), 'default')
    'default'

    We should be able to ``pop`` both original and recently added elements:

    >>> dds.pop(MockImportable(1))
    <(1,) IMPORTED>
    >>> dds.pop(MockImportable(4))
    <(4,) IMPORTED>
    >>> sorted(dds, key=lambda x: x.args)
    [<(0,) IMPORTED>, <(2,) IMPORTED>, <(3,) IMPORTED>, <(5,) IMPORTED>]

    And everything should have been registered:

    >>> sorted(dds.get_added(), key=lambda x: x.args)
    [<(5,) IMPORTED>]
    >>> sorted(dds.get_removed(), key=lambda x: x.args)
    [<(1,) IMPORTED>]
    >>> sorted(dds.get_changed(), key=lambda x: x.args)
    []

    Adding and removing the same
    :py:class:`~.importtools.importables.Importable` shouldn't change anything:

    >>> dds.add(MockImportable(100))
    >>> dds.pop(MockImportable(100))
    <(100,) IMPORTED>
    >>> sorted(dds.get_added(), key=lambda x: x.args)
    [<(5,) IMPORTED>]
    >>> sorted(dds.get_removed(), key=lambda x: x.args)
    [<(1,) IMPORTED>]
    >>> sorted(dds.get_changed(), key=lambda x: x.args)
    []

    But removing an original :py:class:`~.importtools.importables.Importable`
    and adding a new one should be interpreted as a change:

    >>> dds.pop(MockImportable(2))
    <(2,) IMPORTED>
    >>> dds.add(MockImportable(2))
    >>> sorted(dds.get_added(), key=lambda x: x.args)
    [<(5,) IMPORTED>]
    >>> sorted(dds.get_removed(), key=lambda x: x.args)
    [<(1,) IMPORTED>]
    >>> sorted(dds.get_changed(), key=lambda x: x.args)
    [<(2,) IMPORTED>]

    If you change one of the existing elements it will be marked as changed:

    >>> dds.get(MockImportable(3)).register_change()
    >>> sorted(dds.get_added(), key=lambda x: x.args)
    [<(5,) IMPORTED>]
    >>> sorted(dds.get_removed(), key=lambda x: x.args)
    [<(1,) IMPORTED>]
    >>> sorted(dds.get_changed(), key=lambda x: x.args)
    [<(2,) IMPORTED>, <(3,) IMPORTED>]

    Even if the elements are accessed by iterating on over the dataset the
    changes are still tracked:

    >>> list(dds)[0].register_change()
    >>> sorted(dds.get_added(), key=lambda x: x.args)
    [<(5,) IMPORTED>]
    >>> sorted(dds.get_removed(), key=lambda x: x.args)
    [<(1,) IMPORTED>]
    >>> sorted(dds.get_changed(), key=lambda x: x.args)
    [<(0,) IMPORTED>, <(2,) IMPORTED>, <(3,) IMPORTED>]

    When printing the object a list of all the added, removed and created
    elements will be shown, along with the total number for each operation:

    >>> print dds
    To be added: 1
    <(5,) IMPORTED>
    To be removed: 1
    <(1,) IMPORTED>
    To be changed: 3
    <(0,) IMPORTED>
    <(2,) IMPORTED>
    <(3,) IMPORTED>

    >>> dds._added = MemoryDataSet()
    >>> dds._removed = MemoryDataSet()
    >>> print dds
    <BLANKLINE>
    To be changed: 3
    <(0,) IMPORTED>
    <(2,) IMPORTED>
    <(3,) IMPORTED>

    >>> dds._changed = MemoryDataSet()
    >>> print dds
    No elements were affected.

    """
    def __init__(self, data_loader=None):
        self._added = set()
        self._removed = set()
        self._changed = set()
        # XXX for some reason dict subclass methods can't be hashed
        self.rc = lambda x: self.register_change(x)
        super(DiffDataSet, self).__init__(
            element.register_listener(self.rc) for element in data_loader
        )

    def __str__(self):
        result = ''
        added_count = len(self._added)
        if added_count > 0:
            result += 'To be added: %s\n' % added_count
            result += '\n'.join(sorted(repr(x) for x in self.get_added()))
        removed_count = len(self._removed)
        if removed_count > 0:
            result += '\nTo be removed: %s\n' % removed_count
            result += '\n'.join(sorted(repr(x) for x in self.get_removed()))
        changed_count = len(self._changed)
        if changed_count > 0:
            result += '\nTo be changed: %s\n' % changed_count
            result += '\n'.join(sorted(repr(x) for x in self.get_changed()))
        if (added_count + removed_count + changed_count) == 0:
            result += 'No elements were affected.'
        return result

    def add(self, importable):
        if importable in self._removed:
            self._removed.discard(importable)
            self._changed.add(importable)
        else:
            self._added.add(importable)
        super(DiffDataSet, self).add(importable.register_listener(self.rc))

    def pop(self, importable):
        if importable in self._added:
            self._added.discard(importable)
        else:
            self._removed.add(importable)
        return super(DiffDataSet, self).pop(importable)

    def register_change(self, importable):
        """Mark an element in the current dataset as changed.

        This method is registered as a listener for changes in all the
        elements of the wrapped :py:class:`DataSet`.

        """
        assert importable in self
        self._changed.add(importable)

    def get_added(self):
        """A set of all added elements in the wrapped dataset."""
        return self._added

    def get_removed(self):
        """A set of all removed elements in the wrapped dataset."""
        return self._removed

    def get_changed(self):
        """A set of all changed elements in the wrapped dataset."""
        return self._changed
<commit_msg>No more need for sorting key.<commit_after>import abc

class RODataSet(object):
    """
    An :py:mod:`abc` that represents a read-only set of
    :py:class:`~.importtools.importables.Importable` instances.

    """

    @abc.abstractmethod
    def __iter__(self):
        """
        Iterate over all the :py:class:`~.importtools.importables.Importable`
        instances in this set.

        """

    @abc.abstractmethod
    def __contains__(self, importable):
        """
        Check if the dataset contains an equal
        :py:class:`~.importtools.importables.Importable` instance.

        See :py:meth:`~.importtools.importables.Importable.__cmp__` method for
        more details about the equality test between two
        :py:class:`~.importtools.importables.Importable` instances.

        """

    @abc.abstractmethod
    def get(self, importable, default=None):
        """
        Return an equal :py:class:`~.importtools.importables.Importable` from
        the dataset or the default value if no such instances are found.

        """

class DataSet(RODataSet):
    """
    An :py:mod:`abc` that represents a mutable set of
    :py:class:`~.importtools.importables.Importable` instances.

    """

    @abc.abstractmethod
    def add(self, importable):
        """
        Add the :py:class:`~.importtools.importables.Importable` in the dataset
        or replaces an existing and equal one.

        """

    @abc.abstractmethod
    def pop(self, importable):
        """
        Remove and return an equal
        :py:class:`~.importtools.importables.Importable` from the dataset.

        """

class MemoryDataSet(dict, DataSet):
    """
    A simple in-memory :py:class:`dict`-based :py:class:`DataSet`
    implementation.

    Since this particular :py:class:`DataSet` implementation only uses the
    :py:meth:`__hash__` and :py:meth:`__cmp__` methods of an
    :py:class:`~.importtools.importables.Importable` :py:mod:`abc` we can
    exemplify this class functionality using :py:class:`tuple` objects:

    >>> from importtools import MemoryDataSet
    >>> item1 = (1, )
    >>> item2 = (1, )
    >>> item3 = (2, )

    At first, a newly created :py:class:`MemoryDataSet` instance has no
    elements:

    >>> mds = MemoryDataSet()
    >>> list(mds)
    []

    After creation, it can be populated and the elements in the dataset can be
    retrieved using other equal elements.

    >>> mds = MemoryDataSet()
    >>> mds.add(item1)
    >>> mds.get(item1) == item1
    True
    >>> mds.get(item1) is item1
    True
    >>> mds.get(item2) == item2
    True
    >>> mds.get(item2) is item1
    True

    Trying to get an inexistent item should return the default value or
    :py:class:`None`:

    >>> mds = MemoryDataSet()
    >>> mds.add(item1)
    >>> mds.get(item3) is None
    True
    >>> mds.get(item3, 'default')
    'default'

    Initial data can be passed when constructing instances:

    >>> mds = MemoryDataSet((item1, item2, item3))

    """
    def __init__(self, data_loader=None):
        data_iter = tuple()
        if data_loader is not None:
            data_iter = iter(data_loader)
        mapping = ((i, i) for i in data_iter)
        super(MemoryDataSet, self).__init__(mapping)

    def add(self, importable):
        self[importable] = importable

class DiffDataSet(MemoryDataSet):
    """
    A :py:class:`DataSet` implementation that remembers all the changes,
    additions and removals done to it.

    Using a :py:class:`DiffDataSet` as the destination of the import algorithm
    allows optimal persistence of the changes by grouping them in a way suited
    for batch processing.

    In order to exemplify this class functionality we start with a
    instance that we prepopulate:

    An :py:class:`~.importtools.importables.Importable` mock is needed to
    populate the datasets:

    >>> from importtools import (
    ...     Importable, MockImportable, MemoryDataSet, DiffDataSet
    ... )
    >>> dds = DiffDataSet([
    ...     MockImportable(0),
    ...     MockImportable(1),
    ...     MockImportable(2),
    ...     MockImportable(3)
    ... ])

    >>> sorted(dds)
    [<(0,) IMPORTED>, <(1,) IMPORTED>, <(2,) IMPORTED>, <(3,) IMPORTED>]

    Fetching existing elements from ``dds`` should work correctly:

    >>> dds.get(MockImportable(1)) == MockImportable(1)
    True

    New :py:class:`~.importtools.importables.Importable` can still added in
    ``dds``:

    >>> dds.add(MockImportable(4))
    >>> dds.add(MockImportable(5))
    >>> sorted(dds)
    ... # doctest: +NORMALIZE_WHITESPACE
    [<(0,) IMPORTED>, <(1,) IMPORTED>, <(2,) IMPORTED>,
     <(3,) IMPORTED>, <(4,) IMPORTED>, <(5,) IMPORTED>]

    Check if the new elements were added:

    >>> dds.get(MockImportable(4))
    <(4,) IMPORTED>
    >>> dds.get(MockImportable(5))
    <(5,) IMPORTED>

    A lookup with an inexistent element should return the default value:

    >>> dds.get(MockImportable(100), 'default')
    'default'

    We should be able to ``pop`` both original and recently added elements:

    >>> dds.pop(MockImportable(1))
    <(1,) IMPORTED>
    >>> dds.pop(MockImportable(4))
    <(4,) IMPORTED>
    >>> sorted(dds)
    [<(0,) IMPORTED>, <(2,) IMPORTED>, <(3,) IMPORTED>, <(5,) IMPORTED>]

    And everything should have been registered:

    >>> sorted(dds.get_added())
    [<(5,) IMPORTED>]
    >>> sorted(dds.get_removed())
    [<(1,) IMPORTED>]
    >>> sorted(dds.get_changed())
    []

    Adding and removing the same
    :py:class:`~.importtools.importables.Importable` shouldn't change anything:

    >>> dds.add(MockImportable(100))
    >>> dds.pop(MockImportable(100))
    <(100,) IMPORTED>
    >>> sorted(dds.get_added())
    [<(5,) IMPORTED>]
    >>> sorted(dds.get_removed())
    [<(1,) IMPORTED>]
    >>> sorted(dds.get_changed())
    []

    But removing an original :py:class:`~.importtools.importables.Importable`
    and adding a new one should be interpreted as a change:

    >>> dds.pop(MockImportable(2))
    <(2,) IMPORTED>
    >>> dds.add(MockImportable(2))
    >>> sorted(dds.get_added())
    [<(5,) IMPORTED>]
    >>> sorted(dds.get_removed())
    [<(1,) IMPORTED>]
    >>> sorted(dds.get_changed())
    [<(2,) IMPORTED>]

    If you change one of the existing elements it will be marked as changed:

    >>> dds.get(MockImportable(3)).register_change()
    >>> sorted(dds.get_added())
    [<(5,) IMPORTED>]
    >>> sorted(dds.get_removed())
    [<(1,) IMPORTED>]
    >>> sorted(dds.get_changed())
    [<(2,) IMPORTED>, <(3,) IMPORTED>]

    Even if the elements are accessed by iterating on over the dataset the
    changes are still tracked:

    >>> list(dds)[0].register_change()
    >>> sorted(dds.get_added())
    [<(5,) IMPORTED>]
    >>> sorted(dds.get_removed())
    [<(1,) IMPORTED>]
    >>> sorted(dds.get_changed())
    [<(0,) IMPORTED>, <(2,) IMPORTED>, <(3,) IMPORTED>]

    When printing the object a list of all the added, removed and created
    elements will be shown, along with the total number for each operation:

    >>> print dds
    To be added: 1
    <(5,) IMPORTED>
    To be removed: 1
    <(1,) IMPORTED>
    To be changed: 3
    <(0,) IMPORTED>
    <(2,) IMPORTED>
    <(3,) IMPORTED>

    >>> dds._added = MemoryDataSet()
    >>> dds._removed = MemoryDataSet()
    >>> print dds
    <BLANKLINE>
    To be changed: 3
    <(0,) IMPORTED>
    <(2,) IMPORTED>
    <(3,) IMPORTED>

    >>> dds._changed = MemoryDataSet()
    >>> print dds
    No elements were affected.

    """
    def __init__(self, data_loader=None):
        self._added = set()
        self._removed = set()
        self._changed = set()
        # XXX for some reason dict subclass methods can't be hashed
        self.rc = lambda x: self.register_change(x)
        super(DiffDataSet, self).__init__(
            element.register_listener(self.rc) for element in data_loader
        )

    def __str__(self):
        result = ''
        added_count = len(self._added)
        if added_count > 0:
            result += 'To be added: %s\n' % added_count
            result += '\n'.join(sorted(repr(x) for x in self.get_added()))
        removed_count = len(self._removed)
        if removed_count > 0:
            result += '\nTo be removed: %s\n' % removed_count
            result += '\n'.join(sorted(repr(x) for x in self.get_removed()))
        changed_count = len(self._changed)
        if changed_count > 0:
            result += '\nTo be changed: %s\n' % changed_count
            result += '\n'.join(sorted(repr(x) for x in self.get_changed()))
        if (added_count + removed_count + changed_count) == 0:
            result += 'No elements were affected.'
        return result

    def add(self, importable):
        if importable in self._removed:
            self._removed.discard(importable)
            self._changed.add(importable)
        else:
            self._added.add(importable)
        super(DiffDataSet, self).add(importable.register_listener(self.rc))

    def pop(self, importable):
        if importable in self._added:
            self._added.discard(importable)
        else:
            self._removed.add(importable)
        return super(DiffDataSet, self).pop(importable)

    def register_change(self, importable):
        """Mark an element in the current dataset as changed.

        This method is registered as a listener for changes in all the
        elements of the wrapped :py:class:`DataSet`.

        """
        assert importable in self
        self._changed.add(importable)

    def get_added(self):
        """A set of all added elements in the wrapped dataset."""
        return self._added

    def get_removed(self):
        """A set of all removed elements in the wrapped dataset."""
        return self._removed

    def get_changed(self):
        """A set of all changed elements in the wrapped dataset."""
        return self._changed
<|endoftext|>