X-Git-Url: https://git.openstreetmap.org./nominatim.git/blobdiff_plain/681aad7e0dc099658eea15b769fdefd44cd8c484..83054af46fb21ea40c0bebb2c8a528937e5a49a0:/nominatim/db/utils.py diff --git a/nominatim/db/utils.py b/nominatim/db/utils.py index 461cb662..9a7b4f16 100644 --- a/nominatim/db/utils.py +++ b/nominatim/db/utils.py @@ -7,14 +7,14 @@ """ Helper functions for handling DB accesses. """ -from typing import IO, Optional, Union +from typing import IO, Optional, Union, Any, Iterable import subprocess import logging import gzip import io from pathlib import Path -from nominatim.db.connection import get_pg_env +from nominatim.db.connection import get_pg_env, Cursor from nominatim.errors import UsageError LOG = logging.getLogger() @@ -84,20 +84,20 @@ class CopyBuffer: """ Data collector for the copy_from command. """ - def __init__(self): + def __init__(self) -> None: self.buffer = io.StringIO() - def __enter__(self): + def __enter__(self) -> 'CopyBuffer': return self - def __exit__(self, exc_type, exc_value, traceback): + def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None: if self.buffer is not None: self.buffer.close() - def add(self, *data): + def add(self, *data: Any) -> None: """ Add another row of data to the copy buffer. """ first = True @@ -113,9 +113,9 @@ class CopyBuffer: self.buffer.write('\n') - def copy_out(self, cur, table, columns=None): + def copy_out(self, cur: Cursor, table: str, columns: Optional[Iterable[str]] = None) -> None: """ Copy all collected data into the given table. """ if self.buffer.tell() > 0: self.buffer.seek(0) - cur.copy_from(self.buffer, table, columns=columns) + cur.copy_from(self.buffer, table, columns=columns) # type: ignore[no-untyped-call]