Fix PEtab v2 extension config model (#506)

* Fix PEtab v2 extension config model (#474)

ProblemConfig.extensions was typed as `list[ExtensionConfig] | dict`,
but only the dict branch was ever used or supported downstream. On
top of that, the extension config models didn't match the v2 schema:
neither ExtensionConfig nor SciMLConfig had the schema-mandated
`required` field, and the generic ExtensionConfig wrongly nested
extra keys under a `config` field instead of allowing them directly.
Writing a problem with a `sciml` extension to YAML and reading it
back therefore failed schema validation.

- Add a shared `ExtensionConfig` base (version, required, extra
  fields allowed) in petab/v2/extensions/__init__.py, and make
  SciMLConfig subclass it (required defaults to True, since a
  SciML hybrid model is virtually always load-bearing).
- Type `ProblemConfig.extensions` as `dict[str, ExtensionConfig]`
  with `SerializeAsAny` so subclass fields survive serialization.
- Fix a broken import in SciMLConfig.to_yaml() that made it crash
  unconditionally.
- Problem.validate() still only warns about unsupported extensions
  (rejecting based on `required` is left to consumers like
  simulators that actually interpret the extension mathematically).

Closes #474

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

* less verbose

* Apply suggestions from code review

Co-authored-by: Dilan Pathirana <59329744+dilpath@users.noreply.github.com>

* fixup

* Fix CI regression and decouple extension parsing from core.py

- Revert extensions field_validator to raise ValueError instead of
  TypeError: pydantic only converts ValueError/AssertionError raised
  in a validator into a ValidationError, not TypeError, so the
  TypeError variant broke test_problem_config_extensions_rejects_non_dict
  in CI.
- Move the per-extension-ID config dispatch (sciml vs. generic) out
  of ProblemConfig._parse_extensions into
  petab.v2.extensions.parse_extension_config, so core.py no longer
  needs to import SciMLConfig or check C.EXT_ID_SCIML directly.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com>
Co-authored-by: Dilan Pathirana <59329744+dilpath@users.noreply.github.com>

Author: 3 people

petab/v2/core.py

@@ -35,6 +35,7 @@
     BeforeValidator,
     ConfigDict,
     Field,
+    SerializeAsAny,
     ValidationInfo,
     field_serializer,
     field_validator,
@@ -53,6 +54,7 @@
 from ..v1.yaml import get_path_prefix
 from ..versions import parse_version
 from . import C, get_observable_df
+from .extensions import ExtensionConfig, parse_extension_config
 
 if TYPE_CHECKING:
     from ..v2.lint import ValidationResultList, ValidationTask
@@ -311,10 +313,7 @@ def __iadd__(self, other: T) -> BaseTable[T]:
 
 # SciML extension classes — imported after BaseTable is defined to avoid
 # circular imports (sciml.py does not import from core.py).
-from .extensions.sciml import (  # noqa: E402
-    SciMLConfig,
-    SciMLExt,
-)
+from .extensions.sciml import SciMLExt  # noqa: E402
 
 
 class ProblemExtensions:
@@ -2492,13 +2491,6 @@ class ModelFile(BaseModel):
     )
 
 
-class ExtensionConfig(BaseModel):
-    """The configuration of a PEtab extension."""
-
-    version: str
-    config: dict
-
-
 class ProblemConfig(BaseModel):
     """The PEtab problem configuration."""
 
@@ -2541,8 +2533,8 @@ class ProblemConfig(BaseModel):
     # Absolute or relative to `base_path`.
     mapping_files: list[AnyUrl | Path] = []
 
-    #: Extensions used by the problem.
-    extensions: list[ExtensionConfig] | dict = {}
+    #: Extensions used by the problem, keyed by extension ID.
+    extensions: dict[str, SerializeAsAny[ExtensionConfig]] = {}
 
     model_config = ConfigDict(
         validate_assignment=True,
@@ -2553,23 +2545,15 @@ class ProblemConfig(BaseModel):
     def _parse_extensions(cls, v):
         """Parse extensions dict and convert known extensions to their specific
         config classes."""
-        if isinstance(v, dict):
-            parsed_extensions = {}
-            for ext_name, ext_config in v.items():
-                if ext_name == C.EXT_ID_SCIML:
-                    parsed_extensions[ext_name] = (
-                        ext_config
-                        if isinstance(ext_config, SciMLConfig)
-                        else SciMLConfig(**ext_config)
-                    )
-                else:
-                    parsed_extensions[ext_name] = (
-                        ext_config
-                        if isinstance(ext_config, ExtensionConfig)
-                        else ExtensionConfig(**ext_config)
-                    )
-            return parsed_extensions
-        return v
+        if not isinstance(v, dict):
+            raise ValueError(
+                "extensions must be a dict of extension ID to extension "
+                f"config, got {type(v)}."
+            )
+        return {
+            ext_id: parse_extension_config(ext_id, ext_config)
+            for ext_id, ext_config in v.items()
+        }
 
     # convert parameter_file to list
     @field_validator(

petab/v2/extensions/__init__.py

@@ -0,0 +1,46 @@
+from pydantic import BaseModel, ConfigDict
+
+__all__ = ["ExtensionConfig", "parse_extension_config"]
+
+
+class ExtensionConfig(BaseModel):
+    """The configuration of a PEtab extension."""
+
+    #: The extension's semantic version.
+    version: str
+    #: Whether the extension is required for the mathematical
+    #: interpretation of the problem.
+    required: bool
+
+    model_config = ConfigDict(extra="allow", validate_assignment=True)
+
+
+def _extension_config_classes() -> dict[str, type[ExtensionConfig]]:
+    """Registry of extension ID to its specific :class:`ExtensionConfig`
+    subclass, if any.
+
+    Imported lazily (rather than built at module level) to avoid a
+    circular import: extension submodules (e.g. ``sciml``) import
+    :class:`ExtensionConfig` from this package.
+    """
+    from .. import C
+    from .sciml import SciMLConfig
+
+    return {C.EXT_ID_SCIML: SciMLConfig}
+
+
+def parse_extension_config(
+    ext_id: str, config: dict | ExtensionConfig
+) -> ExtensionConfig:
+    """Parse a single extension's configuration.
+
+    Converts ``config`` to the extension-specific :class:`ExtensionConfig`
+    subclass registered for ``ext_id``, or to the generic
+    :class:`ExtensionConfig` if no specific subclass is registered.
+
+    :param ext_id: The extension ID.
+    :param config: The extension's configuration, as a dict or an already
+        parsed :class:`ExtensionConfig` (sub)instance.
+    """
+    cls = _extension_config_classes().get(ext_id, ExtensionConfig)
+    return config if isinstance(config, cls) else cls(**config)

petab/v2/extensions/sciml.py

@@ -25,6 +25,7 @@
     pass
 
 from .. import C
+from . import ExtensionConfig
 
 __all__ = [
     "Hybridization",
@@ -136,11 +137,11 @@ class NeuralNetConfig(BaseModel):
     )
 
 
-class SciMLConfig(BaseModel):
+class SciMLConfig(ExtensionConfig):
     """The extended configuration of a PEtab SciML problem."""
 
-    #: The PEtab SciML format version.
     version: str = "0.1.0"
+    required: bool = True
     #: The paths to the array data files.
     array_files: list[AnyUrl | Path] = []
     #: The paths to the hybridization tables.
@@ -155,7 +156,7 @@ class SciMLConfig(BaseModel):
 
     def to_yaml(self) -> dict:
         """Return a YAML-serializable dict with Paths converted to strings."""
-        from . import C
+        from .. import C
 
         d = self.model_dump(by_alias=True)
         for key in ("array_files", "hybridization_files"):

tests/v2/test_core.py

@@ -29,6 +29,8 @@
     UPPER_BOUND,
 )
 from petab.v2.core import *
+from petab.v2.core import ExtensionConfig
+from petab.v2.lint import ValidationIssueSeverity
 from petab.v2.models.sbml_model import SbmlModel
 from petab.v2.petab1to2 import petab1to2
 
@@ -595,6 +597,83 @@ def test_problem_config_paths():
     # see also https://github.com/pydantic/pydantic/issues/8575
 
 
+def test_problem_config_generic_extension():
+    """A generic (non-sciml) extension is parsed per the PEtab v2 schema:
+    `version` and `required` at the top level, plus arbitrary
+    extension-specific keys alongside them."""
+    pc = ProblemConfig(
+        parameter_files=["parameters.tsv"],
+        measurement_files=["measurements.tsv"],
+        observable_files=["observables.tsv"],
+        extensions={
+            "my_ext": {
+                "version": "1.0.0",
+                "required": False,
+                "some_key": "some_value",
+            }
+        },
+    )
+    ext = pc.extensions["my_ext"]
+    assert isinstance(ext, ExtensionConfig)
+    assert ext.version == "1.0.0"
+    assert ext.required is False
+    assert ext.some_key == "some_value"
+
+    dumped = pc.model_dump(by_alias=True)["extensions"]["my_ext"]
+    assert dumped == {
+        "version": "1.0.0",
+        "required": False,
+        "some_key": "some_value",
+    }
+
+
+def test_problem_config_extensions_rejects_non_dict():
+    """`extensions` must be a dict keyed by extension ID (see #474) -- a
+    list is not a valid PEtab v2 problem configuration."""
+    with pytest.raises(ValidationError):
+        ProblemConfig(
+            parameter_files=["parameters.tsv"],
+            measurement_files=["measurements.tsv"],
+            observable_files=["observables.tsv"],
+            extensions=[{"version": "1.0.0", "required": False}],
+        )
+
+
+def test_validate_unsupported_extension_severity():
+    """libpetab-python doesn't mathematically interpret extensions, so an
+    unsupported extension only ever produces a WARNING (that the problem
+    can't be fully linted) -- regardless of `required`. Rejecting a problem
+    that uses an unsupported `required` extension is up to the consumer
+    (e.g. a simulator) that actually interprets it."""
+    problem = Problem()
+    problem.model = SbmlModel.from_antimony("""
+    model m
+      species A;
+      A = 1;
+      k1 = 1;
+      R1: A -> ; k1 * A;
+    end
+    """)
+    problem.add_observable("obs_A", "A", noise_formula="1")
+    problem.add_parameter(
+        "k1", estimate=True, lb=1e-5, ub=1e5, nominal_value=1
+    )
+    problem.add_measurement("obs_A", time=1, measurement=1, experiment_id="")
+    assert problem.validate() == []
+
+    for required in (False, True):
+        problem.config = ProblemConfig(
+            extensions={"my_ext": {"version": "1.0.0", "required": required}}
+        )
+        results = problem.validate()
+        assert not results.has_errors()
+        assert any(
+            r.level == ValidationIssueSeverity.WARNING
+            and "my_ext" in r.message
+            for r in results
+        )
+
+
 def test_get_changes_for_period():
     """Test getting changes for a specific period."""
     problem = Problem()

tests/v2/test_sciml.py

@@ -157,6 +157,26 @@ def test_lint():
     assert problem.validate() == []
 
 
+def test_sciml_config_yaml_round_trip(tmp_path):
+    """The `sciml` extension config, once written to YAML via
+    `ProblemConfig.to_yaml()`, is schema-valid and can be read back.
+    """
+    from petab.v1.yaml import load_yaml, validate_yaml_syntax
+
+    problem = _get_test_problem()
+    yaml_path = tmp_path / "problem.yaml"
+    problem.config.to_yaml(yaml_path)
+
+    yaml_config = load_yaml(yaml_path)
+    validate_yaml_syntax(yaml_config)
+    assert yaml_config["extensions"]["sciml"]["required"] is True
+
+    reloaded_config = ProblemConfig(**yaml_config, base_path=tmp_path)
+    sciml_config = reloaded_config.extensions["sciml"]
+    assert isinstance(sciml_config, SciMLConfig)
+    assert sciml_config.required is True
+
+
 def test_lint_equinox_network_format():
     """Linter accepts non-YAML formats without reading the network file."""
     problem = _get_test_problem()