Very often when coding in Python (or any language where type is not specified) I run into the following problem:
I want to write a function that can accept either a thing, or a container of such things in general. For example consider this subset function:
def subset(df, months=all_months): return df[pd.DatetimeIndex(df['Timestamp']).month.isin(months)]
I would now like to be able to call this both with a list of months, or a single month. The second case however is only possible by something ugly like:
subset(df, months=["June"]) # Why not just subset(df, "June")?
I can think of 2 solutions, neither of which feels elegant to me:
I can use
isinstanceand treat scalars different that containers. This is not the best because I type check and also there are cases where the code in the function is complex and needs major revision when treating scalar instead of list
I could cast the argument to list:
months = [months]. Again I would be using the
Is there are better looking solution than the ones I thought of? Even better, is there a well known solution that I am not aware of?
Both answers provided acceptable solutions. I decided to stick with the one proposed by @Doc Brown because I can easily isolate the type checking into a decorator and want to avoid introducing unwrapping into the API for workplace specific reasons.