]> git.openstreetmap.org Git - nominatim.git/blob - test/bdd/steps/geometry_factory.py
use context management when processing Tiger data
[nominatim.git] / test / bdd / steps / geometry_factory.py
1 # SPDX-License-Identifier: GPL-2.0-only
2 #
3 # This file is part of Nominatim. (https://nominatim.org)
4 #
5 # Copyright (C) 2022 by the Nominatim developer community.
6 # For a full list of authors see the git log.
7 from pathlib import Path
8 import os
9
10 class GeometryFactory:
11     """ Provides functions to create geometries from scenes and data grids.
12     """
13
14     def __init__(self):
15         defpath = Path(__file__) / '..' / '..' / '..' / 'scenes' / 'data'
16         self.scene_path = os.environ.get('SCENE_PATH', defpath.resolve())
17         self.scene_cache = {}
18         self.grid = {}
19
20     def parse_geometry(self, geom, scene):
21         """ Create a WKT SQL term for the given geometry.
22             The function understands the following formats:
23
24               [<scene>]:<name>
25                  Geometry from a scene. If the scene is omitted, use the
26                  default scene.
27               <P>
28                  Point geometry
29               <P>,...,<P>
30                  Line geometry
31               (<P>,...,<P>)
32                  Polygon geometry
33
34            <P> may either be a coordinate of the form '<x> <y>' or a single
35            number. In the latter case it must refer to a point in
36            a previously defined grid.
37         """
38         if geom.find(':') >= 0:
39             return "ST_SetSRID({}, 4326)".format(self.get_scene_geometry(scene, geom))
40
41         if geom.find(',') < 0:
42             out = "POINT({})".format(self.mk_wkt_point(geom))
43         elif geom.find('(') < 0:
44             out = "LINESTRING({})".format(self.mk_wkt_points(geom))
45         else:
46             out = "POLYGON(({}))".format(self.mk_wkt_points(geom.strip('() ')))
47
48         return "ST_SetSRID('{}'::geometry, 4326)".format(out)
49
50     def mk_wkt_point(self, point):
51         """ Parse a point description.
52             The point may either consist of 'x y' cooordinates or a number
53             that refers to a grid setup.
54         """
55         geom = point.strip()
56         if geom.find(' ') >= 0:
57             return geom
58
59         try:
60             pt = self.grid_node(int(geom))
61         except ValueError:
62             assert False, "Scenario error: Point '{}' is not a number".format(geom)
63
64         assert pt is not None, "Scenario error: Point '{}' not found in grid".format(geom)
65         return "{} {}".format(*pt)
66
67     def mk_wkt_points(self, geom):
68         """ Parse a list of points.
69             The list must be a comma-separated list of points. Points
70             in coordinate and grid format may be mixed.
71         """
72         return ','.join([self.mk_wkt_point(x) for x in geom.split(',')])
73
74     def get_scene_geometry(self, default_scene, name):
75         """ Load the geometry from a scene.
76         """
77         geoms = []
78         for obj in name.split('+'):
79             oname = obj.strip()
80             if oname.startswith(':'):
81                 assert default_scene is not None, "Scenario error: You need to set a scene"
82                 defscene = self.load_scene(default_scene)
83                 wkt = defscene[oname[1:]]
84             else:
85                 scene, obj = oname.split(':', 2)
86                 scene_geoms = self.load_scene(scene)
87                 wkt = scene_geoms[obj]
88
89             geoms.append("'{}'::geometry".format(wkt))
90
91         if len(geoms) == 1:
92             return geoms[0]
93
94         return 'ST_LineMerge(ST_Collect(ARRAY[{}]))'.format(','.join(geoms))
95
96     def load_scene(self, name):
97         """ Load a scene from a file.
98         """
99         if name in self.scene_cache:
100             return self.scene_cache[name]
101
102         scene = {}
103         with open(Path(self.scene_path) / "{}.wkt".format(name), 'r') as fd:
104             for line in fd:
105                 if line.strip():
106                     obj, wkt = line.split('|', 2)
107                     scene[obj.strip()] = wkt.strip()
108             self.scene_cache[name] = scene
109
110         return scene
111
112     def set_grid(self, lines, grid_step):
113         """ Replace the grid with one from the given lines.
114         """
115         self.grid = {}
116         y = 0
117         for line in lines:
118             x = 0
119             for pt_id in line:
120                 if pt_id.isdigit():
121                     self.grid[int(pt_id)] = (x, y)
122                 x += grid_step
123             y += grid_step
124
125     def grid_node(self, nodeid):
126         """ Get the coordinates for the given grid node.
127         """
128         return self.grid.get(nodeid)