Return with argument inside generator

Question

I know that there is a similar question that has already been asked, but does not answer what I need, since mine is a little different.

My code is:

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 true, it means reading it from the json file, and I want to return json.load(f) , otherwise I want to transfer the lines of the file to the generator.

I tried an alternative to converting the generator to json, but it became very dirty, very fast and not very good.

Answer 1

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

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

Here you should write two functions: one that reads the json file and one for the regular file. Then you can write a wrapper that takes no argument which of these two functions to call.

Or simply move the iteration part to another function as follows:

 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 can yield from dictionary loaded using JSON, iterating over key-value pairs in a dict, but that would not be your desired behavior.

 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 the generator, but this will not work, since with the help, 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 can define another function only to get lines from a file and use this in a 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're there, you can simply use two separate functions to read the en-block file as JSON or to iterate over the lines ...