1# encoding: utf-8
2"""A dict subclass that supports attribute style access.
3
4Authors:
5
6* Fernando Perez (original)
7* Brian Granger (refactoring to a dict subclass)
8"""
9
10#-----------------------------------------------------------------------------
11#  Copyright (C) 2008-2011  The IPython Development Team
12#
13#  Distributed under the terms of the BSD License.  The full license is in
14#  the file COPYING, distributed as part of this software.
15#-----------------------------------------------------------------------------
16
17#-----------------------------------------------------------------------------
18# Imports
19#-----------------------------------------------------------------------------
20
21__all__ = ['Struct']
22
23#-----------------------------------------------------------------------------
24# Code
25#-----------------------------------------------------------------------------
26
27
28class Struct(dict):
29    """A dict subclass with attribute style access.
30
31    This dict subclass has a a few extra features:
32
33    * Attribute style access.
34    * Protection of class members (like keys, items) when using attribute
35      style access.
36    * The ability to restrict assignment to only existing keys.
37    * Intelligent merging.
38    * Overloaded operators.
39    """
40    _allownew = True
41    def __init__(self, *args, **kw):
42        """Initialize with a dictionary, another Struct, or data.
43
44        Parameters
45        ----------
46        args : dict, Struct
47            Initialize with one dict or Struct
48        kw : dict
49            Initialize with key, value pairs.
50
51        Examples
52        --------
53
54        >>> s = Struct(a=10,b=30)
55        >>> s.a
56        10
57        >>> s.b
58        30
59        >>> s2 = Struct(s,c=30)
60        >>> sorted(s2.keys())
61        ['a', 'b', 'c']
62        """
63        object.__setattr__(self, '_allownew', True)
64        dict.__init__(self, *args, **kw)
65
66    def __setitem__(self, key, value):
67        """Set an item with check for allownew.
68
69        Examples
70        --------
71
72        >>> s = Struct()
73        >>> s['a'] = 10
74        >>> s.allow_new_attr(False)
75        >>> s['a'] = 10
76        >>> s['a']
77        10
78        >>> try:
79        ...     s['b'] = 20
80        ... except KeyError:
81        ...     print('this is not allowed')
82        ...
83        this is not allowed
84        """
85        if not self._allownew and key not in self:
86            raise KeyError(
87                "can't create new attribute %s when allow_new_attr(False)" % key)
88        dict.__setitem__(self, key, value)
89
90    def __setattr__(self, key, value):
91        """Set an attr with protection of class members.
92
93        This calls :meth:`self.__setitem__` but convert :exc:`KeyError` to
94        :exc:`AttributeError`.
95
96        Examples
97        --------
98
99        >>> s = Struct()
100        >>> s.a = 10
101        >>> s.a
102        10
103        >>> try:
104        ...     s.get = 10
105        ... except AttributeError:
106        ...     print("you can't set a class member")
107        ...
108        you can't set a class member
109        """
110        # If key is an str it might be a class member or instance var
111        if isinstance(key, str):
112            # I can't simply call hasattr here because it calls getattr, which
113            # calls self.__getattr__, which returns True for keys in
114            # self._data.  But I only want keys in the class and in
115            # self.__dict__
116            if key in self.__dict__ or hasattr(Struct, key):
117                raise AttributeError(
118                    'attr %s is a protected member of class Struct.' % key
119                )
120        try:
121            self.__setitem__(key, value)
122        except KeyError as e:
123            raise AttributeError(e)
124
125    def __getattr__(self, key):
126        """Get an attr by calling :meth:`dict.__getitem__`.
127
128        Like :meth:`__setattr__`, this method converts :exc:`KeyError` to
129        :exc:`AttributeError`.
130
131        Examples
132        --------
133
134        >>> s = Struct(a=10)
135        >>> s.a
136        10
137        >>> type(s.get)
138        <... 'builtin_function_or_method'>
139        >>> try:
140        ...     s.b
141        ... except AttributeError:
142        ...     print("I don't have that key")
143        ...
144        I don't have that key
145        """
146        try:
147            result = self[key]
148        except KeyError:
149            raise AttributeError(key)
150        else:
151            return result
152
153    def __iadd__(self, other):
154        """s += s2 is a shorthand for s.merge(s2).
155
156        Examples
157        --------
158
159        >>> s = Struct(a=10,b=30)
160        >>> s2 = Struct(a=20,c=40)
161        >>> s += s2
162        >>> sorted(s.keys())
163        ['a', 'b', 'c']
164        """
165        self.merge(other)
166        return self
167
168    def __add__(self,other):
169        """s + s2 -> New Struct made from s.merge(s2).
170
171        Examples
172        --------
173
174        >>> s1 = Struct(a=10,b=30)
175        >>> s2 = Struct(a=20,c=40)
176        >>> s = s1 + s2
177        >>> sorted(s.keys())
178        ['a', 'b', 'c']
179        """
180        sout = self.copy()
181        sout.merge(other)
182        return sout
183
184    def __sub__(self,other):
185        """s1 - s2 -> remove keys in s2 from s1.
186
187        Examples
188        --------
189
190        >>> s1 = Struct(a=10,b=30)
191        >>> s2 = Struct(a=40)
192        >>> s = s1 - s2
193        >>> s
194        {'b': 30}
195        """
196        sout = self.copy()
197        sout -= other
198        return sout
199
200    def __isub__(self,other):
201        """Inplace remove keys from self that are in other.
202
203        Examples
204        --------
205
206        >>> s1 = Struct(a=10,b=30)
207        >>> s2 = Struct(a=40)
208        >>> s1 -= s2
209        >>> s1
210        {'b': 30}
211        """
212        for k in other.keys():
213            if k in self:
214                del self[k]
215        return self
216
217    def __dict_invert(self, data):
218        """Helper function for merge.
219
220        Takes a dictionary whose values are lists and returns a dict with
221        the elements of each list as keys and the original keys as values.
222        """
223        outdict = {}
224        for k,lst in data.items():
225            if isinstance(lst, str):
226                lst = lst.split()
227            for entry in lst:
228                outdict[entry] = k
229        return outdict
230
231    def dict(self):
232        return self
233
234    def copy(self):
235        """Return a copy as a Struct.
236
237        Examples
238        --------
239
240        >>> s = Struct(a=10,b=30)
241        >>> s2 = s.copy()
242        >>> type(s2) is Struct
243        True
244        """
245        return Struct(dict.copy(self))
246
247    def hasattr(self, key):
248        """hasattr function available as a method.
249
250        Implemented like has_key.
251
252        Examples
253        --------
254
255        >>> s = Struct(a=10)
256        >>> s.hasattr('a')
257        True
258        >>> s.hasattr('b')
259        False
260        >>> s.hasattr('get')
261        False
262        """
263        return key in self
264
265    def allow_new_attr(self, allow = True):
266        """Set whether new attributes can be created in this Struct.
267
268        This can be used to catch typos by verifying that the attribute user
269        tries to change already exists in this Struct.
270        """
271        object.__setattr__(self, '_allownew', allow)
272
273    def merge(self, __loc_data__=None, __conflict_solve=None, **kw):
274        """Merge two Structs with customizable conflict resolution.
275
276        This is similar to :meth:`update`, but much more flexible. First, a
277        dict is made from data+key=value pairs. When merging this dict with
278        the Struct S, the optional dictionary 'conflict' is used to decide
279        what to do.
280
281        If conflict is not given, the default behavior is to preserve any keys
282        with their current value (the opposite of the :meth:`update` method's
283        behavior).
284
285        Parameters
286        ----------
287        __loc_data : dict, Struct
288            The data to merge into self
289        __conflict_solve : dict
290            The conflict policy dict.  The keys are binary functions used to
291            resolve the conflict and the values are lists of strings naming
292            the keys the conflict resolution function applies to.  Instead of
293            a list of strings a space separated string can be used, like
294            'a b c'.
295        kw : dict
296            Additional key, value pairs to merge in
297
298        Notes
299        -----
300
301        The `__conflict_solve` dict is a dictionary of binary functions which will be used to
302        solve key conflicts.  Here is an example::
303
304            __conflict_solve = dict(
305                func1=['a','b','c'],
306                func2=['d','e']
307            )
308
309        In this case, the function :func:`func1` will be used to resolve
310        keys 'a', 'b' and 'c' and the function :func:`func2` will be used for
311        keys 'd' and 'e'.  This could also be written as::
312
313            __conflict_solve = dict(func1='a b c',func2='d e')
314
315        These functions will be called for each key they apply to with the
316        form::
317
318            func1(self['a'], other['a'])
319
320        The return value is used as the final merged value.
321
322        As a convenience, merge() provides five (the most commonly needed)
323        pre-defined policies: preserve, update, add, add_flip and add_s. The
324        easiest explanation is their implementation::
325
326            preserve = lambda old,new: old
327            update   = lambda old,new: new
328            add      = lambda old,new: old + new
329            add_flip = lambda old,new: new + old  # note change of order!
330            add_s    = lambda old,new: old + ' ' + new  # only for str!
331
332        You can use those four words (as strings) as keys instead
333        of defining them as functions, and the merge method will substitute
334        the appropriate functions for you.
335
336        For more complicated conflict resolution policies, you still need to
337        construct your own functions.
338
339        Examples
340        --------
341
342        This show the default policy:
343
344        >>> s = Struct(a=10,b=30)
345        >>> s2 = Struct(a=20,c=40)
346        >>> s.merge(s2)
347        >>> sorted(s.items())
348        [('a', 10), ('b', 30), ('c', 40)]
349
350        Now, show how to specify a conflict dict:
351
352        >>> s = Struct(a=10,b=30)
353        >>> s2 = Struct(a=20,b=40)
354        >>> conflict = {'update':'a','add':'b'}
355        >>> s.merge(s2,conflict)
356        >>> sorted(s.items())
357        [('a', 20), ('b', 70)]
358        """
359
360        data_dict = dict(__loc_data__,**kw)
361
362        # policies for conflict resolution: two argument functions which return
363        # the value that will go in the new struct
364        preserve = lambda old,new: old
365        update   = lambda old,new: new
366        add      = lambda old,new: old + new
367        add_flip = lambda old,new: new + old  # note change of order!
368        add_s    = lambda old,new: old + ' ' + new
369
370        # default policy is to keep current keys when there's a conflict
371        conflict_solve = dict.fromkeys(self, preserve)
372
373        # the conflict_solve dictionary is given by the user 'inverted': we
374        # need a name-function mapping, it comes as a function -> names
375        # dict. Make a local copy (b/c we'll make changes), replace user
376        # strings for the three builtin policies and invert it.
377        if __conflict_solve:
378            inv_conflict_solve_user = __conflict_solve.copy()
379            for name, func in [('preserve',preserve), ('update',update),
380                               ('add',add), ('add_flip',add_flip),
381                               ('add_s',add_s)]:
382                if name in inv_conflict_solve_user.keys():
383                    inv_conflict_solve_user[func] = inv_conflict_solve_user[name]
384                    del inv_conflict_solve_user[name]
385            conflict_solve.update(self.__dict_invert(inv_conflict_solve_user))
386        for key in data_dict:
387            if key not in self:
388                self[key] = data_dict[key]
389            else:
390                self[key] = conflict_solve[key](self[key],data_dict[key])
391
392