Coverage for src/loman/util.py: 100%

1"""Utility functions and classes for loman computation graphs.""" 

2 

3from __future__ import annotations 

4 

5import itertools 

6import types 

7from collections.abc import Callable, Generator, Hashable, Iterable, Mapping, Sequence 

8from dataclasses import dataclass, field 

9from typing import TYPE_CHECKING, Any, Generic, Protocol, TypeVar, cast, overload 

10 

11import numpy as np 

12import pandas as pd 

13 

14if TYPE_CHECKING: 

15 from loman.computeengine import Computation 

16 from loman.nodekey import Name, NodeKey 

17else: 

18 Computation = Any 

19 Name = Any 

20 NodeKey = Any 

21 

22T = TypeVar("T") 

23R = TypeVar("R") 

24K = TypeVar("K", bound=Hashable) 

25 

26 

27_NO_VALUE = object() 

28 

29 

30@dataclass(frozen=True) 

31class PlannedNode: 

32 """One node a feature intends to create, described rather than created. 

33 

34 Features return these instead of changing the computation, so the builder can 

35 check every node and edge of a definition before any of it is applied. A node 

36 with no ``func`` is an input node holding ``value``; otherwise ``func`` is 

37 called with ``args``, where each argument is either a :class:`NodeKey` to 

38 depend on or a :class:`ConstantValue` to pass through unchanged. 

39 """ 

40 

41 node_key: NodeKey 

42 func: Callable[..., Any] | None = None 

43 args: tuple[Any, ...] = () 

44 value: Any = _NO_VALUE 

45 label: Name | None = None 

46 

47 @classmethod 

48 def input_node(cls, node_key: NodeKey, value: Any, label: Name | None = None) -> PlannedNode: 

49 """Plan an input node holding a fixed value.""" 

50 return cls(node_key, value=value, label=label) 

51 

52 @classmethod 

53 def link(cls, node_key: NodeKey, source: NodeKey, label: Name | None = None) -> PlannedNode: 

54 """Plan a node that takes its value from another node unchanged.""" 

55 from loman.computeengine import identity_function 

56 

57 return cls(node_key, identity_function, (source,), label=label) 

58 

59 @classmethod 

60 def calc( 

61 cls, node_key: NodeKey, func: Callable[..., Any], args: Sequence[Any], label: Name | None = None 

62 ) -> PlannedNode: 

63 """Plan a calculation node.""" 

64 return cls(node_key, func, tuple(args), label=label) 

65 

66 @property 

67 def predecessors(self) -> tuple[NodeKey, ...]: 

68 """Return the nodes this planned node would depend on.""" 

69 from loman.computeengine import ConstantValue 

70 

71 return tuple(arg for arg in self.args if not isinstance(arg, ConstantValue)) 

72 

73 def apply_to(self, comp: Computation) -> NodeKey: 

74 """Create this node in a computation.""" 

75 if self.func is None: 

76 comp.add_node(self.node_key, value=self.value) 

77 else: 

78 comp.add_node(self.node_key, self.func, args=list(self.args), inspect=False) 

79 return self.node_key 

80 

81 

82@dataclass(frozen=True) 

83class BlockContext(Generic[K]): 

84 """What a feature is given when it plans its nodes. 

85 

86 ``blocks`` maps each key to the path of its generated block. ``block`` is the 

87 template, so a feature can check that a relative name it was given really is a 

88 node, or an input node, of the block being repeated. 

89 """ 

90 

91 comp: Computation 

92 block: Computation 

93 base_path: NodeKey 

94 blocks: Mapping[K, NodeKey] 

95 definition_object: object = None 

96 ignore_self: bool = False 

97 planned: set[NodeKey] = field(default_factory=set) 

98 

99 def require_block_node(self, name: Name, description: str) -> NodeKey: 

100 """Resolve a relative name that must be a node inside every block. 

101 

102 The name may come from the template, or from a node an earlier feature 

103 already planned, so one feature can build on another's output. Features 

104 are planned in the order they are declared. 

105 """ 

106 node_key = _to_node_name(name, description) 

107 if not self.block.has_node(node_key) and not self._is_planned_in_every_block(node_key): 

108 msg = f"{description} does not exist in the block: {node_key!r}" 

109 raise ValueError(msg) 

