Pivotting in Python

pivot.py

from collections import defaultdict

def pivot(rows, x_getter, y_getter, z_getter, empty_value=''):
  pivotted_data = defaultdict(dict)
  seen_y_values = set()

  for row in rows:
    x_value = x_getter(row)
    y_value = y_getter(row)
    z_value = z_getter(row)

    pivotted_data[x_value][y_value] = z_value
    seen_y_values.add(y_value)
  
  output_data = []
  columns = tuple(sorted(pivotted_data.keys()))
  
  # Write headers
  headers = ['']
  headers.extend(columns)
  output_data.append(headers)
  
  for row in sorted(seen_y_values):
    output_row = [row]
    
    # Write rows
    for column in columns:
      output_row.append(pivotted_data[column].get(row, empty_value))
    
    output_data.append(output_row)
  
  return output_data
  

test_pivot.py

import unittest
from operator import itemgetter

from pivot import pivot

class TestPivot(unittest.TestCase):

    def test_indexed_data(self):
        input_data = [
            ['cat', 'A', 1],
            ['dog', 'B', 2],
            ['rabbit', 'A', 2],
            ['hamster', 'B', 1],
        ]

        expected_output_data = [
            ['', 'A', 'B'],
            [1, 'cat', 'hamster'],
            [2, 'rabbit', 'dog'],
        ]

        output_data = pivot(input_data, itemgetter(1), itemgetter(2), itemgetter(0))
        self.assertEqual(expected_output_data, output_data)

    def test_keyed_data(self):
        input_data = [
            {'x': 'A', 'y': 1, 'z': 'cat'},
            {'x': 'B', 'y': 2, 'z': 'dog'},
            {'x': 'A', 'y': 2, 'z': 'rabbit'},
            {'x': 'B', 'y': 1, 'z': 'hamster'},
        ]

        expected_output_data = [
            ['', 'A', 'B'],
            [1, 'cat', 'hamster'],
            [2, 'rabbit', 'dog'],
        ]

        output_data = pivot(input_data, itemgetter('x'), itemgetter('y'), itemgetter('z'))
        self.assertEqual(expected_output_data, output_data)
    
    def test_empty_value(self):
        input_data = [
            {'x': 'A', 'y': 1, 'z': 'cat'},
            {'x': 'B', 'y': 2, 'z': 'dog'},
        ]

        expected_output_data = [
            ['', 'A', 'B'],
            [1, 'cat', 'N/A'],
            [2, 'N/A', 'dog'],
        ]

        output_data = pivot(input_data, itemgetter('x'), itemgetter('y'), itemgetter('z'), 'N/A')
        self.assertEqual(expected_output_data, output_data)