StatusBase.register decorator obscures the type of decorated classes

[james-garner-canonical]

StatusBase.register is typed as follows:

class StatusBase:
    ...
    @classmethod
    def register(cls, child: Type['StatusBase']):
        ...
        return child

And used like so:

@StatusBase.register
class ActiveStatus(StatusBase):
    ...

Which leads to:

reveal_type(ActiveStatus)   # Type of "ActiveStatus" is "type[StatusBase]"

I believe the solution is to use a type variable with a bound of Type[StatusBase] instead and annotate the return type, e.g.

_StatusBaseType = typing.TypeVar('_StatusBaseType', bound='Type[StatusBase'])

class StatusBase:
    ...
    @classmethod
    def register(cls, child: _StatusBaseType) -> _StatusBaseType:
        ...
        return child

Which results in the expected:

reveal_type(ActiveStatus)   # Type of "ActiveStatus" is "type[ActiveStatus]"

[benhoyt]

This seems reasonable to me -- thanks. Does this also improve the charmer's user experience or autocompletion when using status classes, or is it just a correctness fix?

[james-garner-canonical]

In my case I had a function returning a status, and it was confusing to see that the return type was always StatusBase. However, I only figured out what was wrong when I was investigating a different way to improve ux via types.

I'll create a PR for this later.

[dimaqq]

An alternative approach would be:

class ActiveStatus(StatusBase): ...
class BlockedStatus(StatusBase): ...

StatusBase.register(ActiveStatus)
StatusBase.register(BlockedStatus)

in other words, to avoid the decorator.

[james-garner-canonical]

That's quite an elegant solution, and does avoid needing to do any fancy typing in StatusBase.register, as it no longer interferes with the public classes. The only arguments against it that I can think of right now are the ones in favour of adding decorators to python in the first place -- it's easier to see that something is happening up at the definition line than later on in the code base.

As mentioned in discussion on the PR, we could also remove the decorator by doing the registration in StatusBase.__init_subclass__. What about calling register from there to avoid having StatusBase.register calls elsewhere in the file?

[tonyandrewmeyer]

I think __init_subclass__ probably is the 'correct' solution here, and that the existing code either had to work on Pythons older than 3.6 or the author wasn't familiar with it.

I would be ok with any of:

• leave the typing on the register decorator as it is in main, deprecate (documentation & warnings.warn) it, and use __init_subclass__ (this is still a behaviour change, since previously you could not register a subclass, but less breaking and it's the intended use-case)
• improve the typing (as in the branch now), deprecate it, and use __init_subclass__
• improve the typing (as in the branch now), c'est la vie for the rest

I think removing the register method entirely, even though it's not meant to be used by others, and even though it's probably not actually useful by others, goes too far. I think I may be the most cautious in the team about these things, though :)