77 lines
2.0 KiB
Python
77 lines
2.0 KiB
Python
"""This is taken from https://hakibenita.com/fast-load-data-python-postgresql"""
|
|
|
|
import io
|
|
from typing import Iterator, Optional
|
|
|
|
|
|
class StringIteratorIO(io.TextIOBase):
|
|
def __init__(self, iter: Iterator[str]):
|
|
self._iter = iter
|
|
self._buff = ""
|
|
|
|
def readable(self) -> bool:
|
|
return True
|
|
|
|
def _read1(self, n: Optional[int] = None) -> str:
|
|
while not self._buff:
|
|
try:
|
|
self._buff = next(self._iter)
|
|
except StopIteration:
|
|
break
|
|
ret = self._buff[:n]
|
|
self._buff = self._buff[len(ret) :]
|
|
return ret
|
|
|
|
def read(self, n: Optional[int] = None) -> str:
|
|
line = []
|
|
if n is None or n < 0:
|
|
while True:
|
|
m = self._read1()
|
|
if not m:
|
|
break
|
|
line.append(m)
|
|
else:
|
|
while n > 0:
|
|
m = self._read1(n)
|
|
if not m:
|
|
break
|
|
n -= len(m)
|
|
line.append(m)
|
|
return "".join(line)
|
|
|
|
|
|
class BytesIteratorIO(io.BufferedIOBase):
|
|
def __init__(self, iter: Iterator[bytes]):
|
|
self._iter = iter
|
|
self._buff = b""
|
|
|
|
def readable(self) -> bool:
|
|
return True
|
|
|
|
def _read1(self, n: Optional[int] = None) -> bytes:
|
|
while not self._buff:
|
|
try:
|
|
self._buff = next(self._iter)
|
|
except StopIteration:
|
|
break
|
|
ret = self._buff[:n]
|
|
self._buff = self._buff[len(ret) :]
|
|
return ret
|
|
|
|
def read(self, n: Optional[int] = None) -> bytes:
|
|
line = []
|
|
if n is None or n < 0:
|
|
while True:
|
|
m = self._read1()
|
|
if not m:
|
|
break
|
|
line.append(m)
|
|
else:
|
|
while n > 0:
|
|
m = self._read1(n)
|
|
if not m:
|
|
break
|
|
n -= len(m)
|
|
line.append(m)
|
|
return b"".join(line)
|