diff --git a/zipline/protocol_utils.py b/zipline/protocol_utils.py index 4581d732..8d21b1a0 100644 --- a/zipline/protocol_utils.py +++ b/zipline/protocol_utils.py @@ -1,6 +1,8 @@ import copy import pandas from ctypes import Structure, c_ubyte +from collections import MutableMapping +from itertools import izip def Enum(*options): """ @@ -28,7 +30,7 @@ def FrameExceptionFactory(name): return InvalidFrame -class namedict(object): +class namedict(MutableMapping): """ Namedicts are dict like objects that have fields accessible by attribute lookup @@ -60,7 +62,16 @@ class namedict(object): def __getitem__(self, key): return self.__dict__[key] - + + def __delitem__(self, key): + del self.__dict__[key] + + def __iter__(self): + return self.__dict__.iterkeys() + + def __len__(self): + return len(self.__dict__) + def keys(self): return self.__dict__.keys() @@ -86,8 +97,124 @@ class namedict(object): def has_attr(self, name): return self.__dict__.has_key(name) - + def as_series(self): s = pandas.Series(self.__dict__) s.name = self.sid return s + +class ndict(MutableMapping): + """ + Xtreme Namedicts 2.0 + + Ndicts are dict like objects that have fields accessible by attribute + lookup as well as being indexable and iterable. Done right + this time. + """ + + def __init__(self, dct=None): + self.__internal = dict() + self.cls = frozenset(dir(self)) + + if dct: + self.__internal.update(dct) + + # Abstact Overloads + # ----------------- + + def __setitem__(self, key, value): + """ + Required for use by pymongo as_class parameter to find. + """ + if key == '_id': + self.__internal['id'] = value + else: + self.__internal[key] = value + + + def __getattr__(self, key): + if key in self.cls: + return self.__dict__[key] + else: + return self.__internal[key] + + def __getitem__(self, key): + return self.__internal[key] + + def __delitem__(self, key): + del self.__internal[key] + + def __iter__(self): + return self.__internal.iterkeys() + + def __len__(self): + return len(self.__internal) + + # Compatability with namedicts + # ---------------------------- + + # for compat, not the Python way to do things though... + # Deprecated, use builtin ``del`` operator. + delete = __delitem__ + + def has_attr(self, key): + """ + Deprecated, use builtin ``in`` operator. + """ + return self.__contains__(key) + + def has_key(self, key): + return self.__contains__(key) + + # Custom Methods + # -------------- + + def copy(self): + return ndict(copy.copy(self.__internal)) + + def as_dataframe(self): + """ + Return the representation as a Pandas dataframe. + """ + d = pandas.DataFrame(self.__internal) + return d + + def as_series(self): + """ + Return the representation as a Pandas time series. + """ + s = pandas.Series(self.__internal) + s.name = self.sid + return s + + def as_dict(self): + """ + Return the representation as a vanilla Python dict. + """ + # shallow copy is O(n) + return copy.copy(self.__internal) + + def merge(self, other_nd): + """ + Merge in place with another ndict. + """ + assert isinstance(other_nd, ndict) + self.__internal.update(other_nd.__internal) + + def __repr__(self): + return "namedict: " + str(self.__internal) + + # Faster dictionary comparison? + #def __eq__(self, other): + #assert isinstance(other, ndict) + + #keyeq = set(self.keys()) == set(other.keys()) + + #if not keyeq: + #return False + + #for i, j in izip(self.itervalues(), other.itervalues()): + #if i != j: + #return False + + #return True diff --git a/zipline/test/test_ndict.py b/zipline/test/test_ndict.py new file mode 100644 index 00000000..e2cb8a84 --- /dev/null +++ b/zipline/test/test_ndict.py @@ -0,0 +1,46 @@ +from zipline.protocol_utils import ndict, namedict + +def test_ndict(): + nd = ndict({}) + + # Properties + assert len(nd) == 0 + assert nd.keys() == [] + assert nd.values() == [] + assert list(nd.iteritems()) == [] + + # Accessors + nd['x'] = 1 + assert nd.x == 1 + assert nd['x'] == 1 + assert nd.get('y') == None + assert nd.get('y', 'fizzpop') == 'fizzpop' + assert nd.has_key('x') == True + assert nd.has_key('y') == False + + assert 'x' in nd + assert 'y' not in nd + + # Class isolation + assert '__init__' not in nd + assert '__iter__' not in nd + assert not nd.__dict__.has_key('x') + assert nd.get('__init__') is None + + # Comparison + nd2 = nd.copy() + assert id(nd2) != id(nd) + assert nd2 == nd + nd2['z'] = 3 + assert nd2 != nd + + class ndictlike(object): + x = 1 + + assert { 'x': 1 } == nd + assert ndictlike() != nd + + # Deletion + del nd['x'] + assert not nd.has_key('x') + assert nd.get('x') is None