setattr type conversions

[hmc-cs-mdrissi]

One interesting type behavior I stumbled on with one library is it auto converts certain attributes to a different type when set in that class/subclass. Some example code,

# Some library
class Model:
  ...

class MyModel(Model):
  def __init__(self, x: list[str]):
    self.x = x
    print(type(self.x)) # Not list, but instead lib.ListWrapper

How would I write a stub for Model class that indicates this behavior? Is it possible to describe? Exact behavior is roughly setting attribute to a list becomes ListWrapper, dict to DictWrapper, other types are left untouched. While list/ListWrapper behave similarly there are cases where difference matters (runtime type introspection libraries being one).

[AlexWaygood]

It's hard to say for certain without knowing which library you're talking about, but it seems likely that the library is accomplishing this kind of magic using some kind of descriptor. Luckily, mypy has a pretty good grasp of the descriptor protocol, so you should be able to model this kind of behaviour like this:

### STUB ###

from typing import Any, TypeVar

_T = TypeVar("_T")

class ListWrapper: ...

class MyDescriptor:
    def __get__(self, obj: _T, objtype: type[_T] | None = ...) -> ListWrapper: ...
    def __set__(self, obj: Any, value: list[str]) -> None: ...

class Model:
    x: MyDescriptor
    
### RUNTIME ###

class MyModel(Model):
    def __init__(self, x: list[str]) -> None:
        self.x = x
        reveal_type(self.x)  # Revealed type is "__main__.ListWrapper"
        self.x = 6  # error: Incompatible types in assignment (expression has type "int", variable has type "List[str]")

See it in action on mypy playground here.

[hmc-cs-mdrissi]

The issue is this isn't just for x, but for all attributes defined by Model. The implementation of Model looks like,

class Model:
  def __setattr__(self, name, value):
    if type(value) == list:
     value = ListWrapper(value)
    elif type(value) == dict:
      value = DictWrapper(value)
    
    super().__setattr__(name, value)

The actual library is tensorflow and real implementation is here, but I think that simplified one captures the main idea.

[AlexWaygood]

You could make it more extensible by making the (fictional, as it happens, since the runtime isn't actually using descriptors) descriptor generic:

### STUB ###

from typing import Any, Generic, TypeVar

_T = TypeVar("_T")
_GetType = TypeVar("_GetType")
_SetType = TypeVar("_SetType")

class ListWrapper: ...
class DictWrapper: ...

class _TensorflowDescriptor(Generic[_GetType, _SetType]):
    def __get__(self, obj: _T, objtype: type[_T] | None = ...) -> _GetType: ...
    def __set__(self, obj: Any, value: _SetType) -> None: ...

class Model:
    x: _TensorflowDescriptor[ListWrapper, list[str]]
    y: _TensorflowDescriptor[DictWrapper, dict[Any, Any]]

### RUNTIME ###

class MyModel(Model):
    def __init__(self, x: list[str]) -> None:
        self.x = x
        reveal_type(self.x)  # Revealed type is "__main__.ListWrapper"
        self.x = 6  # error: Incompatible types in assignment (expression has type "int", variable has type "List[str]")
        self.y = {'foo': 'bar'}
        reveal_type(self.y) # Revealed type is "__main__.DictWrapper"

Playground link here.

[hmc-cs-mdrissi]

Hmm, the approach/descriptors is cool and nice to learn about, but the keys here can be any field and having to mark them is a bit awkward. Mainly because this requires the person writing subclass to use descriptor to mark there fields. Model/Layer are basic classes that we currently have subclassed 100 times in our codebase and I'm unsure if idiomatic tensorflow should be always use these descriptors yourself. Other thing is it means any time a teammate/other engineer in department writes a layer they should do same descriptor. I would feel comfortable writing this code for myself, but I feel like I'd be asking everyone else to either learn descriptors or just treat TensorflowDescriptor[...] as a black box that handles this for you. Maybe having others treat it as a blackbox is right way. In practice it would look like

class MyModel1(Model):
  def __init__(self, a, b, c):
    self.a: TensorflowDescriptor = a
    self.b: TensorflowDescriptor = b
    self.c: TensorflowDescriptor = c

class MyModel2(Model):
  def __init__(self, blue, green, red):
    self.blue: TensorflowDescriptor = blue
    self.green: TensorflowDescriptor = green
    self.red: TensorflowDescriptor = red

Mostly because x/y are arbitrary here and each model subclass normally defines it's own unique attributes.

The other weakness is it means most users of tensorflow are unlikely to do this for their own custom attributes.

Although in practice difference between list vs ListWrapper is mild enough most people may never notice differences.

[AlexWaygood]

Right, I see the issue 😕 unfortunately, that kind of behaviour just looks too dynamic to me to be expressible in a generalised way in the current type system. The generic-descriptor hack is the best I've got — but maybe somebody else can think of something better 🙂