Search code examples
pythonpython-3.xdictionary-comprehension

How to reuse an expression in a comprehension expression?


Imagine a theoretical snippet:

# just for this example: `bad_structure` contains a list of dicts with different keys
# for the same semantic
bad_structure = [{'path': '/dir/one'}, {'subdir': '/dir/two'}]

# i want to turn this into
# { '/dir/one': some_func('/dir/one'),
#   '/dir/two': some_func('/dir/two')}

result = {}
for e in bad_structure:
  # calculate a value which we will need more than once (here the key)
  p = next(k for k in ('path', 'subdir') if k in e)
  result[p] = some_func(p)

I want to turn this into a dict comprehension now and my first approach looks like this:

bad_structure = [{'path': '/dir/one'}, {'path': '/dir/two'}]
result = {next(k for k in ('path', 'subdir') if k in e): 
          some_func(next(k for k in ('path', 'subdir') if k in e))
          for e in bad_structure}

which contains the 'calculation' twice which ugly, error prone and slow. I would like to rewrite it to s.th. like

result = {p: some_func(p) 
          for p = next(k for k in ('path', 'subdir') if k in e)
          for e in bad_structure}

which is no valid Python code of course..

Is something like this possible in Python?

For clarification: I don't care about comprehension syntax but reusing a calculation without separate variable declaration (which is not possible in closed expressions)


Solution

  • You can use an intermediate comprehension to bind to a name:

    result = {
        p: some_func(p)
        # bind intermediate result to p
        for p in (  # nested comprehension to produce intermediate result
            next(k for k in ('path', 'subdir') if k in e)
            for e in bad_structure
        )
     }
    

    Instead of mapping directly to two separate expressions, it first maps to a common expression which then is mapped to two separate expressions.

    You can pass along and rename an arbitrary number of values. Create a tuple in the inner comprehension, and unpack it to multiple names in the outer comprehension.

    result = {
        p: some_func(e, p)
        for e, p in (
            (e, next(iter(e)))
            for e in bad_structure
        )
     }