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...
Given a list of these objects:
my_students = [
Student(1048, "John", "Adams"),
Student(2517, "Karen", "Goodman"),
Student(3131, "Anna", "Karenina"),
]
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)
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])
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)
A couple of notes
- the
__iter__performs almost the same task asstudent_to_tuple, but it returns an iterable instead of a tuple. - The
csv.writerwill iterate through thisStudentobject to get the cells in a row, thus trigger the__iter__method - We can now use the
.writerowsmethod 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)
Student('1048', 'John', 'Adams')
Student('2517', 'Karen', 'Goodman')
Student('3131', 'Anna', 'Karenina')
Conclusion
By adding a simple Student.__iter__ method, we allow a csv.writer object to work directly with Student objects.
Top comments (2)
Haha i got a challenge for you, try to save a nested class in this way
This is great thank you. How can I add a header row to the csv file using this
method?