Merge remote-tracking branch 'upstream/master'

[nominatim.git] / nominatim / db / utils.py

diff --git a/nominatim/db/utils.py b/nominatim/db/utils.py
index e2b63e1dd0e360e528724e5019608a7153e31b5f..57048da3f20e0a2e612b496ac99242b44155b4e7 100644 (file)
--- a/nominatim/db/utils.py
+++ b/nominatim/db/utils.py
@@ -7,19 +7,19 @@
 """
 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()
 
-def _pipe_to_proc(proc: subprocess.Popen[bytes],
+def _pipe_to_proc(proc: 'subprocess.Popen[bytes]',
                   fdesc: Union[IO[bytes], gzip.GzipFile]) -> int:
     assert proc.stdin is not None
     chunk = fdesc.read(2048)
@@ -84,20 +84,25 @@ 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 size(self) -> int:
+        """ Return the number of bytes the buffer currently contains.
+        """
+        return self.buffer.tell()
+
+    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 +118,12 @@ 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.
+
+            The buffer is empty and reusable after this operation.
         """
         if self.buffer.tell() > 0:
             self.buffer.seek(0)
             cur.copy_from(self.buffer, table, columns=columns)
+            self.buffer = io.StringIO()