35
loading...
How to Write a Class Object to CSV
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?
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
__iter__ performs almost the same task as student_to_tuple, but it returns an iterable instead of a tuple. csv.writer will iterate through this Student object to get the cells in a row, thus trigger the __iter__ method.writerows method to write all the objects in one step.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')By adding a simple
Student.__iter__ method, we allow a csv.writer object to work directly with Student objects.35