Add utility function to create frozenset with unique values

When used instead of a plain call to “frozenset”, this would have
avoided the issue fixed in commit e2dd6ec. The new function is located
in the “compat” module as it will be used at module load time in most
places and should therefore reside in a place with very few
dependencies.

Signed-off-by: Michael Hanselmann <hansmi@google.com>
Reviewed-by: Iustin Pop <iustin@google.com>

diff --git a/lib/compat.py b/lib/compat.py
index 6515af1..5f1409e 100644 (file)
--- a/lib/compat.py
+++ b/lib/compat.py
@@ -146,6 +146,26 @@ def TryToRoman(val, convert=True):
   else:
     return val
 
+
+def UniqueFrozenset(seq):
+  """Makes C{frozenset} from sequence after checking for duplicate elements.
+
+  @raise ValueError: When there are duplicate elements
+
+  """
+  if isinstance(seq, (list, tuple)):
+    items = seq
+  else:
+    items = list(seq)
+
+  result = frozenset(items)
+
+  if len(items) != len(result):
+    raise ValueError("Duplicate values found")
+
+  return result
+
+
 #: returns the first element of a list-like value
 fst = operator.itemgetter(0)
 

diff --git a/test/ganeti.compat_unittest.py b/test/ganeti.compat_unittest.py
index 54a1c62..8ed54fe 100755 (executable)
--- a/test/ganeti.compat_unittest.py
+++ b/test/ganeti.compat_unittest.py
@@ -99,5 +99,24 @@ class TestTryToRoman(testutils.GanetiTestCase):
     self.assertEquals(compat.TryToRoman("19", convert=False), "19")
 
 
+class TestUniqueFrozenset(unittest.TestCase):
+  def testDuplicates(self):
+    for values in [["", ""], ["Hello", "World", "Hello"]]:
+      self.assertRaises(ValueError, compat.UniqueFrozenset, values)
+
+  def testEmpty(self):
+    self.assertEqual(compat.UniqueFrozenset([]), frozenset([]))
+
+  def testUnique(self):
+    self.assertEqual(compat.UniqueFrozenset([1, 2, 3]), frozenset([1, 2, 3]))
+
+  def testGenerator(self):
+    seq = ("Foo%s" % i for i in range(10))
+    self.assertTrue(callable(seq.next))
+    self.assertFalse(isinstance(seq, (list, tuple)))
+    self.assertEqual(compat.UniqueFrozenset(seq),
+                     frozenset(["Foo%s" % i for i in range(10)]))
+
+
 if __name__ == "__main__":
   testutils.GanetiTestProgram()