String to Literal throws incompatible type error

Question

Running mypy on the below snippet:

from typing import Literal, Final

def extract_literal(d2: Literal["b", "c"]) -> str:
    if d2 == "b":
        return "BA"
    if d2 == "c":
        return "BC"
    
def model(d2_name: str = "b-123") -> None:
    if d2_name[0] not in ["b", "c"]:
        raise AssertionError
    d2: Final = d2_name[0]
    print(extract_literal(d2))

throws:

typing_test.py:17: error: Argument 1 to "extract_literal" has incompatible type "str";
expected "Literal['b', 'c']"  [arg-type]
        print(extract_literal(d2))
                              ^~
Found 1 error in 1 file (checked 1 source file)

For context, d2_name is guaranteed to be either "b-number" or "c-number". Based on its first letter I want to output a different message.

Python version: 3.11
mypy version: 1.9.0

What change would be required to allow mypy to pass?

Answer 1

For context, d2_name is guaranteed to be either "b-number" or "c-number".

How do you know this?

If it's because you are hardcoding d2_name's value somewhere else in your program, then put it in the format you expect to begin with. See mypy playground.

Name: TypeAlias = tuple[Literal["b", "c"], int]

EXAMPLE_D2_NAME_1: Name = ("b", 123)
EXAMPLE_D2_NAME_2: Name = ("c", 456)

def model(d2_name: Name = EXAMPLE_D2_NAME_1) -> None:
    # Note how our assertion is no longer needed, 
    # because we hardcoded the data in our desired format to begin with.
    d2 = d2_name[0]
    print(extract_literal(d2))

model(EXAMPLE_D2_NAME_2)

Otherwise, it sounds like you're obtaining the value externally. In that case,
"Parse, don't validate" is relevant. In order to take advantage of the type system / mypy, it's better to parse data into the format we expect then just check the criteria is met. See mypy playground.

def parse_name(raw: str) -> Name:
    components = raw.split("-")

    if len(components) != 2:
        raise AssertionError("Exactly one hyphen expected.")

    prefix, raw_number = components

    try:
        number = int(components[1])
    except ValueError:   
        raise AssertionError("Expected number after hyphen.")
      
    if prefix in ("b", "c"):
        # Since `mypy` doesn't support literal narrowing yet: https://github.com/python/mypy/issues/16820,
        return prefix, number # type: ignore

    raise AssertionError("Prefix must be b or c.")

external_d2_name = parse_name(os.environ["D2_NAME"])
model(external_d2_name)

Notice how my approach of parsing and your approach of validating both use AssertionErrors, but I additionally return the type-checked data in the happy path, proving to the rest of the code that it is the correct format.

In the comments, TypeGuard was suggested. I would not use this since mypy doesn't actually check that you return what's in the TypeGuard. This mechanism is more intended for variable length containers of items (such as list).