I know this is very similar to a few other questions, but I can't quite get this function to work correctly.

def flatten(*args):
    return list(item for iterable in args for item in iterable)

The output I'm looking for is:

flatten(1) -> [1]
flatten(1,[2]) -> [1, 2]
flatten([1,[2]]) -> [1, 2]

The current function, which I from another SO answer doesn't seem to produce correct results at all:

>>> flatten([1,[2]])
[1, [2]]

I wrote the following function which seems to work for 0 or 1 levels of nesting, but not deeper:

def flatten(*args):
    output = []
    for arg in args:
        if hasattr(arg, '__iter__'):
            output += arg
        else:
            output += [arg]
    return output