From 37fc0ed37d4a29d0b73af036ec0e19fedc7403b3 Mon Sep 17 00:00:00 2001
From: Tal Borenstein <tal@keephq.dev>
Date: Wed, 3 Apr 2024 17:44:44 +0300
Subject: [PATCH] feat(workflows): dict_pop function

---
 keep/functions/__init__.py | 10 ++++++++++
 tests/test_functions.py    | 12 ++++++++++++
 2 files changed, 22 insertions(+)

diff --git a/keep/functions/__init__.py b/keep/functions/__init__.py
index 58686f2453..5bc598b2fa 100644
--- a/keep/functions/__init__.py
+++ b/keep/functions/__init__.py
@@ -1,3 +1,4 @@
+import copy
 import datetime
 import json
 import urllib.parse
@@ -102,3 +103,12 @@ def slice(str_to_slice: str, start: int = 0, end: int = 0) -> str:
     if end == 0 or end == "0":
         return str_to_slice[int(start) :]
     return str_to_slice[int(start) : int(end)]
+
+
+def dict_pop(data: str | dict, *args) -> dict:
+    if isinstance(data, str):
+        data = json.loads(data)
+    dict_copy = copy.deepcopy(data)
+    for arg in args:
+        dict_copy.pop(arg, None)
+    return dict_copy
diff --git a/tests/test_functions.py b/tests/test_functions.py
index 44a113c335..bc9000d747 100644
--- a/tests/test_functions.py
+++ b/tests/test_functions.py
@@ -225,3 +225,15 @@ def test_encode():
 
 def test_dict_to_key_value_list():
     assert functions.dict_to_key_value_list({"a": 1, "b": "test"}) == ["a:1", "b:test"]
+
+
+def test_dict_pop():
+    d = {"a": 1, "b": 2}
+    d2 = functions.dict_pop(d, "a")
+    assert d2 == {"b": 2}
+
+
+def test_dict_pop_str():
+    d = '{"a": 1, "b": 2}'
+    d2 = functions.dict_pop(d, "a")
+    assert d2 == {"b": 2}