Keep original query name for ambiguous result with allow_ambiguous flag

python-package/lets_plot/geo_data/geocodes.py

@@ -43,10 +43,17 @@ class Resolution(enum.Enum):
  world_low = 1
 
 
-def select_request(query: RegionQuery, answer: Answer, feature: GeocodedFeature) -> str:
- # exploding answers (features count > 1) don't have exact request (like us-48, it can't be a proper
- # request for 48 features/states) and so feature name should be used as request.
- return query.request if len(answer.features) <= 1 else feature.name
+def select_request_string(request: Optional[str], name: str) -> str:
+ if request is None:
+ return name
+
+ if len(request) == 0:
+ return name
+
+ if 'us-48' == request.lower():
+ return name
+
+ return request
 
 
 def zip_answers(queries: List, answers: List):
@@ -64,9 +71,9 @@ def __init__(self):
  self._state: List[Optional[str]] = []
  self._country: List[Optional[str]] = []
 
- def append_row(self, request: str, found_name: str, query: Optional[RegionQuery]):
- self._request.append(request)
- self._found_name.append(found_name)
+ def append_row(self, query: RegionQuery, feature: GeocodedFeature):
+ self._request.append(select_request_string(query.request, feature.name))
+ self._found_name.append(feature.name)
 
  if query is None:
  self._county.append(MapRegion.name_or_none(None))
@@ -138,7 +145,7 @@ def to_map_regions(self) -> List[MapRegion]:
  regions: List[MapRegion] = []
  for answer, query in zip_answers(self._answers, self._queries):
  for feature in answer.features:
- regions.append(MapRegion.place(feature.id, select_request(query, answer, feature), self._level_kind))
+ regions.append(MapRegion.place(feature.id, select_request_string(query.request, feature.name), self._level_kind))
  return regions
 
  def as_list(self) -> List['Geocodes']:
@@ -286,7 +293,7 @@ def to_data_frame(self) -> DataFrame:
  # for us-48 queries doesnt' count
  for query, answer in zip_answers(self._queries, self._answers):
  for feature in answer.features:
- places.append_row(select_request(query, answer, feature), feature.name, query)
+ places.append_row(query, feature)
 
  data = {**data, **places.build_dict()}
 

python-package/lets_plot/geo_data/gis/response.py

@@ -38,11 +38,11 @@ def __init__(self, geometry: Union[Multipolygon, Polygon, GeoPoint]):
 class GeocodedFeature:
  def __init__(self,
  id: str, name: str,
- highlights: Optional[List[str]],
- boundary: Optional[Boundary],
- centroid: Optional[GeoPoint],
- limit: Optional[GeoRect],
- position: Optional[GeoRect]):
+ highlights: Optional[List[str]]=None,
+ boundary: Optional[Boundary]=None,
+ centroid: Optional[GeoPoint]=None,
+ limit: Optional[GeoRect]=None,
+ position: Optional[GeoRect]=None):
  assert_type(id, str)
  assert_type(name, str)
  assert_optional_list_type(highlights, str)

python-package/lets_plot/geo_data/to_geo_data_frame.py

@@ -5,9 +5,9 @@
 from pandas import DataFrame
 from shapely.geometry import box
 
-from lets_plot.geo_data import PlacesDataFrameBuilder, zip_answers, select_request, DF_REQUEST, DF_FOUND_NAME, abstractmethod
-from lets_plot.geo_data.gis.response import Answer, GeocodedFeature, GeoRect, Boundary, Multipolygon, Polygon, GeoPoint
+from lets_plot.geo_data import PlacesDataFrameBuilder, zip_answers, select_request_string, abstractmethod
 from lets_plot.geo_data.gis.request import RegionQuery
+from lets_plot.geo_data.gis.response import Answer, GeocodedFeature, GeoRect, Boundary, Multipolygon, Polygon, GeoPoint
 
 ShapelyPoint = shapely.geometry.Point
 ShapelyLinearRing = shapely.geometry.LinearRing
@@ -48,7 +48,7 @@ def to_data_frame(self, answers: List[Answer], queries: List[RegionQuery]) -> Da
  for feature in answer.features:
  rects: List[GeoRect] = self._read_rect(feature)
  for rect in rects:
- places.append_row(request=select_request(query, answer, feature), found_name=feature.name, query=query)
+ places.append_row(query, feature)
  self._lonmin.append(rect.min_lon)
  self._latmin.append(rect.min_lat)
  self._lonmax.append(rect.max_lon)
