From 50c0439f9baf90b0aa6bbf3a309a31088bff2b10 Mon Sep 17 00:00:00 2001
From: Fabio Caccamo <fabio.caccamo@gmail.com>
Date: Fri, 4 Oct 2019 15:51:29 +0200
Subject: [PATCH] Improved remove method to accept a single key or args.

---
 benedict/dicts/__init__.py |  5 ++++-
 tests/test_benedict.py     | 33 ++++++++++++++++++++++++++++++++-
 2 files changed, 36 insertions(+), 2 deletions(-)

diff --git a/benedict/dicts/__init__.py b/benedict/dicts/__init__.py
index 586f4da..d51072b 100644
--- a/benedict/dicts/__init__.py
+++ b/benedict/dicts/__init__.py
@@ -93,7 +93,10 @@ class benedict(IODict, KeypathDict, ParseDict):
         for d in dicts:
             dict_util.merge(self, d)
 
-    def remove(self, keys):
+    def remove(self, keys, *args):
+        if isinstance(keys, string_types):
+            keys = [keys]
+        keys += args
         for key in keys:
             try:
                 del self[key]
diff --git a/tests/test_benedict.py b/tests/test_benedict.py
index c7db383..b78642b 100644
--- a/tests/test_benedict.py
+++ b/tests/test_benedict.py
@@ -911,7 +911,21 @@ class BenedictTestCase(unittest.TestCase):
         self.assertEqual(b.pop('b.c'), 2)
         # self.assertTrue(isinstance(b.pop('b.d'), benedict))
 
-    def test_remove(self):
+    def test_remove_with_key(self):
+        d = {
+            'a': 1,
+            'b': 2,
+            'c': '4',
+        }
+        b = benedict(d)
+        b.remove('c')
+        r = {
+            'a': 1,
+            'b': 2,
+        }
+        self.assertEqual(b, r)
+
+    def test_remove_with_keys_list(self):
         d = {
             'a': 1,
             'b': 2,
@@ -928,6 +942,23 @@ class BenedictTestCase(unittest.TestCase):
         }
         self.assertEqual(b, r)
 
+    def test_remove_with_keys_args(self):
+        d = {
+            'a': 1,
+            'b': 2,
+            'c': '4',
+            'e': '5',
+            'f': 6,
+            'g': 7,
+        }
+        b = benedict(d)
+        b.remove('c', 'e', 'f', 'g', 'x', 'y', 'z')
+        r = {
+            'a': 1,
+            'b': 2,
+        }
+        self.assertEqual(b, r)
+
     def test_remove_with_keypath(self):
         d = {
             'a': {