Coverage for openhcs/debug/pickle_to_python.py: 0.0%

1#!/usr/bin/env python3 

2""" 

3Pickle to Python Converter - Convert OpenHCS debug pickle files to runnable Python scripts 

4""" 

5 

6import sys 

7import dill as pickle 

8import inspect 

9from pathlib import Path 

10from datetime import datetime 

11from collections import defaultdict 

12from enum import Enum 

13import dataclasses 

14from dataclasses import is_dataclass, fields 

15from typing import Callable 

16 

17from openhcs.core.config import GlobalPipelineConfig, PathPlanningConfig, VFSConfig, ZarrConfig 

18from openhcs.core.steps.function_step import FunctionStep 

19 

20def collect_imports_from_data(data_obj): 

21 """Extract function, enum, and dataclass imports by traversing data structure.""" 

22 function_imports = defaultdict(set) 

23 enum_imports = defaultdict(set) 

24 decorated_functions = set() 

25 

26 def register_imports(obj): 

27 if isinstance(obj, Enum): 

28 enum_imports[obj.__class__.__module__].add(obj.__class__.__name__) 

29 elif is_dataclass(obj): 

30 module = obj.__class__.__module__ 

31 name = obj.__class__.__name__ 

32 function_imports[module].add(name) 

33 [register_imports(getattr(obj, f.name)) for f in fields(obj) if getattr(obj, f.name) is not None] 

34 elif callable(obj): 

35 if _is_external_registered_function(obj): 

36 # Use the actual module path but under openhcs namespace 

37 original_module = obj.__module__ 

38 # Convert original module to openhcs namespace: cucim.skimage.filters -> openhcs.cucim.skimage.filters 

39 virtual_module = f'openhcs.{original_module}' 

40 function_imports[virtual_module].add(obj.__name__) 

41 decorated_functions.add(obj.__name__) 

42 else: 

43 function_imports[obj.__module__].add(obj.__name__) 

44 elif isinstance(obj, (list, tuple)): 

45 [register_imports(item) for item in obj] 

46 elif isinstance(obj, dict): 

47 [register_imports(value) for value in obj.values()] 

48 elif hasattr(obj, '__dict__') and obj.__dict__: 

49 [register_imports(value) for value in obj.__dict__.values()] 

50 

51 register_imports(data_obj) 

52 return function_imports, enum_imports, decorated_functions 

53 

54 

55def _is_external_registered_function(func): 

56 """Check if function is an external library function registered with OpenHCS.""" 

57 # External functions have slice_by_slice but not full OpenHCS decorations 

58 return (hasattr(func, 'slice_by_slice') and 

59 not hasattr(func, '__processing_contract__') and 

60 not func.__module__.startswith('openhcs.')) 

61 

62 

63def _get_function_library_name(func): 

64 """Get the library name for an external registered function.""" 

65 from openhcs.processing.backends.lib_registry.registry_service import RegistryService 

66 

67 # Find the function in the registry to get its library name 

68 all_functions = RegistryService.get_all_functions_with_metadata() 

69 for func_name, metadata in all_functions.items(): 

70 if metadata.func is func: 

71 return metadata.registry.library_name 

72 

73 return None 

74 

75 

76def _create_openhcs_library_modules(): 

77 """Create virtual modules that mirror external library structure under openhcs namespace.""" 

78 import sys 

79 import types 

80 from openhcs.processing.backends.lib_registry.registry_service import RegistryService 

81 

82 # Get all registered functions 

83 all_functions = RegistryService.get_all_functions_with_metadata() 

84 

85 # Group functions by their full module path 

86 functions_by_module = {} 

87 for func_name, metadata in all_functions.items(): 

88 if _is_external_registered_function(metadata.func): 

89 original_module = metadata.func.__module__ 

90 virtual_module = f'openhcs.{original_module}' 

91 if virtual_module not in functions_by_module: 

92 functions_by_module[virtual_module] = {} 

93 functions_by_module[virtual_module][metadata.func.__name__] = metadata.func 

94 

95 # Create virtual modules for each module path 

96 created_modules = [] 

97 for virtual_module, functions in functions_by_module.items(): 

98 if virtual_module not in sys.modules: 

99 module = types.ModuleType(virtual_module) 

100 module.__doc__ = f"Virtual module mirroring {virtual_module.replace('openhcs.', '')} with OpenHCS decorations" 

101 sys.modules[virtual_module] = module 

102 

103 # Add all functions from this module 

104 for func_name, func in functions.items(): 

105 setattr(module, func_name, func) 

106 

107 created_modules.append(virtual_module) 

108 

109 return created_modules 

110 

111def format_imports_as_strings(function_imports, enum_imports): 

112 """Convert import dictionaries to list of import strings with collision resolution.""" 

113 # Merge imports 

114 all_imports = function_imports.copy() 

115 for module, names in enum_imports.items(): 

116 all_imports.setdefault(module, set()).update(names) 

117 

118 # Build collision map 

119 name_to_modules = defaultdict(list) 

120 for module, names in all_imports.items(): 

121 for name in names: 

122 name_to_modules[name].append(module) 

123 

124 import_lines, name_mappings = [], {} 

125 for module, names in sorted(all_imports.items()): 

126 if not module or module == 'builtins' or not names: 

127 continue 

128 

129 imports = [] 

130 for name in sorted(names): 

131 if len(name_to_modules[name]) > 1: 

132 qualified = f"{name}_{module.split('.')[-1]}" 

133 imports.append(f"{name} as {qualified}") 

134 name_mappings[(name, module)] = qualified 

135 else: 

136 imports.append(name) 

137 name_mappings[(name, module)] = name 

138 

139 import_lines.append(f"from {module} import {', '.join(imports)}") 

140 

141 return import_lines, name_mappings 

142 

143def generate_complete_function_pattern_code(func_obj, indent=0, clean_mode=False): 

144 """Generate complete Python code for function pattern with imports.""" 

145 # Collect imports from this pattern first to get name mappings 

146 function_imports, enum_imports, decorated_functions = collect_imports_from_data(func_obj) 

147 import_lines, name_mappings = format_imports_as_strings(function_imports, enum_imports) 

148 

149 # Generate pattern representation using the name mappings for collision resolution 

150 pattern_repr = generate_readable_function_repr(func_obj, indent, clean_mode, name_mappings) 

151 

152 # Build complete code 

153 code_lines = ["# Edit this function pattern and save to apply changes", ""] 

154 if import_lines: 

155 code_lines.append("# Dynamic imports") 

156 code_lines.extend(import_lines) 

157 code_lines.append("") 

158 code_lines.append(f"pattern = {pattern_repr}") 

159 

160 return "\n".join(code_lines) 

161 

162def _value_to_repr(value, required_imports=None, name_mappings=None): 

163 """Converts a value to its Python representation string and tracks required imports.""" 

164 if isinstance(value, Enum): 

165 enum_class_name = value.__class__.__name__ 

166 enum_module = value.__class__.__module__ 

167 

168 # Collect import for the enum class 

169 if required_imports is not None and enum_module and enum_class_name: 

170 required_imports[enum_module].add(enum_class_name) 

171 

172 # Use name mapping if available to handle collisions 

173 if name_mappings and (enum_class_name, enum_module) in name_mappings: 

174 mapped_name = name_mappings[(enum_class_name, enum_module)] 

175 return f"{mapped_name}.{value.name}" 

176 else: 

177 return f"{enum_class_name}.{value.name}" 

178 elif isinstance(value, str): 

179 # Use repr() for strings to properly escape newlines and special characters 

180 return repr(value) 

181 elif isinstance(value, Path): 

182 # Track that we need Path import 

183 if required_imports is not None: 

184 required_imports['pathlib'].add('Path') 

185 

186 # Use name mapping if available 

187 path_name = 'Path' 

188 if name_mappings and ('Path', 'pathlib') in name_mappings: 

189 path_name = name_mappings[('Path', 'pathlib')] 

190 

191 return f'{path_name}({repr(str(value))})' 

192 return repr(value) 

193 

194def generate_clean_dataclass_repr(instance, indent_level=0, clean_mode=False, required_imports=None, name_mappings=None): 

195 """ 

196 Generates a clean, readable Python representation of a dataclass instance, 

197 omitting fields that are set to their default values if clean_mode is True. 

198 This function is recursive and handles nested dataclasses. 

199 """ 

200 if not dataclasses.is_dataclass(instance): 

201 return _value_to_repr(instance, required_imports, name_mappings) 

202 

203 lines = [] 

204 indent_str = " " * indent_level 

205 child_indent_str = " " * (indent_level + 1) 

206 

207 # Get a default instance of the same class for comparison 

208 # CRITICAL FIX: For lazy dataclasses, create instance with raw values to preserve None vs concrete distinction 

209 if hasattr(instance, '_resolve_field_value'): 

210 # This is a lazy dataclass - create empty instance without triggering resolution 

211 default_instance = object.__new__(instance.__class__) 

212 

213 # Set all fields to None (their raw default state) using object.__setattr__ 

214 for field in dataclasses.fields(instance): 

215 object.__setattr__(default_instance, field.name, None) 

216 

217 # Initialize any required lazy dataclass attributes 

218 if hasattr(instance.__class__, '_is_lazy_dataclass'): 

219 object.__setattr__(default_instance, '_is_lazy_dataclass', True) 

220 else: 

221 # Regular dataclass - use normal constructor 

222 default_instance = instance.__class__() 

223 

224 for field in dataclasses.fields(instance): 

225 field_name = field.name 

226 

227 # CRITICAL FIX: For lazy dataclasses, use raw stored value to avoid triggering resolution 

228 # This ensures tier 3 code generation only shows explicitly set pipeline config fields 

229 if hasattr(instance, '_resolve_field_value'): 

230 # This is a lazy dataclass - get raw stored value without triggering lazy resolution 

231 current_value = object.__getattribute__(instance, field_name) 

232 default_value = object.__getattribute__(default_instance, field_name) 

233 else: 

234 # Regular dataclass - use normal getattr 

235 current_value = getattr(instance, field_name) 

236 default_value = getattr(default_instance, field_name) 

237 

238 if clean_mode and current_value == default_value: 

239 continue 

240 

241 if dataclasses.is_dataclass(current_value): 

242 # Recursively generate representation for nested dataclasses 

243 nested_repr = generate_clean_dataclass_repr(current_value, indent_level + 1, clean_mode, required_imports, name_mappings) 

244 

245 # Only include nested dataclass if it has non-default content 

246 if nested_repr.strip(): # Has actual content 

247 # Collect import for the nested dataclass 

248 if required_imports is not None: 

249 class_module = current_value.__class__.__module__ 

250 class_name = current_value.__class__.__name__ 

251 if class_module and class_name: 

252 required_imports[class_module].add(class_name) 

253 

254 lines.append(f"{child_indent_str}{field_name}={current_value.__class__.__name__}(\n{nested_repr}\n{child_indent_str})") 

255 elif not clean_mode: 

256 # In non-clean mode, still include empty nested dataclasses 

257 if required_imports is not None: 

258 class_module = current_value.__class__.__module__ 

259 class_name = current_value.__class__.__name__ 

260 if class_module and class_name: 

261 required_imports[class_module].add(class_name) 

262 

263 lines.append(f"{child_indent_str}{field_name}={current_value.__class__.__name__}()") 

264 else: 

265 value_repr = _value_to_repr(current_value, required_imports, name_mappings) 

266 lines.append(f"{child_indent_str}{field_name}={value_repr}") 

267 

268 if not lines: 

269 return "" # Return empty string if all fields were default in clean_mode 

270 

271 return ",\n".join(lines) 

272 

273 

274def convert_pickle_to_python(pickle_path, output_path=None, clean_mode=False): 

275 """Convert an OpenHCS debug pickle file to a runnable Python script.""" 

276 

277 pickle_file = Path(pickle_path) 

278 if not pickle_file.exists(): 

279 print(f"Error: Pickle file not found: {pickle_path}") 

280 return 

281 

282 if output_path is None: 

283 output_path = pickle_file.with_suffix('.py') 

284 

285 print(f"Converting {pickle_file} to {output_path} (Clean Mode: {clean_mode})") 

286 

287 try: 

288 with open(pickle_file, 'rb') as f: 

289 data = pickle.load(f) 

290 

291 # Generate Python script 

292 with open(output_path, 'w') as f: 

293 f.write('#!/usr/bin/env python3\n') 

294 f.write('"""\n') 

295 f.write(f'OpenHCS Pipeline Script - Generated from {pickle_file.name}\n') 

