add typing information to DB properties

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

diff --git a/nominatim/db/utils.py b/nominatim/db/utils.py
index 7d58a6684ef5d11814a02d3d497b384c18538626..461cb662415910844b87042aa76df2754046e13f 100644 (file)
--- a/nominatim/db/utils.py
+++ b/nominatim/db/utils.py
@@ -7,17 +7,21 @@
 """
 Helper functions for handling DB accesses.
 """
 """
 Helper functions for handling DB accesses.
 """

+from typing import IO, Optional, Union

 import subprocess
 import logging
 import gzip
 import io
 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()
 
 
 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:
     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)
 
 
     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
     """ 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


@@ -40,27 +47,28 @@ def execute_file(dsn, fname, ignore_errors=False, pre_code=None, post_code=None)
         cmd.extend(('-v', 'ON_ERROR_STOP=1'))
     if not LOG.isEnabledFor(logging.INFO):
         cmd.append('--quiet')
         cmd.extend(('-v', 'ON_ERROR_STOP=1'))
     if not LOG.isEnabledFor(logging.INFO):
         cmd.append('--quiet')

-    proc = subprocess.Popen(cmd, env=get_pg_env(dsn), stdin=subprocess.PIPE)
-
-    try:
-        if not LOG.isEnabledFor(logging.INFO):
-            proc.stdin.write('set client_min_messages to WARNING;'.encode('utf-8'))
-
-        if pre_code:
-            proc.stdin.write((pre_code + ';').encode('utf-8'))
-
-        if fname.suffix == '.gz':
-            with gzip.open(str(fname), 'rb') as fdesc:
-                remain = _pipe_to_proc(proc, fdesc)
-        else:
-            with fname.open('rb') as fdesc:
-                remain = _pipe_to_proc(proc, fdesc)
-
-        if remain == 0 and post_code:
-            proc.stdin.write((';' + post_code).encode('utf-8'))
-    finally:
-        proc.stdin.close()
-        ret = proc.wait()
+
+    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'))
+
+            if pre_code:
+                proc.stdin.write((pre_code + ';').encode('utf-8'))
+
+            if fname.suffix == '.gz':
+                with gzip.open(str(fname), 'rb') as fdesc:
+                    remain = _pipe_to_proc(proc, fdesc)
+            else:
+                with fname.open('rb') as fdesc:
+                    remain = _pipe_to_proc(proc, fdesc)
+
+            if remain == 0 and post_code:
+                proc.stdin.write((';' + post_code).encode('utf-8'))
+        finally:
+            proc.stdin.close()
+            ret = proc.wait()

 
     if ret != 0 or remain > 0:
         raise UsageError("Failed to execute SQL file.")
 
     if ret != 0 or remain > 0:
         raise UsageError("Failed to execute SQL file.")