@@ -282,48 +282,34 @@ def _identify_start_end(self):
282282 """
283283 Determine the start and end steps.
284284
285- Uses explicit ``@step(start=True)`` / ``@step(end=True)`` annotations
286- if present. Falls back to looking for steps named ``"start"`` /
287- ``"end"`` for backward compatibility.
288-
289- Sets ``self.start_step`` and ``self.end_step`` to step name strings,
290- or ``None`` if the graph is malformed (validated later by lint).
291- Also assigns the ``"start"`` and ``"end"`` node types.
285+ Uses explicit @step(start=True) / @step(end=True) annotations if present.
286+ Falls back to steps named "start" / "end" for backward compatibility.
287+ Sets self.start_step and self.end_step to step name strings, or None
288+ if the graph is malformed (validated later by lint).
292289 """
293- # 1. Look for explicit annotations
294- annotated_start = [
295- name
296- for name , node in self .nodes .items ()
297- if node .is_start_step and not name .startswith ("_" )
298- ]
299- annotated_end = [
300- name
301- for name , node in self .nodes .items ()
302- if node .is_end_step and not name .startswith ("_" )
303- ]
304-
305- # 2. Determine start step (annotation first, then name fallback)
306- if len (annotated_start ) == 1 :
307- self .start_step = annotated_start [0 ]
308- elif len (annotated_start ) == 0 :
309- self .start_step = "start" if "start" in self .nodes else None
310- else :
311- self .start_step = None # Multiple annotated — lint will catch
312290
313- # 3. Determine end step (annotation first, then name fallback)
314- if len (annotated_end ) == 1 :
315- self .end_step = annotated_end [0 ]
316- elif len (annotated_end ) == 0 :
317- self .end_step = "end" if "end" in self .nodes else None
318- else :
319- self .end_step = None # Multiple annotated — lint will catch
291+ def _resolve (attr , fallback_name ):
292+ """Find the unique annotated step, or fall back to a named step."""
293+ annotated = [
294+ name
295+ for name , node in self .nodes .items ()
296+ if getattr (node , attr ) and not name .startswith ("_" )
297+ ]
298+ if len (annotated ) == 1 :
299+ return annotated [0 ]
300+ if len (annotated ) == 0 :
301+ return fallback_name if fallback_name in self .nodes else None
302+ return None # Multiple annotated — lint will catch
320303
321- # 4. Assign types based on identified start/end.
322- # Only upgrade "linear" → "start" for the entry point; do NOT override
304+ self .start_step = _resolve ("is_start_step" , "start" )
305+ self .end_step = _resolve ("is_end_step" , "end" )
306+
307+ # Assign node types for graph traversal.
308+ # Only upgrade "linear" -> "start" for the entry point; do NOT override
323309 # "split", "foreach", etc. since those types are needed for
324310 # split/join balance checking.
325311 if self .start_step and self .start_step == self .end_step :
326- # Single-step flow: terminal node that is also the entry point
312+ # Single-step flow: terminal node that is also the entry point.
327313 self .nodes [self .start_step ].type = "end"
328314 else :
329315 if self .start_step :
@@ -568,6 +554,18 @@ def populate_block(start_name, end_name):
568554 break
569555 return resulting_list
570556
557+ if self .start_step is None or self .end_step is None :
558+ missing = []
559+ if self .start_step is None :
560+ missing .append ("start" )
561+ if self .end_step is None :
562+ missing .append ("end" )
563+ raise ValueError (
564+ "Cannot compute graph structure: no %s step identified. "
565+ "Use @step(start=True)/@step(end=True) or name your steps "
566+ "'start'/'end'." % " or " .join (missing )
567+ )
568+
571569 if self .start_step == self .end_step :
572570 # Single-step flow
573571 graph_structure = []
0 commit comments