From 383740e8ab54008d909de1b7166b9d8840883392 Mon Sep 17 00:00:00 2001
From: asday <sara.and.zuka@gmail.com>
Date: Mon, 19 Aug 2019 22:29:19 +0100
Subject: [PATCH] Added `Choices().subset()`.

---
 model_utils/choices.py | 14 ++++++++++++++
 tests/test_choices.py  | 27 +++++++++++++++++++++++++++
 2 files changed, 41 insertions(+)

diff --git a/model_utils/choices.py b/model_utils/choices.py
index 6339503..31d5aa1 100644
--- a/model_utils/choices.py
+++ b/model_utils/choices.py
@@ -142,3 +142,17 @@ class Choices(object):
 
     def __deepcopy__(self, memo):
         return self.__class__(*copy.deepcopy(self._triples, memo))
+
+    def subset(self, *new_identifiers):
+        identifiers = set(self._identifier_map.keys())
+
+        if not identifiers.issuperset(new_identifiers):
+            raise ValueError(
+                'The following identifiers are not present: %s' %
+                identifiers.symmetric_difference(new_identifiers),
+            )
+
+        return self.__class__(*[
+            choice for choice in self._triples
+            if choice[1] in new_identifiers
+        ])
diff --git a/tests/test_choices.py b/tests/test_choices.py
index 986670c..cb5bec9 100644
--- a/tests/test_choices.py
+++ b/tests/test_choices.py
@@ -279,3 +279,30 @@ class IdentifierChoicesTests(ChoicesTests):
                 ('group b', [(3, 'three')]),
             ],
         )
+
+
+class SubsetChoicesTest(TestCase):
+
+    def setUp(self):
+        self.choices = Choices(
+            (0, 'a', 'A'),
+            (1, 'b', 'B'),
+        )
+
+    def test_nonexistent_identifiers_raise(self):
+        with self.assertRaises(ValueError):
+            self.choices.subset('a', 'c')
+
+    def test_solo_nonexistent_identifiers_raise(self):
+        with self.assertRaises(ValueError):
+            self.choices.subset('c')
+
+    def test_empty_subset_passes(self):
+        subset = self.choices.subset()
+
+        self.assertEqual(subset, Choices())
+
+    def test_subset_returns_correct_subset(self):
+        subset = self.choices.subset('a')
+
+        self.assertEqual(subset, Choices((0, 'a', 'A')))