Add support for cross-source links to entities function (#128)

Author: jasper-elimity

src/elimity_insights_client/api/entities/__init__.py

@@ -1,4 +1,4 @@
 from elimity_insights_client.api.entities._entities import entities
-from elimity_insights_client.api.entities._entity import Entity, Link
+from elimity_insights_client.api.entities._entity import Entity, EntityType, Link
 
-__all__ = ["Entity", "Link", "entities"]
+__all__ = ["Entity", "EntityType", "Link", "entities"]

src/elimity_insights_client/api/entities/_entities.py

@@ -3,14 +3,14 @@
 from elimity_insights_client.api._api import Config
 from elimity_insights_client.api._api import query as api_query
 from elimity_insights_client.api._api import sources
-from elimity_insights_client.api.entities._entity import Entity
+from elimity_insights_client.api.entities._entity import Entity, EntityType
 from elimity_insights_client.api.entities._parse_query_results_page import (
     parse_query_results_page,
 )
 from elimity_insights_client.api.entities._query import query
 
 
-def entities(config: Config, entity_type: str, source_id: int) -> List[Entity]:
+def entities(config: Config, entity_type: EntityType) -> List[Entity]:
     """
     List all entities of the given entity type from the given source.
 
@@ -19,10 +19,8 @@ def entities(config: Config, entity_type: str, source_id: int) -> List[Entity]:
     """
 
     sos = sources(config)
-    schema = next(
-        source.domain_graph_schema for source in sos if source.id == source_id
-    )
-    que = query(entity_type, schema, source_id)
+    schemas = {source.id: source.domain_graph_schema for source in sos}
+    que = query(entity_type, schemas)
     queries = [que]
     (page,) = api_query(config, queries)
-    return parse_query_results_page(entity_type, page, schema)
+    return parse_query_results_page(entity_type, page, schemas)

src/elimity_insights_client/api/entities/_entity.py

@@ -10,10 +10,18 @@ class Entity:
 
     attribute_assignments: Dict[str, Value]
     id: str
-    links: Dict[str, List["Link"]]
+    links: Dict["EntityType", List["Link"]]
     name: str
 
 
+@dataclass(frozen=True)
+class EntityType:
+    """Type of entities for a given source, represents a unique key in dictionaries of links."""
+
+    id: str
+    source_id: int
+
+
 @dataclass
 class Link:
     """Linked entity with assignments for all attributes."""

src/elimity_insights_client/api/entities/_parse_query_results_page.py

@@ -1,5 +1,5 @@
 from dataclasses import dataclass
-from typing import List
+from typing import Dict, List
 
 from more_itertools import chunked
 
@@ -8,7 +8,7 @@
     DomainGraphSchema,
 )
 from elimity_insights_client._util import map_list
-from elimity_insights_client.api.entities._entity import Entity, Link
+from elimity_insights_client.api.entities._entity import Entity, EntityType, Link
 from elimity_insights_client.api.entities._schema import (
     attribute_types as schema_attribute_types,
 )
@@ -32,24 +32,26 @@ class _AttributeAssignment:
 
 @dataclass
 class _LinkGroup:
-    entity_type: str
+    entity_type: EntityType
     links: List[Link]
 
 
 def parse_query_results_page(
-    entity_type: str, page: QueryResultsPage, schema: DomainGraphSchema
+    entity_type: EntityType,
+    page: QueryResultsPage,
+    schemas: Dict[int, DomainGraphSchema],
 ) -> List[Entity]:
     def make_entity(result: QueryResult) -> Entity:
-        entity = _link(entity_type, schema, result)
+        entity = _link(entity_type, result, schemas)
 
-        def make_group(entity_type: str, page: QueryResultsPage) -> _LinkGroup:
+        def make_group(entity_type: EntityType, page: QueryResultsPage) -> _LinkGroup:
             def make_link(result: QueryResult) -> Link:
-                return _link(entity_type, schema, result)
+                return _link(entity_type, result, schemas)
 
             links = map_list(make_link, page.results)
             return _LinkGroup(entity_type, links)
 
-        link_entity_types = schema_link_entity_types(entity_type, schema)
+        link_entity_types = schema_link_entity_types(entity_type, schemas)
         group_iter = map(make_group, link_entity_types, result.link_pages)
         links = {group.entity_type: group.links for group in group_iter}
         return Entity(entity.attribute_assignments, entity.id, links, entity.name)
