Using variables in a for loop to assign legend values in pyplot

Question

I've got a DataFrame with several columns including a name, a date and a statistic. I'm plotting a statistic over a period of time to a graph and trying to make the legend of each line match up with the name.

Each time I try something new, I either get the statistic's column name or one name from the name column repeated twice or truncated.

Can I make both names appear in the legend with the existing for loop?

fig, ax = plt.subplots()
for name, group in statistics.groupby('A name'):
    group.plot(ax=ax, label=name)
    print(name)

plt.legend()

This code, which prints two separate lines on the plot, gives each line the the name of the column for the stat I'm plotting instead the name of the person, which is what I'm trying to do with label=name. The print statement in the for loop prints the name correctly.

How can I set the label to each plotted line equal to the correct name?

Answer 1

You need to grab those line handles:

statistics = pd.DataFrame({'A name':pd.np.random.choice(list('ABCD'),20),'Data Values':np.random.randint(0,100,20)})
fig, ax = plt.subplots()
lh = {}
for name, group in statistics.groupby('A name'):
    lh[name] = group.plot(ax=ax, label=name)
    print(name)

plt.legend(lh)