110 return node_key 

111 

112 def _is_planned_in_every_block(self, node_key: NodeKey) -> bool: 

113 """Return whether an earlier feature planned this node in every block.""" 

114 return all(block_path.join(node_key) in self.planned for block_path in self.blocks.values()) 

115 

116 def require_block_input(self, name: Name, description: str, *, create: bool = False) -> NodeKey: 

117 """Resolve a relative name that must be an input node inside every block. 

118 

119 A node the template declares must be an input: replacing a calculation 

120 would silently discard it. A node an earlier feature planned is accepted 

121 as-is, since the template has nothing to say about it. A name the template 

122 never mentions is rejected, because it is usually a typo that would 

123 otherwise add a dead node to every block — pass ``create=True`` to allow 

124 it deliberately. 

125 """ 

126 from loman.consts import NodeAttributes 

127 

128 node_key = _to_node_name(name, description) 

129 if not self.block.has_node(node_key): 

130 if self._is_planned_in_every_block(node_key) or create: 

131 return node_key 

132 msg = ( 

133 f"{description} does not exist in the block: {node_key!r}. " 

134 "Pass create=True to add it to every block anyway." 

135 ) 

136 raise ValueError(msg) 

137 node = self.block.dag.nodes[node_key] 

138 if node.get(NodeAttributes.FUNC) is not None or next(self.block.dag.predecessors(node_key), None) is not None: 

139 msg = f"{description} must be an input node: {node_key!r}" 

140 raise ValueError(msg) 

141 return node_key 

142 

143 def bind(self, func: Any) -> Any: 

144 """Bind a callback to the class a computation factory was defined from. 

145 

146 Returns anything that is not callable unchanged, so a plain node name 

147 passes through. 

148 """ 

149 from loman.computeengine import _bind_self 

150 

151 return _bind_self(func, self.definition_object, self.ignore_self) 

152 

153 

154class BlockFeature(Protocol): 

155 """One wiring pattern applied to every copy of a repeated block. 

156 

157 A feature never changes the computation itself. It describes the nodes it 

158 wants, and the builder validates every feature's plan together before 

159 applying any of it, so a definition that fails leaves the graph untouched. 

160 Implement this protocol to add a wiring pattern of your own. 

161 """ 

162 

163 def plan(self, ctx: BlockContext[Any]) -> Iterable[PlannedNode]: 

164 """Describe the nodes to create, without changing anything.""" 

165 ... 

166 

167 

168@dataclass(frozen=True) 

169class FanOut(Generic[K]): 

170 """Wire one source node to a relative input in every repeated block. 

171 

172 ``source`` normally names a single outer node feeding every block. Passing a 

173 callable instead resolves a source node per key, as ``source(key)``, so each 

174 block can read from a different outer node. With a ``transform``, each target 

175 is calculated as ``transform(value, key)``. 

176 

177 ``target`` is a name inside each block. It must be something the template 

178 declares or refers to, so a typo does not quietly add a dead node to every 

179 block; set ``create=True`` to feed a name the template never mentions. 

180 """ 

181 

182 source: Name | Callable[[K], Name] 

183 target: Name 

184 transform: Callable[[Any, K], Any] | None = None 

185 create: bool = False 

186 

187 def plan(self, ctx: BlockContext[K]) -> Iterable[PlannedNode]: 

188 """Plan one target node per key, linked or transformed from its source.""" 

189 from loman.computeengine import C 

190 

191 target = ctx.require_block_input(self.target, "Fan-out target", create=self.create) 

192 sources = _resolve_fan_out_sources(ctx.bind(self.source), ctx.blocks) 

193 transform = ctx.bind(self.transform) 

194 for key, block_path in ctx.blocks.items(): 

195 node_key = block_path.join(target) 

196 if transform is None: 

197 yield PlannedNode.link(node_key, sources[key]) 

198 else: 

199 yield PlannedNode.calc(node_key, transform, (sources[key], C(key))) 

200 

201 

202@dataclass(frozen=True) 

203class Positional: 

