diff options
| -rw-r--r-- | rsa/blocks.py | 19 | ||||
| -rw-r--r-- | tests/test_blocks.py | 7 |
2 files changed, 26 insertions, 0 deletions
diff --git a/rsa/blocks.py b/rsa/blocks.py index c38c1f2..0250566 100644 --- a/rsa/blocks.py +++ b/rsa/blocks.py @@ -132,3 +132,22 @@ def yield_varblocks(infile): yield block +def yield_fixedblocks(infile, blocksize): + '''Generator, yields each block of ``blocksize`` bytes in the input file. + + :param infile: file to read and separate in blocks. + :returns: a generator that yields the contents of each block + ''' + + while True: + block = infile.read(blocksize) + + read_bytes = len(block) + if read_bytes == 0: + break + + yield block + + if read_bytes < blocksize: + break + diff --git a/tests/test_blocks.py b/tests/test_blocks.py index 8481c3c..ce5f03a 100644 --- a/tests/test_blocks.py +++ b/tests/test_blocks.py @@ -65,4 +65,11 @@ class VarblockTest(unittest.TestCase): varblocks = list(blocks.yield_varblocks(infile)) self.assertEqual(['12345', 'Sybren'], varblocks) +class FixedblockTest(unittest.TestCase): + def test_yield_fixedblock(self): + + infile = StringIO('123456Sybren') + + fixedblocks = list(blocks.yield_fixedblocks(infile, 6)) + self.assertEqual(['123456', 'Sybren'], fixedblocks) |