Stream sample

stream_sample.py

from builtins import range
from itertools import count
import unittest

def stream_sample(iterable, n):
	"""
	Peek at the first n elements of an iterable, and determine the min and max
	values. Preserve all elements in the iterable!
	
	Return a tuple of minimum, maximum (in sample), and an iterator that yields all
	"""
	
	pass

class StreamSampleTestCase(unittest.TestCase):

	def test_stream_sample(self):

		# sample only the first 10 elements of a range of length 100
		
		stream = range(100)
		
		min_val, max_val, new_stream = stream_sample(iter(stream), 10)
		
		self.assertEqual(0, min_val)
		self.assertEqual(9, max_val)
		# all elements are still there:
		self.assertEqual(list(range(100)), list(new_stream))

		# sample more elements than there are - no error raised
		# now we now the global maximum!
		
		stream = range(100)
		
		min_val, max_val, new_stream = stream_sample(iter(stream), 1000)
		
		self.assertEqual(0, min_val)
		self.assertEqual(99, max_val)
		self.assertEqual(list(range(100)), list(new_stream))

		# and guess what - it also works with infinite iterators

		stream = count(0)

		min_val, max_val, new_stream = stream_sample(iter(stream), 10)
		
		self.assertEqual(0, min_val)
		self.assertEqual(9, max_val)

if __name__ == "__main__":
	unittest.main()