Reputation: 177
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.
Upvotes: 15
Views: 11318
Reputation: 20434
If you don't want to define a dictionary as in @timgeb's answer, then you can call locals()
, which gives you a dict
of the variables available in the local scope.
def get_ext(file_type):
text = ['txt', 'doc']
audio = ['mp3', 'wav']
video = ['mp4', 'mkv']
return locals()[file_type]
and a test to show it works:
>>> get_ext("text")
['txt', 'doc']
Upvotes: 12
Reputation: 78770
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']
Upvotes: 34
Reputation: 209
As per the answer by @timgeb I'd use a dictionary, but if you're accessing a lot, care about speed and don't want to define a class you can use caching.
from functools import lru_cache
def get_ext(file_type):
d = {'text': ['txt', 'doc'],
'audio': ['mp3', 'wav'],
'video': ['mp4', 'mkv']}
return d[file_type]
@lru_cache(maxsize=3, typed=False)
def get_ext_cached(file_type):
d = {'text': ['txt', 'doc'],
'audio': ['mp3', 'wav'],
'video': ['mp4', 'mkv']}
return d[file_type]
from timeit import timeit
# non cached
print(timeit(stmt='get_ext("text")', globals={'get_ext': get_ext}))
# 0.48447531609922706 on my machine
# cached
print(timeit(stmt='get_ext("text")', globals={'get_ext': get_ext_cached}))
# 0.11434909792297276
Though for this particular case it's likely overkill and you can just call get on the dictionary directly (cache just builds its own dictionary and does exactly that) you can use this in future for any pure functions which are effectively a computed lookup.
d = {'text': ['txt', 'doc'],
'audio': ['mp3', 'wav'],
'video': ['mp4', 'mkv']}
# 0.05016115184298542
print(timeit(stmt="d['text']",
globals={'d':d,'c':c}))
Upvotes: 1
Reputation: 1028
You can easily use dict with tuple/list values like so:
def get_ext(file_type):
d = {'text': ['txt', 'doc'],
'audio': ['mp3', 'wav'],
'video': ['mp4', 'mkv']}
return d[file_type]
print(get_ext('audio'))
Upvotes: 4
Reputation: 4538
Use dictionary:
def get_ext(file_type):
d = {'text' : ['txt', 'doc'],
'audio' : ['mp3', 'wav'],
'video' : ['mp4', 'mkv']}
try:
return d[file_type]
except KeyError:
return []
get_ext('audio') # ['mp3', 'wav']
returns empty list in case that key does not exists. how ever this is simplest answer that came in my mind , for better answer see @timgeb answer.
Upvotes: 1