refactor: rewrite the pg_config stub in a more idiomatic way

psycopg · Dec 21, 2024 · a5ae1b6 · a5ae1b6
1 parent c7a3077
commit a5ae1b6
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 32 deletions.
diff --git a/.gitignore b/.gitignore
@@ -22,3 +22,4 @@ wheelhouse/
 /psycopg/build/
 /psycopg_c/build/
 /psycopg_pool/build/
+/tools/build/pg_config_vcpkg_stub/build/
diff --git a/tools/build/pg_config_vcpkg_stub/pg_config_vcpkg_stub/__init__.py b/tools/build/pg_config_vcpkg_stub/pg_config_vcpkg_stub/__init__.py
@@ -1,51 +1,79 @@
+#!/usr/bin/env python
 """
 We use vcpkg in github actions to build psycopg-binary.
 
-This is a stub to work as `pg_config --libdir` or `pg_config --includedir` to make it work with vcpkg.
+This is a stub to work as `pg_config --libdir` or `pg_config --includedir` to
+make it work with vcpkg.
 
-You will need install vcpkg and set `VCPKG_ROOT `env
-and run `vcpkg install libpq:x64-windows-release` before you use this script
+You will need install `vcpkg`, set `VCPKG_ROOT` env, and run `vcpkg install
+libpq:x64-windows-release` before using this script.
 """
 
 import os
-import pathlib
 import sys
 import platform
+from pathlib import Path
+from argparse import ArgumentParser, Namespace, RawDescriptionHelpFormatter
 
 
-def main():
+class ScriptError(Exception):
+    """Controlled exception raised by the script."""
+
+
+def _main() -> None:
     # only x64-windows
     if not (sys.platform == "win32" and platform.machine() == "AMD64"):
-        raise Exception("this script should only be used in x64-windows")
+        raise ScriptError("this script should only be used in x64-windows")
 
-    what = sys.argv[1]
+    vcpkg_root = os.environ.get(
+        "VCPKG_ROOT", os.environ.get("VCPKG_INSTALLATION_ROOT", "")
+    )
+    if not vcpkg_root:
+        raise ScriptError("VCPKG_ROOT/VCPKG_INSTALLATION_ROOT env var not specified")
+    vcpkg_platform_root = (Path(vcpkg_root) / "installed/x64-windows-release").resolve()
 
-    if what == "--help":
-        print(__doc__)
-        return
+    args = parse_cmdline()
 
-    # on github actions it's `VCPKG_INSTALLATION_ROOT`
-    if "VCPKG_INSTALLATION_ROOT" not in os.environ:
-        print("failed to find VCPKG ROOT path", file=sys.stderr)
-        sys.exit(1)
+    if args.libdir:
+        if not (f := vcpkg_platform_root / "lib/libpq.lib").exists():
+            raise ScriptError(f"libpq library not found: {f}")
+        print(vcpkg_platform_root.joinpath("lib"))
+
+    elif args.includedir:
+        if not (d := vcpkg_platform_root / "include/libpq").is_dir():
+            raise ScriptError(f"libpq include directory not found: {d}")
+        print(vcpkg_platform_root.joinpath("include"))
+
+    else:
+        raise ScriptError("command not handled")
 
-    vcpkg_root = pathlib.Path(os.environ["VCPKG_INSTALLATION_ROOT"])
-    vcpkg_platform_root = vcpkg_root.joinpath("installed/x64-windows-release").resolve()
-    if not vcpkg_root.joinpath("packages/libpq_x64-windows-release").exists():
-        print("libpq not installed with vcpkg", file=sys.stderr)
-        sys.exit(1)
 
-    if what == "--libdir":
-        print(str(vcpkg_platform_root.joinpath("lib")))
-        return
-    if what == "--includedir":
-        print(str(vcpkg_platform_root.joinpath("include")))
-        return
-
-    print(
-        "unexpected command: {!r}\n this maybe out-of-sync between 'psycopg_c/setup.py' and 'tools/build/pg_config_vcpkg_stub/pg_config_vcpkg_stub/__init__.py'".format(
-            sys.argv[1:]
-        ),
-        file=sys.stderr,
+def parse_cmdline() -> Namespace:
+    parser = ArgumentParser(
+        description=__doc__, formatter_class=RawDescriptionHelpFormatter
     )
-    sys.exit(1)
+    g = parser.add_mutually_exclusive_group(required=True)
+    g.add_argument(
+        "--libdir",
+        action="store_true",
+        help="show location of object code libraries",
+    )
+    g.add_argument(
+        "--includedir",
+        action="store_true",
+        help="show location of C header files of the client interfaces",
+    )
+    opt = parser.parse_args()
+    return opt
+
+
+def main() -> None:
+    try:
+        _main()
+    except ScriptError as e:
+        print(f"ERROR: {e}.", file=sys.stderr)
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()