Add function for writing domain graphs to CSV files (#132)

Author: jasper-elimity

src/elimity_insights_client/csv.py

@@ -0,0 +1,96 @@
+"""Utilities for writing domain graph schemas to CSV files."""
+
+from collections.abc import Iterable
+from csv import writer
+from json import dumps, loads
+
+from elimity_insights_client._decode_domain_graph_schema import (
+    decode_domain_graph_schema,
+)
+from elimity_insights_client._domain_graph_schema import DomainGraphSchema
+from elimity_insights_client._elimity_insights_client import (
+    BooleanValue,
+    DateTimeValue,
+    DateValue,
+    DomainGraph,
+    Entity,
+    NumberValue,
+    Relationship,
+    StringValue,
+    Value,
+)
+
+
+def write_domain_graph(filename: str, graph: DomainGraph, schema_json: str) -> None:
+    """Serialize the given domain graph to an importable CSV file at the given path."""
+    schema_dict = loads(schema_json)
+    schema = decode_domain_graph_schema(schema_dict)
+    with open(filename, "w", newline="") as file:
+        wri = writer(file)
+        rows = _rows(graph, schema)
+        wri.writerows(rows)
+
+
+def _rows(graph: DomainGraph, schema: DomainGraphSchema) -> Iterable[Iterable[str]]:
+    yield _headers(schema)
+    for entity in graph.entities:
+        yield _entity_cells(entity, schema)
+    for relationship in graph.relationships:
+        yield _relationship_cells(relationship, schema)
+
+
+def _headers(schema: DomainGraphSchema) -> Iterable[str]:
+    for entity_type in schema.entity_types:
+        id = entity_type.id
+        yield from [id + ": id", id + ": name"]
+    for attribute_type in schema.attribute_types:
+        yield f"{attribute_type.entity_type}: {attribute_type.id}"
+
+
+def _entity_cells(entity: Entity, schema: DomainGraphSchema) -> Iterable[str]:
+    type = entity.type
+    for entity_type in schema.entity_types:
+        yield from [entity.id, entity.name] if entity_type.id == type else ["", ""]
+    values = {
+        assignment.attribute_type_id: assignment.value
+        for assignment in entity.attribute_assignments
+    }
+    for attribute_type in schema.attribute_types:
+        value = values.get(attribute_type.id)
+        yield "" if attribute_type.entity_type != type or value is None else _cell(
+            value
+        )
+
+
+def _relationship_cells(
+    relationship: Relationship, schema: DomainGraphSchema
+) -> Iterable[str]:
+    for type in schema.entity_types:
+        dict = {
+            relationship.from_entity_type: relationship.from_entity_id,
+            relationship.to_entity_type: relationship.to_entity_id,
+        }
+        yield dict.get(type.id, "")
+        yield ""
+    yield from [""] * len(schema.attribute_types)
+
+
+def _cell(value: Value) -> str:
+    if isinstance(value, BooleanValue):
+        return "true" if value.value else "false"
+
+    elif isinstance(value, DateValue):
+        return f"{value.year:04}-{value.month:02}-{value.day:02}"
+
+    elif isinstance(value, DateTimeValue):
+        val = value.value
+        return f"{val.year:04}-{val.month:02}-{val.day:02} {val.hour:02}:{val.minute:02}:{val.second:02}.0"
+
+    elif isinstance(value, NumberValue):
+        return dumps(value.value)
+
+    elif isinstance(value, StringValue):
+        return value.value
+
+    else:
+        return f"{value.hour:02}:{value.minute:02}:{value.second:02}.0"

tests/graph.csv

@@ -0,0 +1,15 @@
+user: id,user: name,group: id,group: name,computer: id,computer: name,user: accountExpires,user: accountType,user: badPwdCount,user: cn,user: company,user: countryCode,user: department,user: departmentNumber,user: description,user: disabled,user: distinguishedName,user: displayName,user: division,user: employeeID,user: givenName,user: initials,user: isDeleted,user: l,user: lastLogon,user: lockedOut,user: logonCount,user: mail,user: manager,user: objectClass,user: objectSid,user: passwordExpired,user: passwdNotreqd,user: pwdLastSet,user: sAMAccountName,user: telephoneNumber,user: userPrincipalName,user: whenChanged,user: whenCreated,group: isDeleted,group: description,group: distinguishedName,group: displayName,group: lastLogon,group: logonCount,group: managedBy,group: objectClass,group: objectSid,group: sAMAccountName,group: whenChanged,group: whenCreated,computer: accountExpires,computer: accountType,computer: badPwdCount,computer: cn,computer: description,computer: disabled,computer: distinguishedName,computer: displayName,computer: isDeleted,computer: lockedOut,computer: objectClass,computer: objectSid,computer: operatingSystem,computer: operatingSystemPack,computer: operatingSystemVersion,computer: pwdLastSet,computer: sAMAccountName,computer: whenChanged,computer: whenCreated
+"cn=Amy Wong+sn=Kroker,ou=people,dc=planetexpress,dc=com",Amy Wong,,,,,,,,Amy Wong,,,,,Human,,,,,,Amy,,,,,,,amy@planetexpress.com,,top/user/organizationalPerson/inetOrgPerson,,,,,,,,2021-03-11 13:28:53.0,2021-03-11 13:28:53.0,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
+"cn=Bender Bending Rodríguez,ou=people,dc=planetexpress,dc=com",Bender Bending Rodríguez,,,,,,,,Bender Bending Rodríguez,,,,,Robot,,,Bender,,,Bender,,,,,,,bender@planetexpress.com,,inetOrgPerson/organizationalPerson/user/top,,,,,,,,2021-03-11 13:28:53.0,2021-03-11 13:28:53.0,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
+"cn=Philip J. Fry,ou=people,dc=planetexpress,dc=com",Philip J. Fry,,,,,,,,Philip J. Fry,,,,,Human,,,Fry,,,Philip,,,,,,,fry@planetexpress.com,,inetOrgPerson/organizationalPerson/user/top,,,,,,,,2021-03-11 13:28:53.0,2021-03-11 13:28:53.0,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
+"cn=Hermes Conrad,ou=people,dc=planetexpress,dc=com",Hermes Conrad,,,,,,,,Hermes Conrad,,,,,Human,,,,,,Hermes,,,,,,,hermes@planetexpress.com,,top/user/organizationalPerson/inetOrgPerson,,,,,,,,2021-03-11 13:28:53.0,2021-03-11 13:28:53.0,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
+"cn=Turanga Leela,ou=people,dc=planetexpress,dc=com",Turanga Leela,,,,,,,,Turanga Leela,,,,,Mutant,,,,,,Leela,,,,,,,leela@planetexpress.com,,inetOrgPerson/organizationalPerson/user/top,,,,,,,,2021-03-11 13:28:53.0,2021-03-11 13:28:53.0,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
+"cn=Hubert J. Farnsworth,ou=people,dc=planetexpress,dc=com",Hubert J. Farnsworth,,,,,,,,Hubert J. Farnsworth,,,,,Human,,,Professor Farnsworth,,,Hubert,,,,,,,professor@planetexpress.com/hubert@planetexpress.com,,inetOrgPerson/organizationalPerson/user/top,,,,,,,,2021-03-11 13:28:53.0,2021-03-11 13:28:53.0,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
+"cn=John A. Zoidberg,ou=people,dc=planetexpress,dc=com",John A. Zoidberg,,,,,,,,John A. Zoidberg,,,,,Decapodian,,,Zoidberg,,,John,,,,,,,zoidberg@planetexpress.com,,top/user/organizationalPerson/inetOrgPerson,,,,,,,,2021-03-11 13:28:53.0,2021-03-11 13:28:53.0,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
+,,"cn=admin_staff,ou=people,dc=planetexpress,dc=com",admin_staff,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,group/top,,,2021-03-11 13:28:53.0,2021-03-11 13:28:53.0,,,,,,,,,,,,,,,,,,,
+,,"cn=ship_crew,ou=people,dc=planetexpress,dc=com",ship_crew,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,group/top,,,2021-03-11 13:28:53.0,2021-03-11 13:28:53.0,,,,,,,,,,,,,,,,,,,
+"cn=Hermes Conrad,ou=people,dc=planetexpress,dc=com",,"cn=admin_staff,ou=people,dc=planetexpress,dc=com",,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
+"cn=Hubert J. Farnsworth,ou=people,dc=planetexpress,dc=com",,"cn=admin_staff,ou=people,dc=planetexpress,dc=com",,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
+"cn=Bender Bending Rodríguez,ou=people,dc=planetexpress,dc=com",,"cn=ship_crew,ou=people,dc=planetexpress,dc=com",,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
+"cn=Turanga Leela,ou=people,dc=planetexpress,dc=com",,"cn=ship_crew,ou=people,dc=planetexpress,dc=com",,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,
+"cn=Philip J. Fry,ou=people,dc=planetexpress,dc=com",,"cn=ship_crew,ou=people,dc=planetexpress,dc=com",,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,,