Efficiently read in 2nd column of CSV into List of Lists in Python 3

Question

I created a function to read in arbitrary numbers of .csv files, their 2nd column only, into a list of lists (or if there is a more efficient data structure, then that). I want it to be efficient and scalable/customizable. Is there anything in it that can be tweaked and improved?

import os
FILE_NAMES=["DOCS/1.csv","DOCS/2.csv"]

def FUNC_READ_FILES(file_names): 
    nr_files=len(file_names)
    filedata=[[] for x in range(nr_files)] # efficient list of lists

    for i in range(nr_files):   # read in the files
        if(os.path.isfile(file_names[i])):
           with open(file_names[i],'r') as f:
              filedata[i]=f.readlines()
        else:
             print("ERROR: FILES ARE MISSING!!!!") 
             exit() 

    for i in range(nr_files):  # iterate through the files and only keep the 2nd column, remove \n if it's end of line
       for k in range(len(filedata[0])):
        filedata[i][k]=filedata[i][k].strip().split(',')[1]

    return (filedata)

FUNC_READ_FILES(FILE_NAMES)

Answer 1

Since you are in python 3.x, I'd suggest looking into asyncio for the CPU intensive file I/O operations.

---

In your code, you are first reading each and every line from the csv into memory, and then processing that data. This is highly inefficient. Process those lines as soon as you get to them, so that your memory overhead is minimal.

---

Instead of iterating over the indices of the list of files, iterate over the files themselves:

filedata = []
for file_name in file_names:
    if(os.path.isfile(file_name)):
        with open(file_name, 'r') as f:
            reader = csv.reader(f)
            filedata.append([row[1] for row in reader])
    else:
         print("ERROR: FILES ARE MISSING!!!!") 
         exit()

will replace the following

for i in range(nr_files):   # read in the files
    if(os.path.isfile(file_names[i])):
       with open(file_names[i],'r') as f:
          filedata[i]=f.readlines()
    else:
         print("ERROR: FILES ARE MISSING!!!!") 
         exit() 
for i in range(nr_files):  # iterate through the files and only keep the 2nd column, remove \n if it's end of line
   for k in range(len(filedata[0])):
    filedata[i][k]=filedata[i][k].strip().split(',')[1]

Answer 2

General remarks

pep 8

for your names and code style, try to follow pep-8

• lower_case for variable and function names
• spaces around operators

main guard

put the calling of your functions after a if __name__ == '__main__':, so you can load the script from somewhere else without it executing the code immediately

looping

Don't loop over indices. Code like for i in range(nr_files): is a lot cleaner using enumerate: for i, filename in enumerate(file_names).

I suggest you check out the excellent 'Looping like a Pro' talk by David Baumgold

functions

Instead of having 1 function to load the files, loop over them and pick the correct element, seasiest would be to split if into different functions:

1. takes the list of files, and passes them on one by one to the parse
2. parse a single file

Generators

The most pythonic an efficient approach to do this would be to use generators, pathlib.Path and the built-in csv module

My solution

parse one file

This function takes a filehandle, and parses the requested element from the line

import csv
from pathlib import Path


def parse_file(filehandle, field_name):
    kwargs = {  # https://docs.python.org/3/library/csv.html#csv-fmt-params
        'delimiter': ',',
        'skipinitialspace': True,
        # ...
    }
#     field_name = b  # or the column name 
    reader = csv.DictReader(filehandle, **kwargs)  # or csv.reader if there is no header, and it might be faster
    for line in reader:
        yield line[field_name]

This can be easily tested like this:

from io import StringIO
csv_str = '''a, b, c
0, 1, 2
3, 4, 5'''
with StringIO(csv_str, newline='') as file:
    print(list(parse_file(file, 'b')))

['1', '4']

parse multiple file

def parse_files(files):
    for file in files:
        try:
            with filename.open('r', newline='', ) as csv_file:
                yield list(parse_file(csv_file))
        except FileNotFoundError:
            print("ERROR: FILES ARE MISSING!!!!")
            raise

Now we have a good method to parse the information, we just need to call it with the subsequent files

main

def main(files):
    results = list(parse_files(files))
    return results

if __name__ == '__main__':

    files= [Path("DOCS/1.csv"),Path("DOCS/2.csv")]
    main(files)