ipstruct.py 12 KB

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