@@ -65,9 +67,11 @@ def _attribute_assignment(
     return _AttributeAssignment(assigned, type.id, value_inclusion)
 
 
-def _link(entity_type: str, schema: DomainGraphSchema, result: QueryResult) -> Link:
+def _link(
+    entity_type: EntityType, result: QueryResult, schemas: Dict[int, DomainGraphSchema]
+) -> Link:
     inclusions_iter = chunked(result.inclusions, 2)
-    attribute_types = schema_attribute_types(entity_type, schema)
+    attribute_types = schema_attribute_types(entity_type, schemas)
     assignment_iter = map(_attribute_assignment, inclusions_iter, attribute_types)
     assignments = {
         assignment.attribute_type: assignment.value

src/elimity_insights_client/api/entities/_query.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import Dict, List
 
 from more_itertools import interleave
 
@@ -8,6 +8,7 @@
     Type,
 )
 from elimity_insights_client._util import map_list
+from elimity_insights_client.api.entities._entity import EntityType
 from elimity_insights_client.api.entities._schema import (
     attribute_types as schema_attribute_types,
 )
@@ -47,17 +48,16 @@
 
 
 def query(
-    entity_type: str,
-    schema: DomainGraphSchema,
-    source_id: int,
+    entity_type: EntityType,
+    schemas: Dict[int, DomainGraphSchema],
 ) -> Query:
-    def make_link_query(entity_type: str) -> Query:
+    def make_link_query(entity_type: EntityType) -> Query:
         link_queries: List[Query] = []
-        return _query(entity_type, link_queries, schema, source_id)
+        return _query(entity_type, link_queries, schemas)
 
-    link_entity_types = schema_link_entity_types(entity_type, schema)
+    link_entity_types = schema_link_entity_types(entity_type, schemas)
     link_queries = map_list(make_link_query, link_entity_types)
-    return _query(entity_type, link_queries, schema, source_id)
+    return _query(entity_type, link_queries, schemas)
 
 
 def _assigned_inclusion(type: AttributeType) -> AnyExpression:
@@ -87,37 +87,38 @@ def _attribute_inclusion(type: AttributeType) -> AnyExpression:
     return TimeAnyExpression(time_expr)
 
 
-def _inclusions(entity_type: str, schema: DomainGraphSchema) -> List[AnyExpression]:
-    attribute_types = schema_attribute_types(entity_type, schema)
+def _inclusions(
+    entity_type: EntityType, schemas: Dict[int, DomainGraphSchema]
+) -> List[AnyExpression]:
+    attribute_types = schema_attribute_types(entity_type, schemas)
     assigned_iter = map(_assigned_inclusion, attribute_types)
     attribute_iter = map(_attribute_inclusion, attribute_types)
     inclusion_iter = interleave(assigned_iter, attribute_iter)
     return list(inclusion_iter)
 
 
 def _query(
-    entity_type: str,
+    entity_type: EntityType,
     link_queries: List[Query],
-    schema: DomainGraphSchema,
-    source_id: int,
+    schemas: Dict[int, DomainGraphSchema],
 ) -> Query:
     condition = LiteralBooleanExpression(True)
     direct_link_group_by_queries: List[DirectLinkGroupByQuery] = []
     direct_link_queries: List[DirectLinkQuery] = []
-    inclusions = _inclusions(entity_type, schema)
+    inclusions = _inclusions(entity_type, schemas)
     link_group_by_queries: List[LinkGroupByQuery] = []
     orderings: List[Ordering] = []
     return Query(
         "",
         condition,
         direct_link_group_by_queries,
         direct_link_queries,
-        entity_type,
+        entity_type.id,
         inclusions,
         999999,
         link_group_by_queries,
         link_queries,
         0,
         orderings,
-        source_id,
+        entity_type.source_id,
     )

src/elimity_insights_client/api/entities/_schema.py

@@ -1,18 +1,32 @@
-from typing import List
+from itertools import starmap
+from typing import Dict, Iterable, List
 
 from elimity_insights_client._domain_graph_schema import (
     AttributeType,
     DomainGraphSchema,
 )
+from elimity_insights_client.api.entities._entity import EntityType
 
 
-def attribute_types(entity_type: str, schema: DomainGraphSchema) -> List[AttributeType]:
+def attribute_types(
+    entity_type: EntityType, schemas: Dict[int, DomainGraphSchema]
+) -> List[AttributeType]:
+    schema = schemas[entity_type.source_id]
     return [
         type
         for type in schema.attribute_types
-        if not type.archived and type.entity_type == entity_type
+        if not type.archived and type.entity_type == entity_type.id
     ]
 
 
-def link_entity_types(entity_type: str, schema: DomainGraphSchema) -> List[str]:
-    return [type.id for type in schema.entity_types if type.id != entity_type]
+def link_entity_types(
+    entity_type: EntityType, schemas: Dict[int, DomainGraphSchema]
+) -> List[EntityType]:
+    items = schemas.items()
+    iters = starmap(_entity_types, items)
+    return [typ for iter in iters for typ in iter if typ != entity_type]
+
+
+def _entity_types(source_id: int, schema: DomainGraphSchema) -> Iterable[EntityType]:
+    for type in schema.entity_types:
+        yield EntityType(type.id, source_id)

tests/elimity_insights_client/api/entities/_entities.py

@@ -1,17 +1,35 @@
-from typing import Type, TypeVar
+from typing import List, Tuple, Type, TypedDict, TypeVar
 
 from elimity_insights_client._decode_domain_graph_schema import (
     DomainGraphSchemaDict,
     decode_domain_graph_schema,
 )
+from elimity_insights_client._domain_graph_schema import DomainGraphSchema
+from elimity_insights_client.api.entities._entity import EntityType
 from tests.elimity_insights_client.api._json import decode_file as json_decode_file
 
 _T = TypeVar("_T")
 
+_SourceDict = TypedDict(
+    "_SourceDict",
+    {
+        "schema": DomainGraphSchemaDict,
+        "sourceId": int,
+    },
+)
+
 
 def json_decode_file_local(filename: str, type: Type[_T]) -> _T:
     return json_decode_file(filename, __package__, type)
 
 
-_json = json_decode_file_local("schema.json", DomainGraphSchemaDict)
-schema = decode_domain_graph_schema(_json)
+def _item(dict: _SourceDict) -> Tuple[int, DomainGraphSchema]:
+    schema = dict["schema"]
+    return dict["sourceId"], decode_domain_graph_schema(schema)
+
+
+_json = json_decode_file_local("sources.json", List[_SourceDict])
+_items = map(_item, _json)
+schemas = dict(_items)
+
+type = EntityType("foo", 42)

tests/elimity_insights_client/api/entities/query-results-page.json

@@ -45,6 +45,30 @@
         {
           "count": 0,
           "results": []
+        },
+        {
+          "count": 1,
+          "results": [
+            {
+              "entity": {
+                "active": false,
+                "id": "baz",
+                "name": "lorem"
+              },
+              "inclusions": [
+                {
+                  "type": "boolean",
+                  "value": "true"
+                },
+                {
+                  "type": "number",
+                  "value": "24"
+                }
+              ],
+              "linkGroupByPages": [],
+              "linkPages": []
+            }
+          ]
         }
       ]
     }

tests/elimity_insights_client/api/entities/query.json

@@ -77,6 +77,40 @@
       "offset": 0,
       "orderBy": [],
       "sourceId": 42
+    },
+    {
+      "alias": "",
+      "condition": {
+        "boolean": true,
+        "type": "literal"
+      },
+      "directLinkGroupByQueries": [],
+      "directLinkQueries": [],
+      "entityType": "bar",
+      "include": [
+        {
+          "expr": {
+            "attributeType": "ipsum",
+            "reference": "",
+            "type": "assigned"
+          },
+          "type": "boolean"
+        },
+        {
+          "expr": {
+            "attributeType": "ipsum",
+            "reference": "",
+            "type": "attribute"
+          },
+          "type": "number"
+        }
+      ],
+      "limit": 999999,
+      "linkGroupByQueries": [],
+      "linkQueries": [],
+      "offset": 0,
+      "orderBy": [],
+      "sourceId": 24
     }
   ],
   "offset": 0,