diff --git a/src/linkahead/utils/get_entity.py b/src/linkahead/utils/get_entity.py index ce74089784525b1a38d9c5c5bbf3193a9ba94575..8d280521364a2597419577ef58897aa2d98ccdaa 100644 --- a/src/linkahead/utils/get_entity.py +++ b/src/linkahead/utils/get_entity.py @@ -27,23 +27,27 @@ from ..common.models import Entity, execute_query from .escape import escape_squoted_text -def get_entity_by_name(name: str) -> Entity: +def get_entity_by_name(name: str, role: Optional[str] = None) -> Entity: """Return the result of a unique query that uses the name to find the correct entity. Submits the query "FIND ENTITY WITH name='{name}'". """ name = escape_squoted_text(name) + if role is None: + role = "ENTITY" # type hint can be ignored, it's a unique query, so never Container or int - return execute_query(f"FIND ENTITY WITH name='{name}'", unique=True) # type: ignore + return execute_query(f"FIND {role} WITH name='{name}'", unique=True) # type: ignore -def get_entity_by_id(eid: Union[str, int]) -> Entity: +def get_entity_by_id(eid: Union[str, int], role: Optional[str] = None) -> Entity: """Return the result of a unique query that uses the id to find the correct entity. Submits the query "FIND ENTITY WITH id='{eid}'". """ + if role is None: + role = "ENTITY" # type hint can be ignored, it's a unique query - return execute_query(f"FIND ENTITY WITH id='{eid}'", unique=True) # type: ignore + return execute_query(f"FIND {role} WITH id='{eid}'", unique=True) # type: ignore def get_entity_by_path(path: str) -> Entity: