'How to type hint a function returning a generic type and conformance to a Protocol in Python

If I have a function taking in an instance of type T, and outputting that same instance but modified so it additionally conforms to a Protocol, how should I type hint that? My main goal is to allow my IDE (VSCode) to know that the returned object has all attributes from the original object + all attributes defined in the 'added' protocol.

Specifically, I'm trying to add 'metadata-awareness' to a bunch of objects coming from a 3rd-party library in the following way (code simplified a bit to focus on the type-hinting only):

@runtime_checkable
class HasMetadata(Protocol):
    """Protocol for objects which have a `metadata` property."""

    @property
    def metadata(self) -> Dict[str, Any]: ...


class MetadataMixin(HasMetadata):
    """A mixin addig a `metadata` property to an existing class."""

    def __init__(self, *args, metadata: Optional[Dict[str, Any]] = None, **kwargs):
        """
        Adds a metadata dictionary to this instance and calls the superclass ctor.

        Args:
            metadata (optional): The metadata to associate to this object. Defaults to
                None, meaning an empty dictionary.
        """
        super().__init__(*args, **kwargs)
        self._metadata = metadata or {}

    @property
    def metadata(self) -> Dict[str, Any]:
        return self._metadata


# Here's what I fail to type-hint correctly
T = TypeVar('T')
def inject_metadata(obj: T, metadata: Optional[Dict[str, Any]] = None) -> ??? # In swift I'd write something like `T & HasMetadata`
    """Dynamically injects metadata in the given object."""
    if isinstance(t, HasMetadata):
        raise TypeError('Cannot inject metadata in an object that already has it.')

    # Create a new type which acts like the old but also has `MetadataMixin` in its
    # mro, then change the type of the given object to this new type. This will give
    # the object the `metadata` property.
    old_type = type(t)
    new_type_name = f'_{old_type.__name__}WithMetadata'

    # Retrieve the new type if it was already created, or create it otherwise.
    # Simplified for demo purposes.. just always create the new type.
    # Also don't mind about the Mixin added last as a base class; I know technically it
    # should come first to be a true Mixin but that doesn't matter here.
    new_type = type(new_type_name, (old_type, MetadataMixin), dict(old_type.__dict__))
    obj.__class__ = new_type

    # Now set the actual metadata.
    metadata = metadata or {}
    setattr(obj, '_metadata', metadata)

    return obj

I know I can make this work for specific - known up front - classes by creating helper types, e.g.:

class Bar:
    pass

class Baz(Bar, HasMetadata):
    @property
    def metadata(self) -> Dict[str, Any]: ...

def inject_metadata(obj: Bar, ...) -> Baz:
    ...

but that only works if I know specifically which type(s) I want to inject metadata in. Ideally I'd make this work with any T.



Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source