Search code examples
pythondictionaryenthoughttraitstraitsui

Enthought Traits.HasTraits class as a nested dictionary?


Imagine I have the following Traits objects:

from traits.api import Int, HasTraits, Instance

class Foo(HasTraits):
    a = Int(2)
    b = Int(5)

class Bar(HasTraits):
    c = Int(7)
    foo = Instance(Foo,())

Bar will let me access attribute a on Foo via:

bar = Bar()
bar.foo.a
>>> 2

Is there a standard way to return bar as a nested dictionary of the form:

print bar_as_dict 
>>> {'c':7, 'foo':{'a':2, 'b':5}}

I'm essentially trying to extract all subtraits on an object that are a particular type. Our use case is we have deeply-nested HasTrait objects that have plot traits, and we are trying to dynamically find all the plots on a particular object. We have a function that can return the plots from a nested dictionary, but we need to pass HasTrait objects in, so formatting them into nested dictionaries would be great. If there is another way to dynamically inspect the stack of a HasTraits object and return all traits of a certain type, that would work too.

Here's a link to the HasTraits API... couldn't figure this out directly from that.

Solution Attempt

I've tried using the .traits() method, but it returns these CTrait objects.

print bar.traits()
>>>{'trait_added': <traits.traits.CTrait object at 0x8b7439c>,
 'c': <traits.traits.CTrait object at 0x8b29e9c>, 
'foo': <traits.traits.CTrait object at 0x8b29e44>, 
'trait_modified': <traits.traits.CTrait object at 0x8b74344>}

Which don't evaluate as I'd expect:

isinstance(bar.traits()['c'], int)
>>> False

But after Pieter's suggestion, this works:

print bar.traits()['c'].is_trait_type(Int)
>>> True

Now the question is how to do this recursively.


Solution

  • I figured it out after following a similar question on recursion of a nested dictionary without Trait values

    def flatten_traitobject(traitobject, *types):
        node_map = {}
        node_path = [] 
        def nodeRecursiveMap(traitobject, node_path): 
            for key in traitobject.editable_traits():
                val = traitobject.get(key)[key]
                for type in types:
                    if isinstance(val, types[0]):
                        node_map['.'.join(node_path + [key])] = val 
                try:
                    nodeRecursiveMap(val, node_path + [key])
                except (AttributeError, TypeError):
                    pass
        nodeRecursiveMap(traitobject, node_path)
        return node_map
    

    Testing it out, I can choose to retain only integers and/or only objects of type Foo

    bar=Bar()
    print 'RETAINING ONLY INTEGER TYPES'
    print flatten_traitobject(bar, int)
    print '\nRETAINTING ONLY FOO OBJECT TYPES'
    print flatten_traitobject(bar, Foo)
    
    >>> RETAINING ONLY INTEGER TYPES
    >>> {'c': 7, 'foo.b': 5, 'foo.a': 2}
    
    >>> RETAINTING ONLY FOO OBJECT TYPES
    >>> {'foo': <__main__.Foo object at 0x8f9308c>}
    

    Foo represents the special plotting traits I'm after, so I can effectively return them now.