code.kerkeslager.com Git - fur/blob - util.py

   1 import functools
   2
   3 def force_generator(to_type):
   4     def decorator(generator_function):
   5         @functools.wraps(generator_function)
   6         def forced_generator(*args, **kwargs):
   7             return to_type(generator_function(*args, **kwargs))
   8
   9         return forced_generator
  10
  11     return decorator
  12
  13 if __name__ == '__main__':
  14     import unittest
  15
  16     class ForceGeneratorTests(unittest.TestCase):
  17         def test_forces_generator(self):
  18             forced_range = force_generator(list)(range)
  19
  20             self.assertEqual(
  21                 forced_range(10),
  22                 [0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
  23             )
  24
  25     unittest.main()