CODEWITHSUNDEEP
pythonlist
Ryan Kang
Ryan Kang

Reputation: 61

How to write a python function which returns a list when an input is a list and returns a non-list value when an input is a non-list?

Language: Python 3.x

I want to write a function which returns list type only if input type is list, but in a prettier way hopefully using decorator.

I could do it so as below but how would you do differently if you have to write this pattern many times?

from datetime import datetime

def my_datetime_parse_function(datetime_list, format=None):

    if isinstance(datetime_list, list):
        return [ datetime.strptime(dt, format) for dt in datetime_list ]

    elif isinstance(datetime_list, str):
        return datetime.strptime(datetime_list, format)

    else:
        raise ValueError

Upvotes: 2

Views: 60

Answers (1)

willeM_ Van Onsem
willeM_ Van Onsem

Reputation: 477492

You could make a decorator [PEP-318] that looks like:

from functools import wraps

def map_list(f):
    @wraps(f)
    def decorator(item_list, *args, **kwargs):
        if isinstance(item_list, list):
            return [f(item, *args, **kwargs) for item in item_list]
        else:
            return f(item_list, *args, **kwargs)
    return decorator

then you can apply the decorator on a "simple" function:

@map_list
def my_datetime_parse_function(dt, format=None):
    return datetime.strptime(dt, format)

this will always inspect the first parameter, and in case it is a list, perform a mapping, the other parameters are just passed as extra parameters to the "simple" function. The decorator can thus be applied to all functions you want to "decorate" with this behavior.

The simple function can thus focus on how to "map" a single item, and the decorator will take care of the list case.

For example:

>>> my_datetime_parse_function('1/1/2018', format='%d/%m/%Y')
datetime.datetime(2018, 1, 1, 0, 0)
>>> my_datetime_parse_function(['1/1/2018', '10/1/2018', '11/7/1302'], format='%d/%m/%Y')
[datetime.datetime(2018, 1, 1, 0, 0), datetime.datetime(2018, 1, 10, 0, 0), datetime.datetime(1302, 7, 11, 0, 0)]

You can extend the decorator to work with other collections as well (tuples, sets, etc.).

Upvotes: 2

Related Questions