Namedicts edges fixed, and awesomer ndict

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