Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 39 additions & 4 deletions src/jobflow/core/reference.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__")
74 changes: 72 additions & 2 deletions tests/core/test_reference.py
Original file line number Diff line number Diff line change
Expand Up @@ -153,6 +153,8 @@ def test_set_uuid():


def test_schema():
from functools import cached_property

from pydantic import BaseModel

from jobflow import OutputReference
Expand All @@ -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 == ()

Expand All @@ -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):
Expand Down Expand Up @@ -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")
Loading