added type stub generation for dynamic functions

Author: mbway

.gitignore

@@ -1,3 +1,7 @@
+# IDE
+.vscode
+.idea
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

_generate_type_stubs.py

@@ -0,0 +1,14 @@
+from pathlib import Path
+
+from geoalchemy2._functions_helpers import _generate_stubs
+
+"""
+this script is outside the geoalchemy2 package because the 'geoalchemy2.types'
+package interferes with the 'types' module in the standard library
+"""
+
+script_dir = Path(__file__).resolve().parent
+
+
+if __name__ == "__main__":
+    (script_dir / "geoalchemy2/functions.pyi").write_text(_generate_stubs())

geoalchemy2/_functions.py

@@ -1,9 +1,14 @@
 # -*- coding: utf-8 -*-
 # flake8: noqa
+from typing import List
+from typing import Optional
+from typing import Tuple
+from typing import Union
+
 from geoalchemy2 import types
 
 # fmt: off
-_FUNCTIONS = [
+_FUNCTIONS: List[Tuple[str, Optional[type], Union[None, str, Tuple[str, str]]]] = [
     ('AddGeometryColumn', None,
      '''Adds a geometry column to an existing table.'''),
     ('DropGeometryColumn', None,

geoalchemy2/_functions_helpers.py

@@ -0,0 +1,106 @@
+from typing import Callable
+from typing import Generic
+from typing import Optional
+from typing import Tuple
+from typing import TypeVar
+from typing import Union
+from typing import cast
+
+from sqlalchemy.sql import functions
+from typing_extensions import ParamSpec
+
+
+def _get_docstring(name: str, doc: Union[None, str, Tuple[str, str]], type_: Optional[type]) -> str:
+    doc_string_parts = []
+
+    if isinstance(doc, tuple):
+        doc_string_parts.append(doc[0])
+        doc_string_parts.append("see http://postgis.net/docs/{0}.html".format(doc[1]))
+    elif doc is not None:
+        doc_string_parts.append(doc)
+        doc_string_parts.append("see http://postgis.net/docs/{0}.html".format(name))
+
+    if type_ is not None:
+        return_type_str = "{0}.{1}".format(type_.__module__, type_.__name__)
+        doc_string_parts.append("Return type: :class:`{0}`.".format(return_type_str))
+
+    return "\n\n".join(doc_string_parts)
+
+
+def _replace_indent(text: str, indent: str) -> str:
+    lines = []
+    for i, line in enumerate(text.splitlines()):
+        if i == 0 or not line.strip():
+            lines.append(line)
+        else:
+            lines.append(f"{indent}{line}")
+    return "\n".join(lines)
+
+
+def _generate_stubs() -> str:
+    """Generates type stubs for the dynamic functions described in `geoalchemy2/_functions.py`."""
+    from geoalchemy2._functions import _FUNCTIONS
+    from geoalchemy2.functions import ST_AsGeoJSON
+
+    header = '''\
+# this file is automatically generated
+from typing import Any
+from typing import List
+
+from sqlalchemy.sql import functions
+from sqlalchemy.sql.elements import ColumnElement
+
+import geoalchemy2.types
+from geoalchemy2._functions_helpers import _generic_function
+
+class GenericFunction(functions.GenericFunction): ...
+
+class TableRowElement(ColumnElement):
+    inherit_cache: bool = ...
+    """The cache is disabled for this class."""
+
+    def __init__(self, selectable: bool) -> None: ...
+    @property
+    def _from_objects(self) -> List[bool]: ...  # type: ignore[override]
+'''
+    stub_file_parts = [header]
+
+    functions = _FUNCTIONS.copy()
+    functions.insert(0, ("ST_AsGeoJSON", str, ST_AsGeoJSON.__doc__))
+
+    for name, type_, doc in functions:
+        doc = _replace_indent(_get_docstring(name, doc, type_), "    ")
+
+        if type_ is None:
+            type_str = "None"
+        elif type_.__module__ == "builtins":
+            type_str = type_.__name__
+        else:
+            type_str = f"{type_.__module__}.{type_.__name__}"
+
+        signature = f'''\
+@_generic_function
+def {name}(*args: Any, **kwargs: Any) -> {type_str}:
+    """{doc}"""
+    ...
+'''
+        stub_file_parts.append(signature)
+
+    return "\n".join(stub_file_parts)
+
+
+_P = ParamSpec("_P")
+_R = TypeVar("_R", covariant=True)
+
+
+class _GenericFunction(functions.GenericFunction, Generic[_P, _R]):
+    def __call__(self, *args: _P.args, **kwargs: _P.kwargs) -> _R:  # type: ignore[empty-body]
+        ...
+
+
+def _generic_function(func: Callable[_P, _R]) -> _GenericFunction[_P, _R]:
+    """Take a regular function and extend it with attributes from sqlalchemy GenericFunction.
+
+    based on https://github.com/python/mypy/issues/2087#issuecomment-1194111648
+    """
+    return cast(_GenericFunction[_P, _R], func)

geoalchemy2/functions.py

@@ -67,6 +67,7 @@
 
 """
 import re
+from typing import List
 from typing import Type
 
 from sqlalchemy import inspect
@@ -77,6 +78,7 @@
 
 from geoalchemy2 import elements
 from geoalchemy2._functions import _FUNCTIONS
+from geoalchemy2._functions_helpers import _get_docstring
 
 _GeoFunctionBase: Type[functions.GenericFunction]
 _GeoFunctionParent: Type[functions.GenericFunction]
@@ -131,11 +133,11 @@ class TableRowElement(ColumnElement):
     inherit_cache: bool = False
     """The cache is disabled for this class."""
 
-    def __init__(self, selectable) -> None:
+    def __init__(self, selectable: bool) -> None:
         self.selectable = selectable
 
     @property
-    def _from_objects(self):
+    def _from_objects(self) -> List[bool]:
         return [self.selectable]
 
 
@@ -262,33 +264,24 @@ def __init__(self, *args, **kwargs) -> None:
 ]
 
 
-# Iterate through _FUNCTIONS and create GenericFunction classes dynamically
-for name, type_, doc in _FUNCTIONS:
-    attributes = {
-        "name": name,
-        "inherit_cache": True,
-    }
-    docs = []
+def _create_dynamic_functions() -> None:
+    # Iterate through _FUNCTIONS and create GenericFunction classes dynamically
+    for name, type_, doc in _FUNCTIONS:
+        attributes = {
+            "name": name,
+            "inherit_cache": True,
+            "__doc__": _get_docstring(name, doc, type_),
+        }
 
-    if isinstance(doc, tuple):
-        docs.append(doc[0])
-        docs.append("see http://postgis.net/docs/{0}.html".format(doc[1]))
-    elif doc is not None:
-        docs.append(doc)
-        docs.append("see http://postgis.net/docs/{0}.html".format(name))
+        if type_ is not None:
+            attributes["type"] = type_
 
-    if type_ is not None:
-        attributes["type"] = type_
+        globals()[name] = type(name, (GenericFunction,), attributes)
+        __all__.append(name)
 
-        type_str = "{0}.{1}".format(type_.__module__, type_.__name__)
-        docs.append("Return type: :class:`{0}`.".format(type_str))
 
-    if len(docs) != 0:
-        attributes["__doc__"] = "\n\n".join(docs)
+_create_dynamic_functions()
 
-    globals()[name] = type(name, (GenericFunction,), attributes)
-    __all__.append(name)
 
-
-def __dir__():
+def __dir__() -> list[str]:
     return __all__