Source code for monty.collections

# coding: utf-8
from __future__ import absolute_import, print_function, unicode_literals,\

import collections

__author__ = 'Shyue Ping Ong'
__copyright__ = "Copyright 2014, The Materials Virtual Lab"
__version__ = '0.1'
__maintainer__ = 'Shyue Ping Ong'
__email__ = ''
__date__ = '1/24/14'

[docs]def tree(): """ A tree object, which is effectively a recursive defaultdict that adds tree as members. Usage: x = tree() x['a']['b']['c'] = 1 Returns: A tree. """ return collections.defaultdict(tree)
[docs]def as_set(obj): """ Convert obj into a set, returns None if obj is None. >>> assert as_set(None) is None and as_set(1) == set([1]) and as_set(range(1,3)) == set([1, 2]) """ if obj is None or isinstance(obj, collections.Set): return obj if not isinstance(obj, collections.Iterable): return set((obj,)) else: return set(obj)
[docs]class frozendict(dict): """ A dictionary that does not permit changes. The naming violates PEP8 to be consistent with standard Python's "frozenset" naming. """ def __init__(self, *args, **kwargs): dict.__init__(self, *args, **kwargs) def __setitem__(self, key, val): raise KeyError("Cannot overwrite existing key: %s" % str(key))
[docs] def update(self, *args, **kwargs): raise KeyError("Cannot update a %s" % self.__class__.__name__)
[docs]class Namespace(dict): """A dictionary that does not permit to redefine its keys.""" def __init__(self, *args, **kwargs): self.update(*args, **kwargs) def __setitem__(self, key, val): if key in self: raise KeyError("Cannot overwrite existent key: %s" % str(key)) dict.__setitem__(self, key, val)
[docs] def update(self, *args, **kwargs): for k, v in dict(*args, **kwargs).items(): self[k] = v
[docs]class AttrDict(dict): """ Allows to access dict keys as in addition to the traditional way obj['foo']" Example: >>> d = AttrDict(foo=1, bar=2) >>> assert d["foo"] == >>> = "hello" >>> assert == "hello" """ def __init__(self, *args, **kwargs): super(AttrDict, self).__init__(*args, **kwargs) self.__dict__ = self
[docs] def copy(self): newd = super(AttrDict, self).copy() return self.__class__(**newd)
[docs]class FrozenAttrDict(frozendict): """ A dictionary that: * does not permit changes. * Allows to access dict keys as in addition to the traditional way obj['foo'] """ def __init__(self, *args, **kwargs): super(FrozenAttrDict, self).__init__(*args, **kwargs) def __getattribute__(self, name): try: return super(FrozenAttrDict, self).__getattribute__(name) except AttributeError: try: return self[name] except KeyError as exc: raise AttributeError(str(exc)) def __setattr__(self, name, value): raise KeyError("You cannot modify attribute %s of %s" % (name, self.__class__.__name__))
[docs]class MongoDict(object): """ This dict-like object allows one to access the entries in a nested dict as attributes. Entries (attributes) cannot be modified. It also provides Ipython tab completion hence this object is particularly useful if you need to analyze a nested dict interactively (e.g. documents extracted from a MongoDB database). >>> m = MongoDict({'a': {'b': 1}, 'x': 2}) >>> assert m.a.b == 1 and m.x == 2 >>> assert "a" in m and "b" in m.a >>> m["a"] {'b': 1} .. note:: Cannot inherit from ABC collections.Mapping because otherwise dict.keys and dict.items will pollute the namespace. e.g MongoDict({"keys": 1}).keys would be the ABC dict method. """ def __init__(self, *args, **kwargs): self.__dict__["_mongo_dict_"] = dict(*args, **kwargs) def __repr__(self): return str(self) def __str__(self): return str(self._mongo_dict_) def __setattr__(self, name, value): raise NotImplementedError("You cannot modify attribute %s of %s" % (name, self.__class__.__name__)) def __getattribute__(self, name): try: return super(MongoDict, self).__getattribute__(name) except: #raise try: a = self._mongo_dict_[name] if isinstance(a, collections.Mapping): a = self.__class__(a) return a except Exception as exc: raise AttributeError(str(exc)) def __getitem__(self, slice): return self._mongo_dict_.__getitem__(slice) def __iter__(self): return iter(self._mongo_dict_) def __len__(self): return len(self._mongo_dict_) def __dir__(self): """For Ipython tab completion. See""" return sorted(list(k for k in self._mongo_dict_ if not callable(k)))
[docs]def dict2namedtuple(*args, **kwargs): """ Helper function to create a :class:`namedtuple` from a dictionary. Example: >>> t = dict2namedtuple(foo=1, bar="hello") >>> assert == 1 and == "hello" >>> t = dict2namedtuple([("foo", 1), ("bar", "hello")]) >>> assert t[0] == and t[1] == .. warning:: - The order of the items in the namedtuple is not deterministic if kwargs are used. namedtuples, however, should always be accessed by attribute hence this behaviour should not represent a serious problem. - Don't use this function in code in which memory and performance are crucial since a dict is needed to instantiate the tuple! """ d = collections.OrderedDict(*args) d.update(**kwargs) return collections.namedtuple(typename="dict2namedtuple", field_names=list(d.keys()))(**d)