Function with generic return type causing typing issues in calling functions with specific type

Question

I have a generic lookup function, that mostly returns TypeA, but sometimes can return TypeB:

Types = Union[TypeA,TypeB]
def get_hashed_value(
    key:str, table: Dict[str,Types]
) -> Types:
  return table.get(key)

and I use it in two less-generic functions:

def get_valueA(key: str) -> TypeA:
  return get_hashed_value(key, A_dict)  # A_dict: Dict[str, TypeA]

and

def get_valueB(key: str) -> TypeB:
  return get_hashed_value(key, B_dict)  # B_dict: Dict[str, TypeB]

what is the best way to handle typing on this?

since get_hashed_value can return either TypeA or TypeB, the return statement in the get_* functions throws a typing exception (during my linting)

1. there’s more logic in these methods, and I need the separate get_* functions, so I can’t just collapse all the usages
2. it would be really nice to have explicit return types on the get_* functions
3. it feels like a bad practice to duplicate get_hashed_value, just to get around the typing issue
4. it feels bad to just ignore type everything get_hashed_value is called

Thanks for your help! Also I am sure this has been asked before, but I had trouble finding the answer. :\

Answer 1

Interestingly, this doesn't return a type warning for me (in Pycharm). I'm not sure why it isn't warning on what's comparable to a "downcast", but Pycharm isn't perfect.

Regardless, this seems like a job that's more suited for a TypeVar than a Union:

from typing import TypeVar, Dict

T = TypeVar("T", TypeA, TypeB)  # A generic type that can only be a TypeA or TypeB

# And the T stays consistent from the input to the output
def get_hashed_value(key: str, table: Dict[str, T]) -> T:
    return table.get(key)

# Which means that if you feed it a Dict[str, TypeA], it will expect a TypeA return
def get_valueA(key: str) -> TypeA:
    return get_hashed_value(key, A_dict)

# And if you feed it a Dict[str, TypeB], it will expect an TypeB return
def get_valueB(key: str) -> TypeB:
    return get_hashed_value(key, B_dict)