Python generic that adds typed functionality to extend existing types incorrectly typed

Question

I am interested in creating a generic to add a fixed functionality to different types in python. I have a semi working solution, but it's tying is problematic.

Here's a simplified version of what I'm trying:

from __future__ import annotations
from typing import Generic, Set, TypeVar, Type, Callable

def CreateTaggedType(base_type: type):
    T = TypeVar("T")

    class TaggedTypeGenerator(Generic[T]):
        def __init__(self):
            pass
    
        def tag_provider(self, parent: Type[T]):
            class Imp(parent):
                def __init__(self,  **args):
                    super().__init__(**args)
                    self.tags: Set[str] = set()
    
                def add_tag(self, tag):
                    self.tags.add(tag)
    
                def remove_tag(self, tag):
                    if tag in self.tags:
                        self.tags.remove(tag)
                        return True
                    return False
    
            return Imp
    return TaggedTypeGenerator[base_type]().tag_provider(parent=base_type)

class Counter:
    def __init__(self, count=0):
        self.count = count

    def inc(self):
        self.count += 1
        return self.count


# Create an alias for a new type of tagged Counter
TaggedCounter = CreateTaggedType(Counter)
tc = TaggedCounter(count=2)

print(tc.inc())  # should print "3"
tc.add_tag("foo")
print(tc.remove_tag("bar"))  # should print "False"
print(tc.remove_tag("foo"))  # should print "True"
print(tc.inc())  # should print "4"

The output:

3
False
True
4

Technically, this works, but the type isn't correctly inferred. This causes auto-completion to not work properly. I only get aut-ocompletion for the added functionality, and not for anything related to the base class.

How can I make it so that the type is correct?

I tried this variation using abc.ABCMeta with no effect:

from __future__ import annotations
import abc
from typing import Callable, Generic, Set, Type, TypeVar

T = TypeVar("T")

class TaggedTypeGenerator(Generic[T]):
    def __init__(self, parent: Type[T]):
        self.parent = parent
        
        class _imp_(self.parent):
            def __init__(self,  **args):
                super().__init__(**args)
                self.tags: Set[str] = set()
    
            def add_tag(self, tag):
                self.tags.add(tag)
    
            def remove_tag(self, tag):
                if tag in self.tags:
                    self.tags.remove(tag)
                    return True
                return False
        
        self.Imp = _imp_

class Counter:
    def __init__(self, count=0):
        self.count = count

    def inc(self):
        self.count += 1
        return self.count

def CreateTaggedTypeGenerator(base_type: T) -> TaggedTypeGenerator[T]:  
    return TaggedTypeGenerator[base_type](base_type)

# Create an alias for a new type of tagged Counter
TaggedCounterGererator = CreateTaggedTypeGenerator(Counter)

class counter_meta(metaclass=abc.ABCMeta):
    pass

counter_meta.register(Counter)

@counter_meta.register
class TaggedCounter(TaggedCounterGererator.Imp):
    pass

tc = TaggedCounter(count=2)

print(tc.inc())  # should print "3"
tc.add_tag("foo")
print(tc.remove_tag("bar"))  # should print "False"
print(tc.remove_tag("foo"))  # should print "True"
print(tc.inc())  # should print "4"

This also procudes the expected output. However, not only does this not solve the typing issue, it makes it worse. With this, auto-complete doesn't work for the base class and the added functions.

Any suggestions would are welcome.

NOTE: My editor is VS Code, I don't think it matters much here, but...

Answer 1

The first solution will make autocomplete almost happy:

from __future__ import annotations
from typing import Callable, Generic, Set, Type, TypeVar

T = TypeVar("T")

class TaggedMixin:
    """Extract tagging functionality into separate class to get rid of unnecessary dynamics."""

    def __init__(self,  **args):
        super().__init__(**args)
        self.tags: Set[str] = set()

    def add_tag(self, tag):
        self.tags.add(tag)

    def remove_tag(self, tag):
        if tag in self.tags:
            self.tags.remove(tag)
            return True
        return False


def make_tagged_type(parent: Type[T]) -> Type[T]:
    """Annotate return type to be the same as input.

    This way type checker will think that make_tagged_type(list) is
    class that has all attrs/methods of list (and nothing else).
    """
    return type(f'Tagged{parent.__name__}', (TaggedMixin, parent), {})
    # this is (almost) equivalent to
    # class _Tagged(TaggedMixin, parent):
    #    pass
    # _Tagged.__name__ = f'Tagged{parent.__name__}'
    # return _Tagged


class Counter:
    def __init__(self, count=0):
        self.count = count

    def inc(self):
        self.count += 1
        return self.count

tc = make_tagged_type(Counter)(count=2)

print(tc.inc())  # should print "3"
tc.add_tag("foo")
print(tc.remove_tag("bar"))  # should print "False"
print(tc.remove_tag("foo"))  # should print "True"
print(tc.inc())  # should print "4"

Now this works at runtime, but fails type checking (because add_tag and remove_tag are added dynamically and not visible to the checker - but all other methods are accessible). Probably autocomplete will think the same.

Type you need to declare (which is subtype of both TaggedMixin and Type[T]) is called Intersection[TaggedMixin, Type[T]] from typing theory point of view. There was a large discussion around support for Intersection type which is not implemented neither in mypy nor in other checkers like pylance (and, what's more important, is isn't added to typing by PEP's, so is not approved for python itself).

Related mypy issue

I don't know what type checker does VS Code use (and whether it is related to autocomplete). Looks like they have a thing called "IntelliSense" which doesn't have any plugin capabilities, so this is probably the end of the story about autocomplete and dynamic code together. If autocomplete is really important for you (or you don't want to play with mypy plugins and have just a couple of classes with that functionality), then "now is better than never" and "explicit is better than implicit":

from __future__ import annotations
from typing import Callable, Generic, Set, Type, TypeVar

T = TypeVar("T")

class TaggedMixin:
    """Extract tagging functionality into separate class to get rid of unnecessary dynamics."""

    def __init__(self,  **args):
        super().__init__(**args)
        self.tags: Set[str] = set()

    def add_tag(self, tag):
        self.tags.add(tag)

    def remove_tag(self, tag):
        if tag in self.tags:
            self.tags.remove(tag)
            return True
        return False


class Counter:
    def __init__(self, count=0):
        self.count = count

    def inc(self):
        self.count += 1
        return self.count

class TaggedCounter(TaggedMixin, Counter):
    pass

tc = TaggedCounter(count=2)

print(tc.inc())  # should print "3"
tc.add_tag("foo")
print(tc.remove_tag("bar"))  # should print "False"
print(tc.remove_tag("foo"))  # should print "True"
print(tc.inc())  # should print "4"

This will work both for autocomplete and type checking. Also this is almost DRY: you repeat only class declaration, which is fairly reasonable. I would recommend using this implementation is also the most simple, if you don't have hundreds of classes to modify and your classes are not dynamic.

Please tell me if there is a reason to use dynamic declarations here. Also, don't forget about Any, it is sometimes better just to break type checking strictness in order to make code more readable/efficient. There is (at least) one more solution for mypy I can elaborate on: custom plugin. However, it's really special thing, that is usable mostly for third-party modules that provide access to some types which are not expressible with "normal" typing - you probably don't want to use plugin to remove, say, 5 type errors: ignoring them (# type: ignore is your friend too!) will be much simpler. Just leave a comment explaining ignorance and document class factory to explicitly mention dynamic attributes.