pythonpython-typingclass-variablespyright

Make type checker understand that class and instance attributes share names but have diffent types dynamically


I have a piece of code in Python that in essence looks and works the following way:

from typing import Any

class Column:
    def __init__(self, s: str) -> None:
        self.name = s

    def __repr__(self):
        return f'Column(name={self.name!r})'

class Meta(type):
    def __new__(cls, name: str, bases: tuple[type, ...], namespace: dict[str, Any]) -> type:
        annotations = namespace.get('__annotations__', {})
        namespace.update({k: Column(k) for k in annotations})
        return super().__new__(cls, name, bases, namespace)
        
class Base(metaclass=Meta):
    def __init__(self, data: dict[str, Any]) -> None:
        if self.__class__ == Base:
            raise RuntimeError
        self.__dict__ = data

# Example: inherit from Base and create an object for testing
# (and we assume the data passed is correct).

class Human(Base):
    name: str
    height: int

tom = Human({'name': 'Tom', 'height': 180})

print(Human.height)  # Column(name='height')
print(tom.height)    # 180

Now when we access a field from the class, we get back a Column with the field's name. When we access the field from an instance like tom we get back the actual value for that field of that instance, in our case 180. We also kind of tricked the type checker that these are both valid fields, because it just thinks we are working with class fields. (At least this is the case for Pylance, which is what I use)

Now comes the question: Is there a trick or workaround for the type checker to understand that Human.height is of type Column, without specifying the same field twice like the following I would have to do here:

class Human(Base):
    name: str
    name: ClassVar[Column]
    height: int
    height: ClassVar[Column]

So the goal is to have my type checker think the class has the annotated instance attributes, while all class attributes have the same name but are all of type Column.


Solution

  • I figured out how to achieve this via descriptors:

    from typing import Any, Callable, overload, Never
    
    class Column:
        def __init__(self, name: str) -> None:
            self.name = name
        
        def __repr__(self) -> str:
            return f'Column(name={self.name})'
    
    
    class BaseRow:
        __required_keys__: set[str] = set()
    
        def __init__(self, row: dict[str, Any]) -> None:
            if self.__class__ == BaseRow:
                raise RuntimeError
    
            self.__mapping__ = row
    
    
    class Cell[TRow: BaseRow, TValue: Any]:
        def __init__(self, wrapped: Callable[[TRow], TValue] | Callable[[type[TRow]], TValue]):
            self.__wrapped__ = wrapped
        
        def __set_name__(self, owner: type[TRow], name: str):
            self.name = name
            self.column = Column(name)
    
            if not owner.__required_keys__:
                owner.__required_keys__ = set()
    
            owner.__required_keys__.add(name)
    
        @overload
        def __get__(self, instance: TRow, owner: type[TRow]) -> TValue:
            pass
        
        @overload
        def __get__(self, instance: None, owner: type[TRow]) -> Column:
            pass
    
        def __get__(self, instance: TRow | None, owner: type[TRow]) -> TValue | Column:
            if instance is None:
                return self.column
            
            return instance.__mapping__[self.name] 
    
    # The body of a Cell is never called, so to avoid type checker errors for any type checker we use Never
    def never() -> Never:
        raise RuntimeError()
        
    class Human(BaseRow):
        @Cell 
        def name(self_or_cls) -> str: 
            never()
    
        @Cell
        def height(self_or_cls) -> int:
            never()
    
    tom = Human({'name': 'Tom', 'height': 180})
    
    print(Human.height)  # Column(name='height')
    print(tom.height)    # 180
    

    Through the overloads of __get__ the type checker understands what I want.