🎨 add dict-like behavior to Representation

that and some tests
This commit is contained in:
Ondřej Samohel 2024-11-29 11:28:19 +01:00
parent 32d82e47e6
commit 60f10feeee
No known key found for this signature in database
GPG key ID: 02376E18990A97C6
2 changed files with 80 additions and 2 deletions

View file

@ -1,6 +1,7 @@
"""Defines the base trait model and representation.""" """Defines the base trait model and representation."""
from __future__ import annotations from __future__ import annotations
import contextlib
import inspect import inspect
import re import re
import sys import sys
@ -49,7 +50,7 @@ class Representation:
representation_id (str): Representation ID. representation_id (str): Representation ID.
""" """
_data: dict _data: dict[str, T]
_module_blacklist: ClassVar[list[str]] = [ _module_blacklist: ClassVar[list[str]] = [
"_", "builtins", "pydantic", "_", "builtins", "pydantic",
] ]
@ -60,6 +61,70 @@ class Representation:
"""Return hash of the representation ID.""" """Return hash of the representation ID."""
return hash(self.representation_id) return hash(self.representation_id)
def __getitem__(self, key: str) -> T:
"""Get the trait by ID.
Args:
key (str): Trait ID.
Returns:
TraitBase: Trait instance.
Raises:
MissingTraitError: If the trait is not found.
"""
return self.get_trait_by_id(key)
def __setitem__(self, key: str, value: T) -> None:
"""Set the trait by ID.
Args:
key (str): Trait ID.
value (TraitBase): Trait instance.
"""
with contextlib.suppress(KeyError):
self._data.pop(key)
self.add_trait(value)
def __delitem__(self, key: str) -> None:
"""Remove the trait by ID.
Args:
key (str): Trait ID.
Raises:
ValueError: If the trait is not found.
"""
self.remove_trait_by_id(key)
def __contains__(self, key: str) -> bool:
"""Check if the trait exists by ID.
Args:
key (str): Trait ID.
Returns:
bool: True if the trait exists, False otherwise.
"""
return self.contains_trait_by_id(key)
def __iter__(self):
"""Return the trait ID iterator."""
return iter(self._data)
def __str__(self):
"""Return the representation name."""
return self.name
def items(self) -> dict[str, T]:
"""Return the traits as items."""
return self._data.items()
def add_trait(self, trait: TraitBase, *, exists_ok: bool=False) -> None: def add_trait(self, trait: TraitBase, *, exists_ok: bool=False) -> None:
"""Add a trait to the Representation. """Add a trait to the Representation.

View file

@ -25,15 +25,18 @@ REPRESENTATION_DATA = {
"file_path": Path("/path/to/file"), "file_path": Path("/path/to/file"),
"file_size": 1024, "file_size": 1024,
"file_hash": None, "file_hash": None,
"persitent": True,
}, },
Image.id: {}, Image.id: {"persitent": True},
PixelBased.id: { PixelBased.id: {
"display_window_width": 1920, "display_window_width": 1920,
"display_window_height": 1080, "display_window_height": 1080,
"pixel_aspect_ratio": 1.0, "pixel_aspect_ratio": 1.0,
"persitent": True,
}, },
Planar.id: { Planar.id: {
"planar_configuration": "RGB", "planar_configuration": "RGB",
"persitent": True,
}, },
} }
@ -168,6 +171,16 @@ def test_trait_removing(representation: Representation) -> None:
ValueError, match=f"Trait with ID {Image.id} not found."): ValueError, match=f"Trait with ID {Image.id} not found."):
representation.remove_trait(Image) representation.remove_trait(Image)
def test_representation_dict_properties(representation: Representation) -> None:
"""Test representation as dictionary."""
representation = Representation(name="test")
representation[Image.id] = Image()
assert Image.id in representation
image = representation[Image.id]
assert image == Image()
for trait_id, trait in representation.items():
assert trait_id == Image.id
assert trait == Image()
def test_getting_traits_data(representation: Representation) -> None: def test_getting_traits_data(representation: Representation) -> None: