Search code examples
pythoniteratoryield

call a new function without closing the file


I have a method that opens a csv file and another method that I want to use to process this file -but I also need it to process different types of objects, so I want to keep it as an independent method.

import csv

class Loader():

    def load_from_csv(self, csv_path: str):

        with open(csv_path, mode='r', encoding='utf-8') as file:
            reader = csv.DictReader(file)
            return self._read_and_yield(reader)
    
    @classmethod
    def _read_and_yield(cls,reader):
        for row in reader:
            yield(row)

loader = Loader()
for row in loader.load_from_csv('sample.csv'):
    print(row)

But I get ValueError: I/O operation on closed file, because the file seems to get closed when I call the _read_and_yield method. Here it just yields the rows, but imagine a case in which I have to do more operations: I should repeat this process for each type of file (namely read_from_excel).

How should I handle this?


Solution

  • You could consider rewriting your Loader class as a context manager.

    Something like this:

    from csv import DictReader
    
    class Loader:
        def __init__(self, filename):
            self._filename = filename
        def __enter__(self):
            self._fd = open(self._filename, "r", encoding="utf-8", newline="")
            return self
        def __exit__(self, *_):
            self._fd.close() #type: ignore
        def row(self):
            yield from DictReader(self._fd) #type: ignore
    
    with Loader("foo.csv") as loader:
        for r in loader.row():
            print(r)