Python function that accepts either an integer or a list of integer

Question

I need a function, that takes two arguments a (an integer) and b (either an integer or a list of integer), to return True if a == b or a in b.

This can be easily done by putting a simple or between the two conditions, but I wanted to know if it was possible to do it in a single condition.

I thought of doing:

• a in list(b) but it is not working because integers are not iterable,
• a in [b] but it is not working if b is already a list.

Answer 1

Ask for forgiveness - not permission! Let's assume b is a list, and if it's not then it must be an int:

def foo(a, b):
    try:
        return a in b
    except TypeError:
        return a == b

If you insist on doing it with one expression/condition, we can actually use regular expressions here:

import re

def foo(a, b):
    return bool(re.search(rf"\b{a}\b", str(b)))

---

A different approach would be to change the signature of the function to take a variable amount of variables and instead of passing it a list, it should be unpacked:

def foo(a, *b):
    return a in b

And to call it:

>>> foo(4, 4)
True
>>> foo(4, *[4, 5])
True
>>> foo(4, *[5, 3])
False
>>> foo(4, 3)
False