@@ -85,7 +85,7 @@ def to_data_frame(self, answers: List[Answer], queries: List[RegionQuery]) -> Da
 
  for query, answer in zip_answers(queries, answers):
  for feature in answer.features:
- places.append_row(request=select_request(query, answer, feature), found_name=feature.name, query=query)
+ places.append_row(query, feature)
  self._lons.append(feature.centroid.lon)
  self._lats.append(feature.centroid.lat)
 
@@ -103,7 +103,7 @@ def to_data_frame(self, answers: List[Answer], queries: List[RegionQuery]) -> Da
  geometry = []
  for query, answer in zip_answers(queries, answers):
  for feature in answer.features:
- places.append_row(request=select_request(query, answer, feature), found_name=feature.name, query=query)
+ places.append_row(query, feature)
  geometry.append(self._geo_parse_geometry(feature.boundary))
 
  return _create_geo_data_frame(places.build_dict(), geometry=geometry)

python-package/test/geo_data/test_geocoder.py

@@ -10,9 +10,9 @@
 from lets_plot.geo_data.gis.geometry import GeoRect, GeoPoint
 from lets_plot.geo_data.gis.request import MapRegion, AmbiguityResolver, GeocodingRequest, LevelKind, RegionQuery, \
  IgnoringStrategyKind
-from lets_plot.geo_data.geocoder import geocode_countries, geocode, NamesGeocoder, geocode_cities
+from lets_plot.geo_data.geocoder import geocode_countries, geocode, NamesGeocoder, geocode_cities, geocode_states
 from lets_plot.geo_data.geocodes import Geocodes
-from .geo_data import make_answer
+from .geo_data import make_answer, assert_row
 from .request_assertion import GeocodingRequestAssertion, QueryMatcher, ScopeMatcher, ValueMatcher, eq, empty, \
  eq_map_region_with_name, eq_map_region_with_id
 
@@ -153,11 +153,7 @@ def test_where_closets_to_point():
  GeocodedFeature(
  id='foo_id',
  name='foo',
- highlights=None,
- boundary=None,
- centroid=GeoPoint(1, 2),
- limit=None,
- position=None
+ centroid=GeoPoint(1, 2)
  )
  ])
  ]
@@ -192,6 +188,51 @@ def test_select_all_query_with_empty_result_should_return_empty_dataframe():
  assert 0 == len(centroids)
 
 
+@mock.patch.object(GeocodingService, 'do_request', lambda self, reqest: SuccessResponse(
+ message='',
+ level=LevelKind.state,
+ answers=[
+ Answer(
+ features=[
+ GeocodedFeature(id='foo_id', name='foo'),
+ GeocodedFeature(id='bar_id', name='bar'),
+ GeocodedFeature(id='baz_id', name='baz'),
+ ]
+ )
+ ]
+))
+def test_for_us48_request_should_contain_feature_name():
+ states = geocode_states('us-48')
+
+ assert_row(
+ states.get_geocodes(),
+ names=['foo', 'bar', 'baz'],
+ found_name=['foo', 'bar', 'baz']
+ )
+
+
+@mock.patch.object(GeocodingService, 'do_request', lambda self, reqest: SuccessResponse(
+ message='',
+ level=LevelKind.city,
+ answers=[
+ Answer(
+ features=[
+ GeocodedFeature(id='foo1_id', name='Foo'),
+ GeocodedFeature(id='foo2_id', name='Foo'),
+ GeocodedFeature(id='foo3_id', name='Fooo'),
+ ]
+ )
+ ]
+))
+def test_allow_ambiguous_result_should_keep_request():
+ cities = geocode_cities('foo')
+
+ assert_row(
+ cities.get_geocodes(),
+ names=['foo', 'foo', 'foo'],
+ found_name=['Foo', 'Foo', 'Fooo']
+ )
+
 
 def test_allow_ambiguous():
  request = geocode(names='foo')\

python-package/test/geo_data/test_integration_new_api.py

@@ -278,4 +278,11 @@ def test_counties():
 
  geocoded_counties = geodata.geocode_counties(counties).states(states).scope('usa').get_boundaries('country').request
 
- assert counties == geocoded_counties.tolist()
+ assert counties == geocoded_counties.tolist()
+
+
+def test_request_in_ambiguous_df():
+ us48 = geodata.geocode_states('us-48').get_geocodes()
+ warwick = geodata.geocode_cities('warwick').allow_ambiguous().get_geocodes()
+
+ assert_row(warwick, names='warwick', found_name='Warwick')