Search code examples
pythonlistdictionarylist-comprehensionone-liner

List comprehension to collect a nested dict's items that have given innermost key


Please assume the following dict is given

sources_targets = 
{'front': {'source': [[0.021025050381526675, -0.39686011326197257, 3.328947963092819], [1.0601052368302668, -0.3938359761868055, 3.3247223740425893], [1.0543731204008824, -0.038184154352961984, 2.941639590795943], [0.017868184643970383, -0.0445863307249157, 2.9604912584916665]], 'target': [[-250.0, 0.0, 60.0], [-250.0, 0.0, -60.0], [-190.0, 0.0, -60.0], [-190.0, 0.0, 60.0]]}, 'left': {'source': [[-0.9522471062122733, -1.8444069007997075, 5.372044839796925], [-0.9665739520089994, -1.001259794819009, 5.0057689609608005], [0.9940538769534978, -1.851333804840362, 5.340677879647542], [0.9959517362506759, -1.0049420919111534, 4.942663843894899]], 'target': [[60.0, 0.0, 140.0], [60.0, 0.0, 80.0], [-60.0, 0.0, 140.0], [-60.0, 0.0, 80.0]]}, 'right': {'source': [[-0.8596841529333474, -3.0721166255322663, 4.182871479604773], [-0.8796404109762729, -2.117062488877432, 4.147040556143069], [1.0791152756424247, -2.0436646487085532, 4.08578012939533], [1.0951903113036177, -2.994375693306352, 4.124102127893507]], 'target': [[-60.0, 0.0, -140.0], [-60.0, 0.0, -80.0], [60.0, 0.0, -80.0], [60.0, 0.0, -140.0]]}, 'rear': {'source': [[0.08792816743383122, -0.5260295091566244, 3.0182522276468458], [0.9012540916522604, -0.5012267882763559, 3.0172622143554695], [0.8942115223224005, -0.15635208604951806, 2.6353057539009934], [0.08814313470840558, -0.18017837896764446, 2.6579174137231463]], 'target': [[250.0, 0.0, -40.0], [250.0, 0.0, 60.0], [190.0, 0.0, 60.0], [190.0, 0.0, -40.0]]}}

bundle_names = ['front', 'left', 'right', 'rear']

I want to get a list of all the sources and another list of all the targets.
Also, it needs to be done in a one-liner list comprehension.

My successful attempt is


sources3d = [st[1] for name in self._bundle_names for st in sources_targets[name].items() if st[0] == "source"]
targets3d = [st[1] for name in self._bundle_names for st in sources_targets[name].items() if st[0] == "target"]

with the correct output (for "sources" for example)

[[[0.021025050381526675, -0.39686011326197257, 3.328947963092819], [1.0601052368302668, -0.3938359761868055, 3.3247223740425893], [1.0543731204008824, -0.038184154352961984, 2.941639590795943], [0.017868184643970383, -0.0445863307249157, 2.9604912584916665]], [[-0.9522471062122733, -1.8444069007997075, 5.372044839796925], [-0.9665739520089994, -1.001259794819009, 5.0057689609608005], [0.9940538769534978, -1.851333804840362, 5.340677879647542], [0.9959517362506759, -1.0049420919111534, 4.942663843894899]], [[-0.8596841529333474, -3.0721166255322663, 4.182871479604773], [-0.8796404109762729, -2.117062488877432, 4.147040556143069], [1.0791152756424247, -2.0436646487085532, 4.08578012939533], [1.0951903113036177, -2.994375693306352, 4.124102127893507]], [[0.08792816743383122, -0.5260295091566244, 3.0182522276468458], [0.9012540916522604, -0.5012267882763559, 3.0172622143554695], [0.8942115223224005, -0.15635208604951806, 2.6353057539009934], [0.08814313470840558, -0.18017837896764446, 2.6579174137231463]]]

The way I accessed the inner dict using .items() and accessing the tuple by index seems cumbersome.

I would like to access some way with

[st["source"] for name in self._bundle_names for st in sources_targets[name]] 

which doesn't work but would be much cleaner.

I am sure there is a way to do this correctly.


Solution

  • You iterate all the items and then pick the one (!) item that has the desired key. Instead, access the key directly. This should work, yielding a result equal to your sources3d and targets3d:

    sources3d = [sources_targets[name]["source"] for name in bundle_names]
    targets3d = [sources_targets[name]["target"] for name in bundle_names]
    

    If you need more than one item, you can basically do the same, but also iterate over the list of keys/items you would like to get:

    >>> items = ["source", "target"]                                                                                                                                        
    >>> [sources_targets[name][item] for name in bundle_names for item in items]