Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

How to restrict python overload to not being a literal

Is it possible in python typing to restrict a function argument to be not a literal? My usecase is the following. I have code like this example.

class Test(Enum):
    A = 'A'
    B = 'B'
 
@overload
def test(test: Literal[Test.A], foo: int) -> None: ...
 
@overload
def test(test: Literal[Test.B], foo: str) -> None: ...

def test(test: Test, foo: int | str) -> None:
    print(test)
    print(foo)

This effectively makes the type of the argument foo dependent on the enum value of argument test. Lets assume that this makes sense for some implementation of the function.

The problem arises when I want to supply a variable as argument to the function test. Since a variable is not a literal mypy will complain about the wrong type and a missing overload. One option to solve this is an additional overload like so:

@overload
def test(test: Test, foo: int | str) -> None: ...

Then code giving a variable for parameter test will type check. But wrong code like for example test(Test.A, 'abc') will also type check. If will not match the overload given for the literal Test.A but it will match the additional overload.

Is it possible to restrict the additional overload only to non Literals? My goal is that the type checker will spot an error if args are given as literals, but not for variables since this is out of reach for the type checker. In that case it must be ensured by the caller that only appropriate arguments are passed.

If there is a complete other approach to write that code for the type checker this would also be a solution to my problem.

like image 291
loewexy Avatar asked Dec 05 '25 15:12

loewexy


1 Answers

Functions are first-class objects, so where this kind of single-dispatch overloading makes senses in languages that lack them, a simple mapping of functions makes more sense in Python.

class Test(Enum):
    A = 'A'
    B = 'B'
 
def test_a(foo: int) -> None: ...
 
def test_b(foo: str) -> None: ...


tests = {Test.A: test_a, Test.B: test_b}

Any call like test(something, ...) can be replaced by a dict lookup tests[something](...). But when you already have a statically known first argument, that's really little different from a known function name. (test_a(...) instead of test(Test.A, ...)).

like image 85
chepner Avatar answered Dec 09 '25 00:12

chepner



Donate For Us

If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!