Typehint Union Dictionary branching error-CodePudding

I want to implement a function like the one below, but it throws a type hint warning.

def test(flag: bool)->Dict[str, int]| Dict[str, str]:
    a: Dict[str, str]| Dict[str, int] = {}
    if flag:
        a['a'] = 1
    else:
        a['a'] = 'hello'
    return post_process(a)

With the following warning by Pylance:

Argument of type "Literal[1]" cannot be assigned to parameter "__value" of type "str" in function "__setitem__"
  "Literal[1]" is incompatible with "str"PylancereportGeneralTypeIssues

I know the following is a solution, but it is not very semantically pleasing. Since semantically, a and b are the same, but it appears to be different.

def test()->Dict[str,int]|Dict[str, str]:
    flag = True
    if flag:
        a: Dict[str, int] = {'a': 1}
        return a
    else:
        b: Dict[str, str] = {'a': 'hello'}
        return b
    if flag:
        return post_process(a)
    else:
        return post_process(b)

Is there a solution that looks like the code below, that does not throw type hint warnings?:

def test(flag: bool)->Dict[str, int]| Dict[str, str]:
    a: Dict[str, str]| Dict[str, int] = {}
    if flag:
        a: Dict[str, str]
        a['a'] = 1
    else:
        a: Dict[str, int]
        a['a'] = 'hello'
    return post_process(a)

Note that if a weren't a dictionary and a literal type it could work:

def test(flag: bool)-> int|str:
    a: int| str
    if flag:
        a = 1
    else:
        a = 'hello'
    return post_process(a)

I want an implementation that do not contradict with the type hint system, without sacrificing any of the readability of the actual code. Thank you.

I know type hint warnings are not a part of the python protocol, but I am using Pylance linting system.

CodePudding user response：

Edit: Since your comment outlined that you truly want the return value to be only ever a dict of strings or a dict of ints, and never a mix of both, you can use the @overload decorator to determine return types based on the boolean flag:

from typing import Dict, overload, Literal


@overload
def test(flag: Literal[True]) -> Dict[str, int]:
    ...


@overload
def test(flag: Literal[False]) -> Dict[str, str]:
    ...


def test(flag):
    a = {}
    if flag:
        a["a"] = 1
    else:
        a["a"] = "hello"
    return a


b = test(True)  # Intellisensed as Dict[str, int]

c = test(False)  # Intellisensed as Dict[str, str]

Original Post:

I was able to recreate your error and resolve it with Dict[str, str | int] like so:

def test(flag: bool)->Dict[str, int | str]:
    a: Dict[str, int | str] = {}
    if flag:
        a['a'] = 1
    else:
        a['a'] = 'hello'
    return a

Notice that if you swap the order of the int and str in a: Dict[str, int | str] = {}, then either assignment of a['a'] will fail. This helps indicate that your type hint suggests that the return value is either ONLY a dict of strings, or ONLY a dict of ints, but when intellisense reads your code, it sees that there's a possibility of a variable dict type, that of strings AND integers. If you wanted to type the return to be either a dict ONLY containing strings based on a flag, or ONLY containing integers based on the same flag, then check out the @overload decorator from the typing library.