profiler: fix minor plotting issue

* Plot legends could contain a partial list of the lines plotted if some
  of the lines were partial.

Author: oliver-sanders

cylc/flow/main_loop/log_data_store.py

@@ -107,15 +107,17 @@ def _plot(state, path):
     ax1.set_xlabel('Time (s)')
 
     ax1.set_ylabel('Objects')
-    for key, objects in state['objects'].items():
-        ax1.plot(times, objects, label=key)
+    lines = [
+        ax1.plot(times, objects, label=key)[0]
+        for key, objects in state['objects'].items()
+    ]
 
     ax2 = ax1.twinx()
     ax2.set_ylabel('Size (kb)')
     for sizes in state['size'].values():
         ax2.plot(times, [x / 1000 for x in sizes], linestyle=':')
 
-    ax1.legend(loc=0)
+    ax1.legend(lines, state['objects'], loc=0)
     ax2.legend(
         (ax1.get_children()[0], ax2.get_children()[0]),
         ('objects', 'size'),

cylc/flow/main_loop/log_memory.py

@@ -155,8 +155,8 @@ def _dump(data, path):
 
 def _plot(fields, times, path, title='Objects'):
     if (
-            not PLT
-            or len(times) < 2
+        not PLT
+        or len(times) < 2
     ):
         return False
 
@@ -166,10 +166,12 @@ def _plot(fields, times, path, title='Objects'):
     ax1.set_xlabel('Time (s)')
     ax1.set_ylabel('Memory (kb)')
 
-    for key, sizes in fields.items():
-        ax1.plot(times, [x / 1000 for x in sizes], label=key)
+    lines = [
+        ax1.plot(times, [x / 1000 for x in sizes], label=key)[0]
+        for key, sizes in fields.items()
+    ]
 
-    ax1.legend(loc=0)
+    ax1.legend(lines, fields, loc=0)
 
     # start both axis at 0
     ax1.set_xlim(0, ax1.get_xlim()[1])