Adds Supports with multiple type args

CHANGELOG.md

@@ -3,6 +3,18 @@
 We follow Semantic Versions since the `0.1.0` release.
 
 
+## Version 0.4.0 WIP
+
+### Features
+
+- Adds support for multiple type arguments in `Supports` type
+
+### Bugfixes
+
+- Fixes that types referenced in multiple typeclasses
+  were not handling `Supports` properly #249
+
+
 ## Version 0.3.0
 
 ### Features

classes/_typeclass.py

@@ -140,7 +140,7 @@
 
 _NewInstanceType = TypeVar('_NewInstanceType', bound=Type)
 
-_StrictAssociatedType = TypeVar('_StrictAssociatedType', bound='AssociatedType')
+_AssociatedTypeDef = TypeVar('_AssociatedTypeDef', contravariant=True)
 _TypeClassType = TypeVar('_TypeClassType', bound='_TypeClass')
 _ReturnType = TypeVar('_ReturnType')
 
@@ -240,7 +240,7 @@ def __class_getitem__(cls, type_params) -> type:
 
 
 @final
-class Supports(Generic[_StrictAssociatedType]):
+class Supports(Generic[_AssociatedTypeDef]):
     """
     Used to specify that some value is a part of a typeclass.
 

classes/contrib/mypy/classes_plugin.py

@@ -30,8 +30,9 @@
 from mypy.types import Type as MypyType
 from typing_extensions import Final, final
 
-from classes.contrib.mypy.features import associated_type, typeclass
+from classes.contrib.mypy.features import associated_type, supports, typeclass
 
+_ASSOCIATED_TYPE_FULLNAME: Final = 'classes._typeclass.AssociatedType'
 _TYPECLASS_FULLNAME: Final = 'classes._typeclass._TypeClass'
 _TYPECLASS_DEF_FULLNAME: Final = 'classes._typeclass._TypeClassDef'
 _TYPECLASS_INSTANCE_DEF_FULLNAME: Final = (
@@ -58,8 +59,14 @@ def get_type_analyze_hook(
         fullname: str,
     ) -> Optional[Callable[[AnalyzeTypeContext], MypyType]]:
         """Hook that works on type analyzer phase."""
-        if fullname == 'classes._typeclass.AssociatedType':
+        if fullname == _ASSOCIATED_TYPE_FULLNAME:
             return associated_type.variadic_generic
+        if fullname == 'classes._typeclass.Supports':
+            associated_type_node = self.lookup_fully_qualified(
+                _ASSOCIATED_TYPE_FULLNAME,
+            )
+            assert associated_type_node
+            return supports.VariadicGeneric(associated_type_node)
         return None
 
     def get_function_hook(
@@ -80,7 +87,7 @@ def get_method_hook(
     ) -> Optional[Callable[[MethodContext], MypyType]]:
         """Here we adjust the typeclass with new allowed types."""
         if fullname == '{0}.__call__'.format(_TYPECLASS_DEF_FULLNAME):
-            return typeclass.typeclass_def_return_type
+            return typeclass.TypeClassDefReturnType(_ASSOCIATED_TYPE_FULLNAME)
         if fullname == '{0}.__call__'.format(_TYPECLASS_INSTANCE_DEF_FULLNAME):
             return typeclass.InstanceDefReturnType()
         if fullname == '{0}.instance'.format(_TYPECLASS_FULLNAME):

classes/contrib/mypy/features/associated_type.py

@@ -1,23 +1,11 @@
 from mypy.plugin import AnalyzeTypeContext
-from mypy.types import Instance
 from mypy.types import Type as MypyType
 
+from classes.contrib.mypy.semanal.variadic_generic import (
+    analize_variadic_generic,
+)
 
-def variadic_generic(ctx: AnalyzeTypeContext) -> MypyType:
-    """
-    Variadic generic support.
-
-    What is "variadic generic"?
-    It is a generic type with any amount of type variables.
-    Starting with 0 up to infinity.
 
-    We primarily use it for our ``AssociatedType`` implementation.
-    """
-    sym = ctx.api.lookup_qualified(ctx.type.name, ctx.context)  # type: ignore
-    if not sym or not sym.node:
-        # This will happen if `Supports[IsNotDefined]` will be called.
-        return ctx.type
-    return Instance(
-        sym.node,
-        ctx.api.anal_array(ctx.type.args),  # type: ignore
-    )
+def variadic_generic(ctx: AnalyzeTypeContext) -> MypyType:
+    """Variadic generic support for ``AssociatedType`` type."""
+    return analize_variadic_generic(ctx)

classes/contrib/mypy/features/supports.py

@@ -0,0 +1,46 @@
+from mypy.nodes import SymbolTableNode
+from mypy.plugin import AnalyzeTypeContext
+from mypy.types import Instance
+from mypy.types import Type as MypyType
+from mypy.types import UnionType
+from typing_extensions import final
+
+from classes.contrib.mypy.semanal.variadic_generic import (
+    analize_variadic_generic,
+)
+from classes.contrib.mypy.validation import validate_supports
+
+
+@final
+class VariadicGeneric(object):
+    """
+    Variadic generic support for ``Supports`` type.
+
+    We also need to validate that
+    all type args of ``Supports`` are subtypes of ``AssociatedType``.
+    """
+
+    __slots__ = ('_associated_type_node',)
+
+    def __init__(self, associated_type_node: SymbolTableNode) -> None:
+        """We need ``AssociatedType`` fullname here."""
+        self._associated_type_node = associated_type_node
+
+    def __call__(self, ctx: AnalyzeTypeContext) -> MypyType:
+        """Main entry point."""
+        analyzed_type = analize_variadic_generic(
+            validate_callback=self._validate,
+            ctx=ctx,
+        )
+        if isinstance(analyzed_type, Instance):
+            return analyzed_type.copy_modified(
+                args=[UnionType.make_union(analyzed_type.args)],
+            )
+        return analyzed_type
+
+    def _validate(self, instance: Instance, ctx: AnalyzeTypeContext) -> bool:
+        return validate_supports.check_type(
+            instance,
+            self._associated_type_node,
+            ctx,
+        )

classes/contrib/mypy/features/typeclass.py

@@ -103,35 +103,46 @@ def _process_typeclass_def_return_type(
         return typeclass_intermediate_def
 
 
-def typeclass_def_return_type(ctx: MethodContext) -> MypyType:
+@final
+class TypeClassDefReturnType(object):
     """
     Callback for cases like ``@typeclass(SomeType)``.
 
     What it does? It works with the associated types.
     It checks that ``SomeType`` is correct, modifies the current typeclass.
     And returns it back.
     """
-    assert isinstance(ctx.default_return_type, Instance)
-    assert isinstance(ctx.context, Decorator)
 
-    instance_args.mutate_typeclass_def(
-        typeclass=ctx.default_return_type,
-        definition_fullname=ctx.context.func.fullname,
-        ctx=ctx,
-    )
+    __slots__ = ('_associated_type',)
 
-    validate_typeclass_def.check_type(
-        typeclass=ctx.default_return_type,
-        ctx=ctx,
-    )
-    if isinstance(ctx.default_return_type.args[2], Instance):
-        validate_associated_type.check_type(
-            associated_type=ctx.default_return_type.args[2],
+    def __init__(self, associated_type: str) -> None:
+        """We need ``AssociatedType`` fullname here."""
+        self._associated_type = associated_type
+
+    def __call__(self, ctx: MethodContext) -> MypyType:
+        """Main entry point."""
+        assert isinstance(ctx.default_return_type, Instance)
+        assert isinstance(ctx.context, Decorator)
+
+        instance_args.mutate_typeclass_def(
             typeclass=ctx.default_return_type,
+            definition_fullname=ctx.context.func.fullname,
             ctx=ctx,
         )
 
-    return ctx.default_return_type
+        validate_typeclass_def.check_type(
+            typeclass=ctx.default_return_type,
+            ctx=ctx,
+        )
+        if isinstance(ctx.default_return_type.args[2], Instance):
+            validate_associated_type.check_type(
+                associated_type=ctx.default_return_type.args[2],
+                associated_type_fullname=self._associated_type,
+                typeclass=ctx.default_return_type,
+                ctx=ctx,
+            )
+
+        return ctx.default_return_type
 
 
 def instance_return_type(ctx: MethodContext) -> MypyType:

classes/contrib/mypy/semanal/variadic_generic.py

@@ -0,0 +1,35 @@
+from typing import Callable, Optional
+
+from mypy.plugin import AnalyzeTypeContext
+from mypy.types import Instance
+from mypy.types import Type as MypyType
+
+_ValidateCallback = Callable[[Instance, AnalyzeTypeContext], bool]
+
+
+def analize_variadic_generic(
+    ctx: AnalyzeTypeContext,
+    validate_callback: Optional[_ValidateCallback] = None,
+) -> MypyType:
+    """
+    Variadic generic support.
+
+    What is "variadic generic"?
+    It is a generic type with any amount of type variables.
+    Starting with 0 up to infinity.
+
+    We also conditionally validate types of passed arguments.
+    """
+    sym = ctx.api.lookup_qualified(ctx.type.name, ctx.context)  # type: ignore
+    if not sym or not sym.node:
+        # This will happen if `Supports[IsNotDefined]` will be called.
+        return ctx.type
+
+    instance = Instance(
+        sym.node,
+        ctx.api.anal_array(ctx.type.args),  # type: ignore
+    )
+
+    if validate_callback is not None:
+        validate_callback(instance, ctx)
+    return instance

classes/contrib/mypy/typeops/mro.py

@@ -1,9 +1,10 @@
-from typing import List
+from typing import List, Optional
 
 from mypy.plugin import MethodContext
+from mypy.subtypes import is_equivalent
 from mypy.types import Instance
 from mypy.types import Type as MypyType
-from mypy.types import TypeVarType, union_items
+from mypy.types import TypeVarType, UnionType, union_items
 from typing_extensions import final
 
 from classes.contrib.mypy.typeops import type_loader
@@ -96,10 +97,21 @@ def add_supports_metadata(self) -> None:
                 self._ctx,
             )
 
-            if supports_spec not in instance_type.type.bases:
+            index = self._find_supports_index(instance_type, supports_spec)
+            if index is not None:
+                # We already have `Supports` base class inserted,
+                # it means that we need to unify them:
+                # `Supports[A] + Supports[B] == Supports[Union[A, B]]`
+                self._add_unified_type(instance_type, supports_spec, index)
+            else:
+                # This is the first time this type is referenced in
+                # a typeclass'es instance defintinion.
+                # Just inject `Supports` with no extra steps:
                 instance_type.type.bases.append(supports_spec)
+
             if supports_spec.type not in instance_type.type.mro:
-                instance_type.type.mro.insert(0, supports_spec.type)
+                # We only need to add `Supports` type to `mro` once:
+                instance_type.type.mro.append(supports_spec.type)
 
             self._added_types.append(supports_spec)
 
@@ -110,11 +122,66 @@ def remove_supports_metadata(self) -> None:
 
         for instance_type in self._instance_types:
             assert isinstance(instance_type, Instance)
-
-            for added_type in self._added_types:
-                if added_type in instance_type.type.bases:
-                    instance_type.type.bases.remove(added_type)
-                if added_type.type in instance_type.type.mro:
-                    instance_type.type.mro.remove(added_type.type)
-
+            self._clean_instance_type(instance_type)
         self._added_types = []
+
+    def _clean_instance_type(self, instance_type: Instance) -> None:
+        remove_mro = True
+        for added_type in self._added_types:
+            index = self._find_supports_index(instance_type, added_type)
+            if index is not None:
+                remove_mro = self._remove_unified_type(
+                    instance_type,
+                    added_type,
+                    index,
+                )
+
+            if remove_mro and added_type.type in instance_type.type.mro:
+                # We remove `Supports` type from `mro` only if
+                # there are not associated types left.
+                # For example, `Supports[A, B] - Supports[B] == Supports[A]`
+                # then `Supports[A]` stays.
+                # `Supports[A] - Supports[A] == None`
+                # then `Supports` is removed from `mro` as well.
+                instance_type.type.mro.remove(added_type.type)
+
+    def _find_supports_index(
+        self,
+        instance_type: Instance,
+        supports_spec: Instance,
+    ) -> Optional[int]:
+        for index, base in enumerate(instance_type.type.bases):
+            if is_equivalent(base, supports_spec, ignore_type_params=True):
+                return index
+        return None
+
+    def _add_unified_type(
+        self,
+        instance_type: Instance,
+        supports_spec: Instance,
+        index: int,
+    ) -> None:
+        unified_arg = UnionType.make_union([
+            *supports_spec.args,
+            *instance_type.type.bases[index].args,
+        ])
+        instance_type.type.bases[index] = supports_spec.copy_modified(
+            args=[unified_arg],
+        )
+
+    def _remove_unified_type(
+        self,
+        instance_type: Instance,
+        supports_spec: Instance,
+        index: int,
+    ) -> bool:
+        base = instance_type.type.bases[index]
+        union_types = [
+            type_arg
+            for type_arg in union_items(base.args[0])
+            if type_arg not in supports_spec.args
+        ]
+        instance_type.type.bases[index] = supports_spec.copy_modified(
+            args=[UnionType.make_union(union_types)],
+        )
+        return not bool(union_types)