296 f.write(f'Generated: {datetime.now()}\n') 

297 f.write('"""\n\n') 

298 

299 # Imports 

300 f.write('import sys\n') 

301 f.write('import os\n') 

302 f.write('from pathlib import Path\n\n') 

303 f.write('# Add OpenHCS to path\n') 

304 f.write('sys.path.insert(0, "/home/ts/code/projects/openhcs")\n\n') 

305 

306 f.write('from openhcs.core.orchestrator.orchestrator import PipelineOrchestrator\n') 

307 f.write('from openhcs.core.steps.function_step import FunctionStep\n') 

308 f.write('from openhcs.core.config import (GlobalPipelineConfig, PathPlanningConfig, VFSConfig, ZarrConfig, \n' 

309 ' MaterializationBackend, ZarrCompressor, ZarrChunkStrategy)\n') 

310 f.write('from openhcs.constants.constants import VariableComponents, Backend, Microscope\n\n') 

311 

312 # Use extracted function for orchestrator generation 

313 orchestrator_code = generate_complete_orchestrator_code( 

314 data["plate_paths"], data["pipeline_data"], data['global_config'], clean_mode 

315 ) 

316 

317 # Write orchestrator code (already includes dynamic imports) 

318 f.write(orchestrator_code) 

319 f.write('\n\n') 

320 

321 # ... (rest of the file remains the same for now) ... 

322 f.write('def setup_signal_handlers():\n') 

323 f.write(' """Setup signal handlers to kill all child processes and threads on Ctrl+C."""\n') 

324 f.write(' import signal\n') 

325 f.write(' import os\n') 

326 f.write(' import sys\n\n') 

327 f.write(' def cleanup_and_exit(signum, frame):\n') 

328 f.write(' print(f"\\n🔥 Signal {signum} received! Cleaning up all processes and threads...")\n\n') 

329 f.write(' os._exit(1)\n\n') 

330 f.write(' signal.signal(signal.SIGINT, cleanup_and_exit)\n') 

331 f.write(' signal.signal(signal.SIGTERM, cleanup_and_exit)\n\n') 

332 

333 f.write('def run_pipeline():\n') 

334 f.write(' os.environ["OPENHCS_SUBPROCESS_MODE"] = "1"\n') 

335 f.write(' plate_paths, pipeline_data, global_config = create_pipeline()\n') 

336 f.write(' from openhcs.core.orchestrator.gpu_scheduler import setup_global_gpu_registry\n') 

337 f.write(' setup_global_gpu_registry(global_config=global_config)\n') 

338 f.write(' for plate_path in plate_paths:\n') 

339 f.write(' orchestrator = PipelineOrchestrator(plate_path)\n') 

340 f.write(' orchestrator.initialize()\n') 

341 f.write(' compiled_contexts = orchestrator.compile_pipelines(pipeline_data[plate_path])\n') 

342 f.write(' orchestrator.execute_compiled_plate(\n') 

343 f.write(' pipeline_definition=pipeline_data[plate_path],\n') 

344 f.write(' compiled_contexts=compiled_contexts,\n') 

345 f.write(' max_workers=global_config.num_workers\n') 

346 f.write(' )\n\n') 

347 

348 f.write('if __name__ == "__main__":\n') 

349 f.write(' setup_signal_handlers()\n') 

350 f.write(' run_pipeline()\n') 

351 

352 

353 print(f"✅ Successfully converted to {output_path}") 

354 print(f"You can now run: python {output_path}") 

355 

356 except Exception as e: 

357 print(f"Error converting pickle file: {e}") 

358 import traceback 

359 traceback.print_exc() 

360 

361 

362def generate_readable_function_repr(func_obj, indent=0, clean_mode=False, name_mappings=None): 

363 """Generate readable Python representation with collision-resolved function names.""" 

364 indent_str = " " * indent 

365 next_indent_str = " " * (indent + 1) 

366 name_mappings = name_mappings or {} 

367 

368 # Get qualified function name for collisions (handle both original and virtual modules) 

369 def get_name(f): 

370 if not callable(f): 

371 return str(f) 

372 # Try virtual module first (for external functions), then original module 

373 virtual_module = f'openhcs.{f.__module__}' 

