返回类型取决于 Python 中的输入类型

Question

The following function in Python can returns a list no matter if the input is a list , a numpy.array or a pandas.Series .

What is the pythonic way to write it so that the output type is the same as the input type?

def foo(input):
    output = []

    output.append(input[0])

    for i in range(1, len(input)-1):
        if (some condition):
            output.append(input[i])

    output.append(input[-1])

    return output

Answer 1

In general, you can't do this without making a lot of assumptions about what the input is .

The first step would be to make sure that output is the right type, not a list.

output = type(self)()

However, this assumes that whatever type your input is, you can create an instance by calling it with no arguments.

Next, you have to restrict yourself to operations on output that are supported by all expected inputs. Not all iterables support an append method ( set , for instance, uses add , not append ), and not all iterables support __getitem__ (a generator, for instance). This means that you can't generalize your function too much; you always have to keep in mind which types of input you will support.

Alternatively, if the set of types you want to support can create an instance from a list, you can let output = [] stand, but convert it just before returning:

return type(self)(output)