'Python typing to allow one or a list of values
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?
Solution 1:[1]
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
Solution 2:[2]
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!
Sources
This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.
Source: Stack Overflow
| Solution | Source |
|---|---|
| Solution 1 | JohnAlexINL |
| Solution 2 | Paweł Rubin |
