From 9edd2c6bf3c0d622816fbaa0f3e1600ed97cd5a5 Mon Sep 17 00:00:00 2001
From: Fabio Caccamo <fabio.caccamo@gmail.com>
Date: Fri, 4 Oct 2019 15:52:14 +0200
Subject: [PATCH] Improved subset method to accept a single key or args.

---
 benedict/dicts/__init__.py |  5 ++++-
 tests/test_benedict.py     | 19 +++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/benedict/dicts/__init__.py b/benedict/dicts/__init__.py
index d51072b..dcdf47f 100644
--- a/benedict/dicts/__init__.py
+++ b/benedict/dicts/__init__.py
@@ -104,8 +104,11 @@ class benedict(IODict, KeypathDict, ParseDict):
                 continue
 
     @benediction
-    def subset(self, keys):
+    def subset(self, keys, *args):
         d = self.__class__()
+        if isinstance(keys, string_types):
+            keys = [keys]
+        keys += args
         for key in keys:
             d[key] = self.get(key, None)
         return d
diff --git a/tests/test_benedict.py b/tests/test_benedict.py
index b78642b..c594424 100644
--- a/tests/test_benedict.py
+++ b/tests/test_benedict.py
@@ -1025,6 +1025,25 @@ class BenedictTestCase(unittest.TestCase):
         self.assertEqual(f, r)
         self.assertFalse(f is b)
 
+    def test_subset_with_keys_args(self):
+        d = {
+            'a': 1,
+            'b': 2,
+            'c': '4',
+            'e': '5',
+            'f': 6,
+            'g': 7,
+        }
+        b = benedict(d)
+        f = b.subset('c', 'f', 'x')
+        r = {
+            'c': '4',
+            'f': 6,
+            'x': None,
+        }
+        self.assertEqual(f, r)
+        self.assertFalse(f is b)
+
     def test_subset_with_keypath(self):
         d = {
             'x': {