@@ -93,6 +93,16 @@ def get_labels(
9393 return [graph_name , node_name , field_name ]
9494
9595 def _wrap_field (self , node_name , func ):
96+ """Creates a wrapped field function that
97+ observes time spent in the field.
98+
99+ Args:
100+ node_name: Name of the node.
101+ func: Field function to wrap.
102+
103+ Returns:
104+ Wrapped field function.
105+ """
96106 observe = self ._observe_fields (node_name )
97107 wrapper = self .field_wrapper (observe , func )
98108 if _do_pass_context (func ):
@@ -104,6 +114,16 @@ def _wrap_field(self, node_name, func):
104114 return wrapper
105115
106116 def _wrap_link (self , node_name , link_name , func ):
117+ """Creates a wrapped link function that observes time spent in the link.
118+
119+ Args:
120+ node_name: Name of the node.
121+ link_name: Name of the link.
122+ func: Link function to wrap.
123+
124+ Returns:
125+ Wrapped link function.
126+ """
107127 observe = self ._observe_fields (node_name )
108128 wrapper = self .link_wrapper (observe , func )
109129
@@ -120,6 +140,16 @@ def _wrap_link(self, node_name, link_name, func):
120140 return wrapper
121141
122142 def _wrap_subquery (self , node_name , subquery ):
143+ """Creates a wrapped subquery function that observes
144+ time spent in the subquery.
145+
146+ Args:
147+ node_name: Name of the node.
148+ subquery: Subquery function to wrap.
149+
150+ Returns:
151+ Wrapped subquery function.
152+ """
123153 observe = self ._observe_fields (node_name )
124154 wrapper = self .subquery_wrapper (observe , subquery )
125155 wrapper = _subquery_field_names (wrapper )
@@ -134,22 +164,36 @@ def visit_node(self, obj):
134164 self ._node = None
135165
136166 def visit_field (self , obj ):
167+ """Wrap field function to observe time spent in the field.
168+
169+ If field function if already wrapped and cached in self._wrappers,
170+ it will be used instead of creating a new wrapper.
171+
172+ Field.func is replaced with the wrapped function.
173+
174+ Args:
175+ obj: Field object to wrap.
176+
177+ Returns:
178+ Field with wrapped function.
179+ """
137180 obj = super ().visit_field (obj )
138181 node_name = self .root_name if self ._node is None else self ._node .name
139182 if isinstance (obj .func , CheckedExpr ):
140183 func = obj .func .__subquery__
141184 else :
142185 func = obj .func
143186
144- wrapper = self ._wrappers .get (func )
187+ wrapper_key = (node_name , func )
188+ wrapper = self ._wrappers .get (wrapper_key )
145189 if wrapper is None :
146190 if isinstance (obj .func , CheckedExpr ):
147- wrapper = self ._wrappers [func ] = self ._wrap_subquery (
191+ wrapper = self ._wrappers [wrapper_key ] = self ._wrap_subquery (
148192 node_name ,
149193 func ,
150194 )
151195 else :
152- wrapper = self ._wrappers [func ] = self ._wrap_field (
196+ wrapper = self ._wrappers [wrapper_key ] = self ._wrap_field (
153197 node_name ,
154198 func ,
155199 )
0 commit comments