pythonpython-typing

Map one type to another and make type checker understand


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.


Solution

  • 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]]).