使用 CSV 文件中的列表创建嵌套字典

Question

I have a CSV file ("import.csv") which has the following data where each value is a string

KID, C1, C2, C3
ID1, X1, "X2, X2.5", X3
ID2, "Y1, Y1.5", Y2, Y3

I want to create a nested dictionary using the data like so

'KID' {
    'ID1' {
        'C1': 'X1',
        'C2': ['X2', 'X2.5'],
        'C3': 'X3',
    },
    'ID2' {
        'C1': ['Y1', 'Y1.5'],
        'C2': 'Y2',
        'C3': 'Y3',
    },

I have the following code

import csv
dict = {}

with open("import.csv", mode='r') as data_file:
data = csv.DictReader(data_file)
for row in data:
    item = artists.get(row["artist"], dict())

    item["C1"] = str(row["C1"])
    item["C2"] = str(row["C2"])
    item["C3"] = str(row["C3"])


    dict[row["KID"]] = item

Instead, I am getting a dictionary of

'KID' {
    'ID1' {
        'C1': 'X1',
        'C2': 'X2, X2.5',
        'C3': 'X3',
    },
    'ID2' {
        'C1': 'Y1, Y1.5',
        'C2': 'Y2',
        'C3': 'Y3',
    },

I am looping through the dictionary by finding keys in the dictionary by value, so the string/list issue is messing me up. Is there a way to change my dictionary creation or do I have to change the way my CSV is set up?

Answer 1

Use a wrapper function to look to return list if multiple values are present and string otherwise

def get_str_or_list(value):
    return value.split(",") if "," in value else value


item["C1"] = get_str_or_list(str(row["C1"]))
item["C2"] = get_str_or_list(str(row["C2"]))
item["C3"] = get_str_or_list(str(row["C3"]))