Merge branch 'master' of github.com:quantopian/zipline

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

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