Skip to content

Validate properties attribute in output_schema#839

Merged
utf merged 1 commit intomaterialsproject:mainfrom
gpetretto:validate_properties
Feb 1, 2026
Merged

Validate properties attribute in output_schema#839
utf merged 1 commit intomaterialsproject:mainfrom
gpetretto:validate_properties

Conversation

@gpetretto
Copy link
Contributor

@gpetretto gpetretto commented Jan 15, 2026

As discussed in materialsproject/atomate2#1353 and materialsproject/atomate2#1354, currently a code like

from pydantic import BaseModel
from jobflow import job, Flow, run_locally


class M(BaseModel):
    a: int

    @property
    def a2(self):
        return self.a * 2    


@job(output_schema=M)
def test():
    return M(a=2)

@job
def print_job(x):
    print(f"{x}")

j1 = test()
j2 = print_job(j1.output.a2)

f = Flow([j1,j2])

run_locally(f)

Fails with an error:

AttributeError: M does not have attribute 'a2'.

but if the (output_schema=M) is removed from the job definition the code is executed correctly, meaning that it is just the validation step that fails to identify the existence of the properties.

This PR aims at making the code above correct even when the schema is enforced, by checking the existence of property-like attributes (e.g. @property, functools' @cached_property, monty's @lazy_property) matching the name in the output schema.

The current implementation does not attempt to extract the output type of the object returned by the property to verify if it is itself a subclass of BaseModel. I tried playing a bit with extracting information from the typing of the method, but it seems hard to cover all the possible cases and is probably a less interesting case to cover. It should be easier to cover at least the standard @property, if this is seen as necessary.
Another option would be to implement a @typed_property(output_type=....) decorator, for the cases where it may be important to enforce the checks in the output model.

@utf
Copy link
Member

utf commented Feb 1, 2026

Thanks @gpetretto. Handling the schema validation in an exhaustive way is quite tricky, but this looks like a nice extension.

@utf utf merged commit a33894d into materialsproject:main Feb 1, 2026
7 checks passed
@utf utf added the enhancement New feature or request label Feb 1, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants