2
\$\begingroup\$

As part of the exercise in the book "Automate the boring stuff with Python" chapter 13 on Excel, I had the following exercise:

  1. Reads the data from the Excel spreadsheet
  2. Counts the number of census tracts in each county
  3. Counts the total population of each county
  4. Prints the results

Here is the structure of the spreadsheet enter image description here

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")
\$\endgroup\$
  • \$\begingroup\$ You can use modules like Pandas to help you with that \$\endgroup\$ – Barb Dec 17 '19 at 5:06
1
\$\begingroup\$

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}}
\$\endgroup\$
  • \$\begingroup\$ Thank you very much @Graipher. As I was writing this code, my mind was screaming pandas as it would be so much simpler. However, just to flex my python muscles, I want to go through this exercise. Then, I wanted to rewrite this using Pandas to compare the 2 approaches. \$\endgroup\$ – Siraj Samsudeen Dec 17 '19 at 9:07
  • \$\begingroup\$ @SirajSamsudeen: I agree, which is why I added both possibilities. The groupby function should now work. \$\endgroup\$ – Graipher Dec 17 '19 at 9:12
  • 1
    \$\begingroup\$ @Graiphet, thank you so much. I was simply blown away by your elegant solution, especially the part that you encapsulated it so well. I need to read it multiple times to get it into my head. So beautiful. Thank you so much for the time you put into rewriting my code. \$\endgroup\$ – Siraj Samsudeen Dec 17 '19 at 13:43
  • \$\begingroup\$ can I ask you for a favor? I am trying to reread your code and run it in Jupyter to compare your code and mine to see what changes you have made. Is it possible for you to take my code and comment out the portions you have edited and also provide a quick rationale. For example, you have used defaultdict which as I just looked up helps you with avoiding a check for the existence of the key - is this needed in my program? In my code, I have used a set as the value of the dictionary to ensure that the counties are not duplicated - how have you handled it ?I am sorry to ask you for more \$\endgroup\$ – Siraj Samsudeen Dec 18 '19 at 3:59
  • \$\begingroup\$ any recommendation from you to learn pandas - I want something that helps me to get to an intermediate level fast so that I can replicate the elegant 2-line expression you have written on my own? \$\endgroup\$ – Siraj Samsudeen Dec 18 '19 at 4:09

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.