Skip to content

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant