Skip to content

Commit

Permalink
more robust attribute tests
Browse files Browse the repository at this point in the history
  • Loading branch information
BradyAJohnston committed Dec 20, 2024
1 parent 976fe73 commit 80907c9
Show file tree
Hide file tree
Showing 2 changed files with 57 additions and 5 deletions.
19 changes: 15 additions & 4 deletions databpy/attribute.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,12 @@
COMPATIBLE_TYPES = [bpy.types.Mesh, bpy.types.Curve, bpy.types.PointCloud]


class NamedAttributeError(AttributeError):
def __init__(self, message):
self.message = message
super().__init__(self.message)


def _check_obj_attributes(obj: Object) -> None:
if not isinstance(obj, bpy.types.Object):
raise TypeError(f"Object must be a bpy.types.Object, not {type(obj)}")
Expand Down Expand Up @@ -355,9 +361,14 @@ def store_named_attribute(
if not attribute or not overwrite:
attribute = obj.data.attributes.new(name, atype.value.type_name, str(domain))

target_atype = AttributeTypes[attribute.data_type]
if len(data) != len(attribute.data):
raise AttributeMismatchError(
f"Data length {len(data)}, dimensions {data.shape} does not equal the size of the target domain {domain}, len={len(attribute.data)=}"
raise NamedAttributeError(
f"Data length {len(data)}, dimensions {data.shape} does not equal the size of the target `{domain=}`, `{len(attribute.data)=}`, {target_atype.value.dimensions=}`"
)
if target_atype != atype:
raise NamedAttributeError(
f"Attribute being written to: `{attribute.name}` of type `{target_atype.value.type_name}` does not match the type for the given data: `{atype.value.type_name}`"
)

# the 'foreach_set' requires a 1D array, regardless of the shape of the attribute
Expand Down Expand Up @@ -440,7 +451,7 @@ def named_attribute(
if verbose:
message += f"Possible attributes are: {obj.data.attributes.keys()}"

raise AttributeError(message)
raise NamedAttributeError(message)

return attr.as_array()

Expand Down Expand Up @@ -470,6 +481,6 @@ def remove_named_attribute(
attr = obj.data.attributes[name]
obj.data.attributes.remove(attr)
except KeyError:
raise AttributeError(
raise NamedAttributeError(
f"The selected attribute '{name}' does not exist on the object"
)
43 changes: 42 additions & 1 deletion tests/test_attribute.py
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ def test_named_attribute_custom():
np.testing.assert_array_equal(result, test_data)

db.remove_named_attribute(obj, "test_attr")
with pytest.raises(AttributeError):
with pytest.raises(db.attribute.NamedAttributeError):
db.named_attribute(obj, "test_attr")


Expand All @@ -39,6 +39,47 @@ def test_named_attribute_nonexistent():
db.named_attribute(obj, "nonexistent_attr")


def test_attribute_mismatch():
# Create test object
verts = np.array([[0, 0, 0], [1, 1, 1]])
obj = db.create_object(verts, name="TestObject")
new_data = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
db.store_named_attribute(obj, new_data, "test_attr")

test_data = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 1.0, 0.0]])

with pytest.raises(db.attribute.NamedAttributeError):
db.store_named_attribute(obj, test_data, "test_attr")

with pytest.raises(db.attribute.NamedAttributeError):
db.store_named_attribute(obj, np.repeat(1, 3), "test_attr")


def test_attribute_overwrite():
verts = np.array([[0, 0, 0], [1, 1, 1]])
obj = db.create_object(verts, name="TestObject")
new_data = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]])
db.store_named_attribute(obj, new_data, "test_attr")
# with overwrite = False, the attribute should not be overwritten and a new one will
# be created with a new name instead
new_values = np.repeat(1, 2)
att = db.store_named_attribute(obj, new_values, "test_attr", overwrite=False)

assert new_values.shape != db.named_attribute(obj, "test_attr").shape
assert np.allclose(new_values, db.named_attribute(obj, att.name))

assert db.named_attribute(obj, "test_attr").shape == (2, 3)
with pytest.raises(db.attribute.NamedAttributeError):
db.store_named_attribute(obj, new_values, "test_attr")

db.remove_named_attribute(obj, "test_attr")
with pytest.raises(db.attribute.NamedAttributeError):
db.named_attribute(obj, "test_attr")
db.store_named_attribute(obj, new_values, "test_attr")
assert np.allclose(db.named_attribute(obj, "test_attr"), new_values)
assert db.named_attribute(obj, "test_attr").shape == (2,)


def test_named_attribute_evaluate():
# Create test object with modifier
obj = bpy.data.objects["Cube"]
Expand Down

0 comments on commit 80907c9

Please sign in to comment.