204 """Adapt an aggregator that takes positional arguments to a keyed ``combine``. 

205 

206 ``combine`` receives an ordered mapping so keys stay attached to values, which 

207 is usually what you want. Where an existing function takes the values 

208 positionally, wrap it rather than repeating ``lambda m: fn(*m.values())``:: 

209 

210 FanIn("value", "total", combine=Positional(df_hconcat)) 

211 

212 Keys are discarded. A keyed aggregator can use them instead — for dataframes, 

213 ``lambda m: pd.concat(m, axis=1)`` turns them into column labels — but that is 

214 **a different result, not a drop-in replacement**: it adds an outer level to 

215 the column index, where a flat positional concatenation does not. Choose it 

216 because you want the keys in the output, not as a like-for-like swap. 

217 

218 Like any callable that is not an importable module-level function, this needs 

219 ``use_dill_for_functions=True`` to serialize. 

220 """ 

221 

222 func: Callable[..., Any] 

223 

224 def __call__(self, values: Mapping[Any, Any]) -> Any: 

225 """Call the wrapped function with the mapping's values, in order.""" 

226 return self.func(*values.values()) 

227 

228 

229@dataclass(frozen=True) 

230class FanIn(Generic[K]): 

231 """Collect one relative output from every repeated block into a result node. 

232 

233 ``source`` is a name inside each block. ``result`` is **not** relative to the 

234 definition's ``base_path``: it names a node in the outer computation, so the 

235 aggregate can live wherever it belongs rather than being forced under the 

236 blocks. ``BuiltRepeatedBlocks.named`` reports the key that was created. 

237 """ 

238 

239 source: Name 

240 result: Name 

241 combine: Callable[[Mapping[K, Any]], Any] | None = None 

242 

243 def plan(self, ctx: BlockContext[K]) -> Iterable[PlannedNode]: 

244 """Plan one result node gathering the same relative node from every block.""" 

245 from loman.computeengine import C 

246 

247 source = ctx.require_block_node(self.source, "Fan-in source") 

248 result = _to_node_name(self.result, "Fan-in result") 

249 sources = [block_path.join(source) for block_path in ctx.blocks.values()] 

250 yield PlannedNode.calc( 

251 result, 

252 _combine_keyed_values, 

253 (C(tuple(ctx.blocks)), C(ctx.bind(self.combine)), *sources), 

254 label=self.result, 

255 ) 

256 

257 

258@dataclass(frozen=True) 

259class IdNode: 

260 """Give every block a node holding its own key. 

261 

262 Block functions can then depend on their key by name, to look data up or to 

263 branch on it, without the key being wired in from outside. 

264 

265 ``create`` defaults to ``True``, unlike :class:`FanOut`, because creating the 

266 node is this feature's whole job: a template that never mentions the name is 

267 the ordinary case, not a mistake. The cost is that a misspelled name adds a 

268 node nothing reads and leaves the real one unfilled — `validate()` reports 

269 that as an uninitialized input, but only once the graph is built. Set 

270 ``create=False`` where the template does declare the node, to have the 

271 misspelling rejected at definition time instead. 

272 """ 

273 

274 name: Name 

275 create: bool = True 

276 

277 def plan(self, ctx: BlockContext[K]) -> Iterable[PlannedNode]: 

278 """Plan one value node per key, holding that key.""" 

279 name = ctx.require_block_input(self.name, "Identifier node", create=self.create) 

280 for key, block_path in ctx.blocks.items(): 

281 yield PlannedNode.input_node(block_path.join(name), key) 

282 

283 

284@dataclass(frozen=True) 

285class InputValue: 

286 """Give every block the same constant value for one relative input. 

287 

288 Unlike :class:`FanIn`'s ``result``, the shared node this creates **is** 

289 relative to the definition's ``base_path``, landing at ``<base_path>/<name>`` 

290 so that two definitions with different base paths do not collide. 

291 

292 ``create`` behaves as it does on :class:`FanOut`, and defaults the same way: 

293 seeding a value into a name the template never mentions is usually a typo, 

294 so it must be asked for. 

295 """ 

296 

297 name: Name 

298 value: Any 

299 create: bool = False 

300 

301 def plan(self, ctx: BlockContext[K]) -> Iterable[PlannedNode]: 

302 """Plan one shared outer node, linked into every block.""" 

303 name = ctx.require_block_input(self.name, "Input value", create=self.create) 

304 shared = ctx.base_path.join(name) 

305 yield PlannedNode.input_node(shared, self.value, label=self.name) 

