Return with argument inside generator

Question

I know there's a similar question already asked, but doesn't answer what I need, as mine is a little different.

My code:

def tFileRead(fileName, JSON=False):
    with open(fileName) as f:
        if JSON:
            return json.load(f)
        for line in f:
            yield line.rstrip('\n')

What I want to do: if JSON is true, it means its reading from a json file and I want to return json.load(f), otherwise, I want to yield the lines of the file into a generator.

I've tried the alternative of converting the generator into json, but that got very messy, very fast, and doesn't work very well.

Answer 1

The first solution that came to my mind was to explicitly return a generator object which would provide the exact behavior you tried to achieve.

The problem is: if you explicitly returned a generator object like this return (line.rstrip('\n') for line in f) the file would be closed after returning and any further reading from the file would raise an exception.

You should write two functions here: one that reads a json file and one for the normal file. Then you can write a wrapper that desides on an argument which of these two functions to call.

Or just move the iteration part into another function like this:

def iterate_file(file_name):
    with open(file_name) as fin:
        for line in fin:
            yield line.rstrip("\n")


def file_read(file_name, as_json=False):
    if as_json:
        with open(file_name) as fin:
            return json.load(fin)
    else:
        return iterate_file(file_name)

Answer 2

You could yield from the dictionary loaded with JSON, thus iterating over the key-value-pairs in the dict, but this would not be your desired behaviour.

def tFileRead(fileName, JSON=False):
    with open(fileName) as f:
        if JSON:
            yield from json.load(f).items()  # works, but differently
        for line in f:
            yield line.rstrip('\n')

It would be nice if you could just return a generator, but this will not work, since using with, the file is closed as soon as the function returns, i.e. before the generator is consumed.

def tFileRead(fileName, JSON=False):
    with open(fileName) as f:
        if JSON:
            return json.load(f)
        else:
            return (line.rstrip('\n') for line in f)  # won't work

Alternatively, you could define another function just for yielding the lines from the file and use that in the generator:

def tFileRead(fileName, JSON=False):
    if JSON:
        with open(fileName) as f:
            return json.load(f)
    else:
        def withopen(fileName):
            with open(fileName) as f:
                yield from f
        return (line.rstrip('\n') for line in withopen(fileName))

But once you are there, you can really just use two separate functions for reading the file en-block as JSON or for iterating the lines...