diff --git a/test/test_util.py b/test/test_util.py index 7da9be2f5..93c4b7ee6 100644 --- a/test/test_util.py +++ b/test/test_util.py @@ -14,37 +14,64 @@ class TestGeneratorFile(unittest.TestCase): def test_basic_generator(self): with GeneratorFile(self.sample_generator()) as f: + self.assertEquals(0, f.tell()) self.assertEquals("thisisatest", f.read()) + self.assertEquals(len("thisisatest"), f.tell()) def test_same_lengths(self): with GeneratorFile(self.sample_generator()) as f: self.assertEquals("this", f.read(4)) + self.assertEquals(4, f.tell()) + self.assertEquals("is", f.read(2)) + self.assertEquals(6, f.tell()) + self.assertEquals("a", f.read(1)) + self.assertEquals(7, f.tell()) + self.assertEquals("test", f.read(4)) + self.assertEquals(11, f.tell()) def test_indexed_lengths(self): with GeneratorFile(self.sample_generator()) as f: self.assertEquals("thisis", f.read(6)) + self.assertEquals(6, f.tell()) + self.assertEquals("atest", f.read(5)) + self.assertEquals(11, f.tell()) def test_misindexed_lengths(self): with GeneratorFile(self.sample_generator()) as f: self.assertEquals("thisis", f.read(6)) + self.assertEquals(6, f.tell()) + self.assertEquals("ate", f.read(3)) + self.assertEquals(9, f.tell()) + self.assertEquals("st", f.read(2)) + self.assertEquals(11, f.tell()) + self.assertEquals("", f.read(2)) + self.assertEquals(11, f.tell()) def test_misindexed_lengths_2(self): with GeneratorFile(self.sample_generator()) as f: self.assertEquals("thisisat", f.read(8)) + self.assertEquals(8, f.tell()) + self.assertEquals("e", f.read(1)) + self.assertEquals(9, f.tell()) + self.assertEquals("st", f.read(2)) + self.assertEquals(11, f.tell()) + self.assertEquals("", f.read(2)) + self.assertEquals(11, f.tell()) def test_overly_long(self): with GeneratorFile(self.sample_generator()) as f: self.assertEquals("thisisatest", f.read(60)) + self.assertEquals(11, f.tell()) class TestUsernameGenerator(unittest.TestCase): diff --git a/util/registry/generatorfile.py b/util/registry/generatorfile.py index 36544cc28..e7d2bb657 100644 --- a/util/registry/generatorfile.py +++ b/util/registry/generatorfile.py @@ -10,10 +10,16 @@ class GeneratorFile(object): self._generator = generator self._closed = False self._buf = '' + self._position = 0 def __iter__(self): return self + def tell(self): + """Return the file's current position, like stdio's ftell().""" + _complain_ifclosed(self._closed) + return self._position + def next(self): """A file object is its own iterator, for example iter(f) returns f (unless f is closed). When a file is used as an iterator, typically @@ -62,8 +68,11 @@ class GeneratorFile(object): self._buf = '' returned = buf + self._position = self._position + len(returned) return returned + def seek(self): + raise NotImplementedError def close(self): self._closed = True