diff --git a/mergic/mergic.py b/mergic/mergic.py index 66332f9..f4801eb 100755 --- a/mergic/mergic.py +++ b/mergic/mergic.py @@ -226,17 +226,37 @@ def apply_diff_(args): print_json(partition) +def table(partition): + """Generate 'merge table' rows from a partition. + + Parameters + ---------- + partition : dict + A partition dictionary where the values are lists. Items appear + exactly once through all the value lists (they are "assigned to" + their key value.) + + Yields + ------ + tuple + Pairs of original and new names, as specified in the partition, + encoded in UTF-8. + + """ + for key, values in partition.items(): + for value in values: + yield (value, key) + + def table_(args): """Print out a two-column 'merge table' at the command line.""" - data = json.loads(args.partition.read()) - check(data) + partition = json.loads(args.partition.read()) + check(partition) writer = csv.writer(sys.stdout) writer.writerow(["original", "mergic"]) - for key, values in data.items(): - key = key.encode('utf-8') - for value in values: - value = value.encode('utf-8') - writer.writerow([value, key]) + for row in table(partition): + row = [str(element).encode('utf-8') for element in row] + writer.writerow(row) def _calc_(self, args): diff --git a/mergic/test.py b/mergic/test.py index 05021ef..df4e09f 100644 --- a/mergic/test.py +++ b/mergic/test.py @@ -91,5 +91,20 @@ def test_slightly_interesting_equality(self): {2: ['a'], 1: [3, 2, 1]})) +class TestTable(unittest.TestCase): + + def test_nothing_for_empty(self): + self.assertEqual(len(list(mergic.table({}))), 0) + + def test_right_number_of_rows(self): + p = {1: [2, 3], 'b': [4, 5, 6]} + self.assertEqual(len(list(mergic.table(p))), 5) + + def test_correct_rows_appear(self): + p = {1: [2, 3], 'b': [4, 5, 6]} + rowset = set([(2, 1), (3, 1), (4, 'b'), (5, 'b'), (6, 'b')]) + self.assertEqual(set(mergic.table(p)), rowset) + + if __name__ == '__main__': unittest.main()