306 for block_path in ctx.blocks.values(): 

307 yield PlannedNode.link(block_path.join(name), shared) 

308 

309 

310@dataclass(frozen=True) 

311class BuiltRepeatedBlocks(Generic[K]): 

312 """Node paths created by :meth:`RepeatedBlocks.add_to`. 

313 

314 ``named`` collects the nodes that features chose to label, so a fan-in result 

315 can be looked up by the name it was declared with. 

316 """ 

317 

318 blocks: dict[K, NodeKey] 

319 nodes: tuple[NodeKey, ...] 

320 named: dict[Name, NodeKey] 

321 

322 

323@dataclass(frozen=True) 

324class RepeatedBlocks(Generic[K]): 

325 """Reusable definition for keyed copies of a computation block. 

326 

327 ``features`` describe how data flows in and out of every copy: see 

328 :class:`FanOut`, :class:`FanIn`, :class:`IdNode` and :class:`InputValue`, or 

329 write your own against :class:`BlockFeature`. Features are applied in the 

330 order given, and every feature's plan is validated before any of it is 

331 applied. 

332 

333 ``keep_values`` defaults to ``False``, so only the structure of ``block`` is 

334 copied. This differs from :meth:`Computation.add_block`, which copies values 

335 by default: that call adds one specific block, which may be a sub-model that 

336 has already been populated, whereas this one stamps out many copies of a 

337 template. To give every copy the same value, use an :class:`InputValue`, or a 

338 :class:`FanOut` with no ``transform`` to broadcast an existing node. 

339 """ 

340 

341 block: Computation 

342 keys: Sequence[K] 

343 base_path: Name 

344 features: Sequence[BlockFeature] = () 

345 keep_values: bool = False 

346 

347 def __post_init__(self) -> None: 

348 """Freeze collection inputs as tuples for reusable definitions.""" 

349 object.__setattr__(self, "keys", tuple(self.keys)) 

350 object.__setattr__(self, "features", tuple(self.features)) 

351 

352 def add_to( 

353 self, comp: Computation, definition_object: object = None, ignore_self: bool = False 

354 ) -> BuiltRepeatedBlocks[K]: 

355 """Add this repeated-block definition to a computation.""" 

356 return _add_repeated_blocks_definition(comp, self, definition_object, ignore_self) 

357 

358 

359def _combine_keyed_values( 

360 keys: tuple[Hashable, ...], combine: Callable[[Mapping[Hashable, Any]], Any] | None, *values: Any 

361) -> Any: 

362 """Build a keyed value mapping and optionally combine it.""" 

363 keyed_values = dict(zip(keys, values, strict=True)) 

364 if combine is None: 

365 return keyed_values 

366 return combine(keyed_values) 

367 

368 

369def _to_node_name(name: Name, description: str) -> NodeKey: 

370 """Convert a node name to a key, rejecting a callable given by mistake. 

371 

372 ``Name`` admits any hashable, so a callable would otherwise silently become a 

373 node whose key is the function object itself. Only a fan-out source may be a 

374 callable, where it resolves a different source node for each key. 

375 """ 

376 from loman.nodekey import to_nodekey 

377 

378 if callable(name): 

379 msg = f"{description} must be a node name, not a callable: {name!r}" 

380 raise TypeError(msg) 

381 return to_nodekey(name) 

382 

383 

384def _resolve_fan_out_sources(source: Name | Callable[[K], Name], keys: Iterable[K]) -> dict[K, NodeKey]: 

385 """Resolve a fan-out source to one source node per key. 

386 

387 A callable is applied to each key, so every target can read from a different 

388 node. Any other value names one node broadcast to every key. 

389 """ 

390 if not callable(source): 

391 return dict.fromkeys(keys, _to_node_name(source, "Fan-out source")) 

392 resolve = cast("Callable[[K], Name]", source) 

393 return {key: _to_node_name(resolve(key), f"Fan-out source for key {key!r}") for key in keys} 

394 

395 

396def _repeated_block_paths(keys: Iterable[K], base_path: Name) -> dict[K, NodeKey]: 

397 """Build and validate the paths for repeated block keys.""" 

398 from loman.nodekey import to_nodekey 

399 

400 base_path_node_key = to_nodekey(base_path) 

401 blocks: dict[K, NodeKey] = {} 

