From 07c6668278dc2d063b2c2646f70a8cdcda80c9a3 Mon Sep 17 00:00:00 2001
From: Leonardo Schwarz <leonardo.schwarz@fgcz.ethz.ch>
Date: Mon, 12 Aug 2024 09:37:30 +0200
Subject: [PATCH] add Entity.__getitem__

---
 bfabric/entities/core/entity.py                 | 4 ++++
 bfabric/tests/unit/entities/core/test_entity.py | 8 ++++++--
 docs/changelog.md                               | 3 ++-
 3 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/bfabric/entities/core/entity.py b/bfabric/entities/core/entity.py
index 40c044f1..0069142f 100644
--- a/bfabric/entities/core/entity.py
+++ b/bfabric/entities/core/entity.py
@@ -51,6 +51,10 @@ def find_all(cls, ids: list[int], client: Bfabric) -> dict[int, Self]:
             logger.warning(f"Only found {len(results)} out of {len(ids)}.")
         return results
 
+    def __getitem__(self, key: str) -> Any:
+        """Returns the value of a key in the data dictionary."""
+        return self.__data_dict[key]
+
     def __repr__(self) -> str:
         """Returns the string representation of the workunit."""
         return f"{self.__class__.__name__}({repr(self.__data_dict)}, client={repr(self.__client)})"
diff --git a/bfabric/tests/unit/entities/core/test_entity.py b/bfabric/tests/unit/entities/core/test_entity.py
index ce1b3cea..47b4ca6c 100644
--- a/bfabric/tests/unit/entities/core/test_entity.py
+++ b/bfabric/tests/unit/entities/core/test_entity.py
@@ -6,8 +6,7 @@
 
 @pytest.fixture
 def mock_client(mocker):
-    client = mocker.Mock(spec=Bfabric)
-    return client
+    return mocker.Mock(spec=Bfabric)
 
 
 @pytest.fixture
@@ -70,6 +69,11 @@ def test_find_all_when_not_all_found(mocker, mock_client) -> None:
     mock_client.read.assert_called_once_with("test_endpoint", obj={"id": [1, 5]})
 
 
+def test_get_item(mock_entity) -> None:
+    assert mock_entity["id"] == 1
+    assert mock_entity["name"] == "Test Entity"
+
+
 def test_repr(mock_entity, mock_data_dict) -> None:
     entity = Entity(mock_data_dict, None)
     assert repr(entity) == "Entity({'id': 1, 'name': 'Test Entity'}, client=None)"
diff --git a/docs/changelog.md b/docs/changelog.md
index 6e515203..4f3ee80f 100644
--- a/docs/changelog.md
+++ b/docs/changelog.md
@@ -14,9 +14,10 @@ Versioning currently follows `X.Y.Z` where
 
 - The `Bfabric` instance is now pickleable.
 - Entities mapping:
+  - Add `Entity.id` and `Entity.web_url` properties.
+  - Add `Entity.__getitem__` to access fields from the data dictionary directly.
   - More types and relationships
   - Relationships defer imports to descriptor call, i.e. circular relationships are possible now.
-  - Add `Entity.id` and `Entity.web_url` properties.
   - `HasOne` and `HasMany` allow defining `optional=True` to indicate fields which can be missing under some circumstances.
 
 ## \[1.13.4\] - 2024-08-05