From 5968dc31ddb55ce2ba47fb2fee50641f95fee498 Mon Sep 17 00:00:00 2001
From: Matthew <git@matthewgamble.net>
Date: Sun, 19 Feb 2017 00:58:23 +1100
Subject: [PATCH] Support for a nicer structure for ID maps (#12)

This adds support for a more expressive (albeit more verbose) IdentMap
structure. It makes it easier to annotate the structure with additional
data (such as names to associate with the IDs), to assist with
maintaining the IdentMap.
---
 docs/source/identmap.rst        | 35 ++++++++++++++++++++++++++-------
 tsstats/__main__.py             |  3 +++
 tsstats/tests/test_ident_map.py | 23 ++++++++++++++++++++++
 tsstats/utils.py                | 19 ++++++++++++++++++
 4 files changed, 73 insertions(+), 7 deletions(-)

diff --git a/docs/source/identmap.rst b/docs/source/identmap.rst
index e69da79..578dbfc 100644
--- a/docs/source/identmap.rst
+++ b/docs/source/identmap.rst
@@ -2,19 +2,40 @@ IdentMap
 ********
 An IdentMap is used to map multiple (U)ID's of one client to one client.
 This can be useful, if a user creates multiple identities and you want to summarize all actions from all identities.
+To pass an IdentMap to TeamspeakStats, create your IdentMap as shown above and pass it to the cli::
+
+  tsstats --idmap <path to idmap.json>
 
 TeamspeakStats' IdentMap-file is saved in json-format::
 
+  [
+    {
+      "primary_id": "1",
+      "alternate_ids": ["2", "3", "4"]
+    }
+  ]
+
+If you would pass this IdentMap to TeamspeakStats and your log would contain entries for clients with id's 1, 2, 3 and 4,
+your output will just show data for one client (1).
+
+The format is flexible enough to support other arbitrary data to assist you in maintaining your IdentMap::
+
+  [
+    {
+      "name": "Friend 1",
+      "primary_id": "1",
+      "alternate_ids": ["2", "3", "4"]
+    }
+  ]
+
+The parser will ignore all nodes other than the "primary_id" and "alternate_ids" nodes, allowing you to use them for whatever you want.
+
+The original IdentMap format is still supported:
+
   {
     '2': '1',
     '3': '1',
     '4': '1'
   }
 
-
-If you would pass this IdentMap to TeamspeakStats and your log would contain entries for clients with id's 1, 2, 3 and 4,
-your output will just show data for one client (1).
-
-To pass an IdentMap to TeamspeakStats, create your IdentMap as shown above and pass it to the cli::
-
-  tsstats --idmap <path to idmap.json>
+While it is less expressive, it is also less verbose.
diff --git a/tsstats/__main__.py b/tsstats/__main__.py
index 549a5ff..5c40a8c 100644
--- a/tsstats/__main__.py
+++ b/tsstats/__main__.py
@@ -9,6 +9,7 @@ from tsstats import config
 from tsstats.exceptions import InvalidConfiguration
 from tsstats.log import parse_logs
 from tsstats.template import render_servers
+from tsstats.utils import transform_pretty_identmap
 
 logger = logging.getLogger('tsstats')
 
@@ -77,6 +78,8 @@ def main(configuration):
         identmap = json.load(open(idmap))
     else:
         identmap = None
+    if isinstance(identmap, list):
+        identmap = transform_pretty_identmap(identmap)
 
     log = configuration.get('General', 'log')
     if not log:
diff --git a/tsstats/tests/test_ident_map.py b/tsstats/tests/test_ident_map.py
index ce0c8c4..26e20a3 100644
--- a/tsstats/tests/test_ident_map.py
+++ b/tsstats/tests/test_ident_map.py
@@ -1,6 +1,7 @@
 import pytest
 
 from tsstats.client import Client, Clients
+from tsstats.utils import transform_pretty_identmap
 
 
 @pytest.fixture(scope='module')
@@ -24,3 +25,25 @@ def test_ident_map(identmap_clients):
     assert clients['5'] == cl
     assert clients['UID1'] == uidcl
     assert clients['UID5'] == uidcl
+
+
+@pytest.mark.parametrize('test_input,expected', [
+    (
+        [
+            {'primary_id': '1', 'alternate_ids': ['3', '6']},
+            {'primary_id': '4', 'alternate_ids': ['9', '42', '23']}
+        ],
+        (('3', '1'), ('6', '1'), ('9', '4'), ('42', '4'), ('23', '4'))
+    ),
+    (
+        [
+            {'name': 'Friend 1', 'primary_id': '2', 'alternate_ids': ['4']},
+            {'name': 'Friend 3', 'primary_id': '8', 'alternate_ids': ['9', '14']}
+        ],
+        (('4', '2'), ('9', '8'), ('14', '8'))
+    )
+])
+def test_transform_pretty_identmap(test_input, expected):
+    transformed_identmap = transform_pretty_identmap(test_input)
+    for alternate, primary in expected:
+        assert transformed_identmap[alternate] == primary
diff --git a/tsstats/utils.py b/tsstats/utils.py
index d11f233..12b8459 100644
--- a/tsstats/utils.py
+++ b/tsstats/utils.py
@@ -78,3 +78,22 @@ def tz_aware_datime(datetime, timezone=UTC()):
     :type timezone: datetime.timezone
     '''
     return datetime.replace(tzinfo=timezone)
+
+
+def transform_pretty_identmap(pretty_identmap):
+    '''
+    Transforms a list of client ID mappings from a more descriptive format
+    to the traditional format of alternative IDs to actual ID.
+
+    :param pretty_identmap: ID mapping in "nice" form
+    :type pretty_identmap: list
+
+    :return: ID mapping in simple key/value pairs
+    :rtype: dict
+    '''
+
+    final_identmap = {}
+    for mapping in pretty_identmap:
+        for alt_id in mapping['alternate_ids']:
+            final_identmap[alt_id] = mapping['primary_id']
+    return final_identmap