Skip to content

Instantly share code, notes, and snippets.

@euangoddard

euangoddard/pivot.py

Last active March 15, 2018 08:51
Show Gist options
  • Save euangoddard/fbd22c4e5f15b5d8dab342ee992a93c7 to your computer and use it in GitHub Desktop.
Save euangoddard/fbd22c4e5f15b5d8dab342ee992a93c7 to your computer and use it in GitHub Desktop.
Download ZIP
Pivotting in Python
Raw
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
Raw
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment