Python is powerful enough for this:
with Node('root', parent=None) as root:
Node.current.is_red = True
Node('child')
child = Node('named child')
with Node('complex child'):
@Node.event_handler
def on_walk(self, distance_km):
if distance_km > 10:
print(f'{self.name!r} is tired!')
Node('grand')
Node('grand')
Node('new child')
pprint(root)
assert root.is_red
root.children[2].on_walk(42)
output:
Node(name='root',
children=[Node(name='child', children=[]),
Node(name='named child', children=[]),
Node(name='complex child',
children=[Node(name='grand', children=[]),
Node(name='grand', children=[])]),
Node(name='new child', children=[])])
'complex child' is tired!
There’s some magic involved, but its effect is simple to explain:
Node.current
depends on the enclosing with
.
Node.__init__
takes a parent argument. If you omit it, Node.current
is used.
Node.event_handler
defines a method on Node.current
.
IMO, the only thing that’s substantially less ergonomic than the OP example is that the attribute is defined with Node.current.attr = ...
rather than simply attr = ...
.
Defining attributes by variable assignment wouldn’t be right. For example for i in range(n)
sets i
, but you probably don’t want it set to an i
attribute. (And there’s a lot of other cases of variable assignment, not all of which are as clear-cut as for
or attr=...
)
This could be simplified to CURRENT.attr = ...
(with a global CURRENT
), at the cost of significantly more magic in the implementation. Not worth it, IMO. (See flask.g
– popular, but full of sharp edge cases.)
(Of course in this particular case you could use root.attr = ...
, since that node has a name.)
If I was making a library like this I’d like to be a bit more explicit and avoid magic, and make you always name nodes when you use the with
statement – but that doesn’t work if you reuse names across levels:
with Node('root') as current:
with current.add(Node('child')) as current:
...
with current.add(Node('child')) as current:
...
current.is_red = True # oops! current now refers to the child!
So, I can’t find a way to avoid context – like wxWize from Andreas’ example. I used contextvars
, and tried to ensure the magic remains contained to the Node
class…
The other bit of magic is using a metaclass, which is needed to have Node.current
rather than Node.get_current()
, and also ensures instance namespace isn’t polluted unnecessarily (there’s no root.current
or child.event_handler
).
click for the magic
from dataclasses import dataclass
from functools import partial
from pprint import pprint
import contextvars
_USE_CURRENT_ROOT = object()
class NodeMeta(type):
@property
def current(cls):
try:
return cls._root_context.get()
except LookupError:
raise LookupError(
f"no current {cls.__name__!r}, use a with statement")
def event_handler(cls, func=None, name=None, node=None):
# this gimmick is not necessary for the main idea
if func is None:
return partial(cls.event_handler, name=name, node=node)
if node is None:
node = cls.current
try:
descr_get = func.__get__
except AttributeError:
pass
else:
func = func.__get__(node, type(node))
return setattr(node, name or func.__name__, func)
@dataclass
class Node(metaclass=NodeMeta):
name: str
children: list
_root_context = contextvars.ContextVar('_root_context')
def __init__(self, name, *, parent=_USE_CURRENT_ROOT):
self.name = name
self.children = []
if parent is _USE_CURRENT_ROOT:
parent = type(self).current
if parent is not None:
parent.children.append(self)
self._reset_tokens = []
def __enter__(self):
self._reset_tokens.append(self._root_context.set(self))
return self
def __exit__(self, *exc_info):
self._root_context.reset(self._reset_tokens.pop())