402 for key in keys: 

403 if key in blocks: 

404 msg = f"Duplicate repeated block key: {key!r}" 

405 raise ValueError(msg) 

406 blocks[key] = base_path_node_key.join_parts(key) 

407 return blocks 

408 

409 

410def _validate_block_template(comp: Computation, block: Computation) -> None: 

411 """Ensure a block template is not the computation being added to.""" 

412 if block is comp: 

413 msg = "Repeated block template must be a different computation" 

414 raise ValueError(msg) 

415 

416 

417def _is_placeholder(comp: Computation, node_key: NodeKey) -> bool: 

418 """Return whether a node exists only as an unfulfilled forward reference.""" 

419 from loman.consts import NodeAttributes, States 

420 

421 return comp.dag.nodes[node_key][NodeAttributes.STATE] == States.PLACEHOLDER 

422 

423 

424def _is_defined(comp: Computation, node_key: NodeKey) -> bool: 

425 """Return whether a node already has a definition or a value in a computation. 

426 

427 Placeholder nodes do not count as defined: they record that another node 

428 refers to a name that has not been defined yet, so generated nodes are free 

429 to supply that definition. 

430 """ 

431 return comp.has_node(node_key) and not _is_placeholder(comp, node_key) 

432 

433 

434def _validate_repeated_block_nodes(comp: Computation, block: Computation, blocks: Mapping[K, NodeKey]) -> None: 

435 """Ensure repeated blocks will not replace existing nodes.""" 

436 generated_nodes = [block_path.join(node_name) for block_path in blocks.values() for node_name in block.nodes()] 

437 collisions = [node_key for node_key in generated_nodes if _is_defined(comp, node_key)] 

438 if collisions: 

439 msg = f"Repeated blocks would replace existing nodes: {collisions!r}" 

440 raise ValueError(msg) 

441 

442 

443def _validate_acyclic_edges(comp: Computation, edges: Iterable[tuple[NodeKey, NodeKey]]) -> None: 

444 """Ensure proposed dependency edges preserve the computation DAG.""" 

445 import networkx as nx 

446 

447 proposed_graph = nx.DiGraph(comp.dag) 

448 proposed_graph.add_edges_from(edges) 

449 if not nx.is_directed_acyclic_graph(proposed_graph): 

450 msg = "Generated computation utilities would create a cycle" 

451 raise ValueError(msg) 

452 

453 

454def add_repeated_blocks( 

455 comp: Computation, 

456 block: Computation, 

457 keys: Iterable[K], 

458 *, 

459 base_path: Name, 

460 keep_values: bool = False, 

461) -> dict[K, NodeKey]: 

462 """Add one copy of a computation block for each key. 

463 

464 Each key becomes one path segment below ``base_path``. Block values are not 

465 copied by default, making the supplied block a reusable calculation 

466 template. This differs from :meth:`Computation.add_block`, which copies 

467 values by default; see :class:`RepeatedBlocks` for why the defaults differ. 

468 

469 Args: 

470 comp: Computation to add the blocks to. 

471 block: Computation used as the block template. 

472 keys: Unique keys identifying the block instances. 

473 base_path: Parent path for all generated blocks. 

474 keep_values: Whether to copy current values from the template block. 

475 Prefer a broadcast :func:`add_fan_out` over ``True`` when every copy 

476 needs the same value. 

477 

478 Returns: 

479 A mapping from each key to its generated block path. 

480 

481 Raises: 

482 ValueError: If ``keys`` contains a duplicate, or ``block`` is the 

483 computation being added to. 

484 """ 

485 _validate_block_template(comp, block) 

486 blocks = _repeated_block_paths(keys, base_path) 

487 _validate_repeated_block_nodes(comp, block, blocks) 

488 for block_path in blocks.values(): 

489 comp.add_block(block_path, block, keep_values=keep_values) 

490 return blocks 

491 

492 

493def add_fan_out( 

494 comp: Computation, 

495 source: Name | Callable[[K], Name], 

496 targets: Mapping[K, Name], 

497 *, 

498 transform: Callable[[Any, K], Any] | None = None, 

499) -> dict[K, NodeKey]: 

