Merge pull request #839 from gpetretto/validate_properties

Validate properties attribute in output_schema

Author: utf

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__")

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")