DEV Community

Hai Vu
Hai Vu

Posted on • Edited on

1 1

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.

AWS Q Developer image

Your AI Code Assistant

Automate your code reviews. Catch bugs before your coworkers. Fix security issues in your code. Built to handle large projects, Amazon Q Developer works alongside you from idea to production code.

Get started free in your IDE

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?

Sentry image

See why 4M developers consider Sentry, “not bad.”

Fixing code doesn’t have to be the worst part of your day. Learn how Sentry can help.

Learn more

👋 Kindness is contagious

Please leave a ❤️ or a friendly comment on this post if you found it helpful!

Okay