Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_dump_group_nested(query_factory):
"""Tests dumping a query with nested entity groups"""
query_text = "Order one large Tesora with medium cream and medium sugar"
query = query_factory.create_query(query_text)
entities = [
QueryEntity.from_query(query, Span(6, 8), entity_type="quantity"),
QueryEntity.from_query(query, Span(10, 14), entity_type="size"),
QueryEntity.from_query(query, Span(16, 21), entity_type="product"),
QueryEntity.from_query(query, Span(28, 33), entity_type="size"),
QueryEntity.from_query(query, Span(35, 39), entity_type="option"),
QueryEntity.from_query(query, Span(45, 50), entity_type="size"),
QueryEntity.from_query(query, Span(52, 56), entity_type="option"),
]
entities[4] = entities[4].with_children((entities[3],))
entities[6] = entities[6].with_children((entities[5],))
entities[2] = entities[2].with_children(
(entities[0], entities[1], entities[4], entities[6])
)
processed_query = ProcessedQuery(query, entities=entities)
markup_text = (
"Order [{one|quantity} {large|size} {Tesora|product} with [{medium|size} "
"{cream|option}|option] and [{medium|size} {sugar|option}|option]|product]"
)
def test_dump_group_nested(query_factory):
"""Tests dumping a query with nested entity groups"""
query_text = "Order one large Tesora with medium cream and medium sugar"
query = query_factory.create_query(query_text)
entities = [
QueryEntity.from_query(query, Span(6, 8), entity_type="quantity"),
QueryEntity.from_query(query, Span(10, 14), entity_type="size"),
QueryEntity.from_query(query, Span(16, 21), entity_type="product"),
QueryEntity.from_query(query, Span(28, 33), entity_type="size"),
QueryEntity.from_query(query, Span(35, 39), entity_type="option"),
QueryEntity.from_query(query, Span(45, 50), entity_type="size"),
QueryEntity.from_query(query, Span(52, 56), entity_type="option"),
]
entities[4] = entities[4].with_children((entities[3],))
entities[6] = entities[6].with_children((entities[5],))
entities[2] = entities[2].with_children(
(entities[0], entities[1], entities[4], entities[6])
)
processed_query = ProcessedQuery(query, entities=entities)
markup_text = (
"Order [{one|quantity} {large|size} {Tesora|product} with [{medium|size} "
def test_dump_group(query_factory):
"""Tests dumping a query with an entity group"""
query_text = "a large latte with nonfat milk please"
query = query_factory.create_query(query_text)
size = QueryEntity.from_query(query, Span(2, 6), entity_type="size")
option = QueryEntity.from_query(query, Span(19, 29), entity_type="option")
product = QueryEntity.from_query(
query, Span(8, 12), entity_type="product", children=(size, option)
)
processed_query = ProcessedQuery(query, entities=[size, product, option])
markup_text = (
"a [{large|size} {latte|product} with {nonfat milk|option}|product] please"
)
entity_text = "a {large|size} {latte|product} with {nonfat milk|option} please"
group_text = "a [large latte with nonfat milk|product] please"
assert markup.dump_query(processed_query) == markup_text
assert markup.dump_query(processed_query, no_group=True) == entity_text
assert markup.dump_query(processed_query, no_entity=True) == group_text
assert (
def test_dump_role(query_factory):
"""Tests dumping a basic query with an entity with a role"""
query_text = "What stores are open between 3 and 5"
query = query_factory.create_query(query_text)
entities = [
QueryEntity.from_query(
query, Span(29, 29), entity_type="sys_time", role="open_hours"
),
QueryEntity.from_query(
query, Span(35, 35), entity_type="sys_time", role="close_hours"
),
]
processed_query = ProcessedQuery(query, entities=entities)
markup_text = (
"What stores are open between {3|sys_time|open_hours} and "
"{5|sys_time|close_hours}"
)
entity_text = "What stores are open between {3|sys_time} and {5|sys_time}"
assert markup.dump_query(processed_query) == markup_text
assert markup.dump_query(processed_query, no_role=True) == entity_text
assert (
markup.dump_query(processed_query, no_role=True, no_entity=True) == query_text
)
def test_dump_group_with_role(query_factory):
"""Tests dumping a query with an entity group with role type"""
query_text = "a large latte with nonfat milk please"
query = query_factory.create_query(query_text)
size = QueryEntity.from_query(query, Span(2, 6), entity_type="size")
option = QueryEntity.from_query(
query, Span(19, 29), entity_type="option", role="beverage"
)
product = QueryEntity.from_query(
query,
Span(8, 12),
entity_type="dish-type",
role="beverage",
children=(size, option),
)
processed_query = ProcessedQuery(query, entities=[size, product, option])
markup_text = (
"a [{large|size} {latte|dish-type|beverage} with "
"{nonfat milk|option|beverage}|dish-type] please"
)
def test_dump_group_with_role(query_factory):
"""Tests dumping a query with an entity group with role type"""
query_text = "a large latte with nonfat milk please"
query = query_factory.create_query(query_text)
size = QueryEntity.from_query(query, Span(2, 6), entity_type="size")
option = QueryEntity.from_query(
query, Span(19, 29), entity_type="option", role="beverage"
)
product = QueryEntity.from_query(
query,
Span(8, 12),
entity_type="dish-type",
role="beverage",
children=(size, option),
)
processed_query = ProcessedQuery(query, entities=[size, product, option])
markup_text = (
"a [{large|size} {latte|dish-type|beverage} with "
"{nonfat milk|option|beverage}|dish-type] please"
)
entity_text = (
"a {large|size} {latte|dish-type|beverage} with "
"{nonfat milk|option|beverage} please"
)
def test_dump_group_with_role(query_factory):
"""Tests dumping a query with an entity group with role type"""
query_text = "a large latte with nonfat milk please"
query = query_factory.create_query(query_text)
size = QueryEntity.from_query(query, Span(2, 6), entity_type="size")
option = QueryEntity.from_query(
query, Span(19, 29), entity_type="option", role="beverage"
)
product = QueryEntity.from_query(
query,
Span(8, 12),
entity_type="dish-type",
role="beverage",
children=(size, option),
)
processed_query = ProcessedQuery(query, entities=[size, product, option])
markup_text = (
"a [{large|size} {latte|dish-type|beverage} with "
"{nonfat milk|option|beverage}|dish-type] please"
)
entity_text = (
Args:
query (Query): The query to construct the QueryEntity from
item (dict): The duckling item
offset (int, optional): The offset into the query that the item's
indexing begins
Returns:
QueryEntity: The query entity described by the duckling item or \
None if no item is present
"""
if item:
start = int(item["start"]) + offset
end = int(item["end"]) - 1 + offset
entity = _duckling_item_to_entity(item)
return QueryEntity.from_query(query, Span(start, end), entity=entity)
else:
return
def _append_entity(token_start, entity_type, tokens):
prefix = " ".join(normalized_tokens[:token_start])
# If there is a prefix, we have to add one for the whitespace
start = len(prefix) + 1 if len(prefix) else 0
end = start - 1 + len(" ".join(tokens))
norm_span = Span(start, end)
entity = QueryEntity.from_query(
query, normalized_span=norm_span, entity_type=entity_type
)
entities.append(entity)
logger.debug("Appended %s.", entity)
if ann.get("parent"):
parent = ann.get("parent")
if parent["ann_type"] == "entity":
children = parent.get("children", [])
children.append(
NestedEntity.from_query(
query,
span.shift(-parent["start"]),
entity=raw_entity,
parent_offset=parent["start"],
)
)
parent["children"] = children
if parent["ann_type"] == "group":
entity = QueryEntity.from_query(query, span, entity=raw_entity)
entities.append(entity)
if parent["type"] == ann["type"]:
# this is the head
parent["head"] = entity
else:
children = parent.get("children", [])
children.append(entity)
parent["children"] = children
else:
entities.append(QueryEntity.from_query(query, span, entity=raw_entity))