Coverage for backend/ahuora-unit-ops/src/ahuora_unit_ops/wrapper_metadata.py: 14%
158 statements
« prev ^ index » next coverage.py v7.10.7, created at 2026-07-22 05:22 +0000
1from __future__ import annotations
3import inspect
4import sys
5from dataclasses import dataclass
6from importlib import import_module
7from pathlib import Path
8from typing import Any
10from .json_config import JsonUnitOpConfig, config_filename_from_unit_key
13@dataclass(frozen=True, slots=True)
14class AhuoraUnitRegistration:
15 """One unit-operation JSON config exposed by an IDAES wrapper module."""
17 metadata: JsonUnitOpConfig
18 constructor: type | str | None = None
21class ConfigSpecError(ValueError):
22 """Raised when unit-operation metadata cannot be discovered safely."""
25@dataclass(frozen=True, slots=True)
26class WrapperDiscoveryResult:
27 """Wrapper metadata discovery result with soft import failures."""
29 wrappers_by_key: dict[str, type | str]
30 registrations_by_key: dict[str, AhuoraUnitRegistration]
31 import_errors: dict[str, str]
34def import_object(path: str) -> Any:
35 """Import an object from a dotted Python path."""
37 module_path, _, object_name = path.rpartition(".")
38 if not module_path or not object_name:
39 raise ValueError(f"expected dotted object path, got {path!r}")
40 module = import_module(module_path)
41 return getattr(module, object_name)
44def get_wrapper_metadata(wrapper_cls: type) -> JsonUnitOpConfig:
45 """Read the first JSON config exposed by a wrapper or its data class."""
47 registrations = _registrations_from_metadata(_read_ahuora_metadata(wrapper_cls))
48 if not registrations:
49 data_cls = _matching_data_class_for_wrapper(wrapper_cls)
50 if data_cls is not None:
51 registrations = _registrations_from_metadata(_read_ahuora_metadata(data_cls))
52 if not registrations:
53 raise TypeError(
54 "Wrapper imports, but has no ahuora_metadata() metadata factory. "
55 "Attach an ahuora_metadata function to the data class. "
56 f"Wrapper: {wrapper_cls!r}"
57 )
58 return registrations[0].metadata
61def registration_constructor_path(registration: AhuoraUnitRegistration) -> str:
62 """Return the dotted runtime constructor path for display purposes."""
64 constructor = registration.constructor
65 if constructor is None:
66 adapter = registration.metadata.idaes_adapter
67 if adapter is None:
68 return ""
69 return adapter.constructor
70 if isinstance(constructor, str):
71 return constructor
72 return constructor_path_from_class(constructor)
75def constructor_path_from_class(wrapper_cls: type) -> str:
76 """Return the import path used to instantiate a wrapper class."""
78 return f"{wrapper_cls.__module__}.{wrapper_cls.__name__}"
81def _matching_data_class_for_wrapper(wrapper_cls: type) -> type | None:
82 """Return a same-module data class whose config targets the wrapper path."""
84 module = sys.modules.get(wrapper_cls.__module__)
85 if module is None:
86 return None
87 wrapper_path = constructor_path_from_class(wrapper_cls)
88 for class_name, data_cls in inspect.getmembers(module, inspect.isclass):
89 if not class_name.endswith("Data"):
90 continue
91 if not callable(getattr(data_cls, "ahuora_metadata", None)):
92 continue
93 for registration in _registrations_from_metadata(_read_ahuora_metadata(data_cls)):
94 adapter = registration.metadata.idaes_adapter
95 if adapter is not None and adapter.constructor == wrapper_path:
96 return data_cls
97 return None
100def django_config_output_path_from_unit_key(key: str) -> Path:
101 """Return the conventional Django config output path without writing it."""
103 repo_root = Path(__file__).resolve().parents[4]
104 return (
105 repo_root
106 / "backend"
107 / "django"
108 / "flowsheetInternals"
109 / "unitops"
110 / "config"
111 / "objects"
112 / config_filename_from_unit_key(key)
113 )
116def discover_wrapper_metadata() -> WrapperDiscoveryResult:
117 """Import custom modules and collect data classes with Ahuora metadata."""
119 import ahuora_builder.custom as custom_package
121 wrappers_by_key: dict[str, type | str] = {}
122 registrations_by_key: dict[str, AhuoraUnitRegistration] = {}
123 import_errors: dict[str, str] = {}
124 duplicate_keys: dict[str, list[str]] = {}
126 def add_registration(
127 registration: AhuoraUnitRegistration,
128 *,
129 owner: type | None,
130 source: str,
131 ) -> None:
132 resolved = _resolve_registration_constructor(
133 registration,
134 owner=owner,
135 )
136 key = resolved.metadata.key
137 constructor = resolved.constructor
138 if key in registrations_by_key:
139 if _registration_source(registrations_by_key[key]) == _registration_source(
140 resolved
141 ):
142 return
143 duplicate_keys.setdefault(key, []).extend(
144 [
145 _registration_source(registrations_by_key[key]),
146 _registration_source(resolved),
147 ]
148 )
149 return
150 registrations_by_key[key] = resolved
151 if constructor is not None:
152 wrappers_by_key[key] = constructor
154 for module in _iter_imported_custom_modules(custom_package):
155 for class_name, wrapper_cls in inspect.getmembers(module, inspect.isclass):
156 if not callable(getattr(wrapper_cls, "ahuora_metadata", None)):
157 continue
158 metadata_result = _read_ahuora_metadata(wrapper_cls)
159 for index, registration in enumerate(
160 _registrations_from_metadata(metadata_result)
161 ):
162 add_registration(
163 registration,
164 owner=wrapper_cls,
165 source=(
166 f"{constructor_path_from_class(wrapper_cls)}"
167 f".ahuora_metadata[{index}]"
168 ),
169 )
171 if duplicate_keys:
172 details = "; ".join(
173 f"{key}: {sorted(set(paths))}" for key, paths in duplicate_keys.items()
174 )
175 raise ConfigSpecError(f"Duplicate wrapper metadata keys found: {details}")
177 return WrapperDiscoveryResult(
178 wrappers_by_key=wrappers_by_key,
179 registrations_by_key=registrations_by_key,
180 import_errors=import_errors,
181 )
184def find_registration_by_key(key: str) -> AhuoraUnitRegistration:
185 """Discover wrapper metadata and return the registration for one exact key."""
187 discovery = discover_wrapper_metadata()
188 try:
189 return discovery.registrations_by_key[key]
190 except KeyError as exc:
191 available = ", ".join(sorted(discovery.registrations_by_key))
192 raise ConfigSpecError(
193 "ERROR [unit-key-not-found]\n"
194 f"No wrapper metadata found for unit key {key!r}.\n"
195 f"Available keys: {available}"
196 ) from exc
199def validate_variant_families(configs: list[JsonUnitOpConfig]) -> None:
200 """Ensure every variant family has exactly one default config."""
202 families: dict[str, list[JsonUnitOpConfig]] = {}
203 for config in configs:
204 variant = config.frontend.variant if config.frontend is not None else None
205 if variant is None:
206 continue
207 families.setdefault(variant.familyKey, []).append(config)
209 for family_key, family_configs in sorted(families.items()):
210 defaults = [
211 config.key
212 for config in family_configs
213 if config.frontend is not None
214 and config.frontend.variant is not None
215 and config.frontend.variant.default
216 ]
217 if len(defaults) != 1:
218 members = [config.key for config in family_configs]
219 raise ConfigSpecError(
220 "ERROR [variant-family-default]\n"
221 f"Variant family {family_key!r} must have exactly one default. "
222 f"Defaults: {defaults}; members: {members}."
223 )
226def _read_ahuora_metadata(owner: type) -> object:
227 metadata_factory = getattr(owner, "ahuora_metadata", None)
228 if not callable(metadata_factory):
229 return None
230 return metadata_factory()
233def _registrations_from_metadata(value: object) -> tuple[AhuoraUnitRegistration, ...]:
234 if isinstance(value, JsonUnitOpConfig):
235 return (AhuoraUnitRegistration(metadata=value),)
236 if isinstance(value, AhuoraUnitRegistration):
237 return (value,)
238 if value is None:
239 return ()
240 try:
241 values = tuple(value) # type: ignore[arg-type]
242 except TypeError as exc:
243 raise ConfigSpecError(
244 "ERROR [ahuora-metadata]\n"
245 "ahuora_metadata() must return JsonUnitOpConfig, "
246 "AhuoraUnitRegistration, or an iterable of those objects."
247 ) from exc
249 registrations: list[AhuoraUnitRegistration] = []
250 for item in values:
251 if isinstance(item, JsonUnitOpConfig):
252 registrations.append(AhuoraUnitRegistration(metadata=item))
253 continue
254 if isinstance(item, AhuoraUnitRegistration):
255 registrations.append(item)
256 continue
257 raise ConfigSpecError(
258 "ERROR [ahuora-metadata-entry]\n"
259 "ahuora_metadata() iterables may only contain "
260 "AhuoraUnitRegistration or JsonUnitOpConfig entries."
261 )
262 return tuple(registrations)
265def _resolve_registration_constructor(
266 registration: AhuoraUnitRegistration,
267 *,
268 owner: type | None,
269) -> AhuoraUnitRegistration:
270 if registration.constructor is not None:
271 return registration
272 adapter = registration.metadata.idaes_adapter
273 if adapter is not None:
274 return AhuoraUnitRegistration(
275 metadata=registration.metadata,
276 constructor=adapter.constructor,
277 )
278 return AhuoraUnitRegistration(metadata=registration.metadata, constructor=owner)
281def _iter_imported_custom_modules(custom_package) -> list[Any]:
282 prefix = custom_package.__name__ + "."
283 modules = [custom_package]
284 for name, module in sorted(sys.modules.items()):
285 if name.startswith(prefix) and module is not None:
286 modules.append(module)
287 return modules
290def _registration_source(registration: AhuoraUnitRegistration) -> str:
291 constructor = registration.constructor
292 if constructor is None:
293 return registration.metadata.key
294 if isinstance(constructor, str):
295 return constructor
296 return constructor_path_from_class(constructor)