diff --git a/zipline/protocol_utils.py b/zipline/protocol_utils.py index 4581d732..9553f843 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, key): + return self.__dict__.iterkeys() + + def __len__(self, key): + return len(self.__dict__) + def keys(self): return self.__dict__.keys() @@ -86,8 +97,118 @@ 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, key): + 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) + + # Custom Methods + # -------------- + + 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