mirror of
https://github.com/OPM/ResInsight.git
synced 2026-10-24 22:27:55 -05:00
232 lines
7.8 KiB
Python
232 lines
7.8 KiB
Python
"""Introspect a taskmaestro workflow directory and emit a JSON schema.
|
|
|
|
The schema is consumed by ResInsight to auto-generate property-editor
|
|
fields for each task's config inputs. Only fields listed in the workflow's
|
|
`config_fields` are inspected, so models that also reference non-Pydantic
|
|
types (e.g. live `rips` objects passed via `depends_on`) do not break
|
|
introspection.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import datetime
|
|
import enum
|
|
import json
|
|
import pathlib
|
|
import sys
|
|
from pathlib import Path
|
|
from types import UnionType
|
|
from typing import Annotated, Any, Union, get_args, get_origin
|
|
|
|
from pydantic import TypeAdapter
|
|
from pydantic.fields import FieldInfo
|
|
|
|
from .refs import annotation_to_resinsight_type
|
|
|
|
|
|
_SCALAR_TYPE_MAP: dict[type, str] = {
|
|
str: "string",
|
|
bool: "boolean",
|
|
int: "integer",
|
|
float: "number",
|
|
}
|
|
|
|
|
|
def _format_for(field_info: FieldInfo) -> str | None:
|
|
"""Ask Pydantic for the JSON-schema `format` of a field, if any.
|
|
|
|
Pydantic stores e.g. `pydantic.DirectoryPath` (which is
|
|
`Annotated[Path, PathType('dir')]`) as `annotation=Path` and
|
|
`metadata=[PathType('dir')]`. We reassemble both before asking
|
|
TypeAdapter, so format markers like `directory-path` survive.
|
|
"""
|
|
annotation = field_info.annotation
|
|
if field_info.metadata:
|
|
annotation = Annotated[tuple([annotation, *field_info.metadata])]
|
|
try:
|
|
return TypeAdapter(annotation).json_schema().get("format")
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def _annotation_type(annotation: object) -> str:
|
|
"""Map a Pydantic field annotation to a JSON-schema-style type label."""
|
|
if isinstance(annotation, type) and annotation in _SCALAR_TYPE_MAP:
|
|
return _SCALAR_TYPE_MAP[annotation]
|
|
|
|
# Optional[T] / T | None — peel the union and map the non-None arm
|
|
origin = get_origin(annotation)
|
|
if origin in (Union, UnionType):
|
|
non_none = [a for a in get_args(annotation) if a is not type(None)]
|
|
if len(non_none) == 1:
|
|
return _annotation_type(non_none[0])
|
|
|
|
if origin in (list, tuple, set, frozenset):
|
|
return "array"
|
|
|
|
if isinstance(annotation, type):
|
|
# Recognised rips object types are reported as "object" with a separate marker.
|
|
if annotation_to_resinsight_type(annotation) is not None:
|
|
return "object"
|
|
|
|
return "string"
|
|
|
|
|
|
_SENTINEL = object()
|
|
|
|
|
|
def _serialize_default(value: Any) -> Any:
|
|
"""Convert a Python value into a JSON-friendly form for the schema's `default`."""
|
|
if isinstance(value, (datetime.date, datetime.datetime)):
|
|
return value.isoformat()
|
|
if isinstance(value, pathlib.PurePath):
|
|
return str(value)
|
|
if isinstance(value, enum.Enum):
|
|
return _serialize_default(value.value)
|
|
if isinstance(value, dict):
|
|
return {key: _serialize_default(item) for key, item in value.items()}
|
|
if isinstance(value, (list, tuple, set, frozenset)):
|
|
return [_serialize_default(item) for item in value]
|
|
return value
|
|
|
|
|
|
def _field_schema(
|
|
field_name: str, field_info: FieldInfo, input_value: Any = _SENTINEL
|
|
) -> dict[str, Any]:
|
|
entry: dict[str, Any] = {
|
|
"name": field_name,
|
|
"type": _annotation_type(field_info.annotation),
|
|
"required": field_info.is_required(),
|
|
}
|
|
if field_info.description:
|
|
entry["description"] = field_info.description
|
|
|
|
# Prefer the value the user already has in input.yaml; fall back to the
|
|
# Pydantic-side default so the UI is pre-filled in either case.
|
|
if input_value is not _SENTINEL:
|
|
entry["default"] = _serialize_default(input_value)
|
|
elif not field_info.is_required():
|
|
default = field_info.get_default(call_default_factory=True)
|
|
entry["default"] = _serialize_default(default)
|
|
|
|
fmt = _format_for(field_info)
|
|
if fmt is not None:
|
|
entry["format"] = fmt
|
|
ri_type = annotation_to_resinsight_type(field_info.annotation)
|
|
if ri_type is not None:
|
|
entry["resinsight_type"] = ri_type
|
|
return entry
|
|
|
|
|
|
def _dependency_edges(task_name: str, deps: Any) -> list[dict[str, str]]:
|
|
"""Flatten taskmaestro's resolved dependency forms into directed edges."""
|
|
if deps is None:
|
|
return []
|
|
if isinstance(deps, str):
|
|
return [{"from": deps, "to": task_name}]
|
|
if isinstance(deps, tuple):
|
|
return [{"from": deps[0], "to": task_name, "output": deps[1]}]
|
|
|
|
edges = []
|
|
for input_name, upstream in sorted(deps.items()):
|
|
edge = {
|
|
"from": upstream[0] if isinstance(upstream, tuple) else upstream,
|
|
"to": task_name,
|
|
"input": input_name,
|
|
}
|
|
if isinstance(upstream, tuple):
|
|
edge["output"] = upstream[1]
|
|
edges.append(edge)
|
|
return edges
|
|
|
|
|
|
def _output_fields(output_type: type) -> list[str]:
|
|
"""Expose output fields, but not ObjectModel's internal wrapped value."""
|
|
from taskmaestro import ObjectModel
|
|
|
|
fields = set(output_type.model_fields)
|
|
if issubclass(output_type, ObjectModel):
|
|
fields.discard("value")
|
|
return sorted(fields)
|
|
|
|
|
|
def collect_schema(workflow_dir: Path) -> dict[str, Any]:
|
|
"""Load the workflow at `workflow_dir` and produce its UI schema."""
|
|
workflow_yaml = workflow_dir / "workflow.yaml"
|
|
input_yaml = workflow_dir / "input.yaml"
|
|
if not workflow_yaml.is_file():
|
|
raise FileNotFoundError(f"Missing workflow.yaml in {workflow_dir}")
|
|
|
|
# Workflow source files (e.g. pipeline.py) live next to workflow.yaml and
|
|
# are referenced as plain dotted paths in workflow.yaml — make them importable.
|
|
workflow_dir_str = str(workflow_dir)
|
|
if workflow_dir_str not in sys.path:
|
|
sys.path.insert(0, workflow_dir_str)
|
|
|
|
from taskmaestro.task import get_input_type, get_output_type
|
|
from taskmaestro.yaml_config import _load_workflow_only
|
|
|
|
wf, jc = _load_workflow_only(
|
|
workflow_yaml,
|
|
input_yaml if input_yaml.is_file() else None,
|
|
)
|
|
|
|
tasks: list[dict[str, Any]] = []
|
|
edges: list[dict[str, str]] = []
|
|
for task_name, task_cls in wf.topological_order():
|
|
edges.extend(_dependency_edges(task_name, wf.get_dependencies(task_name)))
|
|
config_field_names = wf.get_config_fields(task_name)
|
|
task_config = jc.get_config_for_task(task_name) if jc is not None else {}
|
|
|
|
input_type = get_input_type(task_cls)
|
|
config_fields: list[dict[str, Any]] = []
|
|
for field_name in sorted(config_field_names):
|
|
field_info = input_type.model_fields.get(field_name)
|
|
if field_info is None:
|
|
config_fields.append(
|
|
{
|
|
"name": field_name,
|
|
"type": "string",
|
|
"required": True,
|
|
"error": "field not found in input model",
|
|
}
|
|
)
|
|
continue
|
|
input_value = task_config.get(field_name, _SENTINEL)
|
|
config_fields.append(
|
|
_field_schema(field_name, field_info, input_value=input_value)
|
|
)
|
|
|
|
tasks.append(
|
|
{
|
|
"name": task_name,
|
|
"inputs": sorted(input_type.model_fields),
|
|
"outputs": _output_fields(get_output_type(task_cls)),
|
|
"config_fields": config_fields,
|
|
}
|
|
)
|
|
|
|
return {
|
|
"name": wf.name,
|
|
"description": "",
|
|
"tasks": tasks,
|
|
"edges": edges,
|
|
}
|
|
|
|
|
|
def main(argv: list[str]) -> int:
|
|
if len(argv) != 1:
|
|
print("usage: introspect <workflow_dir>", file=sys.stderr)
|
|
return 2
|
|
workflow_dir = Path(argv[0]).resolve()
|
|
try:
|
|
schema = collect_schema(workflow_dir)
|
|
except Exception as exc:
|
|
print(
|
|
json.dumps({"error": str(exc), "type": type(exc).__name__}), file=sys.stderr
|
|
)
|
|
return 1
|
|
json.dump(schema, sys.stdout, indent=2)
|
|
sys.stdout.write("\n")
|
|
return 0
|