Merge branch 'master' of github.com:wesselb/plum

Author: wesselb

plum/parametric.py

@@ -1,4 +1,4 @@
-from typing import Union
+from typing import Type, TypeVar, Union
 
 import beartype.door
 from beartype.roar import BeartypeDoorNonpepException
@@ -13,11 +13,14 @@
     "CovariantMeta",
     "parametric",
     "type_parameter",
+    "type_unparametrized",
     "kind",
     "Kind",
     "Val",
 ]
 
+T = TypeVar("T")
+
 
 _dispatch = Dispatcher()
 
@@ -274,11 +277,82 @@ def class_new(cls, *args, **kw_args):
             cls.__new__ = class_new
         super(original_class, cls).__init_subclass__(**kw_args)
 
+    def __class_nonparametric__(cls):
+        """Return the non-parametric type of an object.
+
+        :mod:`plum.parametric` produces parametric subtypes of classes. This
+        method can be used to get the non-parametric type of an object.
+
+        Examples
+        --------
+        >>> from plum import parametric
+        >>> @parametric
+        ... class Obj:
+        ...     @classmethod
+        ...     def __infer_type_parameter__(cls, *arg):
+        ...         return type(arg[0])
+        ...     def __init__(self, x):
+        ...         self.x = x
+        ...     def __repr__(self):
+        ...         return f"Obj({self.x})"
+
+        >>> obj = Obj(1)
+        >>> obj
+        Obj(1)
+
+        >>> type(obj).mro()
+        [Obj[int], Obj, object]
+
+        >>> obj.__class_nonparametric__().mro()
+        [Obj, object]
+        """
+        return original_class
+
+    def __class_unparametrized__(cls):
+        """Return the unparametrized type of an object.
+
+        :mod:`plum.parametric` produces parametric subtypes of classes. This
+        method can be used to get the un-parametrized type of an object.
+
+        Examples
+        --------
+        >>> from plum import parametric
+        >>> @parametric
+        ... class Obj:
+        ...     @classmethod
+        ...     def __infer_type_parameter__(cls, *arg):
+        ...         return type(arg[0])
+        ...     def __init__(self, x):
+        ...         self.x = x
+        ...     def __repr__(self):
+        ...         return f"Obj({self.x})"
+
+        >>> obj = Obj(1)
+        >>> obj
+        Obj(1)
+
+        >>> type(obj).__name__
+        Obj[int]
+
+        >>> obj.__class_unparametrized__().mro()
+        [Obj, Obj, object]
+
+        Note that this is still NOT the 'original' non-`parametric`-wrapped
+        type.  This is the type that is wrapped by :mod:`plum.parametric`, but
+        without the inferred type parameter(s).
+        """
+        return parametric_class
+
     # Create parametric class.
     parametric_class = meta(
         original_class.__name__,
         (original_class,),
-        {"__new__": __new__, "__init_subclass__": __init_subclass__},
+        {
+            "__new__": __new__,
+            "__init_subclass__": __init_subclass__,
+            "__class_nonparametric__": __class_nonparametric__,
+            "__class_unparametrized__": __class_unparametrized__,
+        },
     )
     parametric_class._parametric = True
     parametric_class._concrete = False
@@ -356,6 +430,44 @@ def type_parameter(x):
     )
 
 
+def type_unparametrized(q: T) -> Type[T]:
+    """Return the unparametrized type of an object.
+
+    :mod:`plum.parametric` produces parametric subtypes of classes.  This
+    function can be used to get the un-parametrized type of an object.
+    This function also works for normal, :mod:`plum.parametric`-wrapped classes.
+
+    Examples
+    --------
+    >>> from plum import parametric
+    >>> @parametric
+    ... class Obj:
+    ...     @classmethod
+    ...     def __infer_type_parameter__(cls, *arg):
+    ...         return type(arg[0])
+    ...     def __init__(self, x):
+    ...         self.x = x
+    ...     def __repr__(self):
+    ...         return f"Obj({self.x})"
+
+    >>> obj = Obj(1)
+    >>> obj
+    Obj(1)
+
+    >>> type(obj).__name__
+    Obj[int]
+
+    >>> type_unparametrized(obj).__name__
+    Obj
+
+    Note that this is still NOT the 'original' non-`parametric`-wrapped type.
+    This is the type that is wrapped by :mod:`plum.parametric`, but without the
+    inferred type parameter(s).
+    """
+    typ = type(q)
+    return q.__class_unparametrized__() if isinstance(typ, ParametricTypeMeta) else typ
+
+
 def kind(SuperClass=object):
     """Create a parametric wrapper type for dispatch purposes.
 

tests/test_parametric.py

@@ -15,7 +15,7 @@
     parametric,
     type_parameter,
 )
-from plum.parametric import CovariantMeta, is_concrete, is_type
+from plum.parametric import CovariantMeta, is_concrete, is_type, type_unparametrized
 
 
 def test_covariantmeta():
@@ -60,6 +60,15 @@ class A(Base1, metaclass=metaclass):
     assert issubclass(type(a1), Base1)
     assert not issubclass(type(a1), Base2)
 
+    assert a1.__class_unparametrized__() is A
+    assert a2.__class_unparametrized__() is A
+
+    # Here we are testing that the class returned by `__class_nonparametric__`
+    # is the 'original' class that the @parametric decorator was applied to.
+    assert a1.__class_nonparametric__() is A.mro()[1]
+    assert issubclass(a2.__class_nonparametric__(), Base1)
+    assert a2.__class_nonparametric__() is not Base1
+
     # Test multiple type parameters.
     assert A[1, 2] == A[1, 2]
 
@@ -575,3 +584,25 @@ class Wrapper(Pytree):
 
     Wrapper[int]
     assert Wrapper[int] in register
+
+
+def test_type_unparametrized():
+    """Test the `type_unparametrized` function."""
+
+    @parametric
+    class Obj:
+        @classmethod
+        def __infer_type_parameter__(cls, *arg):
+            return type(arg[0])
+
+        def __init__(self, x):
+            self.x = x
+
+        def __repr__(self):
+            return f"Obj({self.x})"
+
+    pobj = Obj(1)
+
+    assert type(pobj) is Obj[int]
+    assert type_unparametrized(pobj) is not Obj[int]
+    assert type_unparametrized(pobj) is Obj