from __future__ import annotations
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Mapping, Optional, Tuple
from copy import deepcopy
from ClearMap.config.compound_keys import PairKey
from ClearMap.config.config_adjusters.type_hints import SampleManagerProtocol
from ClearMap.config.defaults_provider import DefaultsProvider
from ClearMap.config.config_adjusters.policy_specs import ContainerSpecRegistry, TemplateKind, InstanceContainerSpec
[docs]
def require_dict(parent: Mapping[str, Any], key: str, *, path: tuple[str, ...] = (), allow_missing: bool = False) -> dict[str, Any]:
"""
Extract a dict-valued child from `parent[key]`.
Parameters
----------
parent:
Mapping to extract from.
key:
Key to extract.
path:
Cosmetic only; dotted path prefix used for error reporting.
allow_missing:
If True, missing key returns {}. If False (default), missing key raises.
Returns
-------
dict[str, Any]
The dict value at parent[key].
Raises
------
KeyError
If key is missing and allow_missing is False.
ValueError
If key exists but is not a dict.
"""
dotted = '.'.join(path + (key,)) if path else key
if key not in parent:
if allow_missing:
return {}
raise KeyError(f'Defaults invariant violated: missing key {dotted}')
v = parent.get(key)
if not isinstance(v, dict):
raise ValueError(f'Defaults invariant violated: {dotted} must be a dict,'
f' got {type(v).__name__} with value: {v!r}')
return v
[docs]
@dataclass(frozen=True)
class RunFacts:
"""
Snapshot of SM-derived facts used for binding templates in a *single* resolver instance.
Notes
-----
- renamed_map provides ergonomic lookup old->new.
- renamed_items provides deterministic iteration and hash-friendly representation (if needed).
"""
sample_channels: Tuple[str, ...]
alignment_reference: Optional[str]
renamed_map: Mapping[str, str]
renamed_items: Tuple[Tuple[str, str], ...]
[docs]
@dataclass(frozen=True)
class BindPolicy:
"""
Central binding policy knobs for the resolver.
This is intentionally narrow: "binding" policy, not "reconcile" policy.
"""
keep_templates_out_of_working_config: bool = True
allow_entry_overwrite: bool = False # templates are generally missing-only in adjusters
# ---------------------------------------------------------------------
# Resolver
# ---------------------------------------------------------------------
[docs]
class TemplateResolverBase:
def __init__(self, *, defaults: DefaultsProvider, policy: BindPolicy = BindPolicy()) -> None:
self._defaults = defaults
self._policy = policy
@property
def defaults(self) -> DefaultsProvider:
return self._defaults
@property
def policy(self) -> BindPolicy:
return self._policy
# Static missing only
# FIXME: do we exclude templates here, or leave to caller?
[docs]
def section_missing_only_defaults(self, section_name: str) -> Dict[str, Any]:
"""
Returns the raw defaults dict for a section (caller performs missing-only merge).
If keep_templates_out_of_working_config is True, this function *does not* strip
templates here (because the correct strip scope is section-specific).
The adjuster can do template skipping in the deep-merge implementation.
"""
d = self._defaults.get(section_name)
return deepcopy(d) if isinstance(d, dict) else {}
[docs]
class GroupTemplatesResolver(TemplateResolverBase):
def __init__(self, *, defaults: DefaultsProvider, group_base_dir: Optional[Path],
policy: BindPolicy = BindPolicy()) -> None:
super().__init__(defaults=defaults, policy=policy)
self.group_base_dir = group_base_dir
[docs]
class ExperimentTemplatesResolver(TemplateResolverBase):
"""
Binds defaults templates to a specific run context (sm + defaults) and offers
*pure* template selection/expansion utilities.
Design intent
-------------
- This object should NOT mutate config.
- It should NOT decide reconciliation strategy (preserve vs overwrite vs prune).
- It should provide:
* target membership lists
* fully-expanded template dicts for a given entry key
* canonicalization helpers for compound keys
* stable per-run facts (RunFacts)
About "alignment_reference_channel changes"
-------------------------------------------
If you allow earlier steps to rename channels, then the "reference channel" may
need to track those renames. This resolver addresses that in two ways:
1) It snapshots the SM-reported reference and then normalizes it through the
rename map (resolve_renamed_channel). This prevents a stale reference if SM
reports the "old" name but the run applies renames.
2) If your SM is itself config-dependent and can change during the run, you
should instantiate a new resolver after the rename/reconcile step (or at least
rebuild facts). This keeps binding coherent with the current working state.
"""
PAIRS_ORIENTED = True
_TEMPLATE_BINDERS = {
("sample", TemplateKind.CHANNELS): 'sample_channel_entry_template',
("registration", TemplateKind.CHANNELS): 'registration_entry_template',
("stitching", TemplateKind.CHANNELS): 'stitching_entry_template',
("cell_map", TemplateKind.CHANNELS): 'cell_map_entry_template',
('cell_map', TemplateKind.PERF): 'cell_map_perf_entry_template',
("tract_map", TemplateKind.CHANNELS): 'tract_map_entry_template',
("vasculature", TemplateKind.CHANNELS): 'vasculature_binarization_template',
("vasculature", TemplateKind.PERF): 'vasculature_perf_binarization_template',
("colocalization", TemplateKind.PAIRS): 'colocalization_entry_template',
}
def __init__(self, defaults: DefaultsProvider, sm: SampleManagerProtocol, *,
specs_registry: ContainerSpecRegistry, policy: BindPolicy = BindPolicy()) -> None:
super().__init__(defaults=defaults, policy=policy)
self._sm = sm
self._specs = specs_registry.by_template_kind()
self._validate_template_binders()
# Snapshot, but normalize reference through rename map to avoid stale references.
renamed_map: Mapping[str, str] = dict(sm.renamed_channels or {})
renamed_items: Tuple[Tuple[str, str], ...] = tuple(renamed_map.items())
sample_channels = tuple(sm.channels or ())
ref = sm.alignment_reference_channel
ref = self.resolve_renamed_channel(ref, renamed_map=renamed_map)
self._facts = RunFacts( sample_channels=sample_channels, alignment_reference=ref,
renamed_map=renamed_map, renamed_items=renamed_items)
def _validate_template_binders(self):
for key, method_name in self._TEMPLATE_BINDERS.items():
# if key not in self._specs: # We're allowed not to have specs for all bindings
# raise KeyError(f"No InstanceContainerSpec for template binding {key}")
_ = self.get_template_method(key)
for (section, kind), spec in self._specs.items():
if (section, kind) not in self._TEMPLATE_BINDERS:
raise KeyError(f"No template binding for InstanceContainerSpec {(section, kind)}")
[docs]
def get_template_method(self, key):
method_name = self._TEMPLATE_BINDERS.get(key)
if method_name is None:
raise KeyError(f"No template binding for {key}")
fn = getattr(self, method_name, None)
if not callable(fn):
raise RuntimeError(f"Template binding {method_name!r} not found or not callable")
return fn
# def _templates_and_base(self, *, section: str, kind: TemplateKind, base_key: str) -> tuple[
# dict[str, Any], dict[str, Any]]:
# root = self._templates_root(section=section, kind=kind)
# spec = self._specs[(section, kind)]
# path = spec.abs_templates_path()
# base = require_dict(root, base_key, path=*path)
# return root, base
@property
def sm(self) -> SampleManagerProtocol:
return self._sm
[docs]
def run_facts(self) -> RunFacts:
return self._facts
# ---- common utilities
[docs]
@staticmethod
def resolve_renamed_channel(name: Optional[str], *, renamed_map: Mapping[str, str],
max_hops: int = 32) -> Optional[str]:
"""
Follow rename chains old->new transitively to avoid stale references.
Protects against cycles via max_hops.
"""
if name is None:
return None
cur = name
hops = 0
while cur in renamed_map and hops < max_hops:
nxt = renamed_map[cur]
if nxt == cur:
break
cur = nxt
hops += 1
return cur
[docs]
def expand(self, obj: Any, *, channel: str, reference: Optional[str] = None) -> Any:
"""Recursively substitute ${channel} and ${reference} in strings within obj."""
def sub_one(s: str) -> str:
out = s.replace('${channel}', channel)
if reference is not None:
out = out.replace('${reference}', reference)
return out
if isinstance(obj, str):
return sub_one(obj)
if isinstance(obj, list):
return [self.expand(x, channel=channel, reference=reference) for x in obj]
if isinstance(obj, tuple):
return tuple(self.expand(x, channel=channel, reference=reference) for x in obj)
if isinstance(obj, dict):
return {k: self.expand(v, channel=channel, reference=reference) for k, v in obj.items()}
return obj
[docs]
def entry_template(self, *, section: str, kind: 'TemplateKind', key: str, **ctx: Any) -> Dict[str, Any]:
fn = self.get_template_method((section, kind))
# Most bindings accept (key) or (channel), stitching accepts root_layout via ctx.
return fn(key, **ctx)
def _get_nested(self, root: Any, path: tuple[str, ...]) -> Any:
node = root
for p in path:
if not isinstance(node, dict):
return None
node = node.get(p)
return node
def _templates_root_and_spec(self, *, section: str, kind: 'TemplateKind') -> tuple[dict[str, Any], InstanceContainerSpec]:
spec = self._specs.get((section, kind))
if spec is None:
raise KeyError(f'No InstanceContainerSpec for {(section, kind)}')
rel = spec.rel_templates_path()
if rel is None:
raise KeyError(f'{(section, kind)} has no defaults_templates_path')
# defaults are section-rooted, and rel_templates_path is section-relative
sec_defaults = self._defaults.get(section)
node = self._get_nested(sec_defaults, rel)
root = deepcopy(node) if isinstance(node, dict) else {}
return root, spec
def _require_template_dict(self, *, section: str, kind: TemplateKind, key: str) -> dict[str, Any]:
root, spec = self._templates_root_and_spec(section=section, kind=kind)
prefix = spec.abs_templates_path() or (section,) # defensive fallback
return require_dict(root, key, path=prefix)
# -----------------------------------------------------------------
# Sample-level defaults
# -----------------------------------------------------------------
[docs]
def sample_channel_entry_template(self, channel: str) -> dict[str, Any]:
return self._require_template_dict(section='sample', kind=TemplateKind.CHANNELS, key='channel')
# -----------------------------------------------------------------
# Stitching
# -----------------------------------------------------------------
[docs]
def stitching_target_channels(self) -> list[str]:
return self._sm.stitchable_channels
[docs]
def stitching_entry_template(self, channel: str, *, root_layout: Optional[str]) -> Dict[str, Any]:
"""
Returns an expanded stitching template for `channel`.
The caller decides role (layout vs derived) and any canonicalization.
"""
templates, spec = self._templates_root_and_spec(section='stitching', kind=TemplateKind.CHANNELS)
err_path = spec.abs_templates_path()
layout_tpl = require_dict(templates, 'layout_template', path=err_path)
derived_tpl = require_dict(templates, 'derived_template', path=err_path)
# If root_layout is None or equals channel, we consider it a layout role template bind.
if root_layout is None or root_layout == channel:
return deepcopy(self.expand(layout_tpl, channel=channel, reference=channel))
return deepcopy(self.expand(derived_tpl, channel=channel, reference=root_layout))
@property
def stitching_template_keys(self):
return 'layout_template', 'derived_template'
# -----------------------------------------------------------------
# Registration
# -----------------------------------------------------------------
[docs]
def registration_target_channels(self) -> list[str]:
return list(self._facts.sample_channels)
[docs]
def registration_entry_template(self, channel: str) -> Dict[str, Any]:
if self._sm.data_type(channel) == 'autofluorescence':
tpl_channel_key = 'autofluorescence'
reference = 'atlas'
else:
tpl_channel_key = 'regular'
reference = self._facts.alignment_reference or channel
base = self._require_template_dict(section='registration', kind=TemplateKind.CHANNELS, key=tpl_channel_key)
return deepcopy(self.expand(base, channel=channel, reference=reference))
# -----------------------------------------------------------------
# Cell map / tract map
# -----------------------------------------------------------------
[docs]
def cell_map_target_channels(self) -> list[str]:
return self._sm.get_channels_by_pipeline('CellMap', as_list=True)
[docs]
def cell_map_entry_template(self, channel: str) -> Dict[str, Any]:
base = self._require_template_dict(section='cell_map', kind=TemplateKind.CHANNELS, key='channel')
return deepcopy(self.expand(base, channel=channel, reference=self._facts.alignment_reference))
[docs]
def cell_map_perf_entry_template(self, channel: str) -> Dict[str, Any]:
base = self._require_template_dict(section='cell_map', kind=TemplateKind.PERF, key='channel')
return deepcopy(self.expand(base, channel=channel, reference=self._facts.alignment_reference))
[docs]
def tract_map_target_channels(self) -> list[str]:
return self._sm.get_channels_by_pipeline('TractMap', as_list=True)
[docs]
def tract_map_entry_template(self, channel: str) -> Dict[str, Any]:
base = self._require_template_dict(section='tract_map', kind=TemplateKind.CHANNELS, key='channel')
return deepcopy(self.expand(base, channel=channel, reference=self._facts.alignment_reference))
# -----------------------------------------------------------------
# Colocalization (compound)
# -----------------------------------------------------------------
[docs]
def colocalization_entry_template(self, canonical_pair_key: str) -> Dict[str, Any]:
# templates_root = self._templates_root_and_spec(section='colocalization', kind=TemplateKind.PAIRS)
# base = require_dict(templates_root, 'channel', path=('colocalization', 'templates'))
base = self._require_template_dict(section='colocalization', kind=TemplateKind.PAIRS, key='channel')
pk = PairKey.from_string(canonical_pair_key, oriented=False)
expanded = self.expand(base, channel=str(pk), reference=self._facts.alignment_reference)
if isinstance(expanded, dict):
expanded = {**expanded, 'channels': [pk.a, pk.b]}
return deepcopy(expanded)
# -----------------------------------------------------------------
# Vasculature
# -----------------------------------------------------------------
[docs]
def vasculature_target_channels(self) -> list[str]:
return self._sm.get_channels_by_pipeline('TubeMap', as_list=True)
[docs]
def vasculature_binarization_template(self, channel: str) -> Dict[str, Any]:
# Select template by data type
tpl_chan_k = 'vessels_template' if self._sm.data_type(channel) == 'vessels' else 'large_vessels_template'
return self._require_template_dict(section='vasculature', kind=TemplateKind.CHANNELS, key=tpl_chan_k)
[docs]
def vasculature_perf_binarization_template(self, channel: str) -> Dict[str, Any]:
# Select template by data type
tpl_chan_k = 'vessels_template' if self._sm.data_type(channel) == 'vessels' else 'large_vessels_template'
return self._require_template_dict(section='vasculature', kind=TemplateKind.PERF, key=tpl_chan_k)
# ----------------- Runner-managed resolver instance -----------------
_CURRENT_RESOLVER: Optional[TemplateResolverBase] = None
[docs]
def set_current_resolver(resolver: TemplateResolverBase) -> Optional[TemplateResolverBase]:
global _CURRENT_RESOLVER
old = _CURRENT_RESOLVER
_CURRENT_RESOLVER = resolver
return old
[docs]
def get_current_resolver() -> TemplateResolverBase:
if _CURRENT_RESOLVER is None:
raise RuntimeError('No TemplateResolverBase derived class has been installed for this run')
return _CURRENT_RESOLVER