1# !/usr/bin/env python3
2###############################################################################
3# $Id: osr_util.py 368b3b607f79055191feb2d8674361ce8fb53bce 2021-05-12 19:00:06 +0200 Even Rouault $
4#
5#  Project:  GDAL utils.auxiliary
6#  Purpose:  OSR utility functions
7#  Author:   Idan Miara <idan@miara.com>
8#
9# ******************************************************************************
10#  Copyright (c) 2021, Idan Miara <idan@miara.com>
11#
12# Permission is hereby granted, free of charge, to any person obtaining a
13# copy of this software and associated documentation files (the "Software"),
14# to deal in the Software without restriction, including without limitation
15# the rights to use, copy, modify, merge, publish, distribute, sublicense,
16# and/or sell copies of the Software, and to permit persons to whom the
17# Software is furnished to do so, subject to the following conditions:
18#
19# The above copyright notice and this permission notice shall be included
20# in all copies or substantial portions of the Software.
21#
22# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
23# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
24# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
25# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
26# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
27# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
28# DEALINGS IN THE SOFTWARE.
29###############################################################################
30from typing import Union, Tuple, Optional
31
32import osgeo
33from osgeo import osr, ogr, gdal
34
35from osgeo_utils.auxiliary.array_util import ArrayLike
36
37AnySRS = Union[str, int, osr.SpatialReference, gdal.Dataset]
38OAMS_AXIS_ORDER = int
39_default_axis_order: Optional[OAMS_AXIS_ORDER] = None
40
41
42def get_srs(srs_like: AnySRS,
43            axis_order: Optional[OAMS_AXIS_ORDER] = None) -> osr.SpatialReference:
44    """
45    returns an SRS object from epsg, pj_string or DataSet or SRS object
46
47    gis_order and axis_order are mutually exclusive
48    if gis_order is not None then axis_order is set according to gis_order
49    if axis_order == None -> set axis order to default_axis_order iff the input is not an osr.SpatialRefernce reference,
50    otherwise set requested axis_order
51    """
52    if isinstance(srs_like, osr.SpatialReference):
53        srs = srs_like
54    elif isinstance(srs_like, gdal.Dataset):
55        srs = osr.SpatialReference()
56        srs.ImportFromWkt(srs_like.GetProjection())
57    elif isinstance(srs_like, int):
58        srs = osr.SpatialReference()
59        if srs.ImportFromEPSG(srs_like) != ogr.OGRERR_NONE:
60            raise Exception(f'ogr error when parsing srs from epsg:{srs_like}')
61    elif isinstance(srs_like, str):
62        srs = osr.SpatialReference()
63        if srs.SetFromUserInput(srs_like) != ogr.OGRERR_NONE:  # accept PROJ string, WKT, PROJJSON, etc.
64            raise Exception(f'ogr error when parsing srs from user input: {srs_like}')
65    else:
66        raise Exception(f'Unknown SRS: {srs_like}')
67
68    if axis_order is None and srs != srs_like:
69        axis_order = _default_axis_order
70    if axis_order is not None and int(osgeo.__version__[0]) >= 3:
71        # GDAL 3 changes axis order: https://github.com/OSGeo/gdal/issues/1546
72        if not isinstance(axis_order, OAMS_AXIS_ORDER):
73            raise Exception(f'Unexpected axis_order: {axis_order}')
74        srs_axis_order = srs.GetAxisMappingStrategy()
75        if axis_order != srs_axis_order:
76            if srs == srs_like:
77                # we don't want to change the input srs, thus create a copy
78                srs = srs.Clone()
79            srs.SetAxisMappingStrategy(axis_order)
80    return srs
81
82
83def get_axis_order_from_gis_order(gis_order: Optional[bool]):
84    return None if gis_order is None \
85        else osr.OAMS_TRADITIONAL_GIS_ORDER if gis_order \
86        else osr.OAMS_AUTHORITY_COMPLIANT
87
88
89def get_gis_order_from_axis_order(axis_order: Optional[OAMS_AXIS_ORDER]):
90    return None if axis_order is None else axis_order == osr.OAMS_TRADITIONAL_GIS_ORDER
91
92
93def set_default_axis_order(axis_order: Optional[OAMS_AXIS_ORDER] = None) -> Optional[OAMS_AXIS_ORDER]:
94    global _default_axis_order
95    _default_axis_order = axis_order
96    return _default_axis_order
97
98
99def get_default_axis_order() -> Optional[OAMS_AXIS_ORDER]:
100    global _default_axis_order
101    return _default_axis_order
102
103
104def get_srs_pj(srs: AnySRS) -> str:
105    srs = get_srs(srs)
106    srs_pj4 = srs.ExportToProj4()
107    return srs_pj4
108
109
110def are_srs_equivalent(srs1: AnySRS, srs2: AnySRS) -> bool:
111    if srs1 == srs2:
112        return True
113    srs1 = get_srs(srs1)
114    srs2 = get_srs(srs2)
115    return srs1.IsSame(srs2)
116
117
118def get_transform(src_srs: AnySRS, tgt_srs: AnySRS) -> Optional[osr.CoordinateTransformation]:
119    src_srs = get_srs(src_srs)
120    tgt_srs = get_srs(tgt_srs)
121    if src_srs.IsSame(tgt_srs):
122        return None
123    else:
124        return osr.CoordinateTransformation(src_srs, tgt_srs)
125
126
127def transform_points(ct: Optional[osr.CoordinateTransformation],
128                     x: ArrayLike, y: ArrayLike, z: Optional[ArrayLike] = None) -> None:
129    if ct is not None:
130        if z is None:
131            for idx, (x0, y0) in enumerate(zip(x, y)):
132                x[idx], y[idx], _z = ct.TransformPoint(x0, y0)
133        else:
134            for idx, (x0, y0, z0) in enumerate(zip(x, y, z)):
135                x[idx], y[idx], z[idx] = ct.TransformPoint(x0, y0, z0)
136