Write using an iterator

This commit is contained in:
Luke Van Seters 2022-01-03 14:50:27 -05:00
parent 6b1c469a10
commit ada540c1d4
2 changed files with 44 additions and 8 deletions

View File

@ -1,10 +1,11 @@
import io
import os
from typing import Any, Iterable, Optional
from sqlalchemy import create_engine, orm
from sqlalchemy.orm import sessionmaker
from mev_inspect.string_io import StringIteratorIO
def get_trace_database_uri() -> Optional[str]:
username = os.getenv("TRACE_DB_USER")
@ -77,12 +78,9 @@ def write_as_csv(
table_name: str,
items: Iterable[Iterable[Any]],
) -> None:
csv_file_like_object = io.StringIO()
for item in items:
csv_file_like_object.write("|".join(map(_clean_csv_value, item)) + "\n")
csv_file_like_object.seek(0)
csv_iterator = StringIteratorIO(
("|".join(map(_clean_csv_value, item)) + "\n" for item in items)
)
with db_session.connection().connection.cursor() as cursor:
cursor.copy_from(csv_file_like_object, table_name, sep="|")
cursor.copy_from(csv_iterator, table_name, sep="|")

38
mev_inspect/string_io.py Normal file
View File

@ -0,0 +1,38 @@
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)