Skip to content

Commit

Permalink
added type stub generation for dynamic functions
Browse files Browse the repository at this point in the history
  • Loading branch information
mbway committed Oct 10, 2023
1 parent 86d807a commit bdaf3b6
Show file tree
Hide file tree
Showing 9 changed files with 3,506 additions and 27 deletions.
4 changes: 4 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
@@ -1,3 +1,7 @@
# IDE
.vscode
.idea

# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
Expand Down
14 changes: 14 additions & 0 deletions _generate_type_stubs.py
Original file line number Diff line number Diff line change
@@ -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())
7 changes: 6 additions & 1 deletion geoalchemy2/_functions.py
Original file line number Diff line number Diff line change
@@ -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,
Expand Down
106 changes: 106 additions & 0 deletions geoalchemy2/_functions_helpers.py
Original file line number Diff line number Diff line change
@@ -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)
43 changes: 18 additions & 25 deletions geoalchemy2/functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@
"""
import re
from typing import List
from typing import Type

from sqlalchemy import inspect
Expand All @@ -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]
Expand Down Expand Up @@ -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]


Expand Down Expand Up @@ -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__
Loading

0 comments on commit bdaf3b6

Please sign in to comment.