@@ -234,13 +234,27 @@ def write_report(node, report_type=None, is_mapnode=False):
234
234
235
235
236
236
def _identify_collapses (hastraits ):
237
+ """ Identify traits that will collapse when being set to themselves.
238
+
239
+ ``OutputMultiObject``s automatically unwrap a list of length 1 to directly
240
+ reference the element of that list.
241
+ If that element is itself a list of length 1, then the following will
242
+ result in modified values.
243
+
244
+ hastraits.trait_set(**hastraits.trait_get())
245
+
246
+ Cloning performs this operation on a copy of the original traited object,
247
+ allowing us to identify traits that will be affected.
248
+ """
237
249
raw = hastraits .trait_get ()
238
250
cloned = hastraits .clone_traits ().trait_get ()
239
251
240
252
collapsed = set ()
241
253
for key in cloned :
242
254
orig = raw [key ]
243
255
new = cloned [key ]
256
+ # Allow numpy to handle the equality checks, as mixed lists and arrays
257
+ # can be problematic.
244
258
if isinstance (orig , list ) and len (orig ) == 1 and (
245
259
not np .array_equal (orig , new ) and np .array_equal (orig [0 ], new )):
246
260
collapsed .add (key )
@@ -249,13 +263,30 @@ def _identify_collapses(hastraits):
249
263
250
264
251
265
def _uncollapse (indexable , collapsed ):
266
+ """ Wrap collapsible values in a list to prevent double-collapsing.
267
+
268
+ Should be used with _identify_collapses to provide the following
269
+ idempotent operation:
270
+
271
+ collapsed = _identify_collapses(hastraits)
272
+ hastraits.trait_set(**_uncollapse(hastraits.trait_get(), collapsed))
273
+
274
+ NOTE: Modifies object in-place, in addition to returning it.
275
+ """
276
+
252
277
for key in indexable :
253
278
if key in collapsed :
254
279
indexable [key ] = [indexable [key ]]
255
280
return indexable
256
281
257
282
258
283
def _protect_collapses (hastraits ):
284
+ """ A collapse-protected replacement for hastraits.trait_get()
285
+
286
+ May be used as follows to provide an idempotent trait_set:
287
+
288
+ hastraits.trait_set(**_protect_collapses(hastraits))
289
+ """
259
290
collapsed = _identify_collapses (hastraits )
260
291
return _uncollapse (hastraits .trait_get (), collapsed )
261
292
@@ -264,14 +295,16 @@ def save_resultfile(result, cwd, name):
264
295
"""Save a result pklz file to ``cwd``"""
265
296
resultsfile = os .path .join (cwd , 'result_%s.pklz' % name )
266
297
if result .outputs :
267
- collapsed = set ()
268
298
try :
269
299
collapsed = _identify_collapses (result .outputs )
270
300
outputs = _uncollapse (result .outputs .trait_get (), collapsed )
301
+ # Double-protect tosave so that the original, uncollapsed trait
302
+ # is saved in the pickle file. Thus, when the loading process
303
+ # collapses, the original correct value is loaded.
304
+ tosave = _uncollapse (outputs .copy (), collapsed )
271
305
except AttributeError :
272
- outputs = result .outputs .dictcopy () # outputs was a bunch
273
- outputs = modify_paths (outputs , relative = True , basedir = cwd )
274
- result .outputs .set (** _uncollapse (outputs , collapsed ))
306
+ tosave = outputs = result .outputs .dictcopy () # outputs was a bunch
307
+ result .outputs .set (** modify_paths (tosave , relative = True , basedir = cwd ))
275
308
276
309
savepkl (resultsfile , result )
277
310
logger .debug ('saved results in %s' , resultsfile )
0 commit comments