python - Making objects iterable?

Question

I'm trying to iterate over each row in a list of lists, append an element from each row to a new list, then find the unique elements in the new list.

I understand that I can do this easily with a for loop. I'm trying a different route because I want to learn more about classes and functions.

Here's an example of the list of lists. The first row is the header:

legislators = [
 ['last_name', 'first_name', 'birthday', 'gender', 'type', 'state', 'party'],
 ['Bassett', 'Richard', '1745-04-02', 'M', 'sen', 'DE', 'Anti-Administration'],
 ['Bland', 'Theodorick', '1742-03-21', '', 'rep', 'VA', ''],
 ['Burke', 'Aedanus', '1743-06-16', '', 'rep', 'SC', ''],
 ['Carroll', 'Daniel', '1730-07-22', 'M', 'rep', 'MD', ''],
 ['Clymer', 'George', '1739-03-16', 'M', 'rep', 'PA', ''],
 ['Contee', 'Benjamin', '', 'M', 'rep', 'MD', ''],...]

Here's my code:

import csv
f = open("legislators.csv")
csvreader = csv.reader(f)
legislators = list(csvreader)

class Dataset:
    def __init__(self, data):
        self.header = data[0] #Isolate header from CSV file
        self.data = data[1:] #Subset CSV data to remove header

legislators_dataset = Dataset(legislators)

def the_set_maker(dataset):
    gender = []
    for each in dataset:
        gender.append(each[3])
    return set(gender)

t=the_set_maker(legislators_dataset)
print(t)

I get the following error:

TypeErrorTraceback (most recent call last)
<ipython-input-1-d65cb459931b> in <module>()
     20     return set(gender)
     21
---> 22 t=the_set_maker(legislators_dataset)
     23 print(t)

<ipython-input-1-d65cb459931b> in the_set_maker(dataset)
     16 def the_set_maker(dataset):
     17     gender = []
---> 18     for each in dataset:
     19         gender.append(each[3])
     20     return set(gender)

TypeError: 'Dataset' object is not iterable

I think the answer is to try to create a method using def __iter__(self) in my Dataset class, but I haven't been able to get this to work. Is this the right track? If not, what's a better one?

Answer 1

According to the documentation for __iter__:

This method should return a new iterator object that can iterate over all the objects in the container.

You might try the following class definition:

class Dataset:
    def __init__(self, data):
        self.header = data[0] #Isolate header from CSV file
        self.data = data[1:] #Subset CSV data to remove header

    def __iter__(self):
        return iter(self.data)

If you're open to trying new options, consider using Pandas:

import pandas as pd
df = pd.read_csv('legislators.csv')
t=df['gender']

Or, if you really want to read in the CSV yourself,

df = pd.DataFrame(legislators[1:], columns=legislators[0])