From: Sarah Hoffmann <lonvia@denofr.de>
Date: Sat, 2 Jul 2022 08:18:10 +0000 (+0200)
Subject: type annotations for DB utils
X-Git-Tag: v4.1.0~10^2~36
X-Git-Url: https://git.openstreetmap.org./nominatim.git/commitdiff_plain/992e6f72cf26088097bd40110e201cd63b55b871

type annotations for DB utils
---

diff --git a/nominatim/db/connection.py b/nominatim/db/connection.py
index cbdbbfef..2cc0fef8 100644
--- a/nominatim/db/connection.py
+++ b/nominatim/db/connection.py
@@ -226,7 +226,7 @@ _PG_CONNECTION_STRINGS = {
 
 
 def get_pg_env(dsn: str,
-               base_env: Optional[Mapping[str, Optional[str]]] = None) -> Mapping[str, Optional[str]]:
+               base_env: Optional[Mapping[str, str]] = None) -> Mapping[str, str]:
     """ Return a copy of `base_env` with the environment variables for
         PostgresSQL set up from the given database connection string.
         If `base_env` is None, then the OS environment is used as a base
diff --git a/nominatim/db/utils.py b/nominatim/db/utils.py
index b859afa8..e2b63e1d 100644
--- 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'))