_helpers.py 2.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182
  1. from functools import wraps
  2. import six
  3. from pyrsistent._pmap import PMap, pmap
  4. from pyrsistent._pset import PSet, pset
  5. from pyrsistent._pvector import PVector, pvector
  6. def freeze(o):
  7. """
  8. Recursively convert simple Python containers into pyrsistent versions
  9. of those containers.
  10. - list is converted to pvector, recursively
  11. - dict is converted to pmap, recursively on values (but not keys)
  12. - set is converted to pset, but not recursively
  13. - tuple is converted to tuple, recursively.
  14. Sets and dict keys are not recursively frozen because they do not contain
  15. mutable data by convention. The main exception to this rule is that
  16. dict keys and set elements are often instances of mutable objects that
  17. support hash-by-id, which this function can't convert anyway.
  18. >>> freeze(set([1, 2]))
  19. pset([1, 2])
  20. >>> freeze([1, {'a': 3}])
  21. pvector([1, pmap({'a': 3})])
  22. >>> freeze((1, []))
  23. (1, pvector([]))
  24. """
  25. typ = type(o)
  26. if typ is dict:
  27. return pmap(dict((k, freeze(v)) for k, v in six.iteritems(o)))
  28. if typ is list:
  29. return pvector(map(freeze, o))
  30. if typ is tuple:
  31. return tuple(map(freeze, o))
  32. if typ is set:
  33. return pset(o)
  34. return o
  35. def thaw(o):
  36. """
  37. Recursively convert pyrsistent containers into simple Python containers.
  38. - pvector is converted to list, recursively
  39. - pmap is converted to dict, recursively on values (but not keys)
  40. - pset is converted to set, but not recursively
  41. - tuple is converted to tuple, recursively.
  42. >>> from pyrsistent import s, m, v
  43. >>> thaw(s(1, 2))
  44. {1, 2}
  45. >>> thaw(v(1, m(a=3)))
  46. [1, {'a': 3}]
  47. >>> thaw((1, v()))
  48. (1, [])
  49. """
  50. if isinstance(o, PVector):
  51. return list(map(thaw, o))
  52. if isinstance(o, PMap):
  53. return dict((k, thaw(v)) for k, v in o.iteritems())
  54. if isinstance(o, PSet):
  55. return set(o)
  56. if type(o) is tuple:
  57. return tuple(map(thaw, o))
  58. return o
  59. def mutant(fn):
  60. """
  61. Convenience decorator to isolate mutation to within the decorated function (with respect
  62. to the input arguments).
  63. All arguments to the decorated function will be frozen so that they are guaranteed not to change.
  64. The return value is also frozen.
  65. """
  66. @wraps(fn)
  67. def inner_f(*args, **kwargs):
  68. return freeze(fn(*[freeze(e) for e in args], **dict(freeze(item) for item in kwargs.items())))
  69. return inner_f