374 return (name_mappings.get((f.__name__, virtual_module), None) or 

375 name_mappings.get((f.__name__, f.__module__), f.__name__)) 

376 

377 if callable(func_obj): 

378 return get_name(func_obj) 

379 

380 elif isinstance(func_obj, tuple) and len(func_obj) == 2 and callable(func_obj[0]): 

381 func, args = func_obj 

382 

383 if not args and clean_mode: 

384 return get_name(func) 

385 

386 # Filter out defaults in clean mode 

387 try: 

388 defaults = {k: v.default for k, v in inspect.signature(func).parameters.items() 

389 if v.default is not inspect.Parameter.empty} 

390 except (ValueError, TypeError): 

391 defaults = {} 

392 

393 final_args = {k: v for k, v in args.items() 

394 if not clean_mode or k not in defaults or v != defaults[k]} 

395 

396 if not final_args: 

397 return get_name(func) if clean_mode else f"({get_name(func)}, {{}})" 

398 

399 args_items = [f"{next_indent_str} '{k}': {generate_readable_function_repr(v, indent + 2, clean_mode, name_mappings)}" 

400 for k, v in final_args.items()] 

401 args_str = "{\n" + ",\n".join(args_items) + f"\n{next_indent_str}}}" 

402 return f"({get_name(func)}, {args_str})" 

403 

404 elif isinstance(func_obj, list): 

405 if clean_mode and len(func_obj) == 1: 

406 return generate_readable_function_repr(func_obj[0], indent, clean_mode, name_mappings) 

407 if not func_obj: 

408 return "[]" 

409 items = [generate_readable_function_repr(item, indent, clean_mode, name_mappings) for item in func_obj] 

410 return f"[\n{next_indent_str}{f',\n{next_indent_str}'.join(items)}\n{indent_str}]" 

411 

412 elif isinstance(func_obj, dict): 

413 if not func_obj: 

414 return "{}" 

415 items = [f"{next_indent_str}'{k}': {generate_readable_function_repr(v, indent, clean_mode, name_mappings)}" 

416 for k, v in func_obj.items()] 

417 return f"{{{',\n'.join(items)}\n{indent_str}}}" 

418 

419 else: 

420 return _value_to_repr(func_obj) 

421 

422 

423def _format_parameter_value(param_name, value, name_mappings=None): 

424 """Format parameter values with lazy dataclass preservation.""" 

425 if isinstance(value, Enum): 

426 enum_class_name = value.__class__.__name__ 

427 enum_module = value.__class__.__module__ 

428 

429 # Use name mapping if available to handle collisions 

430 if name_mappings and (enum_class_name, enum_module) in name_mappings: 

431 mapped_name = name_mappings[(enum_class_name, enum_module)] 

432 return f"{mapped_name}.{value.name}" 

433 else: 

434 return f"{enum_class_name}.{value.name}" 

435 elif isinstance(value, str): 

436 return f'"{value}"' 

437 elif isinstance(value, list) and value and isinstance(value[0], Enum): 

438 formatted_items = [] 

439 for item in value: 

440 enum_class_name = item.__class__.__name__ 

441 enum_module = item.__class__.__module__ 

442 

443 # Use name mapping if available to handle collisions 

444 if name_mappings and (enum_class_name, enum_module) in name_mappings: 

445 mapped_name = name_mappings[(enum_class_name, enum_module)] 

446 formatted_items.append(f"{mapped_name}.{item.name}") 

447 else: 

448 formatted_items.append(f"{enum_class_name}.{item.name}") 

449 

450 return f"[{', '.join(formatted_items)}]" 

451 elif is_dataclass(value) and 'Lazy' in value.__class__.__name__: 

452 # Preserve lazy behavior by only including explicitly set fields 

453 class_name = value.__class__.__name__ 

454 explicit_args = [ 

455 f"{f.name}={_format_parameter_value(f.name, object.__getattribute__(value, f.name), name_mappings)}" 

456 for f in fields(value) 

457 if object.__getattribute__(value, f.name) is not None 

458 ] 

459 return f"{class_name}({', '.join(explicit_args)})" if explicit_args else f"{class_name}()" 

460 else: 

461 return repr(value) 

462 

463 

464 

465 

466 

