ENH: allow to provide the role when using get_entity

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: