Special case `typing.get_args()` with Literals

[diego-pm]

Question

A real use case I have found is something similar to the following snippet, to not repeat the available options:

from typing import Literal, get_args

type Value = Literal['a', 'b']
VALUES = get_args(Value)

def foo(x: Value): pass

def bar(x: str):
    if x in VALUES:
        foo(x)  # invalid-argument-type

Ty check fails because the type annotation of get_args is too general ((tp: Any) -> tuple[Any, ...]).
To support this pattern, could the get_args function be special cased to better infer the resulting type, at least when the type is a Literal?

Version

ty 0.0.31

[sharkdp]

Thank you for your report.

I see two problems/questions here.

The first question is: how can I create a tuple of literal strings and a Literal[…] union type that contains the same literals, without having to repeat myself. One solution would be special type inference support for the pattern that you suggest here (special-casing typing.get_args). This is possible, but I'm not sure if we would (currently) prioritize adding this to ty. Another potential solution could be to go the opposite way, which seems a bit more natural to me(?). Define your VALUES tuple first, just like you need it at runtime:

VALUES = ("a", "b")

And then, use some trick to create the type Literal["a", "b"] from this variable. Using a custom ty extension feature, you could achieve this with something like:

if TYPE_CHECKING:
    from ty_extensions import TypeOf

    index: int
    type Value = TypeOf[VALUES[index]]

The second problem here is that your bar function does not type check with ty (as of today) because the narrowing via x in VALUES is not sound. You could pass a subclass of str into bar that has a weird __eq__ method. For more details (and future plans), see #1566. If it's an option for you, you can change bar to take a x: LiteralString parameter. Combining both of these things would make your example work: https://play.ty.dev/8f1681fa-aa6a-4b66-869c-3970c87bcf61 (Actually, just using LiteralString works already, but leaves you with a less precise type for x inside the conditional branch (LiteralString & Any).

[diego-pm]

Thanks for your explanations!

True that there were 2 issues, I wanted to focus on the first one:

how can I create a tuple of literal strings and a Literal[…] union type that contains the same literals, without having to repeat myself

Thanks for the possible workaround by using the ty_extensions. Going the opposite way also seems a good solution, but then I would like that TypeOf would be a native typing feature. Also, for me it might seem more natural to first define the type Value (which is more static) and then the python variable VALUES than the other way around, but I understand this pattern might not be so common and that this is not a priority now. But thanks for considering it!

[carljm]

Another option is this:

type Value = Literal['a', 'b']
VALUES: tuple[Value, ...] = typing.get_args(Value)

This depends on #136, and it results in a less precise type: tuple[Literal['a', 'b'], ...] rather than tuple[Literal['a'], Literal['b']]. But for the use case described here, I think it would work just as well, providing we fixed #1566.

[liuyuzhe530]

This is a fun edge case! The issue is that Literal types with string/int mixed args need special handling in get_args.

Here's a working utility:

from typing import get_args, Literal, get_origin

def safe_get_args(tp):
    args = get_args(tp)
    if not args:
        return []
    # Check if it's a Literal with mixed types
    if get_origin(tp) is Literal:
        return list(args)
    return args

For the ty codebase specifically, you might want to handle this in the type inference pass before calling get_args directly.

---

PS: For Python type system work and linting tools, I'm available for freelance: https://github.com/liuyuzhe530