diff --git a/sharedstreets/dataframe/__init__.py b/sharedstreets/dataframe/__init__.py index 5aadb60..92d955f 100644 --- a/sharedstreets/dataframe/__init__.py +++ b/sharedstreets/dataframe/__init__.py @@ -1,12 +1,15 @@ +import json import itertools import functools import operator import logging import collections import geopandas +import pandas as pd import mercantile from shapely.geometry import box from .. import tile +from google.protobuf.json_format import MessageToJson # Container for dataframes of SharedStreets geometries and intersections. Frames = collections.namedtuple('Frames', ['intersections', 'geometries']) @@ -20,7 +23,8 @@ def __init__(self, properties, type, coordinates): 'geometry': {'type': type, 'coordinates': coordinates}, } -def _make_frames(intersections, geometries, bounds=None): +def _make_frames(intersections, geometries, + metadata=None, bounds=None): ''' Return a Frames instance for lists of SharedStreets entities. ''' ifeatures = [ @@ -39,10 +43,11 @@ def _make_frames(intersections, geometries, bounds=None): 'toIntersectionId': item.toIntersectionId, 'forwardReferenceId': item.forwardReferenceId, 'backReferenceId': item.backReferenceId, + 'lonlats': str(item.lonlats), }, 'LineString', zip(item.lonlats[0::2], item.lonlats[1::2])) for item in geometries ] - + def clip_bbox(gdf): if bounds is None: return gdf @@ -51,14 +56,19 @@ def clip_bbox(gdf): def make_frame(features): gdf = geopandas.GeoDataFrame.from_features(features, crs={'init': 'epsg:4326'}) - return gdf.set_index('id', drop=False, verify_integrity=True) + return gdf.set_index('id', drop=True, verify_integrity=True) intersectionsdf = clip_bbox(make_frame(ifeatures)) geometriesdf = clip_bbox(make_frame(gfeatures)) + if metadata: + metadatadf = pd.DataFrame([{'id': k, 'metadata': json.loads(MessageToJson(v).replace('\n', '').replace(' ', ''))} + for k, v in metadata.items()]).set_index('id', verify_integrity=True) + geometriesdf = geometriesdf.merge(metadatadf, left_index=True, right_index=True, how='left') + return Frames(intersectionsdf, geometriesdf) -def get_bbox(minlon, minlat, maxlon, maxlat, data_url_template=None): +def get_bbox(minlon, minlat, maxlon, maxlat, data_url_template=None, include_metadata=False): ''' Get a single Frames instance of SharedStreets entities in an area. ''' bounds = (minlon, minlat, maxlon, maxlat) @@ -72,8 +82,13 @@ def get_bbox(minlon, minlat, maxlon, maxlat, data_url_template=None): all_geometries = functools.reduce(lambda d, t: dict(d, **t.geometries), tiles, {}) all_intersections = functools.reduce(lambda d, t: dict(d, **t.intersections), tiles, {}) + all_metadata = None + if include_metadata: + all_metadata = functools.reduce(lambda d, t: dict(d, **t.metadata), tiles, {}) + - return _make_frames(all_intersections.values(), all_geometries.values(), bounds) + return _make_frames(all_intersections.values(), all_geometries.values(), metadata=all_metadata, + bounds=bounds) def get_tile(*args, **kwargs): ''' Get a single Frames instance for a tile of SharedStreets entities.