Python CSV reader

Question

I have a function that pulls data from a file or string literal in my case and returns them in either a list or a dict.

My concerns

1. I'm not sure if a function should have two different return types.

2. I feel like I'm repeating myself a lot.

3. If it is ok to return two different datatypes, how do you document that in the docstring?

---

def pull_csv(data, delimiter=',', field_names=None, perserve_line_nums=False):

    """
    Pulls data out of a CSV file and returns as either a list or dict.

    :param data:
    :param delimiter: str
    :param field_names: bool
    :param perserve_line_nums: bool
    :return:
    """

    if field_names:

        if perserve_line_nums:

            lines = {}
            for line_num, line in enumerate(data.split('\n')):
                lines.update({line_num: {}})
                for index, value in enumerate(line.split(delimiter)):
                    lines[line_num].update({field_names[index]: value})

        else:

            lines = []
            for line_num, line in enumerate(data.split('\n')):
                lines.append({})
                for index, value in enumerate(line.split(delimiter)):
                    lines[line_num].update({field_names[index]: value})

    else:

        if perserve_line_nums:

            lines = {}
            for line_num, line in enumerate(data.split('\n')):
                lines.update({line_num: []})
                for value in line.split(delimiter):
                    lines[line_num].append(value)

        else:

            lines = []
            for line_num, line in enumerate(data.split('\n')):
                lines.append([])
                for value in line.split(delimiter):
                    lines[line_num].append(value)

    return lines

Example

data = """ \
Bill, 1, Jan, 28, 93
Joe, 7, Dec, 02, 67 \
"""
print(pull_csv(data, ',', ['Name', 'Fav_num', 'Birth_month', 'Birth_date', 'Birth_year'], True))

Output

[{'Birth_date': ' 28', 'Birth_year': ' 03', 'Name': ' Jason', 'Birth_month': ' Jan', 'Fav_num': ' 1'}, {'Birth_date': ' 02', 'Birth_year': ' 67 ', 'Name': 'George', 'Birth_month': ' Dec', 'Fav_num': ' 7'}]

Answer 1

I would like to offer some suggestions that will organize your code a little better and give you some tips that to make things less confusing.

I should first point out that there is already a Python CSV module that already accomplishes what you are setting out to to: https://docs.python.org/2/library/csv.html However I will be constructive with some advice in your current code.

Doc Strings

Your doc string is misleading and not complete. According to your description you would expect to give it a file path to a CSV file. Your parameter descriptions are missing what the argument means and the expected value. For example your "field_names" in the doc says "bool" when this is not accurate as it is really a list of column names in the data. Do be as descriptive as possible if your arguments are unclear.

Returning different types in a function

The flexibility in Python allows you to do this and it can be a blessing and also abused. I would think carefully before committing to do this. Ask yourself if a function can return two different types of data, are they actually two different functions? Can one type of data return be easily converted or used to the other if the user so desired? Is your function trying to be too clever? Could there be a separate function that can convert one type to the other so the user can easily understand and control what is going to happen?

In your example, I don't see a need for using a dictionary for preserving line numbers as you can use the index of a list as the line number.

Over complicated loops and dictionary calls

There is a lot of unneeded use of enumerate and dictionary updates in your code. This is both inefficient and doesn't take advantage of Python's strengths.

For example:

lines = []
for line_num, line in enumerate(data.split('\n')):
    lines.append([])
    for value in line.split(delimiter):
        lines[line_num].append(value)

Could be written like this which is more efficient and easier to read:

lines = []
for dataLine in data.split('\n'):
    lines.append(dataLine.split(delimiter))

If using field names, instead of this:

lines = []
for line_num, line in enumerate(data.split('\n')):
    lines.append({})
    for index, value in enumerate(line.split(delimiter)):
        lines[line_num].update({field_names[index]: value})

Do this:

lines = []
for dataLine in data.split('\n'):
    lines.append(dict([(name, value) for name, value in zip(field_names, dataLine.split(delimiter))]))

Better implementation

As I said I would probably break this up into two functions, one to parse and break up the data into lists and return it. Then another function that will accept this data and convert it into a dictionary of dictionaries when using field names:

def parseCSVData(dataString, delimiter=','):
    lines = []
    for data in dataString.split('\n'):
        lines.append(data.split(delimiter))
    return lines

def buildFieldNameDict(csvListData, fieldNames):
    result = []
    for line in csvListData:
        result.append(dict([(name, value) for name, value in zip(fieldNames, line)]))
    return result

csvListData = parseCSVData(myCSVString)
csvFieldDataDict = buildFieldNameDict(csvListData, ['Name', 'Fav_num', 'Birth_month', 'Birth_date', 'Birth_year'])

Overall I would not try to reinvent the wheel if Python provides what you need since it's already written and maintained for you. But this can be a good learning exercise.