467def _collect_dataclass_classes_from_object(obj, visited=None): 

468 """Recursively collect dataclass classes that will be referenced in generated code.""" 

469 if visited is None: 

470 visited = set() 

471 

472 if id(obj) in visited: 

473 return set(), set() 

474 visited.add(id(obj)) 

475 

476 dataclass_classes = set() 

477 enum_classes = set() 

478 

479 if is_dataclass(obj): 

480 dataclass_classes.add(obj.__class__) 

481 for field in fields(obj): 

482 nested_dataclasses, nested_enums = _collect_dataclass_classes_from_object(getattr(obj, field.name), visited) 

483 dataclass_classes.update(nested_dataclasses) 

484 enum_classes.update(nested_enums) 

485 elif isinstance(obj, Enum): 

486 enum_classes.add(obj.__class__) 

487 elif isinstance(obj, (list, tuple)): 

488 for item in obj: 

489 nested_dataclasses, nested_enums = _collect_dataclass_classes_from_object(item, visited) 

490 dataclass_classes.update(nested_dataclasses) 

491 enum_classes.update(nested_enums) 

492 elif isinstance(obj, dict): 

493 for value in obj.values(): 

494 nested_dataclasses, nested_enums = _collect_dataclass_classes_from_object(value, visited) 

495 dataclass_classes.update(nested_dataclasses) 

496 enum_classes.update(nested_enums) 

497 

498 return dataclass_classes, enum_classes 

499 

500 

501def _collect_enum_classes_from_step(step): 

502 """Collect enum classes referenced in step parameters for import generation.""" 

503 from openhcs.core.steps.function_step import FunctionStep 

504 import inspect 

505 from enum import Enum 

506 

507 enum_classes = set() 

508 sig = inspect.signature(FunctionStep.__init__) 

509 

510 for param_name, param in sig.parameters.items(): 

511 # Skip constructor-specific parameters and **kwargs 

512 if param_name in ['self', 'func'] or param.kind == inspect.Parameter.VAR_KEYWORD: 

513 continue 

514 

515 value = getattr(step, param_name, param.default) 

516 if isinstance(value, Enum): 

517 enum_classes.add(type(value)) 

518 elif isinstance(value, (list, tuple)): 

519 # Check for lists/tuples of enums 

520 for item in value: 

521 if isinstance(item, Enum): 

522 enum_classes.add(type(item)) 

523 

524 return enum_classes 

525 

526 

527def _generate_step_parameters(step, default_step, clean_mode=False, name_mappings=None): 

528 """Generate FunctionStep constructor parameters using functional introspection.""" 

529 from openhcs.core.steps.abstract import AbstractStep 

530 

531 signatures = [(name, param) for name, param in inspect.signature(FunctionStep.__init__).parameters.items() 

532 if name != 'self' and param.kind != inspect.Parameter.VAR_KEYWORD] + \ 

533 [(name, param) for name, param in inspect.signature(AbstractStep.__init__).parameters.items() 

534 if name != 'self'] 

535 

536 return [f"{name}={generate_readable_function_repr(getattr(step, name, param.default), 1, clean_mode, name_mappings) if name == 'func' else _format_parameter_value(name, getattr(step, name, param.default), name_mappings)}" 

537 for name, param in signatures 

538 if not clean_mode or getattr(step, name, param.default) != getattr(default_step, name, param.default)] 

539 

540 

541def generate_complete_pipeline_steps_code(pipeline_steps, clean_mode=False): 

542 """Generate complete Python code for pipeline steps with imports.""" 

543 # Build code with imports and steps 

544 code_lines = ["# Edit this pipeline and save to apply changes", ""] 

545 

546 # Collect imports from ALL data in pipeline steps (functions AND parameters) 

547 all_function_imports = defaultdict(set) 

548 all_enum_imports = defaultdict(set) 

549 all_decorated_functions = set() 

550 

551 for step in pipeline_steps: 

552 # Collect all imports from step (functions, enums, dataclasses) 

553 func_imports, enum_imports, func_decorated = collect_imports_from_data(step.func) 

554 param_imports, param_enums, param_decorated = collect_imports_from_data(step) 

555 

556 # Merge imports 

557 for module, names in func_imports.items(): 

558 all_function_imports[module].update(names) 

