1"""
2
3This module defines the :class:`GenericFunction` class, which is the base for
4the implementation of spatial functions in GeoAlchemy.  This module is also
5where actual spatial functions are defined. Spatial functions supported by
6GeoAlchemy are defined in this module. See :class:`GenericFunction` to know how
7to create new spatial functions.
8
9.. note::
10
11    By convention the names of spatial functions are prefixed by ``ST_``.  This
12    is to be consistent with PostGIS', which itself is based on the ``SQL-MM``
13    standard.
14
15Functions created by subclassing :class:`GenericFunction` can be called
16in several ways:
17
18* By using the ``func`` object, which is the SQLAlchemy standard way of calling
19  a function. For example, without the ORM::
20
21      select([func.ST_Area(lake_table.c.geom)])
22
23  and with the ORM::
24
25      Session.query(func.ST_Area(Lake.geom))
26
27* By applying the function to a geometry column. For example, without the
28  ORM::
29
30      select([lake_table.c.geom.ST_Area()])
31
32  and with the ORM::
33
34      Session.query(Lake.geom.ST_Area())
35
36* By applying the function to a :class:`geoalchemy2.elements.WKBElement`
37  object (:class:`geoalchemy2.elements.WKBElement` is the type into
38  which GeoAlchemy converts geometry values read from the database), or
39  to a :class:`geoalchemy2.elements.WKTElement` object. For example,
40  without the ORM::
41
42      conn.scalar(lake['geom'].ST_Area())
43
44  and with the ORM::
45
46      session.scalar(lake.geom.ST_Area())
47
48Reference
49---------
50
51"""
52import re
53
54from sqlalchemy import inspect
55from sqlalchemy.sql import functions
56from sqlalchemy.sql.elements import ColumnElement
57from sqlalchemy.ext.compiler import compiles
58from sqlalchemy.util import with_metaclass
59
60from . import elements
61from ._functions import _FUNCTIONS
62
63
64class _GenericMeta(functions._GenericMeta):
65    """Extend the metaclass mechanism of sqlalchemy to register the functions in
66    a specific registry for geoalchemy2"""
67
68    _register = False
69
70    def __init__(cls, clsname, bases, clsdict):
71        # Register the function
72        elements.function_registry.add(clsname.lower())
73
74        super(_GenericMeta, cls).__init__(clsname, bases, clsdict)
75
76
77class TableRowElement(ColumnElement):
78    def __init__(self, selectable):
79        self.selectable = selectable
80
81    @property
82    def _from_objects(self):
83        return [self.selectable]
84
85
86class ST_AsGeoJSON(with_metaclass(_GenericMeta, functions.GenericFunction)):
87    """Special process for the ST_AsGeoJSON() function to be able to work with its
88    feature version introduced in PostGIS 3."""
89
90    name = "ST_AsGeoJSON"
91
92    def __init__(self, *args, **kwargs):
93        expr = kwargs.pop('expr', None)
94        args = list(args)
95        if expr is not None:
96            args = [expr] + args
97        for idx, element in enumerate(args):
98            if isinstance(element, functions.Function):
99                continue
100            elif isinstance(element, elements.HasFunction):
101                if element.extended:
102                    func_name = element.geom_from_extended_version
103                    func_args = [element.data]
104                else:
105                    func_name = element.geom_from
106                    func_args = [element.data, element.srid]
107                args[idx] = getattr(functions.func, func_name)(*func_args)
108            else:
109                try:
110                    insp = inspect(element)
111                    if hasattr(insp, "selectable"):
112                        args[idx] = TableRowElement(insp.selectable)
113                except Exception:
114                    continue
115
116        functions.GenericFunction.__init__(self, *args, **kwargs)
117
118    __doc__ = (
119        'Return the geometry as a GeoJSON "geometry" object, or the row as a '
120        'GeoJSON feature" object (PostGIS 3 only). (Cf GeoJSON specifications RFC '
121        '7946). 2D and 3D Geometries are both supported. GeoJSON only support SFS '
122        '1.1 geometry types (no curve support for example). '
123        'See https://postgis.net/docs/ST_AsGeoJSON.html')
124
125
126@compiles(TableRowElement)
127def _compile_table_row_thing(element, compiler, **kw):
128    # In order to get a name as reliably as possible, noting that some
129    # SQL compilers don't say "table AS name" and might not have the "AS",
130    # table and alias names can have spaces in them, etc., get it from
131    # a column instead because that's what we want to be showing here anyway.
132
133    compiled = compiler.process(list(element.selectable.columns)[0], **kw)
134
135    # 1. check for exact name of the selectable is here, use that.
136    # This way if it has dots and spaces and anything else in it, we
137    # can get it w/ correct quoting
138    schema = getattr(element.selectable, "schema", "")
139    name = element.selectable.name
140    pattern = r"(.?%s.?\.)?(.?%s.?)\." % (schema, name)
141    m = re.match(pattern, compiled)
142    if m:
143        return m.group(2)
144
145    # 2. just split on the dot, assume anonymized name
146    return compiled.split(".")[0]
147
148
149class GenericFunction(with_metaclass(_GenericMeta, functions.GenericFunction)):
150    """
151    The base class for GeoAlchemy functions.
152
153    This class inherits from ``sqlalchemy.sql.functions.GenericFunction``, so
154    functions defined by subclassing this class can be given a fixed return
155    type. For example, functions like :class:`ST_Buffer` and
156    :class:`ST_Envelope` have their ``type`` attributes set to
157    :class:`geoalchemy2.types.Geometry`.
158
159    This class allows constructs like ``Lake.geom.ST_Buffer(2)``. In that
160    case the ``Function`` instance is bound to an expression (``Lake.geom``
161    here), and that expression is passed to the function when the function
162    is actually called.
163
164    If you need to use a function that GeoAlchemy does not provide you will
165    certainly want to subclass this class. For example, if you need the
166    ``ST_TransScale`` spatial function, which isn't (currently) natively
167    supported by GeoAlchemy, you will write this::
168
169        from geoalchemy2 import Geometry
170        from geoalchemy2.functions import GenericFunction
171
172        class ST_TransScale(GenericFunction):
173            name = 'ST_TransScale'
174            type = Geometry
175    """
176
177    # Set _register to False in order not to register this class in
178    # sqlalchemy.sql.functions._registry. Only its children will be registered.
179    _register = False
180
181    def __init__(self, *args, **kwargs):
182        expr = kwargs.pop('expr', None)
183        args = list(args)
184        if expr is not None:
185            args = [expr] + args
186        for idx, elem in enumerate(args):
187            if isinstance(elem, elements.HasFunction):
188                if elem.extended:
189                    func_name = elem.geom_from_extended_version
190                    func_args = [elem.data]
191                else:
192                    func_name = elem.geom_from
193                    func_args = [elem.data, elem.srid]
194                args[idx] = getattr(functions.func, func_name)(*func_args)
195        functions.GenericFunction.__init__(self, *args, **kwargs)
196
197
198# Iterate through _FUNCTIONS and create GenericFunction classes dynamically
199for name, type_, doc in _FUNCTIONS:
200    attributes = {'name': name}
201    docs = []
202
203    if isinstance(doc, tuple):
204        docs.append(doc[0])
205        docs.append('see http://postgis.net/docs/{0}.html'.format(doc[1]))
206    elif doc is not None:
207        docs.append(doc)
208        docs.append('see http://postgis.net/docs/{0}.html'.format(name))
209
210    if type_ is not None:
211        attributes['type'] = type_
212
213        type_str = '{0}.{1}'.format(type_.__module__, type_.__name__)
214        docs.append('Return type: :class:`{0}`.'.format(type_str))
215
216    if len(docs) != 0:
217        attributes['__doc__'] = '\n\n'.join(docs)
218
219    globals()[name] = type(name, (GenericFunction,), attributes)
220
221
222#
223# Define compiled versions for functions in SpatiaLite whose names don't have
224# the ST_ prefix.
225#
226
227
228_SQLITE_FUNCTIONS = {
229    "ST_GeomFromEWKT": "GeomFromEWKT",
230    "ST_GeomFromEWKB": "GeomFromEWKB",
231    "ST_AsBinary": "AsBinary",
232    "ST_AsEWKB": "AsEWKB",
233    "ST_AsGeoJSON": "AsGeoJSON",
234}
235
236
237# Default handlers are required for SQLAlchemy < 1.1
238# See more details in https://github.com/geoalchemy/geoalchemy2/issues/213
239def _compiles_default(cls):
240    def _compile_default(element, compiler, **kw):
241        return "{}({})".format(cls, compiler.process(element.clauses, **kw))
242    compiles(globals()[cls])(_compile_default)
243
244
245def _compiles_sqlite(cls, fn):
246    def _compile_sqlite(element, compiler, **kw):
247        return "{}({})".format(fn, compiler.process(element.clauses, **kw))
248    compiles(globals()[cls], "sqlite")(_compile_sqlite)
249
250
251for cls, fn in _SQLITE_FUNCTIONS.items():
252    _compiles_default(cls)
253    _compiles_sqlite(cls, fn)
254