read an excel that contains population data of counties and compute summaries by state, county As part of the exercise in the book "Automate the boring stuff with Python" chapter 13 on Excel, I had the following exercise:


Reads the data from the Excel spreadsheet
Counts the number of census tracts in each county
Counts the total population of each county
Prints the results


Here is the structure of the spreadsheet


Here is my code and I would like to get reviews on how I can improve this code as I am quite new to Python.

from collections import Counter
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 = set()
for cell in list(sheet.columns)[1][2:]:
    states.add(cell.value)

#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 ={}
for state in states:
    counties_in_state[state]=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

# Sort the counties in each state
for state,county_set in counties_in_state.items():
    counties_in_state[state] = sorted(county_set)

print("State","County", "Count", "Pop", sep="\t")
for state,county_set in counties_in_state.items():
    for county in county_set:
        print(state, county, counter[county], counter_pop[county], sep="\t")
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 &gt;= 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}}