500 """Connect one or more source nodes to a keyed collection of target nodes. 

501 

502 With no ``transform``, each target receives its source value unchanged. If 

503 supplied, ``transform(value, key)`` is evaluated independently for each 

504 target when the target is computed. 

505 

506 ``source`` normally names one node broadcast to every target. Passing a 

507 callable instead resolves a source node per key, as ``source(key)``, so each 

508 target can read from a different node. 

509 

510 Args: 

511 comp: Computation to add the fan-out nodes to. 

512 source: Source node to broadcast, or a function of the key returning one. 

513 targets: Mapping from target keys to target node names. 

514 transform: Optional keyed transformation applied at computation time. 

515 

516 Returns: 

517 A mapping from each key to its target node key. 

518 

519 Raises: 

520 ValueError: If targets are repeated, replace calculation nodes, or a 

521 transformed target is also its own source. 

522 TypeError: If a node name is a callable, or ``source`` resolves one. 

523 """ 

524 from loman.computeengine import C 

525 from loman.consts import NodeAttributes 

526 

527 source_node_keys = _resolve_fan_out_sources(source, targets) 

528 target_node_keys = {key: _to_node_name(target, "Fan-out target") for key, target in targets.items()} 

529 if len(set(target_node_keys.values())) != len(target_node_keys): 

530 msg = "Fan-out targets must be unique" 

531 raise ValueError(msg) 

532 if transform is not None and any(source_node_keys[key] == target_node_keys[key] for key in target_node_keys): 

533 msg = "A transformed fan-out target cannot also be the source node" 

534 raise ValueError(msg) 

535 for target_node_key in target_node_keys.values(): 

536 if comp.has_node(target_node_key) and ( 

537 comp.dag.nodes[target_node_key].get(NodeAttributes.FUNC) is not None 

538 or next(comp.dag.predecessors(target_node_key), None) is not None 

539 ): 

540 msg = f"Fan-out target must be an input or placeholder node: {target_node_key!r}" 

541 raise ValueError(msg) 

542 _validate_acyclic_edges(comp, ((source_node_keys[key], target) for key, target in target_node_keys.items())) 

543 

544 for key, target_node_key in target_node_keys.items(): 

545 if transform is None: 

546 comp.link(target_node_key, source_node_keys[key]) 

547 else: 

548 comp.add_node(target_node_key, transform, args=[source_node_keys[key], C(key)], inspect=False) 

549 return target_node_keys 

550 

551 

552def add_fan_in( 

553 comp: Computation, 

554 result: Name, 

555 sources: Mapping[K, Name], 

556 *, 

557 combine: Callable[[Mapping[K, Any]], Any] | None = None, 

558) -> NodeKey: 

559 """Collect keyed source nodes into one result node. 

560 

561 Source values are assembled into an insertion-ordered mapping when the 

562 result is computed. With no ``combine`` function, that mapping is the 

563 result. Otherwise, ``combine(mapping)`` produces the result value. 

564 

565 Args: 

566 comp: Computation to add the fan-in node to. 

567 result: Name of the generated result node. 

568 sources: Mapping from source keys to source node names. 

569 combine: Optional function that combines the keyed values. 

570 

571 Returns: 

572 The generated result node key. 

573 

574 Raises: 

575 ValueError: If source nodes are repeated, the result already exists, the 

576 result is also a source, or a source already depends on the result. 

577 TypeError: If a node name is a callable. 

578 """ 

579 from loman.computeengine import C 

580 

581 result_node_key = _to_node_name(result, "Fan-in result") 

582 source_node_keys = [_to_node_name(source, "Fan-in source") for source in sources.values()] 

583 if len(set(source_node_keys)) != len(source_node_keys): 

584 msg = "Fan-in source nodes must be unique" 

585 raise ValueError(msg) 

586 if result_node_key in source_node_keys: 

587 msg = "A fan-in result cannot also be a source node" 

588 raise ValueError(msg) 

589 if _is_defined(comp, result_node_key): 

590 msg = f"Fan-in result node already exists: {result_node_key!r}" 

591 raise ValueError(msg) 

592 _validate_acyclic_edges(comp, ((source, result_node_key) for source in source_node_keys)) 

593 

594 comp.add_node( 

595 result_node_key, 

596 _combine_keyed_values, 

597 args=[C(tuple(sources)), C(combine), *source_node_keys], 

598 inspect=False, 

599 ) 

