DEV Community

Hai Vu
Hai Vu

Posted on • Edited on

How to Write a Class Object to CSV

The Problem

Given a class such as

class Student:
    def __init__(self, student_id, first_name, last_name):
        self.student_id = student_id
        self.first_name = first_name
        self.last_name = last_name

    def __repr__(self):
        return f"Student({self.student_id!r}, {self.first_name!r}, {self.last_name!r})"

    # Other methods...
Enter fullscreen mode Exit fullscreen mode

Given a list of these objects:

my_students = [
    Student(1048, "John", "Adams"),
    Student(2517, "Karen", "Goodman"),
    Student(3131, "Anna", "Karenina"),
]
Enter fullscreen mode Exit fullscreen mode

Our goal is to write them into a CSV file with 3 columns: ID, first, last name. What is a Pythonic way to do that?

The Solution

At first, we might want to create a function which converts a Student object into a tuple and use that:

def student_to_tuple(student):
    return (student.student_id, student.first_name, student.last_name)

with open("students.csv", "w") as stream:
    writer = csv.writer(stream)
    for student in my_students:
        row = student_to_tuple(student)
        writer.writerow(row)
Enter fullscreen mode Exit fullscreen mode

This method works, but if we can modify the Student class, we can do better. By creating an Student.__iter__ method, we can work directly with csv.writer:

class Student:
    def __init__(self, student_id, first_name, last_name):
        self.student_id = student_id
        self.first_name = first_name
        self.last_name = last_name

    def __repr__(self):
        return f"Student({self.student_id!r}, {self.first_name!r}, {self.last_name!r})"

    def __iter__(self):
        return iter([self.student_id, self.first_name, self.last_name])
Enter fullscreen mode Exit fullscreen mode
my_students = [
    Student(1048, "John", "Adams"),
    Student(2517, "Karen", "Goodman"),
    Student(3131, "Anna", "Karenina"),
]

with open("students.csv", "w") as stream:
    writer = csv.writer(stream)
    writer.writerows(my_students)
Enter fullscreen mode Exit fullscreen mode

A couple of notes

  • the __iter__ performs almost the same task as student_to_tuple, but it returns an iterable instead of a tuple.
  • The csv.writer will iterate through this Student object to get the cells in a row, thus trigger the __iter__ method
  • We can now use the .writerows method to write all the objects in one step.

Reading Them Back

Now that we successfully wrote the list of students to a file, how do we read them back?

with open("students.csv") as stream:
    reader = csv.reader(stream)
    for row in reader:
        student = Student(*row)
        print(student)    
Enter fullscreen mode Exit fullscreen mode
Student('1048', 'John', 'Adams')
Student('2517', 'Karen', 'Goodman')
Student('3131', 'Anna', 'Karenina')
Enter fullscreen mode Exit fullscreen mode




Conclusion

By adding a simple Student.__iter__ method, we allow a csv.writer object to work directly with Student objects.

Top comments (2)

Collapse
 
tutorfx profile image
Gabriel Serejo

Haha i got a challenge for you, try to save a nested class in this way

Collapse
 
glenndawg profile image
Glenn MacRae

This is great thank you. How can I add a header row to the csv file using this
method?