Coverage for src/objectstate/context_manager.py: 69%
310 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-17 07:22 +0000
1"""
2Generic contextvars-based context management system for lazy configuration.
4This module provides explicit context scoping using Python's contextvars to enable
5hierarchical configuration resolution without explicit parameter passing.
7Key features:
81. Explicit context scoping with config_context() manager
92. Config extraction from functions, dataclasses, and objects
103. Config merging for context hierarchy
114. Clean separation between UI windows and contexts
13Key components:
14- current_temp_global: ContextVar holding current merged global config
15- config_context(): Context manager for creating context scopes
16- merge_configs(): Merge overrides into base config
17"""
19import contextvars
20import functools
21import logging
22import threading
23from collections.abc import Sequence
24from contextlib import ExitStack, contextmanager
25from typing import Any, Dict, get_type_hints
26from dataclasses import fields, is_dataclass
27from python_introspect import resolve_annotated, resolve_optional
29logger = logging.getLogger(__name__)
32class NoActiveContextError(RuntimeError):
33 """A field was resolved against the active context while no context was active."""
36# Core contextvar for current merged global config
37# This holds the current context state that resolution functions can access
38current_temp_global = contextvars.ContextVar('current_temp_global')
40# Context type stack - tracks the types of objects pushed via config_context()
41# This enables generic hierarchy inference without hardcoding specific types
42# The stack is a tuple of types, ordered from outermost to innermost context
43context_type_stack = contextvars.ContextVar('context_type_stack', default=())
45class ContextLayer:
46 """One object entered with ``config_context``, with the scope it was entered under.
48 ``scope_id`` is ``None`` for a layer entered without one. ``masks_outer_layers``
49 is True for a layer entered with ``mask_with_none``: a ``None`` in it is final,
50 so no layer outside it supplies a value.
51 """
53 __slots__ = ("scope_id", "obj", "masks_outer_layers", "_configs")
55 def __init__(self, scope_id: str | None, obj: Any, masks_outer_layers: bool = False):
56 self.scope_id = scope_id
57 self.obj = obj
58 self.masks_outer_layers = masks_outer_layers
59 self._configs: dict[type, Any] | None = None
61 @property
62 def configs(self) -> dict[type, Any]:
63 """The config instances this layer holds, by type. Extracted once per layer."""
64 configs = self._configs
65 if configs is None:
66 configs = self._configs = extract_all_configs(self.obj)
67 return configs
69 def __repr__(self) -> str:
70 return f"ContextLayer({self.scope_id!r}, {type(self.obj).__name__})"
73class ContextLayerStack:
74 """The layers of one active context, outermost first.
76 This is the value field resolution reads. Every ``config_context`` entry makes a
77 new stack, and ``resolved_fields`` holds the answers computed from that stack, so
78 the answers are dropped together with the stack they belong to.
79 """
81 __slots__ = ("layers", "resolved_fields", "__weakref__")
83 def __init__(self, layers: tuple[ContextLayer, ...] = ()):
84 self.layers = layers
85 self.resolved_fields: dict[tuple[type, str], Any] = {}
87 def with_layer(self, layer: ContextLayer) -> "ContextLayerStack":
88 return ContextLayerStack(self.layers + (layer,))
90 def __repr__(self) -> str:
91 return f"ContextLayerStack({self.layers!r})"
94NO_LAYERS = ContextLayerStack()
96# The layer stack of the active context. NO_LAYERS means no context is active.
97context_layer_stack: contextvars.ContextVar[ContextLayerStack] = contextvars.ContextVar(
98 'context_layer_stack', default=NO_LAYERS
99)
101def _merge_nested_dataclass(base, override, mask_with_none: bool = False):
102 """
103 Recursively merge nested dataclass fields.
105 For each field in override:
106 - If value is None and mask_with_none=False: skip (don't override base)
107 - If value is None and mask_with_none=True: override with None (mask base)
108 - If value is dataclass: recursively merge with base's value
109 - Otherwise: use override value
111 Args:
112 base: Base dataclass instance
113 override: Override dataclass instance
114 mask_with_none: If True, None values override base values
116 Returns:
117 Merged dataclass instance
118 """
119 if not is_dataclass(base) or not is_dataclass(override):
120 return override
122 from objectstate.lazy_factory import get_base_type_for_lazy, replace_raw
124 base_type = get_base_type_for_lazy(type(base)) or type(base)
125 override_type = get_base_type_for_lazy(type(override)) or type(override)
126 # Lazy overlays of the same owner inherit on the base instance. An authored
127 # nominal change instead owns the result, including subtype-only fields.
128 result_owner = base if base_type is override_type else override
129 base_fields = {field_info.name for field_info in fields(base)}
130 merge_values = {}
131 for field_info in fields(override):
132 field_name = field_info.name
133 override_value = object.__getattribute__(override, field_name)
134 if field_name not in base_fields:
135 # This field belongs only to the authored nominal owner.
136 continue
137 base_value = object.__getattribute__(base, field_name)
139 if override_value is None:
140 if mask_with_none:
141 # None overrides base value (masking mode)
142 merge_values[field_name] = None
143 else:
144 # None means inherit, including when the result owner changes.
145 if result_owner is override:
146 merge_values[field_name] = base_value
147 elif is_dataclass(override_value):
148 # Recursively merge nested dataclass
149 if base_value is not None and is_dataclass(base_value):
150 merge_values[field_name] = _merge_nested_dataclass(base_value, override_value, mask_with_none)
151 else:
152 merge_values[field_name] = override_value
153 else:
154 # Concrete value - use override
155 merge_values[field_name] = override_value
157 # Merge on the nominal owner using replace_raw to preserve None values
158 # (dataclasses.replace triggers lazy resolution, baking in resolved values)
159 if merge_values:
160 return replace_raw(result_owner, **merge_values)
161 else:
162 return result_owner
165@contextmanager
166def config_context(
167 obj,
168 mask_with_none: bool = False,
169 use_live_global: bool = True,
170 scope_id: str | None = None,
171):
172 """
173 Create new context scope with obj's matching fields merged into base config.
175 This is the universal context manager for all config context needs. It works by:
176 1. Finding fields that exist on both obj and the base config type
177 2. Using matching field values to create a temporary merged config
178 3. Setting that as the current context
180 Args:
181 obj: Object with config fields (pipeline_config, step, etc.)
182 mask_with_none: If True, None values override/mask base config values, and
183 field resolution does not look at any layer outside this one.
184 If False (default), None values are ignored (normal inheritance).
185 Use True when editing GlobalPipelineConfig to mask the
186 loaded instance with static class defaults.
187 use_live_global: If True (default), use LIVE global config (UI sees unsaved edits).
188 If False, use SAVED global config (compiler sees saved values).
189 scope_id: The scope this layer belongs to (e.g., "plate_123::step_0"). Field
190 resolution reports it as the source scope of every value the layer
191 supplies. A layer entered without one supplies values all the same and
192 is reported with source scope ``None``.
194 Usage:
195 # UI operations (default - uses LIVE)
196 with config_context(orchestrator.pipeline_config):
197 # UI sees unsaved global config edits
199 # Compilation (explicit - uses SAVED)
200 with config_context(orchestrator.pipeline_config, use_live_global=False):
201 # Compiler uses saved global config only
203 # With a scope id
204 with config_context(step, scope_id="plate_123::step_0"):
205 # Values this layer supplies are reported as coming from that scope
206 """
207 # Get current context as base for nested contexts, or fall back to base global config
208 current_context = get_current_temp_global()
209 base_config = current_context if current_context is not None else get_base_global_config(use_live=use_live_global)
211 # Find matching fields between obj and base config type
212 overrides = {}
213 if obj is not None:
214 from objectstate.config import get_base_config_type
216 base_config_type = get_base_config_type()
218 for field_info in fields(base_config_type):
219 field_name = field_info.name
220 expected_type = field_info.type
222 # Check if obj has this field
223 try:
224 # Use object.__getattribute__ to avoid triggering lazy resolution
225 if hasattr(obj, field_name):
226 value = object.__getattribute__(obj, field_name)
227 # CRITICAL: When mask_with_none=True, None values override base config
228 # This allows static defaults to mask loaded instance values
229 if value is not None or mask_with_none:
230 # When masking with None, always include the value (even if None)
231 if mask_with_none:
232 # For nested dataclasses, merge with mask_with_none=True
233 if is_dataclass(value):
234 base_value = getattr(base_config, field_name, None)
235 if base_value is not None and is_dataclass(base_value):
236 merged_nested = _merge_nested_dataclass(base_value, value, mask_with_none=True)
237 overrides[field_name] = merged_nested
238 else:
239 overrides[field_name] = value
240 else:
241 overrides[field_name] = value
242 # Normal mode: only include non-None values
243 elif value is not None:
244 # Check if value is compatible (handles lazy-to-base type mapping)
245 if _is_compatible_config_type(value, expected_type):
246 # CRITICAL FIX: Do NOT call to_base_config() here!
247 # to_base_config() passes None values to base class constructor,
248 # but frozen dataclasses with non-Optional defaults substitute the
249 # default value for None. This breaks cross-hierarchy inheritance.
250 # Example: LazyPartitionFilterConfig(partition_filter_mode=None)
251 # becomes PartitionFilterConfig(partition_filter_mode=INCLUDE)
252 # instead of keeping None.
253 #
254 # Instead, pass lazy dataclass directly to _merge_nested_dataclass,
255 # which uses object.__getattribute__ to get raw values and correctly
256 # skips None overrides.
258 # Recursively merge nested dataclass fields
259 if is_dataclass(value):
260 base_value = getattr(base_config, field_name, None)
261 if base_value is not None and is_dataclass(base_value):
262 merged_nested = _merge_nested_dataclass(base_value, value, mask_with_none=False)
263 overrides[field_name] = merged_nested
264 else:
265 # No base value to merge with - convert if needed
266 if hasattr(value, 'to_base_config'):
267 value = value.to_base_config()
268 overrides[field_name] = value
269 else:
270 # Non-dataclass field, use override as-is
271 overrides[field_name] = value
272 except AttributeError:
273 continue
275 # Create merged config if we have overrides
276 # Use replace_raw to preserve None values (dataclasses.replace triggers lazy resolution)
277 if overrides:
278 from objectstate.lazy_factory import replace_raw
279 merged_config = replace_raw(base_config, **overrides)
280 else:
281 merged_config = base_config
283 # Track the type in the context type stack
284 current_types = context_type_stack.get()
285 new_types = current_types + (type(obj),) if obj is not None else current_types
287 # A context entered while none is active starts from the global config, so the
288 # global config is its outermost layer.
289 if current_context is None:
290 current_stack = ContextLayerStack((ContextLayer("", base_config),))
291 else:
292 current_stack = context_layer_stack.get()
293 new_stack = (
294 current_stack.with_layer(ContextLayer(scope_id, obj, mask_with_none))
295 if obj is not None
296 else current_stack
297 )
299 merged_token = current_temp_global.set(merged_config)
300 type_token = context_type_stack.set(new_types)
301 layer_token = context_layer_stack.set(new_stack)
302 try:
303 yield
304 finally:
305 current_temp_global.reset(merged_token)
306 context_type_stack.reset(type_token)
307 context_layer_stack.reset(layer_token)
310def get_context_type_stack():
311 """
312 Get the current stack of context types (outermost to innermost).
314 This enables generic hierarchy inference without hardcoding specific types.
315 The stack represents the order in which config_context() calls were nested.
317 Returns:
318 Tuple of types representing the context hierarchy.
319 Empty tuple if no context is active.
321 Example:
322 with config_context(global_config): # stack = (GlobalPipelineConfig,)
323 with config_context(pipeline_config): # stack = (GlobalPipelineConfig, PipelineConfig)
324 with config_context(step): # stack = (GlobalPipelineConfig, PipelineConfig, Step)
325 types = get_context_type_stack()
326 # types == (GlobalPipelineConfig, PipelineConfig, Step)
327 """
328 return context_type_stack.get()
331def get_context_layer_stack() -> ContextLayerStack:
332 """
333 Get the layer stack of the active context: the value field resolution reads.
335 ``.layers`` holds one ``ContextLayer`` per object entered with ``config_context``,
336 outermost first, below the global config the context started from. It is empty
337 when no context is active.
339 Example:
340 with config_context(plate, scope_id="plate_123"):
341 with config_context(step, scope_id="plate_123::step_0"):
342 scopes = [layer.scope_id for layer in get_context_layer_stack().layers]
343 # scopes == ["", "plate_123", "plate_123::step_0"]
344 """
345 return context_layer_stack.get()
348@functools.lru_cache(maxsize=None)
349def _normalize_type(t):
350 """Normalize a type by getting its base type if it's a lazy variant.
352 This function is defined here to avoid circular imports with lazy_factory.
353 The actual get_base_type_for_lazy is imported lazily when needed.
355 PERFORMANCE: Results are cached with lru_cache since type->base_type
356 mappings are immutable after class creation.
357 """
358 from objectstate.lazy_factory import get_base_type_for_lazy
359 return get_base_type_for_lazy(t) or t
362def _is_global_type(t):
363 """Check if a type is a global config type.
365 This function is defined here to avoid circular imports with lazy_factory.
366 """
367 from objectstate.lazy_factory import is_global_config_type
368 return is_global_config_type(t)
371# Hierarchy registry - built from active form managers
372# Maps: child_type -> parent_type (normalized base types)
373# This is populated by ParameterFormManager when it registers
374_known_hierarchy: dict = {}
377def _normalize_type_for_hierarchy(t):
378 """Normalize a type for hierarchy registry, but preserve lazy-global distinction.
380 Normalization maps lazy types to their base types (e.g., LazyPathPlanningConfig → PathPlanningConfig).
381 However, PipelineConfig → GlobalPipelineConfig is a special case: they represent DISTINCT
382 hierarchy levels (pipeline overrides vs global defaults), so we must NOT collapse them.
384 Rule: If normalization would produce a global type from a non-global type, keep the original.
385 """
386 normalized = _normalize_type(t)
387 # Don't collapse lazy-to-global - they're distinct hierarchy levels
388 if _is_global_type(normalized) and not _is_global_type(t):
389 return t
390 return normalized
393def register_hierarchy_relationship(context_obj_type, object_instance_type):
394 """Register that context_obj_type is the parent of object_instance_type in the hierarchy.
396 Called by ParameterFormManager when a root manager registers.
397 This builds up the known hierarchy from actual usage patterns.
399 Args:
400 context_obj_type: The parent/context type (e.g., PipelineConfig for Step editor)
401 object_instance_type: The child type being edited (e.g., Step)
403 Note:
404 Types are normalized to base types (e.g., LazyPathPlanningConfig → PathPlanningConfig),
405 but lazy-to-global normalization is prevented to preserve distinct hierarchy levels
406 (PipelineConfig stays as PipelineConfig, not collapsed to GlobalPipelineConfig).
407 """
408 if context_obj_type is None or object_instance_type is None:
409 return
411 parent_base = _normalize_type_for_hierarchy(context_obj_type)
412 child_base = _normalize_type_for_hierarchy(object_instance_type)
414 if parent_base != child_base:
415 _known_hierarchy[child_base] = parent_base
416 logger.debug(f"Registered hierarchy: {parent_base.__name__} -> {child_base.__name__}")
419def unregister_hierarchy_relationship(object_instance_type):
420 """Remove a type from the hierarchy registry when its editor closes.
422 Args:
423 object_instance_type: The type to remove from the registry
424 """
425 child_base = _normalize_type_for_hierarchy(object_instance_type)
426 if child_base in _known_hierarchy:
427 del _known_hierarchy[child_base]
428 logger.debug(f"Unregistered hierarchy for: {child_base.__name__}")
431def get_ancestors_from_hierarchy(target_type):
432 """Get all ancestor types for target_type by walking up the known hierarchy.
434 Returns ancestors in order from outermost to innermost (excluding target_type itself).
436 Args:
437 target_type: The type to find ancestors for
439 Returns:
440 List of ancestor types in hierarchy order (grandparent, parent, ...)
441 """
442 target_base = _normalize_type_for_hierarchy(target_type)
443 ancestors = []
445 # Walk up the hierarchy
446 current = target_base
447 visited = set()
448 while current in _known_hierarchy:
449 if current in visited:
450 # Cycle detected - stop
451 break
452 visited.add(current)
453 parent = _known_hierarchy[current]
454 ancestors.append(parent)
455 current = parent
457 # Reverse so outermost is first
458 ancestors.reverse()
459 return ancestors
462def get_normalized_stack():
463 """
464 Get the context type stack with normalized (base) types, excluding global configs.
466 Returns:
467 List of base types in hierarchy order (outermost to innermost),
468 with global config types filtered out.
469 """
470 canonical_stack = get_context_type_stack()
471 normalized = []
472 for t in canonical_stack:
473 base_t = _normalize_type(t)
474 if not _is_global_type(base_t):
475 normalized.append(base_t)
476 return normalized
479def get_types_before_in_stack(target_type):
480 """
481 Get all non-global types that come before target_type in the hierarchy.
483 First tries the active context_type_stack (for resolution during config_context),
484 then falls back to the known hierarchy registry (for cross-window updates).
486 Args:
487 target_type: The type to find ancestors for (will be normalized)
489 Returns:
490 List of base types that come before target_type in the hierarchy.
491 Empty list if no ancestors found.
492 """
493 # First try active context stack
494 normalized_stack = get_normalized_stack()
495 if normalized_stack:
496 target_base = _normalize_type(target_type)
498 # Find target's position
499 target_index = -1
500 for i, base_t in enumerate(normalized_stack):
501 if base_t == target_base:
502 target_index = i
503 break
505 if target_index > 0:
506 return normalized_stack[:target_index]
508 # Fall back to known hierarchy registry
509 return get_ancestors_from_hierarchy(target_type)
512def is_ancestor_in_context(ancestor_type, descendant_type):
513 """
514 Check if ancestor_type comes before descendant_type in the context hierarchy.
516 This determines whether changes to ancestor_type should affect descendant_type.
518 Args:
519 ancestor_type: The potential ancestor type
520 descendant_type: The potential descendant type
522 Returns:
523 True if ancestor_type is an ancestor of descendant_type,
524 False otherwise.
525 """
526 from objectstate.lazy_factory import get_base_type_for_lazy
528 # Check 1: Is ancestor_type the lazy base of descendant_type?
529 # This handles GlobalPipelineConfig → PipelineConfig relationship
530 # PipelineConfig is a lazy version of GlobalPipelineConfig
531 descendant_base = get_base_type_for_lazy(descendant_type)
532 if descendant_base is not None and descendant_base == ancestor_type:
533 return True
535 # Check 2: Normalize for comparison (handles nested lazy configs like LazyPathPlanningConfig)
536 ancestor_base = _normalize_type(ancestor_type)
537 descendant_normalized = _normalize_type(descendant_type)
539 # Check 3: Active context stack (uses normalized types)
540 normalized_stack = get_normalized_stack()
541 if normalized_stack:
542 ancestor_index = -1
543 descendant_index = -1
544 for i, base_t in enumerate(normalized_stack):
545 if base_t == ancestor_base:
546 ancestor_index = i
547 if base_t == descendant_normalized:
548 descendant_index = i
550 if ancestor_index >= 0 and descendant_index >= 0:
551 return ancestor_index < descendant_index
553 # Check 4: Known hierarchy registry (uses actual types, not normalized)
554 ancestors = get_ancestors_from_hierarchy(descendant_type)
555 # Check if ancestor_type OR its normalized form is in the ancestor list
556 return ancestor_type in ancestors or ancestor_base in [_normalize_type(a) for a in ancestors]
559def is_same_type_in_context(type_a, type_b):
560 """
561 Check if two types are the same when normalized.
563 Handles lazy vs base type equivalence.
565 Args:
566 type_a: First type to compare
567 type_b: Second type to compare
569 Returns:
570 True if both types normalize to the same base type.
571 """
572 return _normalize_type(type_a) == _normalize_type(type_b)
575# ============================================================================
576# Context Stack Building (for UI placeholder resolution)
577# ============================================================================
580@contextmanager
581def build_context_stack(
582 object_instance: object,
583 ancestor_objects_with_scopes: Sequence[tuple[str, object]] = (),
584 current_scope_id: str | None = None,
585 use_live: bool = True,
586):
587 """
588 Enter the complete layer stack that resolves ``object_instance``; leave it on exit.
590 The stack is built from the arguments and the global config alone. A context the
591 caller has entered is set aside for the duration of the block and is active again
592 after it: none of its layers is below, between or above the layers entered here.
593 This is how an ``ObjectState`` resolves, so a state's resolved values depend on
594 its object, the states registered above it and the global config, and never on
595 where the read was made. A lazy object that no state owns is resolved by plain
596 attribute access against whatever ``config_context`` is entered.
598 Layers, outermost first:
599 1. The global config, under scope ``""``. When ``object_instance`` is itself a
600 global config, a fresh instance of its type holding the declared defaults is
601 entered above it, also under scope ``""``, and masks it.
602 2. Each ancestor, under its own scope id. Global-config ancestors are left out
603 because layer 1 already stands for them.
604 3. ``object_instance``, under ``current_scope_id``.
606 Nothing is set until the ``with`` block is entered. Any layer that cannot be
607 built raises, and the layers entered before it are left again.
609 Args:
610 object_instance: The object whose fields are being resolved.
611 ancestor_objects_with_scopes: ``(scope_id, object)`` pairs, least specific first.
612 current_scope_id: Scope id of ``object_instance``.
613 use_live: Read the live global config (True) or the saved one (False).
614 """
615 from objectstate.lazy_factory import GlobalConfigBase
617 def enter(stack: ExitStack, obj: object, scope_id: str | None, mask: bool = False):
618 stack.enter_context(
619 config_context(obj, mask_with_none=mask, use_live_global=use_live, scope_id=scope_id)
620 )
622 with ExitStack() as stack:
623 # With no context active, the first layer entered starts from the global config.
624 for variable, no_context in (
625 (current_temp_global, None),
626 (context_type_stack, ()),
627 (context_layer_stack, NO_LAYERS),
628 ):
629 stack.callback(variable.reset, variable.set(no_context))
631 if isinstance(object_instance, GlobalConfigBase):
632 enter(stack, type(object_instance)(), "", mask=True)
634 for scope_id, ancestor_obj in ancestor_objects_with_scopes:
635 if isinstance(ancestor_obj, GlobalConfigBase):
636 continue
637 enter(stack, ancestor_obj, scope_id)
639 enter(stack, object_instance, current_scope_id)
640 yield
643def merge_configs(base, overrides: Dict[str, Any]):
644 """
645 Merge overrides into base config, creating new immutable instance.
647 A ``None`` override is kept: it means "inherit this field from the parent context".
648 Raises whatever the config's constructor raises when the overrides are not accepted.
650 Args:
651 base: Base config instance (base config type)
652 overrides: Dict of field_name -> value to override
654 Returns:
655 New config instance with overrides applied
656 """
657 if not base or not overrides:
658 return base
660 from objectstate.lazy_factory import replace_raw
661 return replace_raw(base, **overrides)
664def get_base_global_config(use_live: bool = True):
665 """
666 Get the global config a context starts from when none is active.
668 This is the global config of the base config type (see ``objectstate.global_config``),
669 or an instance of that type holding its declared defaults when there is none.
670 Raises when no base config type is set or its defaults cannot be constructed.
672 Args:
673 use_live: If True (default), return LIVE config (UI sees unsaved edits).
674 If False, return SAVED config (compiler sees saved values).
675 """
676 from objectstate.config import get_base_config_type
677 from objectstate.global_config import get_live_global_config, get_saved_global_config
679 base_config_type = get_base_config_type()
681 read = get_live_global_config if use_live else get_saved_global_config
682 current_global = read(base_config_type)
683 if current_global is not None:
684 return current_global
686 return base_config_type()
689def get_current_temp_global():
690 """
691 Get current context or None.
693 This is the primary interface for resolution functions to access
694 the current context. Returns None if no context is active.
696 Returns:
697 Current merged global config or None
698 """
699 return current_temp_global.get(None)
702# Utility functions for debugging and introspection
704def get_context_info() -> Dict[str, Any]:
705 """
706 Get information about current context for debugging.
708 Returns:
709 Dict with context information including type, field count, etc.
710 """
711 current = get_current_temp_global()
712 if current is None:
713 return {"active": False}
715 return {
716 "active": True,
717 "type": type(current).__name__,
718 "field_count": len(fields(current)) if is_dataclass(current) else 0,
719 "non_none_fields": sum(1 for f in fields(current)
720 if getattr(current, f.name) is not None) if is_dataclass(current) else 0
721 }
724def extract_all_configs_from_context() -> Dict[type, Any]:
725 """
726 Extract all *_config attributes from current context.
728 This is used by the resolution system to get all available configs
729 for cross-dataclass inheritance resolution.
731 Returns:
732 Dict of config type -> config instance.
733 """
734 current = get_current_temp_global()
735 if current is None:
736 return {}
738 return extract_all_configs(current)
741@functools.cache
742def _declared_config_fields(dataclass_type: type) -> tuple[tuple[str, type], ...]:
743 """The fields of a dataclass type that are declared as dataclass types, with those types.
745 A fact about the class, so it is kept per class.
746 """
747 annotations = get_type_hints(dataclass_type, include_extras=True)
748 declared = []
749 for field_info in fields(dataclass_type):
750 field_type = annotations.get(field_info.name, field_info.type)
751 actual_type = resolve_optional(resolve_annotated(field_type))
752 if is_dataclass(actual_type):
753 declared.append((field_info.name, actual_type))
754 return tuple(declared)
757def extract_all_configs(context_obj) -> Dict[type, Any]:
758 """
759 Extract all config instances from a context object using type-driven approach.
761 A context layer keeps its own result (``ContextLayer.configs``); this function
762 itself keeps nothing.
764 This function leverages dataclass field type annotations to efficiently extract
765 config instances, avoiding string matching and runtime attribute scanning.
767 Args:
768 context_obj: Object to extract configs from (orchestrator, merged config, etc.)
770 Returns:
771 Dict mapping nominal config types to config instances.
772 """
773 if context_obj is None:
774 return {}
776 configs: Dict[type, Any] = {}
778 # Include the context object itself if it's a dataclass
779 if is_dataclass(context_obj):
780 _record_config_instance(configs, context_obj, source="context object")
782 # Type-driven extraction: Use dataclass field annotations to find config fields
783 if is_dataclass(type(context_obj)):
784 for field_name, declared_type in _declared_config_fields(type(context_obj)):
785 # Read the stored value: a None field holds no config in this object,
786 # and resolving it here would ask the context that is being read.
787 field_value = object.__getattribute__(context_obj, field_name)
788 if field_value is None:
789 continue
790 if not _is_compatible_config_type(field_value, declared_type):
791 raise TypeError(
792 f"{type(context_obj).__name__}.{field_name} declares "
793 f"{declared_type.__name__}, got {type(field_value).__name__}."
794 )
795 _record_config_instance(
796 configs,
797 field_value,
798 source=f"{type(context_obj).__name__}.{field_name}",
799 )
801 # For non-dataclass objects (orchestrators, etc.), extract dataclass attributes
802 else:
803 _extract_from_object_attributes_typed(context_obj, configs)
805 if logger.isEnabledFor(logging.DEBUG):
806 logger.debug(
807 "Extracted %d configs: %s",
808 len(configs),
809 [config_type.__name__ for config_type in configs],
810 )
811 return configs
814def _record_config_instance(
815 configs: Dict[type, Any],
816 config_instance: Any,
817 *,
818 source: str,
819) -> None:
820 """Record one nominal config identity and reject ambiguous duplicates."""
822 from objectstate.lazy_factory import get_base_type_for_lazy
824 instance_type = type(config_instance)
825 base_type = get_base_type_for_lazy(instance_type) or instance_type
826 existing = configs.get(base_type)
827 if existing is not None and existing is not config_instance:
828 raise ValueError(
829 f"{source} introduces a second {base_type.__name__} instance into "
830 "one configuration context."
831 )
832 configs[base_type] = config_instance
835def _extract_from_object_attributes_typed(obj, configs: Dict[type, Any]) -> None:
836 """
837 Type-safe extraction from object attributes for non-dataclass objects.
839 This is used for orchestrators and other objects that aren't dataclasses
840 but have config attributes. Uses type checking instead of string matching.
841 """
842 for attr_name in dir(obj):
843 if attr_name.startswith('_'):
844 continue
846 try:
847 attr_value = getattr(obj, attr_name)
848 except (AttributeError, TypeError):
849 continue
850 if attr_value is not None and is_dataclass(attr_value):
851 _record_config_instance(
852 configs,
853 attr_value,
854 source=f"{type(obj).__name__}.{attr_name}",
855 )
856 logger.debug(
857 "Extracted config %s from attribute %s",
858 type(attr_value).__name__,
859 attr_name,
860 )
863def _is_compatible_config_type(value, expected_type) -> bool:
864 """
865 Check if value is compatible with expected_type, handling lazy-to-base type mapping.
867 This handles cases where:
868 - value is LazyStepMaterializationConfig, expected_type is StepMaterializationConfig
869 - value is a subclass of the expected type
870 - value is exactly the expected type
871 """
872 from objectstate.lazy_factory import get_base_type_for_lazy
874 expected_type = resolve_optional(resolve_annotated(expected_type))
875 if not isinstance(expected_type, type):
876 return False
877 value_type = type(value)
878 value_base = get_base_type_for_lazy(value_type) or value_type
879 expected_base = get_base_type_for_lazy(expected_type) or expected_type
880 return issubclass(value_base, expected_base)
883def spawn_thread_with_context(
884 target,
885 daemon: bool = True,
886 name: str = None,
887 *args,
888 **kwargs,
889):
890 """Spawn a thread with current context propagated.
892 A new thread starts with no ``config_context`` entered. This starts one with the
893 layers the caller has entered (``contextvars.copy_context()``), so plain lazy
894 objects read in it what they read in the caller. The global config and the
895 registry are one per process and need no copying.
897 Args:
898 target: The function to run in thread
899 daemon: Whether thread should be a daemon thread (default: True)
900 name: Optional thread name
901 *args: Positional arguments to pass to target
902 **kwargs: Keyword arguments to pass to target
904 Returns:
905 The spawned Thread object
907 Example:
908 >>> spawn_thread_with_context(my_function, arg1, kwarg=value)
909 >>> # Equivalent to:
910 >>> ctx = contextvars.copy_context()
911 >>> threading.Thread(target=ctx.run, args=(my_function, arg1), kwargs={'kwarg': value}).start()
912 """
913 ctx = contextvars.copy_context()
914 thread = threading.Thread(
915 target=lambda: ctx.run(target, *args, **kwargs),
916 daemon=daemon,
917 name=name,
918 )
919 thread.start()
920 return thread