Map one type to another and make type checker understand

Question

I have the following piece of code:

from typing import Any

class RawA:
  pass

class A:
  pass

class RawB:
  pass

class B:
  pass

# Example of a kind of mapping (that doesn't work)
mapping = {
  RawA: A,
  RawB: B
}


def unraw[TRaw: Any](obj: TRaw) -> Any  # -> the unraw type
  ...

Now how would I create some sort of mapping (of course a dictionary like here won't work) that the type checker then uses to find the appropriate return type?

For example, if I pass an instance of RawA into my function, the type checker understands the return value will be an instance of the type RawA is mapped to, in this case A.

I know could just exhaust all possibilities via overloads, but that is way to cumbersome and error prone for my case.

Answer 1

You could set a common attribute on all "raw" classes and refer to them using a protocol:

_{(playgrounds: Mypy, Pyright, Pyrefly)}

class A: ...
class B: ...

class RawA:
    unraw = A  # ClassVar
class RawB:
    unraw = B  # ClassVar

class Raw[T](Protocol):
    @property
    def unraw(self) -> type[T]: ...

def unraw[T](obj: Raw[T]) -> T: ...

reveal_type(unraw(RawA()))  # A
reveal_type(unraw(RawB()))  # B

unraw cannot be explicitly marked a class variable because, as far as I'm aware, there is no way to express the "instance to class-level protocol" relationship (i.e., Instance[Raw[T]]).