diff --git a/src/jobflow/core/reference.py b/src/jobflow/core/reference.py index 35780d9d..0bbf2205 100644 --- a/src/jobflow/core/reference.py +++ b/src/jobflow/core/reference.py @@ -514,12 +514,47 @@ def validate_schema_access( The BaseModel class associated with the item, if any. """ schema_dict = schema.model_json_schema() - if item not in schema_dict["properties"]: + item_in_schema = item in schema_dict["properties"] + property_like = isinstance(item, str) and has_property_like(schema, item) + if not item_in_schema and not property_like: raise AttributeError(f"{schema.__name__} does not have attribute '{item}'.") subschema = None - item_type = schema.model_fields[item].annotation - if lenient_issubclass(item_type, BaseModel): - subschema = item_type + if item_in_schema: + item_type = schema.model_fields[item].annotation + if lenient_issubclass(item_type, BaseModel): + subschema = item_type return True, subschema + + +def has_property_like(obj_type: type, name: str) -> bool: + """ + Check if a class has an attribute and if it is property-like. + + Parameters + ---------- + obj_type + The class that needs to be checked + name + The name of the attribute to be verified + + Returns + ------- + bool + True if the property corresponding to the name is property-like. + """ + if not hasattr(obj_type, name): + return False + + attr = getattr(obj_type, name) + + if isinstance(attr, property): + return True + + if callable(attr): + return False + + # Check for custom property-like descriptors with __get__ but not callable + # If not, is not property-like. + return hasattr(attr, "__get__") diff --git a/tests/core/test_reference.py b/tests/core/test_reference.py index 353a300b..b91b6fc7 100644 --- a/tests/core/test_reference.py +++ b/tests/core/test_reference.py @@ -153,6 +153,8 @@ def test_set_uuid(): def test_schema(): + from functools import cached_property + from pydantic import BaseModel from jobflow import OutputReference @@ -173,6 +175,17 @@ class MySchema(BaseModel): name: str nested: MediumSchema + @property + def b(self) -> int: + return self.number + + @cached_property + def c(self) -> int: + return self.number + + def some_method(self) -> int: + return self.number + ref = OutputReference("123", output_schema=MySchema) assert ref.attributes == () @@ -185,15 +198,26 @@ class MySchema(BaseModel): assert new_ref.uuid == "123" assert new_ref.output_schema is None - with pytest.raises(AttributeError): + with pytest.raises(AttributeError, match="does not have attribute 'a'"): _ = ref.a.uuid - with pytest.raises(AttributeError): + with pytest.raises(AttributeError, match="does not have attribute 'a'"): _ = ref["a"].uuid with pytest.raises(AttributeError): _ = ref[1].uuid + assert ref.b + assert ref["b"] + assert ref.c + assert ref["c"] + + with pytest.raises(AttributeError, match="does not have attribute 'some_method'"): + _ = ref.some_method + + with pytest.raises(AttributeError, match="does not have attribute 'some_method'"): + _ = ref["some_method"] + # check valid nested schemas assert ref.nested.s.uuid == "123" with pytest.raises(AttributeError): @@ -518,3 +542,49 @@ def test_not_iterable(): with pytest.raises(TypeError): for _ in ref: pass + + +def test_has_property_like(): + from functools import cached_property + + from monty.functools import lazy_property + from pydantic import BaseModel, ConfigDict + + from jobflow.core.reference import has_property_like + + class TestModel(BaseModel): + model_config = ConfigDict(ignored_types=(lazy_property,)) + + x: str = "x" + y: int = 1 + + def method(self): + return self.x + + @staticmethod + def static_method(): + return 1 + + @classmethod + def class_method(cls): + return cls.y + + @property + def standard_property(self): + return self.x + + @lazy_property + def monty_lazy_property(self): + return self.x + + @cached_property + def functools_cached_property(self): + return self.x + + assert not has_property_like(TestModel, "method") + assert not has_property_like(TestModel, "static_method") + assert not has_property_like(TestModel, "class_method") + assert has_property_like(TestModel, "standard_property") + assert has_property_like(TestModel, "monty_lazy_property") + assert has_property_like(TestModel, "functools_cached_property") + assert not has_property_like(TestModel, "not_existing_prop")