test/bdd/steps/geometry_factory.py

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