Manually set legend colors

Question

I try to plot the following line plot, but I have difficulties to manually set legend colors. Currently the legend colors did not match the line colors. Any help would be very helpful. Thank you.

import random
import matplotlib.pyplot as plt
random.seed(10)

data=[(i, i+random.randint(1,20), random.choice(list("ABC"))) for i in range(2000,2025)]

plt.figure(figsize=(14,8))
for x, y,z in data:
    a=(x,x+y)
    b=(y+random.random(),y+random.random())
    if z=="A":
        a=(x,x)
        plt.plot(a,b,"bo",linestyle="-", linewidth=0.4, color="blue")
    elif z=="B":
        plt.plot(a,b,"bo",linestyle="-", linewidth=0.4, color="green")
    else:
        plt.plot(a,b,"bo",linestyle="-", linewidth=0.4, color="red")
ax = plt.gca()
plt.legend(['A', 'B',"C"])

Answer 1

A simple approach would be to save handles to each type of element. In the code below, handleA, = plt.plot(..., label='A') stored the line element created by plt.plot into a variable named handleA. The handle will keep its label to automatically use in the legend. (A comma is needed because plt.plot always returns a tuple, even if only one line element is created.)

import random
import matplotlib.pyplot as plt

random.seed(10)
data = [(i, i + random.randint(1, 20), random.choice(list("ABC"))) for i in range(2000, 2025)]

plt.figure(figsize=(14, 8))
for x, y, z in data:
    a = (x, x + y)
    b = (y + random.random(), y + random.random())
    if z == "A":
        a = (x, x)
        handleA, = plt.plot(a, b, '-o', linewidth=0.4, color="blue", label='A')
    elif z == "B":
        handleB, = plt.plot(a, b, '-o', linewidth=0.4, color="green", label='B')
    else:
        handleC, = plt.plot(a, b, '-o', linewidth=0.4, color="red", label='C')

plt.legend(handles=[handleA, handleB, handleC], bbox_to_anchor=(1.01, 1.01), loc='upper left')
plt.tight_layout()
plt.show()