From e9093fae60e4e995c73a7b984da2f9f3f52cf2b0 Mon Sep 17 00:00:00 2001
From: Marcelo Prates <marceloorp@gmail.com>
Date: Wed, 25 Aug 2021 12:15:47 -0300
Subject: [PATCH] Added support for waterways and fetching by OSM Id

---
 prettymaps/draw.py  | 10 ++++++----
 prettymaps/fetch.py | 13 ++++++++-----
 2 files changed, 14 insertions(+), 9 deletions(-)

diff --git a/prettymaps/draw.py b/prettymaps/draw.py
index b417f85..456644d 100644
--- a/prettymaps/draw.py
+++ b/prettymaps/draw.py
@@ -18,6 +18,7 @@ from shapely.affinity import *
 from geopandas import GeoDataFrame
 
 # etc
+import re
 import pandas as pd
 from functools import reduce
 from tabulate import tabulate
@@ -105,9 +106,11 @@ def plot_shapes(shapes, ax, vsketch = None, palette = None, **kwargs):
 
 # Parse query (by coordinates, OSMId or name)
 def parse_query(query):
-    if type(query) == tuple:
+    if type(query) in([Polygon, MultiPolygon]):
+        return 'polygon'
+    elif type(query) == tuple:
         return 'coordinates'
-    elif False:
+    elif re.match('''[A-Z][0-9]+''', query):
         return 'osmid'
     else:
         return 'address'
@@ -173,9 +176,8 @@ def plot(
                 'radius': radius
             }
         else:
-            by_osmid = False
             base_kwargs = {
-                'perimeter': get_perimeter(query, by_osmid = by_osmid)
+                'perimeter': query if query_mode == 'polygon' else get_perimeter(query, by_osmid = query_mode == 'osmid')
             }
 
         # Fetch layers
diff --git a/prettymaps/fetch.py b/prettymaps/fetch.py
index 8a62982..c88fd7b 100644
--- a/prettymaps/fetch.py
+++ b/prettymaps/fetch.py
@@ -80,7 +80,10 @@ def get_geometries(perimeter = None, point = None, radius = None, tags = {}, per
     return geometries
 
 # Get streets
-def get_streets(perimeter = None, point = None, radius = None, width = 6, custom_filter = None, circle = True, dilate = 0):
+def get_streets(perimeter = None, point = None, radius = None, layer = 'streets', width = 6, custom_filter = None, circle = True, dilate = 0):
+
+    if layer == 'streets':
+        layer = 'highway'
 
     # Boundary defined by polygon (perimeter)
     if perimeter is not None:
@@ -106,10 +109,10 @@ def get_streets(perimeter = None, point = None, radius = None, width = 6, custom
         streets = unary_union([
             # Dilate streets of each highway type == 'highway' using width 'w'
             MultiLineString(
-                streets[(streets.highway == highway) & (streets.geometry.type == 'LineString')].geometry.tolist() +
+                streets[(streets[layer] == highway) & (streets.geometry.type == 'LineString')].geometry.tolist() +
                 list(reduce(lambda x, y: x+y, [
                     list(lines)
-                    for lines in streets[(streets.highway == highway) & (streets.geometry.type == 'MultiLineString')].geometry
+                    for lines in streets[(streets[layer] == highway) & (streets.geometry.type == 'MultiLineString')].geometry
                 ], []))
             ).buffer(w)
             for highway, w in width.items()
@@ -139,8 +142,8 @@ def get_layer(layer, **kwargs):
         else:
             raise Exception("Either 'perimeter' or 'point' & 'radius' must be provided")
     # Fetch streets or railway
-    if layer in ['streets', 'railway']:
-        return get_streets(**kwargs)
+    if layer in ['streets', 'railway', 'waterway']:
+        return get_streets(**kwargs, layer = layer)
     # Fetch geometries
     else:
         return get_geometries(**kwargs)
\ No newline at end of file