type annotations for DB utils

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

diff --git a/nominatim/db/utils.py b/nominatim/db/utils.py
index b859afa8137e9254b4a4a88ce3d8236fb1383d8c..e2b63e1dd0e360e528724e5019608a7153e31b5f 100644 (file)
--- a/nominatim/db/utils.py
+++ b/nominatim/db/utils.py
@@ -7,17 +7,21 @@
 """
 Helper functions for handling DB accesses.
 """
+from typing import IO, Optional, Union
 import subprocess
 import logging
 import gzip
 import io
+from pathlib import Path
 
 from nominatim.db.connection import get_pg_env
 from nominatim.errors import UsageError
 
 LOG = logging.getLogger()
 
-def _pipe_to_proc(proc, fdesc):
+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)
     while chunk and proc.poll() is None:
         try:
@@ -28,7 +32,10 @@ def _pipe_to_proc(proc, fdesc):
 
     return len(chunk)
 
-def execute_file(dsn, fname, ignore_errors=False, pre_code=None, post_code=None):
+def execute_file(dsn: str, fname: Path,
+                 ignore_errors: bool = False,
+                 pre_code: Optional[str] = None,
+                 post_code: Optional[str] = None) -> None:
     """ Read an SQL file and run its contents against the given database
         using psql. Use `pre_code` and `post_code` to run extra commands
         before or after executing the file. The commands are run within the
@@ -42,6 +49,7 @@ def execute_file(dsn, fname, ignore_errors=False, pre_code=None, post_code=None)
         cmd.append('--quiet')
 
     with subprocess.Popen(cmd, env=get_pg_env(dsn), stdin=subprocess.PIPE) as proc:
+        assert proc.stdin is not None
         try:
             if not LOG.isEnabledFor(logging.INFO):
                 proc.stdin.write('set client_min_messages to WARNING;'.encode('utf-8'))