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.
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']