visitor.py 5.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158
  1. """Generic visitor pattern implementation for Python objects."""
  2. import enum
  3. import weakref
  4. class Visitor(object):
  5. defaultStop = False
  6. _visitors = {
  7. # By default we skip visiting weak references to avoid recursion
  8. # issues. Users can override this by registering a visit
  9. # function for weakref.ProxyType.
  10. weakref.ProxyType: {None: lambda self, obj, *args, **kwargs: False}
  11. }
  12. @classmethod
  13. def _register(celf, clazzes_attrs):
  14. assert celf != Visitor, "Subclass Visitor instead."
  15. if "_visitors" not in celf.__dict__:
  16. celf._visitors = {}
  17. def wrapper(method):
  18. assert method.__name__ == "visit"
  19. for clazzes, attrs in clazzes_attrs:
  20. if type(clazzes) != tuple:
  21. clazzes = (clazzes,)
  22. if type(attrs) == str:
  23. attrs = (attrs,)
  24. for clazz in clazzes:
  25. _visitors = celf._visitors.setdefault(clazz, {})
  26. for attr in attrs:
  27. assert attr not in _visitors, (
  28. "Oops, class '%s' has visitor function for '%s' defined already."
  29. % (clazz.__name__, attr)
  30. )
  31. _visitors[attr] = method
  32. return None
  33. return wrapper
  34. @classmethod
  35. def register(celf, clazzes):
  36. if type(clazzes) != tuple:
  37. clazzes = (clazzes,)
  38. return celf._register([(clazzes, (None,))])
  39. @classmethod
  40. def register_attr(celf, clazzes, attrs):
  41. clazzes_attrs = []
  42. if type(clazzes) != tuple:
  43. clazzes = (clazzes,)
  44. if type(attrs) == str:
  45. attrs = (attrs,)
  46. for clazz in clazzes:
  47. clazzes_attrs.append((clazz, attrs))
  48. return celf._register(clazzes_attrs)
  49. @classmethod
  50. def register_attrs(celf, clazzes_attrs):
  51. return celf._register(clazzes_attrs)
  52. @classmethod
  53. def _visitorsFor(celf, thing, _default={}):
  54. typ = type(thing)
  55. for celf in celf.mro():
  56. _visitors = getattr(celf, "_visitors", None)
  57. if _visitors is None:
  58. break
  59. for base in typ.mro():
  60. m = celf._visitors.get(base, None)
  61. if m is not None:
  62. return m
  63. return _default
  64. def visitObject(self, obj, *args, **kwargs):
  65. """Called to visit an object. This function loops over all non-private
  66. attributes of the objects and calls any user-registered (via
  67. ``@register_attr()`` or ``@register_attrs()``) ``visit()`` functions.
  68. The visitor will proceed to call ``self.visitAttr()``, unless there is a
  69. user-registered visit function and:
  70. * It returns ``False``; or
  71. * It returns ``None`` (or doesn't return anything) and
  72. ``visitor.defaultStop`` is ``True`` (non-default).
  73. """
  74. keys = sorted(vars(obj).keys())
  75. _visitors = self._visitorsFor(obj)
  76. defaultVisitor = _visitors.get("*", None)
  77. for key in keys:
  78. if key[0] == "_":
  79. continue
  80. value = getattr(obj, key)
  81. visitorFunc = _visitors.get(key, defaultVisitor)
  82. if visitorFunc is not None:
  83. ret = visitorFunc(self, obj, key, value, *args, **kwargs)
  84. if ret == False or (ret is None and self.defaultStop):
  85. continue
  86. self.visitAttr(obj, key, value, *args, **kwargs)
  87. def visitAttr(self, obj, attr, value, *args, **kwargs):
  88. """Called to visit an attribute of an object."""
  89. self.visit(value, *args, **kwargs)
  90. def visitList(self, obj, *args, **kwargs):
  91. """Called to visit any value that is a list."""
  92. for value in obj:
  93. self.visit(value, *args, **kwargs)
  94. def visitDict(self, obj, *args, **kwargs):
  95. """Called to visit any value that is a dictionary."""
  96. for value in obj.values():
  97. self.visit(value, *args, **kwargs)
  98. def visitLeaf(self, obj, *args, **kwargs):
  99. """Called to visit any value that is not an object, list,
  100. or dictionary."""
  101. pass
  102. def visit(self, obj, *args, **kwargs):
  103. """This is the main entry to the visitor. The visitor will visit object
  104. ``obj``.
  105. The visitor will first determine if there is a registered (via
  106. ``@register()``) visit function for the type of object. If there is, it
  107. will be called, and ``(visitor, obj, *args, **kwargs)`` will be passed
  108. to the user visit function.
  109. The visitor will not recurse if there is a user-registered visit
  110. function and:
  111. * It returns ``False``; or
  112. * It returns ``None`` (or doesn't return anything) and
  113. ``visitor.defaultStop`` is ``True`` (non-default)
  114. Otherwise, the visitor will proceed to dispatch to one of
  115. ``self.visitObject()``, ``self.visitList()``, ``self.visitDict()``, or
  116. ``self.visitLeaf()`` (any of which can be overriden in a subclass).
  117. """
  118. visitorFunc = self._visitorsFor(obj).get(None, None)
  119. if visitorFunc is not None:
  120. ret = visitorFunc(self, obj, *args, **kwargs)
  121. if ret == False or (ret is None and self.defaultStop):
  122. return
  123. if hasattr(obj, "__dict__") and not isinstance(obj, enum.Enum):
  124. self.visitObject(obj, *args, **kwargs)
  125. elif isinstance(obj, list):
  126. self.visitList(obj, *args, **kwargs)
  127. elif isinstance(obj, dict):
  128. self.visitDict(obj, *args, **kwargs)
  129. else:
  130. self.visitLeaf(obj, *args, **kwargs)