From c515d64868b7f79fc25ded68f27dad2f4ce50d82 Mon Sep 17 00:00:00 2001
From: Danny Meijer <10511979+dannymeijer@users.noreply.github.com>
Date: Wed, 29 Jan 2025 21:07:00 +0100
Subject: [PATCH] [BUGFIX] Enhance key access in Context class to support
 direct and nested dotted notation

---
 src/koheesio/context.py    | 18 ++++++++++--------
 tests/core/test_context.py |  5 +++++
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/src/koheesio/context.py b/src/koheesio/context.py
index 0d1ee69..13c072a 100644
--- a/src/koheesio/context.py
+++ b/src/koheesio/context.py
@@ -337,18 +337,20 @@ def get(self, key: str, default: Any = None, safe: bool = True) -> Any:
         Returns `c`
         """
         try:
-            # try full dotted natation first
+            # in case key is directly available, or is written in dotted notation
             try:
                 return self.__dict__[key]
             except KeyError:
                 pass
-
-            # handle nested keys
-            nested_keys = key.split(".")
-            value = self  # parent object
-            for k in nested_keys:
-                value = value[k]  # iterate through nested values
-            return value
+            if "." in key:
+                # handle nested keys
+                nested_keys = key.split(".")
+                value = self  # parent object
+                for k in nested_keys:
+                    value = value[k]  # iterate through nested values
+                return value
+            
+            raise KeyError
 
         except (AttributeError, KeyError, TypeError) as e:
             if not safe:
diff --git a/tests/core/test_context.py b/tests/core/test_context.py
index 71c6958..e1c2d01 100644
--- a/tests/core/test_context.py
+++ b/tests/core/test_context.py
@@ -234,6 +234,11 @@ def test_from_dict():
             Context(foo="bar", baz="qux"),
             {"foo": "bar", "baz": "qux"},
         ),
+        (
+            # Test with dotted keys
+            Context({"a.b.c": 1}),
+            {"a.b.c": 1},
+        )
     ],
 )
 def test_to_dict(context, expected):