There are two ways in which you can change your code. The easiest one is to use a more powerful library, that can do more than just reading excel worksheets. For this I would recommend pandas, in which this task is very few lines:
import pandas as pd
df = pd.read_excel("support files/censuspopdata.xlsx")
df.groupby("County").agg({"State": "last", "CensusTract": "count", "POP2010": "sum"})
# State CensusTract POP2010
# County
# San Francisco CA 7 31060
But you can also improve your use of vanilla Python, by using set comprehensions, the standard library module collections and following Python's official style-guide, PEP8:
from collections import Counter, defaultdict
import openpyxl
wb = xl.load_workbook("support files/censuspopdata.xlsx")
sheet = wb.active
#Create a unique list of all state names from the state column (which is 2nd column)
states = {cell.value for cell in list(sheet.columns)[1][2:]}
#Initiate the countie_in_state dict to hold a set for each state key as it would be populated when the Excel is read row-by-row
counties_in_state = defaultdict(set)
# counter counts the number of census tracts(rows) and counter_pop holds the cumulative population
counter = Counter()
counter_pop = Counter()
#Skip the first row as it contains column names
for row in range(2, sheet.max_row+1):
#Each row represents a census tract
state = sheet.cell(row, 2).value
county = sheet.cell(row, 3).value
pop = sheet.cell(row, 4).value
counties_in_state[state].add(county)
counter[county] += 1
counter_pop[county] += pop
print("State","County", "Count", "Pop", sep="\t")
for state,county_set in counties_in_state.items():
# Sort the counties in each state
for county in sorted(county_set):
print(state, county, counter[county], counter_pop[county], sep="\t")
You could put this code under a if __name__ == "__main__": guard to allow importing from this script and generalize this code so it gains an interface similar to the pandas one:
from collections import Counter, defaultdict
import openpyxl
def update_sum(old, x):
return old + x
def update_count(old, x):
return old + 1
def update_last(old, x):
return x
FUNCS = {"sum": update_sum, "count": update_count, "last": update_last}
def groupby(sheet, key, **agg):
assert agg, "Must give an aggregation function"
rows = sheet.rows
headers = [cell.value for cell in next(rows)]
funcs = {name: FUNCS[func_name] for name, func_name in agg.items()}
counters = defaultdict(lambda: defaultdict(int))
key_index = headers.index(key)
assert key_index >= 0, "Key not found in data"
for row in rows:
row_key = row[key_index].value
for i, (name, cell) in enumerate(zip(headers, row)):
if name in agg:
counters[row_key][name] = funcs[name](counters[row_key][name], cell.value)
return dict(counters)
if __name__ == "__main__":
wb = xl.load_workbook("support files/censuspopdata.xlsx")
sheet = wb.active
agg = {"State": "last", "POP2010": "sum", "CensusTract": "count"}
for county, data in groupby(sheet, "County", **agg).items():
print(data["State"], county, data["CensusTract"], data["POP2010"], sep="\t")
# CA San Francisco 7 31060
This is not the most flexible code, extending it to support e.g. a "mean" is quite hard. For this you would probably have to read the table into a different data structure, sort by the key and then use itertools.groupby. This way you can use any function that acts on the whole group:
from itertools import groupby
from operator import itemgetter
def first(x):
return next(x)
def last(x):
return list(x)[-1]
def count(x):
return len(list(x))
def groupby_agg(sheet, *key, **agg):
rows = sheet.rows
headers = [cell.value for cell in next(rows)]
data = [{col: cell.value for col, cell in zip(headers, row)}
for row in rows]
data.sort(key=itemgetter(*key))
d = {}
for name, group in groupby(data, key=itemgetter(*key)):
group = list(group)
d[name] = {col_name: func(map(itemgetter(col_name), group))
for col_name, func in agg.items()}
return d
...
groupby_agg(sheet, "County", State=first, CensusTract=count, POP2010=sum)
# {'San Francisco': {'CensusTract': 7, 'POP2010': 31060, 'State': 'CA'}}
This has the advantage that you can also use the built-in sum, min, max, statistics.mean, .... It also allows using multiple keys, although it does not give you a multiply nested dictionary in this case:
groupby_agg(sheet, "State", "County", CensusTract=count, POP2010=sum)
# {('CA', 'San Francisco'): {'CensusTract': 7, 'POP2010': 31060}}