559 for module, names in enum_imports.items(): 

560 all_enum_imports[module].update(names) 

561 for module, names in param_imports.items(): 

562 all_function_imports[module].update(names) 

563 for module, names in param_enums.items(): 

564 all_enum_imports[module].update(names) 

565 all_decorated_functions.update(func_decorated) 

566 all_decorated_functions.update(param_decorated) 

567 

568 # Add FunctionStep import (always needed for generated code) 

569 all_function_imports['openhcs.core.steps.function_step'].add('FunctionStep') 

570 

571 # Virtual modules are now automatically created during OpenHCS import 

572 # No need to generate runtime virtual module creation code 

573 

574 # Format and add all collected imports 

575 import_lines, name_mappings = format_imports_as_strings(all_function_imports, all_enum_imports) 

576 if import_lines: 

577 code_lines.append("# Automatically collected imports") 

578 code_lines.extend(import_lines) 

579 code_lines.append("") 

580 

581 # Generate pipeline steps (extract exact logic from lines 164-198) 

582 code_lines.append("# Pipeline steps") 

583 code_lines.append("pipeline_steps = []") 

584 code_lines.append("") 

585 

586 default_step = FunctionStep(func=lambda: None) 

587 for i, step in enumerate(pipeline_steps): 

588 code_lines.append(f"# Step {i+1}: {step.name}") 

589 

590 # Generate all FunctionStep parameters automatically using introspection 

591 step_args = _generate_step_parameters(step, default_step, clean_mode, name_mappings) 

592 

593 args_str = ",\n ".join(step_args) 

594 code_lines.append(f"step_{i+1} = FunctionStep(\n {args_str}\n)") 

595 code_lines.append(f"pipeline_steps.append(step_{i+1})") 

596 code_lines.append("") 

597 

598 return "\n".join(code_lines) 

599 

600 

601def generate_complete_orchestrator_code(plate_paths, pipeline_data, global_config, clean_mode=False, pipeline_config=None): 

602 """Generate complete Python code for orchestrator config with imports.""" 

603 # Build complete code (extract exact logic from lines 150-200) 

604 code_lines = ["# Edit this orchestrator configuration and save to apply changes", ""] 

605 

606 # Collect imports from ALL data in orchestrator (functions, parameters, config) 

607 all_function_imports = defaultdict(set) 

608 all_enum_imports = defaultdict(set) 

609 all_decorated_functions = set() 

610 

611 # Collect from pipeline steps 

612 for plate_path, steps in pipeline_data.items(): 

613 for step in steps: 

614 # Get imports from function patterns 

615 func_imports, enum_imports, func_decorated = collect_imports_from_data(step.func) 

616 # Get imports from step parameters 

617 param_imports, param_enums, param_decorated = collect_imports_from_data(step) 

618 

619 # Merge all imports 

620 for module, names in func_imports.items(): 

621 all_function_imports[module].update(names) 

622 for module, names in enum_imports.items(): 

623 all_enum_imports[module].update(names) 

624 for module, names in param_imports.items(): 

625 all_function_imports[module].update(names) 

626 for module, names in param_enums.items(): 

627 all_enum_imports[module].update(names) 

628 all_decorated_functions.update(func_decorated) 

629 all_decorated_functions.update(param_decorated) 

630 

631 # Don't collect imports from entire global config upfront - only collect what's actually used 

632 # This prevents importing unused classes and keeps the generated code clean 

633 

634 # First pass: Collect imports needed for config representation (e.g., Path) BEFORE formatting imports 

635 config_repr_imports = defaultdict(set) 

636 temp_config_repr = generate_clean_dataclass_repr(global_config, indent_level=0, clean_mode=clean_mode, required_imports=config_repr_imports) 

637 

638 # Merge config representation imports with main imports 

639 for module, names in config_repr_imports.items(): 

640 all_function_imports[module].update(names) 

641 

642 # Don't collect imports from entire pipeline config upfront - let representation generation handle it 

643 # This ensures only actually used imports are collected 

644 

645 # Add always-needed imports for generated code structure 

646 all_function_imports['openhcs.core.steps.function_step'].add('FunctionStep') 

647 all_function_imports['openhcs.core.config'].add('PipelineConfig') 

648 all_function_imports['openhcs.core.orchestrator.orchestrator'].add('PipelineOrchestrator') 

649 all_function_imports['openhcs.core.config'].add('GlobalPipelineConfig') # Always needed for global_config constructor 

650 

651 # Virtual modules are now automatically created during OpenHCS import 

652 # No need for runtime virtual module creation 

653 

654 # First pass: Generate name mappings for collision resolution (don't add imports yet) 

655 import_lines, name_mappings = format_imports_as_strings(all_function_imports, all_enum_imports) 

656 

657 # Generate config representation and collect only the imports it actually needs 

658 config_repr_imports = defaultdict(set) 

659 config_repr = generate_clean_dataclass_repr(global_config, indent_level=0, clean_mode=clean_mode, required_imports=config_repr_imports, name_mappings=name_mappings) 

660 

661 # Add only the imports that are actually used in the config representation 

662 for module, names in config_repr_imports.items(): 

663 all_function_imports[module].update(names) 

664 

665 code_lines.extend([ 

666 "# Plate paths", 

667 f"plate_paths = {repr(plate_paths)}", 

668 "", 

669 "# Global configuration", 

670 ]) 

671 

672 code_lines.append(f"global_config = GlobalPipelineConfig(\n{config_repr}\n)") 

673 code_lines.append("") 

674 

675 # Add PipelineConfig creation with actual values (if any) 

676 if pipeline_config is not None: 

677 # Collect imports needed for pipeline config representation 

678 pipeline_config_imports = defaultdict(set) 

679 pipeline_config_repr = generate_clean_dataclass_repr( 

680 pipeline_config, 

681 indent_level=0, 

682 clean_mode=clean_mode, 

683 required_imports=pipeline_config_imports, 

684 name_mappings=name_mappings 

685 ) 

686 

687 # Add the collected imports to the main import collection 

688 for module, names in pipeline_config_imports.items(): 

689 all_function_imports[module].update(names) 

690 

691 # Regenerate import lines with the new imports 

692 import_lines, name_mappings = format_imports_as_strings(all_function_imports, all_enum_imports) 

693 

694 code_lines.extend([ 

695 "# Pipeline configuration (lazy GlobalPipelineConfig)", 

696 f"pipeline_config = PipelineConfig(\n{pipeline_config_repr}\n)", 

697 "" 

698 ]) 

699 else: 

700 # No pipeline config overrides 

701 code_lines.extend([ 

702 "# Pipeline configuration (lazy GlobalPipelineConfig)", 

703 "pipeline_config = PipelineConfig()", 

704 "" 

705 ]) 

706 

707 # Generate pipeline data (exact logic from lines 164-198) 

708 code_lines.extend(["# Pipeline steps", "pipeline_data = {}", ""]) 

709 

710 default_step = FunctionStep(func=lambda: None) 

711 for plate_path, steps in pipeline_data.items(): 

712 # Extract plate name without using Path in generated code 

713 plate_name = str(plate_path).split('/')[-1] if '/' in str(plate_path) else str(plate_path) 

714 code_lines.append(f'# Steps for plate: {plate_name}') 

715 code_lines.append("steps = []") 

716 code_lines.append("") 

717 

718 for i, step in enumerate(steps): 

719 code_lines.append(f"# Step {i+1}: {step.name}") 

720 

721 # Generate all FunctionStep parameters automatically using introspection with name mappings 

722 step_args = _generate_step_parameters(step, default_step, clean_mode, name_mappings) 

723 

724 args_str = ",\n ".join(step_args) 

725 code_lines.append(f"step_{i+1} = FunctionStep(\n {args_str}\n)") 

726 code_lines.append(f"steps.append(step_{i+1})") 

727 code_lines.append("") 

728 

729 code_lines.append(f'pipeline_data["{plate_path}"] = steps') 

730 code_lines.append("") 

731 

732 # Add orchestrator creation example 

733 code_lines.extend([ 

734 "# Example: Create orchestrators with PipelineConfig", 

735 "# orchestrators = {}", 

736 "# for plate_path in plate_paths:", 

737 "# orchestrator = PipelineOrchestrator(", 

738 "# plate_path=plate_path,", 

739 "# pipeline_config=pipeline_config", 

740 "# )", 

741 "# orchestrators[plate_path] = orchestrator",