600 return result_node_key 

601 

602 

603def add_id_nodes(comp: Computation, blocks: Mapping[K, NodeKey], name: Name) -> dict[K, NodeKey]: 

604 """Give each block a node holding its own key. 

605 

606 Block functions can then depend on their key by name, to look data up or to 

607 branch on it, without the key being wired in from outside. Unlike a fan-out, 

608 the generated nodes have no predecessors: each simply holds its key as a 

609 value. 

610 

611 Args: 

612 comp: Computation holding the blocks. 

613 blocks: Mapping from each key to its block path. 

614 name: Relative node name to create inside every block. 

615 

616 Returns: 

617 A mapping from each key to its generated identifier node key. 

618 

619 Raises: 

620 ValueError: If a generated node would replace a calculation node. 

621 TypeError: If ``name`` is a callable. 

622 """ 

623 from loman.consts import NodeAttributes 

624 

625 id_node_key = _to_node_name(name, "Identifier node name") 

626 id_nodes = {key: block_path.join(id_node_key) for key, block_path in blocks.items()} 

627 for node_key in id_nodes.values(): 

628 if comp.has_node(node_key) and ( 

629 comp.dag.nodes[node_key].get(NodeAttributes.FUNC) is not None 

630 or next(comp.dag.predecessors(node_key), None) is not None 

631 ): 

632 msg = f"Identifier node must be an input or placeholder node: {node_key!r}" 

633 raise ValueError(msg) 

634 for key, node_key in id_nodes.items(): 

635 comp.add_node(node_key, value=key) 

636 return id_nodes 

637 

638 

639def _validate_planned_nodes(comp: Computation, planned: Sequence[PlannedNode], block_nodes: set[NodeKey]) -> None: 

640 """Ensure a set of planned nodes can all be created without conflict.""" 

641 seen: set[NodeKey] = set() 

642 for planned_node in planned: 

643 node_key = planned_node.node_key 

644 if node_key in seen: 

645 msg = f"Repeated block features would write the same node twice: {node_key!r}" 

646 raise ValueError(msg) 

647 seen.add(node_key) 

648 if node_key not in block_nodes and _is_defined(comp, node_key): 

649 msg = f"Repeated block feature node already exists: {node_key!r}" 

650 raise ValueError(msg) 

651 

652 

653def _add_repeated_blocks_definition( 

654 comp: Computation, 

655 definition: RepeatedBlocks[K], 

656 definition_object: object = None, 

657 ignore_self: bool = False, 

658) -> BuiltRepeatedBlocks[K]: 

659 """Validate and add a :class:`RepeatedBlocks` definition atomically.""" 

660 _validate_block_template(comp, definition.block) 

661 base_path = _to_node_name(definition.base_path, "Repeated block base_path") 

662 blocks = _repeated_block_paths(definition.keys, base_path) 

663 _validate_repeated_block_nodes(comp, definition.block, blocks) 

664 block_nodes = { 

665 block_path.join(node_name) for block_path in blocks.values() for node_name in definition.block.nodes() 

666 } 

667 

668 ctx: BlockContext[K] = BlockContext(comp, definition.block, base_path, blocks, definition_object, ignore_self) 

669 planned: list[PlannedNode] = [] 

670 for feature in definition.features: 

671 feature_nodes = list(feature.plan(ctx)) 

672 planned.extend(feature_nodes) 

673 ctx.planned.update(planned_node.node_key for planned_node in feature_nodes) 

674 _validate_planned_nodes(comp, planned, block_nodes) 

675 

676 generated_edges = [ 

677 (block_path.join(source_node), block_path.join(target_node)) 

678 for block_path in blocks.values() 

679 for source_node, target_node in definition.block.dag.edges() 

680 ] 

681 for planned_node in planned: 

682 generated_edges.extend((predecessor, planned_node.node_key) for predecessor in planned_node.predecessors) 

683 _validate_acyclic_edges(comp, generated_edges) 

684 

685 for block_path in blocks.values(): 

686 comp.add_block(block_path, definition.block, keep_values=definition.keep_values) 

687 

688 named: dict[Name, NodeKey] = {} 

689 for planned_node in planned: 

690 planned_node.apply_to(comp) 

691 if planned_node.label is not None: 

692 named[planned_node.label] = planned_node.node_key 

693 return BuiltRepeatedBlocks(blocks, tuple(planned_node.node_key for planned_node in planned), named) 

694 

695 

696@overload 

697def apply1(f: Callable[..., R], xs: list[T], *args: Any, **kwds: Any) -> list[R]: ... 

698 

699 

700@overload 

701def apply1(f: Callable[..., R], xs: T, *args: Any, **kwds: Any) -> R: ... 

702 

703 

704@overload 

705def apply1(f: Callable[..., R], xs: Generator[T, None, None], *args: Any, **kwds: Any) -> Generator[R, None, None]: ... 

706 

707 

708def apply1( 

709 f: Callable[..., R], xs: T | list[T] | Generator[T, None, None], *args: Any, **kwds: Any 

710) -> R | list[R] | Generator[R, None, None]: 

711 """Apply function f to xs, handling generators, lists, and single values.""" 

712 if isinstance(xs, types.GeneratorType): 

713 return (f(x, *args, **kwds) for x in xs) 

714 if isinstance(xs, list): 

715 return [f(x, *args, **kwds) for x in xs] 

716 return f(xs, *args, **kwds) 

717 

718 

719def as_iterable(xs: T | Iterable[T]) -> Iterable[T]: 

720 """Convert input to iterable form if not already iterable.""" 

721 if isinstance(xs, (types.GeneratorType, list, set)): 

722 return xs # type: ignore[return-value] 

723 return (xs,) # type: ignore[return-value] 

724 

725 

726def apply_n(f: Callable[..., Any], *xs: Any, **kwds: Any) -> None: 

727 """Apply function f to the cartesian product of iterables xs.""" 

728 for p in itertools.product(*[as_iterable(x) for x in xs]): 

729 f(*p, **kwds) 

730 

731 

732class AttributeView: 

733 """Provides attribute-style access to dynamic collections.""" 

734 

735 def __init__( 

736 self, 

737 get_attribute_list: Callable[[], Iterable[str]], 

738 get_attribute: Callable[[str], Any], 

739 get_item: Callable[[Any], Any] | None = None, 

740 ) -> None: 

741 """Initialize with functions to get attribute list and individual attributes. 

742 

743 Args: 

744 get_attribute_list: Function that returns list of available attributes 

745 get_attribute: Function that takes an attribute name and returns its value 

746 get_item: Optional function for item access, defaults to get_attribute 

747 """ 

748 self.get_attribute_list = get_attribute_list 

749 self.get_attribute = get_attribute 

750 self.get_item: Callable[[Any], Any] = get_item if get_item is not None else get_attribute 

751 

752 def __dir__(self) -> list[str]: 

753 """Return list of available attributes.""" 

754 return list(self.get_attribute_list()) 

755 

756 def __getattr__(self, attr: str) -> Any: 

757 """Get attribute by name, raising AttributeError if not found.""" 

758 try: 

759 return self.get_attribute(attr) 

760 except KeyError as e: 

761 raise AttributeError(attr) from e 

762 

763 def __getitem__(self, key: Any) -> Any: 

764 """Get item by key.""" 

765 return self.get_item(key) 

766 

767 def __getstate__(self) -> dict[str, Any]: 

768 """Prepare object for serialization.""" 

769 return { 

770 "get_attribute_list": self.get_attribute_list, 

771 "get_attribute": self.get_attribute, 

772 "get_item": self.get_item, 

773 } 

774 

775 def __setstate__(self, state: dict[str, Any]) -> None: 

776 """Restore object from serialized state.""" 

777 self.get_attribute_list = state["get_attribute_list"] 

778 self.get_attribute = state["get_attribute"] 

779 self.get_item = state["get_item"] 

780 if self.get_item is None: 

781 self.get_item = self.get_attribute 

782 

783 @staticmethod 

784 def from_dict(d: dict[Any, Any], use_apply1: bool = True) -> AttributeView: 

785 """Create an AttributeView from a dictionary.""" 

786 if use_apply1: 

787 

788 def get_attribute(xs: Any) -> Any: 

789 """Get attribute value from dictionary with apply1 support.""" 

790 return apply1(d.get, xs) 

791 else: 

792 get_attribute = d.get 

793 return AttributeView(d.keys, get_attribute) 

794