Need python dictionary loop assistance

Question

I have not done a lot of Python programming and am attempting to read in a basic csv then create a nested dictionary from it. Here is what I have so far and I seem to have some issues with looping or overwriting my dict. I know it's not very efficient.

import csv

reader = csv.DictReader(open("fruit.csv"))

fruit_dict = {}
color_dict = {}
for row in reader:
    info_list = []
    count = row.pop('count')
    info_list.append(count)
    year = row.pop('year')
    info_list.append(year)
    info = row.pop('info')
    info_list.append(info)
    if row['color'] not in color_dict:
        #print row['color']
        color_dict['color'] = row['color']
            #print fruit_dict  
        if row['fruit'] not in fruit_dict:
            fruit_dict['name'] = row['fruit']
            #print fruit_dict
            #print info_list
            list_of_info_lists =[]              
            list_of_info_lists.append(info_list)
            fruit_dict['fruitInfo'] = list_of_info_lists
            color_dict['fruit'] = fruit_dict
            #print color_dict
        else:
            list_of_info_lists.append(info_list)
            fruit_dict['fruitInfo'] = list_of_info_lists
            color_dict['fruit'] = fruit_dict
            #print color_dict
    else:
        if row['color'] in color_dict:
            if row['fruit'] not in fruit_dict:
                fruit_dict['name'] = row['fruit']
                #print fruit_dict
                #print info_list
                list_of_info_lists =[]              
                list_of_info_lists.append(info_list)
                fruit_dict['fruitInfo'] = list_of_info_lists
                color_dict['fruit'] = fruit_dict
                #print color_dict
            else:
                list_of_info_lists.append(info_list)
                fruit_dict['fruitInfo'] = list_of_info_lists
                color_dict['fruit'] = fruit_dict
                #print color_dict

#print color_dict

Here is the csv:

color,fruit,year,count,info
red,apple,1970,3,good
red,apple,1922,5,okay
orange,orange,1935,2,okay
green,celery,2001,22,marginal
red,cherries,1999,5,outstanding
orange,carrot,1952,7,okay
green,celery,2014,2,good
green,grapes,2001,12,good

What I'm getting is this:

{'color': 'green', 'fruit': {'name': 'grapes', 'fruitInfo': [['12', '2001', 'good']]}}

Which is lovely except that I am expecting a few more lines than that and am expecting a list of lists when the 'name' already exists for example:

{'color': 'red', 'fruit': {'name': 'apple', 'fruitInfo': [['5', '1922', 'okay'],['3', '1970', 'good']]}}

Any advice would be greatly appreciated. The eventual goal is generate a json file.

Thanks, susan

Here is the format I'd like to have in the end:

[{'color': 'red', 'fruit': {'name': 'apple', 'fruitInfo': [['5', '1922', 'okay'],['3', '1970', 'good']]}},
{'color': 'red', 'fruit': {'name': 'cherries', 'fruitInfo': [['5', '1999', 'outstanding']]}},
{'color': 'orange', 'fruit': {'name': 'orange', 'fruitInfo': [['2', '1935', 'okay']]}},
{'color': 'orange', 'fruit': {'name': 'carrot', 'fruitInfo': [['7', '1952', 'okay']]}},
{'color': 'green', 'fruit': {'name': 'celery', 'fruitInfo': [['2', '2014', 'good'],['22', '2001', 'marginal']]}},
{'color': 'green', 'fruit': {'name': 'grapes', 'fruitInfo': [['12', '2001', 'good']]}}]

Answer 1

Jon Clements' answer is the optimal solution. If you wanted something a little more along the lines of how you originally started to help learn where you might have gone wrong, have a look at this:

results_list = []
colorFruitTuple_set = set()
for row in reader:
    info_list = [row['count'], row['year'],row['info']]
    if (row['color'], row['fruit']) not in colorFruitTuple_set:
        color_dict = {}
        fruit_dict = {}
        color_dict['color'] = row['color']
        fruit_dict['name'] = row['fruit']

        list_of_info_lists = [info_list]

        fruit_dict['fruitInfo'] = list_of_info_lists
        color_dict['fruit'] = fruit_dict
        results_list.append(color_dict)
        colorFruitTuple_set.add((row['color'], row['fruit']))
    else:
        for color_dict in results_list:
            if color_dict["color"] == row['color'] and color_dict["fruit"]["name"] == row["fruit"]:
                color_dict["fruit"]["fruitInfo"].append(info_list)

I think that's along the lines of what you were going for. You were trying to use the same color_dict and fruit_dict when you needed to create several - which also means you can't use either to keep track of duplicates. This is just for learning purposes though - Jon's way is the right way to do it.

Hope this helps!

Answer 2

You can use a defaultdict here with a list to keep your fruitInfo and a 2-tuple as your key (color and fruit) then reformat after, eg:

import csv
from collections import defaultdict

dd = defaultdict(list)
with open('yourfile.csv') as fin:
    csvin = csv.DictReader(fin)
    for row in csvin:
        dd[row['color'], row['fruit']].append([row['count'], row['year'], row['info']])

Then a slight reformat of dd using:

reformatted = [{'color': c, 'fruit': {'name': f, 'fruitInfo': v}} for (c, f), v in dd.items()]

Gives you:

[{'color': 'orange',
  'fruit': {'fruitInfo': [['7', '1952', 'okay']], 'name': 'carrot'}},
 {'color': 'green',
  'fruit': {'fruitInfo': [['12', '2001', 'good']], 'name': 'grapes'}},
 {'color': 'orange',
  'fruit': {'fruitInfo': [['2', '1935', 'okay']], 'name': 'orange'}},
 {'color': 'red',
  'fruit': {'fruitInfo': [['3', '1970', 'good'], ['5', '1922', 'okay']],
            'name': 'apple'}},
 {'color': 'red',
  'fruit': {'fruitInfo': [['5', '1999', 'outstanding']], 'name': 'cherries'}},
 {'color': 'green',
  'fruit': {'fruitInfo': [['22', '2001', 'marginal'], ['2', '2014', 'good']],
            'name': 'celery'}}]

Answer 3

When dealing with dictionaries of dictionaries, my pattern is like this:

sub_dict = main_dict.get(key, {})
sub_dict[sub_key] = sub_value
main_dict[key] = sub_dict

This gets the sub-dictionary, or {} if it doesn't exist. Then it assigns a value to the sub-dictionary, and puts the sub-dictionary back into the main dictionary.

fruit_dict = {}
for row in reader:
    # make the info_list
    info_list = [row['count'], row['year'], row['info']]
    # extract color and fruit into variables
    color = row['color']
    fruit = row['fruit']
    # unpack the dictionaries and list
    colors = fruit_dict.get(color, {})
    fruits = colors.get(fruit, {})
    info = fruits.get('info', [])
    # reassemble the list and dictionaries
    info.append(info_list)
    fruits['info'] = info
    colors[fruit] = fruits
    fruit_dict[color] = colors

The result is a little different than your example, but it's necessary to change it to use the color and fruit as keys.

{'orange': {'orange': {'info': [['2', '1935', 'okay']]}, 'carrot': {'info': [['7', '1952', 'okay']]}}, 'green': {'celery': {'info': [['22', '2001', 'marginal'], ['2', '2014', 'good']]}, 'grapes': {'info': [['12', '2001', 'good']]}}, 'red': {'cherries': {'info': [['5', '1999', 'outstanding']]}, 'apple': {'info': [['3', '1970', 'good'], ['5', '1922', 'okay']]}}}