How can I select a variable by (string) name?

I want to return a pre-determined list from my function, based on the string input.

def get_ext(file_type):
    text = ['txt', 'doc']
    audio = ['mp3', 'wav']
    video = ['mp4', 'mkv']
    return # what do I return here?

get_ext('audio')  #should return the list ['mp3', 'wav']

What is the easiest way to do it?

_{For the related problem of trying to use strings to assign or create variables, see How do I create variable variables?. This question is about looking them up.}

_{For lookup on an existing object (rather than in the current local variables), see How to access object attribute given string corresponding to name of that attribute.}

Solution

In most cases like this, an ordinary dictionary will do the job just fine.

>>> get_ext = {'text': ['txt', 'doc'],
...            'audio': ['mp3', 'wav'],
...            'video': ['mp4', 'mkv']
... }
>>> 
>>> get_ext['video']
['mp4', 'mkv']

If you really want or need a function (for which there can be valid reasons) you have a couple of options. One of the easiest is to assign to the get method of the dictionary. You can even re-assign the name get_ext if you don't have use for the dictionary behind the curtain.

>>> get_ext = get_ext.get
>>> get_ext('video')
['mp4', 'mkv']

This function will return None per default if you enter an unknown key:

>>> x = get_ext('binary')
>>> x is None
True

If you want a KeyError instead for unknown keys, assign to get_ext.__getitem__ instead of get_ext.get.

If you want a custom default-value one approach is to wrap the dictionary inside a function. This example uses an empty list as the default value.

def get_ext(file_type):
    types = {'text': ['txt', 'doc'],
             'audio': ['mp3', 'wav'],
             'video': ['mp4', 'mkv']
    }

    return types.get(file_type, [])

However, @omri_saadon gave the valid remark that the types = ... assignment is performed every function call. Here's what you can do to get around that if this bothers you.

class get_ext(object):
    def __init__(self):
        self.types = {'text': ['txt', 'doc'],
                      'audio': ['mp3', 'wav'],
                      'video': ['mp4', 'mkv']
        }

    def __call__(self, file_type):
        return self.types.get(file_type, [])

get_ext = get_ext()

You can use get_ext like a regular function from here on, because in the end callables are callables. :)

Note that this approach - besides the fact that self.types is only created once - has the considerable advantage that you can still easily change the file types your function recognizes.

>>> get_ext.types['binary'] = ['bin', 'exe']
>>> get_ext('binary')
['bin', 'exe']