From c48ba5c09cbe1404262fd8159722bb16409d1042 Mon Sep 17 00:00:00 2001
From: fubuloubu <3859395+fubuloubu@users.noreply.github.com>
Date: Sat, 4 May 2024 13:11:00 -0400
Subject: [PATCH] fix: add dict methods to Datapoints

---
 silverback/types.py | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/silverback/types.py b/silverback/types.py
index ad0ad002..613f1aa3 100644
--- a/silverback/types.py
+++ b/silverback/types.py
@@ -1,7 +1,7 @@
 from datetime import datetime, timezone
 from decimal import Decimal
 from enum import Enum  # NOTE: `enum.StrEnum` only in Python 3.11+
-from typing import Literal
+from typing import Iterator, Literal
 
 from ape.logging import get_logger
 from pydantic import BaseModel, Field, RootModel, ValidationError, model_validator
@@ -88,3 +88,19 @@ def parse_datapoints(cls, datapoints: dict) -> dict:
             )
 
         return datapoints
+
+    # Add dict methods
+    def get(self, key: str, default: Datapoint | None = None) -> Datapoint | None:
+        if key in self:
+            return self[key]
+
+        return default
+
+    def __iter__(self):
+        return iter(self.root)
+
+    def __getitem__(self, item):
+        return self.root[item]
+
+    def items(self):
+        return self.root.items()