Last active
June 16, 2024 15:27
-
-
Save jsheedy/ed81cdf18190183b3b7d to your computer and use it in GitHub Desktop.
Sometimes you need a file-like object when all you have is an iterator, for instance when using psycopg2's cursor.copy_from. This class will handle the impedance mismatch.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import io | |
import sys | |
class IteratorFile(io.TextIOBase): | |
""" given an iterator which yields strings, | |
return a file like object for reading those strings """ | |
def __init__(self, it): | |
self._it = it | |
self._f = io.StringIO() | |
def read(self, length=sys.maxsize): | |
try: | |
while self._f.tell() < length: | |
self._f.write(next(self._it) + "\n") | |
except StopIteration as e: | |
# soak up StopIteration. this block is not necessary because | |
# of finally, but just to be explicit | |
pass | |
except Exception as e: | |
print("uncaught exception: {}".format(e)) | |
finally: | |
self._f.seek(0) | |
data = self._f.read(length) | |
# save the remainder for next read | |
remainder = self._f.read() | |
self._f.seek(0) | |
self._f.truncate(0) | |
self._f.write(remainder) | |
return data | |
def readline(self): | |
return next(self._it) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Thanks for asking!