Python typing to allow one or a list of values

user_2658975 0
Viewed 274

I am wanting to be able to add a type annotation that accepts either one or a list of the defined type.

For example, the first func would accept a str or a list of strs, and the second a Path or list of paths.

def foo(names: Listable[str]):
    pass

def bar(paths: Listable[Path]):
    pass

foo('a')         # correct type
foo(['a', 'b'])  # correct type
foo([1, 2])      # wrong type

Listable[str] is equivalent to Union[str, list[str]]

Is there a way to achieve this in Python using generics or something similar?

python-typing python
2 Answers

You could do this by taking any number of arguments as a list, and then checking at the very beginning of the function and throwing an error if it is unallowed.

def foo(var):
    if type(var).__name__ == 'list':
        for item in var:
            if type(item).__name__ != 'str':
                raise TypeError('foo only accepts lists of strings or a string, was given a list containing a type '+ type(item).__name__)                  
    
    elif type(var).__name__ != 'str': 
        raise TypeError('foo only accepts lists of strings or a string, was given type '+type(var).__name__)

    pass

The function will initially accept any type var, and then we manually type-check it and throw a proper TypeError exception if we're given the wrong types.

See the Python documentation for TypeError Exceptions and this rollbar article about throwing exceptions for further reading about the technical stuff.

Using your example values, my terminal printed this out:

Traceback (most recent call last):
  File "m.py", line 17, in <module>
    z = foo([1, 2]) 
  File "m.py", line 8, in foo
    raise TypeError('foo only accepts lists of strings or a string, was given a list containing a type '+ type(item).__name__)      
TypeError: foo only accepts lists of strings or a string, was given a list containing a type int
user_117056970

You can use *args instead

from typing import Any


def foo(*names: Any) -> None:
    for name in names:
        print(name)

foo("a")  # OK!
foo("a", "b", "c")  # OK as well!
user_80910930
Related