def _test_suggestions(self, **kwargs):
raise NotImplementedError()
""" Test basic suggested entity searches """
suggestedEntities = MongoSuggestedEntities()
suggested = suggestedEntities.getSuggestedEntities(**kwargs)
utils.log()
utils.log("-" * 80)
utils.log("TEST %s" % pformat(kwargs))
for section in suggested:
utils.log()
utils.log("SECTION %s:" % section["name"])
for i in xrange(len(section["entities"])):
entity = section["entities"][i]
types = list(entity.types)[0] if len(entity.types) == 1 else entity.types
utils.log(" %d) %s (%s) (id=%s)" % (i + 1, entity.title, types, entity.search_id))
utils.log("-" * 80)
utils.log()
self.assertTrue(len(suggested) > 0)
count = 0
seen = defaultdict(set)
for section in suggested:
title = section["name"]
entities = section["entities"]
# ensure the entities and title are valid for this section
self.assertIsInstance(title, basestring)
self.assertTrue(len(title) > 0)
self.assertTrue(len(entities) > 0)
count += len(entities)
for entity in entities:
self.assertIsInstance(entity, BasicEntity)
self.assertTrue(entity.search_id is not None)
self.assertTrue(len(entity.search_id) > 0)
self.assertTrue(entity.title is not None)
self.assertTrue(len(entity.title) > 0)
# ensure the sections contain no obvious duplicates
entities2 = Entity.fast_id_dedupe(entities, seen)
self.assertEqual(len(entities), len(entities2))
# if we set a limit on the number of results returned, ensure that
# that limit was satisfied.
limit = kwargs.get("limit", None)
if limit is not None:
self.assertTrue(count <= limit)
def searchEntities(self,
query,
coords = None,
full = True,
local = False,
category = None,
subcategory = None,
offset = 0,
limit = 10):
results = []
kinds, types = _convertCategorySubcategory(category, subcategory)
try:
if coords.lat is not None and coords.lng is not None:
coords = (coords.lat, coords.lng)
else:
coords = None
except Exception as e:
logs.warning("Exception (%s) - setting coords to None" % e)
coords = None
search = self.search(query,
coords = coords,
full = full,
local = local,
offset = offset,
limit = limit,
kinds = kinds,
types = types)
for item in search:
source = item[1].target.source
if source not in self._sources_map:
source = 'stamped'
entity = EntityProxyContainer().addProxy(item[1].target).buildEntity()
# Hack to make sure entity_id is set (since it's not a part of a group)
if source == 'stamped':
entity.entity_id = item[1].target.key
results.append(entity)
return Entity.fast_id_dedupe(results)
def __init__(self, name, types, maxQueueSize=None):
Greenlet.__init__(self)
self._name = name
self._id = EntitySources.getSourceID(name)
self._output = Queue(maxQueueSize)
self._started = False
self._maxQueueSize = maxQueueSize
self.subcategories = types
self.categories = set()
# validate the types
for t in types:
if not t in Constants.subcategories:
raise AttributeError("Source subcategory '%s' not supported" % t)
self.categories.add(Entity.mapSubcategoryToCategory(t))
def buildEntity(self):
if not self.__proxies:
raise Exception('No proxies added so far, cannot build entity')
primaryProxy = self.__proxies[0]
entity = Entity.buildEntity(kind=primaryProxy.kind)
entity.kind = primaryProxy.kind
entity.types = primaryProxy.types
if entity.isType('book'):
entity.title = self.__chooseBookTitle()
else:
entity.title = self.__chooseBestTitle()
sourceContainer = BasicSourceContainer()
for proxy in self.__proxies:
sourceContainer.addSource(EntityProxySource(proxy))
sourceContainer.enrichEntity(entity, {})
return entity
def getSuggestedEntities(self, userId=None, coords=None, category=None, subcategory=None, limit=None):
"""
Returns a list of suggested entities (separated into sections), restricted
to the given category / subcategory, and possibly personalized with respect
to the given userId.
Each section is a dict, with at least two required attributes, name, and entities.
"""
if category is not None:
category = category.lower().strip()
if subcategory is not None:
subcategory = subcategory.lower().strip()
if category is None and subcategory is None:
# TODO: what is the expected behavior here?
raise NotImplementedError
else:
if category is None:
try:
category = Constants.subcategories[subcategory]
except KeyError:
raise StampedInvalidSubcategoryError("invalid subcategory '%s'" % subcategory)
if category != 'place':
coords = None
if coords is not None:
# Convert from Coordinates object
coords = (coords.lat, coords.lng)
suggested = self._getGlobalSuggestedEntities(coords, category, subcategory)
num_sections = len(suggested)
if num_sections > 0:
out_suggested = []
seen = defaultdict(set)
# filter suggested entities already stamped by this user
if userId is not None:
stampIds = self._collection_collection.getUserStampIds(userId)
stamps = self._stamp_collection.getStamps(stampIds, limit=1000, sort='modified')
for stamp in stamps:
seen['entity_id'].add(stamp.entity.entity_id)
def _get_section_limit(i):
if limit:
return (limit / num_sections) + (1 if i + 1 <= (limit % num_sections) else 0)
return None
# process each section, removing obvious duplicates and enforcing per section limits
for i in xrange(num_sections):
section_limit = _get_section_limit(i)
section = suggested[i]
entities = Entity.fast_id_dedupe(section['entities'], seen)
entities = entities[:section_limit]
if len(entities) > 0:
section['entities'] = entities
out_suggested.append(section)
suggested = out_suggested
return suggested
