]> git.openstreetmap.org Git - nominatim.git/commitdiff
type annotations for DB utils
authorSarah Hoffmann <lonvia@denofr.de>
Sat, 2 Jul 2022 08:18:10 +0000 (10:18 +0200)
committerSarah Hoffmann <lonvia@denofr.de>
Mon, 18 Jul 2022 07:47:57 +0000 (09:47 +0200)
nominatim/db/connection.py
nominatim/db/utils.py

index cbdbbfef323ccf94abc9917ddbc6322739dceb81..2cc0fef8596db5ae4d1f259476f3da477ec20f7a 100644 (file)
@@ -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
index b859afa8137e9254b4a4a88ce3d8236fb1383d8c..e2b63e1dd0e360e528724e5019608a7153e31b5f 100644 (file)
@@ -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'))