# Test cases related to the functools.singledispatch decorator

# Most of these tests are marked as xfails because mypyc doesn't support singledispatch yet

# (These tests will be re-enabled when mypyc supports singledispatch)

[case testSpecializedImplementationUsed]

from functools import singledispatch

from typing import Any

@singledispatch

def fun(arg: Any) -> bool:

return False

@fun.register

def fun_specialized(arg: str) -> bool:

return True

def test_specialize() -> None:

assert fun('a')

assert not fun(3)

[case testSubclassesOfExpectedTypeUseSpecialized]

from functools import singledispatch

from typing import Any

class A: pass

class B(A): pass

@singledispatch

def fun(arg: Any) -> bool:

return False

@fun.register

def fun_specialized(arg: A) -> bool:

return True

def test_specialize() -> None:

assert fun(B())

assert fun(A())

[case testSuperclassImplementationNotUsedWhenSubclassHasImplementation]

from functools import singledispatch

from typing import Any

class A: pass

class B(A): pass

@singledispatch

def fun(arg: Any) -> bool:

# shouldn't be using this

assert False

@fun.register

def fun_specialized(arg: A) -> bool:

return False

@fun.register

def fun_specialized2(arg: B) -> bool:

return True

def test_specialize() -> None:

assert fun(B())

assert not fun(A())

[case testMultipleUnderscoreFunctionsIsntError]

from functools import singledispatch

from typing import Any

@singledispatch

def fun(arg: Any) -> str:

return 'default'

@fun.register

def _(arg: str) -> str:

return 'str'

@fun.register

def _(arg: int) -> str:

return 'int'

# extra function to make sure all 3 underscore functions aren't treated as one OverloadedFuncDef

def a(b: Any) -> Any: pass

@fun.register

def _(arg: list) -> str:

return 'list'

def test_singledispatch() -> None:

assert fun(0) == 'int'

assert fun('a') == 'str'

assert fun([1, 2]) == 'list'

assert fun({'a': 'b'}) == 'default'

[case testCanRegisterCompiledClasses]

from functools import singledispatch

from typing import Any

class A: pass

@singledispatch

def fun(arg: Any) -> bool:

return False

@fun.register

def fun_specialized(arg: A) -> bool:

return True

def test_singledispatch() -> None:

assert fun(A())

assert not fun(1)

[case testTypeUsedAsArgumentToRegister]

from functools import singledispatch

from typing import Any

@singledispatch

def fun(arg: Any) -> bool:

return False

@fun.register(int)

def fun_specialized(arg: Any) -> bool:

return True

def test_singledispatch() -> None:

assert fun(1)

assert not fun('a')

[case testUseRegisterAsAFunction]

from functools import singledispatch

from typing import Any

@singledispatch

def fun(arg: Any) -> bool:

return False

def fun_specialized_impl(arg: Any) -> bool:

return True

fun.register(int, fun_specialized_impl)

def test_singledispatch() -> None:

assert fun(0)

assert not fun('a')

[case testRegisterDoesntChangeFunction]

from functools import singledispatch

from typing import Any

@singledispatch

def fun(arg: Any) -> bool:

return False

@fun.register(int)

def fun_specialized(arg: Any) -> bool:

return True

def test_singledispatch() -> None:

assert fun_specialized('a')

# TODO: turn this into a mypy error

[case testNoneIsntATypeWhenUsedAsArgumentToRegister]

from functools import singledispatch

from typing import Any

@singledispatch

def fun(arg: Any) -> bool:

return False

def test_argument() -> None:

try:

@fun.register

def fun_specialized(arg: None) -> bool:

return True

assert False, "expected to raise an exception"

except TypeError:

pass

[case testRegisteringTheSameFunctionSeveralTimes]

from functools import singledispatch

from typing import Any

@singledispatch

def fun(arg: Any) -> bool:

return False

@fun.register(int)

@fun.register(str)

def fun_specialized(arg: Any) -> bool:

return True

def test_singledispatch() -> None:

assert fun(0)

assert fun('a')

assert not fun([1, 2])

[case testTypeIsAnABC]

# mypy: allow-untyped-defs

from functools import singledispatch

from collections.abc import Mapping

@singledispatch

def fun(arg) -> bool:

# TODO: Adding an Any parameter annotation breaks the test case

return False

@fun.register

def fun_specialized(arg: Mapping) -> bool:

return True

def test_singledispatch() -> None:

assert not fun(1)

assert fun({'a': 'b'})

[case testSingleDispatchMethod-xfail]

from functools import singledispatchmethod

class A:

@singledispatchmethod

def fun(self, arg) -> str:

return 'default'

@fun.register

def fun_int(self, arg: int) -> str:

return 'int'

@fun.register

def fun_str(self, arg: str) -> str:

return 'str'

def test_singledispatchmethod() -> None:

x = A()

assert x.fun(5) == 'int'

assert x.fun('a') == 'str'

assert x.fun([1, 2]) == 'default'

[case testSingleDispatchMethodWithOtherDecorator-xfail]

from functools import singledispatchmethod

class A:

@singledispatchmethod

@staticmethod

def fun(arg) -> str:

return 'default'

@fun.register

@staticmethod

def fun_int(arg: int) -> str:

return 'int'

@fun.register

@staticmethod

def fun_str(arg: str) -> str:

return 'str'

def test_singledispatchmethod() -> None:

x = A()

assert x.fun(5) == 'int'

assert x.fun('a') == 'str'

assert x.fun([1, 2]) == 'default'

[case testSingledispatchTreeSumAndEqual]

from functools import singledispatch

from typing import cast

class Tree:

pass

class Leaf(Tree):

pass

class Node(Tree):

def __init__(self, value: int, left: Tree, right: Tree) -> None:

self.value = value

self.left = left

self.right = right

@singledispatch

def calc_sum(x: Tree) -> int:

raise TypeError('invalid type for x')

@calc_sum.register

def _(x: Leaf) -> int:

return 0

@calc_sum.register

def _(x: Node) -> int:

return x.value + calc_sum(x.left) + calc_sum(x.right)

@singledispatch

def equal(to_compare: Tree, known: Tree) -> bool:

raise TypeError('invalid type for x')

@equal.register

def _(to_compare: Leaf, known: Tree) -> bool:

return isinstance(known, Leaf)

@equal.register

def _(to_compare: Node, known: Tree) -> bool:

if isinstance(known, Node):

if to_compare.value != known.value:

return False

else:

return equal(to_compare.left, known.left) and equal(to_compare.right, known.right)

return False

def build(n: int) -> Tree:

if n == 0:

return Leaf()

return Node(n, build(n - 1), build(n - 1))

def test_sum_and_equal() -> None:

tree = build(5)

tree2 = build(5)

cast(Node, cast(Node, cast(Node, cast(Node, tree2).right).right).right).value = 10

assert calc_sum(tree) == 57

assert calc_sum(tree2) == 65

assert equal(tree, tree)

assert not equal(tree, tree2)

tree3 = build(4)

assert not equal(tree, tree3)

[case testSimulateMypySingledispatch]

from functools import singledispatch

from mypy_extensions import trait

from typing import Iterator, Union, TypeVar, Any, List, Type

# based on use of singledispatch in stubtest.py

class Error:

def __init__(self, msg: str) -> None:

self.msg = msg

@trait

class Node: pass

class MypyFile(Node): pass

class TypeInfo(Node): pass

@trait

class SymbolNode(Node): pass

@trait

class Expression(Node): pass

class TypeVarLikeExpr(SymbolNode, Expression): pass

class TypeVarExpr(TypeVarLikeExpr): pass

class TypeAlias(SymbolNode): pass

class Missing: pass

MISSING = Missing()

T = TypeVar("T")

MaybeMissing = Union[T, Missing]

@singledispatch

def verify(stub: Node, a: MaybeMissing[Any], b: List[str]) -> Iterator[Error]:

yield Error('unknown node type')

@verify.register(MypyFile)

def verify_mypyfile(stub: MypyFile, a: MaybeMissing[int], b: List[str]) -> Iterator[Error]:

if isinstance(a, Missing):

yield Error("shouldn't be missing")

return

if not isinstance(a, int):

# this check should be unnecessary because of the type signature and the previous check,

# but stubtest.py has this check

yield Error("should be an int")

return

yield from verify(TypeInfo(), str, ['abc', 'def'])

@verify.register(TypeInfo)

def verify_typeinfo(stub: TypeInfo, a: MaybeMissing[Type[Any]], b: List[str]) -> Iterator[Error]:

yield Error('in TypeInfo')

yield Error('hello')

@verify.register(TypeVarExpr)

def verify_typevarexpr(stub: TypeVarExpr, a: MaybeMissing[Any], b: List[str]) -> Iterator[Error]:

if False:

yield None

def verify_list(stub: Any, a: Any, b: Any) -> List[str]:

"""Helper function that converts iterator of errors to list of messages"""

return list(err.msg for err in verify(stub, a, b))

def test_verify() -> None:

assert verify_list(TypeAlias(), 'a', ['a', 'b']) == ['unknown node type']

assert verify_list(MypyFile(), MISSING, ['a', 'b']) == ["shouldn't be missing"]

assert verify_list(MypyFile(), 5, ['a', 'b']) == ['in TypeInfo', 'hello']

assert verify_list(TypeInfo(), str, ['a', 'b']) == ['in TypeInfo', 'hello']

assert verify_list(TypeVarExpr(), 'a', ['x', 'y']) == []

[case testArgsInRegisteredImplNamedDifferentlyFromMainFunction]

from functools import singledispatch

from typing import Any

@singledispatch

def f(a: Any) -> bool:

return False

@f.register

def g(b: int) -> bool:

return True

def test_singledispatch() -> None:

assert f(5)

assert not f('a')

[case testKeywordArguments]

from functools import singledispatch

from typing import Any

@singledispatch

def f(arg: Any, *, kwarg: int = 0) -> int:

return kwarg + 10

@f.register

def g(arg: int, *, kwarg: int = 5) -> int:

return kwarg - 10

def test_keywords() -> None:

assert f('a') == 10

assert f('a', kwarg=3) == 13

assert f('a', kwarg=7) == 17

assert f(1) == -5

assert f(1, kwarg=4) == -6

assert f(1, kwarg=6) == -4

[case testGeneratorAndMultipleTypesOfIterable]

from functools import singledispatch

from typing import *

@singledispatch

def f(arg: Any) -> Iterable[int]:

yield 1

@f.register

def g(arg: str) -> Iterable[int]:

return [0]

def test_iterables() -> None:

assert f(1) != [1]

assert list(f(1)) == [1]

assert f('a') == [0]

[case testRegisterUsedAtSameTimeAsOtherDecorators]

from functools import singledispatch

from typing import TypeVar, Any

class A: pass

class B: pass

T = TypeVar('T')

def decorator(f: T) -> T:

return f

@singledispatch

def f(arg: Any) -> int:

return 0

@f.register

@decorator

def h(arg: str) -> int:

return 2

def test_singledispatch() -> None:

assert f(1) == 0

assert f('a') == 2

[case testDecoratorModifiesFunction]

from functools import singledispatch

from typing import Callable, Any

class A: pass

def decorator(f: Callable[[Any], int]) -> Callable[[Any], int]:

def wrapper(x: Any) -> int:

return f(x) * 7

return wrapper

@singledispatch

def f(arg: Any) -> int:

return 10

@f.register

@decorator

def h(arg: str) -> int:

return 5

def test_singledispatch() -> None:

assert f('a') == 35

assert f(A()) == 10

[case testMoreSpecificTypeBeforeLessSpecificType]

from functools import singledispatch

from typing import Any

class A: pass

class B(A): pass

@singledispatch

def f(arg: Any) -> str:

return 'default'

@f.register

def g(arg: B) -> str:

return 'b'

@f.register

def h(arg: A) -> str:

return 'a'

def test_singledispatch() -> None:

assert f(B()) == 'b'

assert f(A()) == 'a'

assert f(5) == 'default'

[case testMultipleRelatedClassesBeingRegistered]

from functools import singledispatch

from typing import Any

class A: pass

class B(A): pass

class C(B): pass

@singledispatch

def f(arg: Any) -> str: return 'default'

@f.register

def _(arg: A) -> str: return 'a'

@f.register

def _(arg: C) -> str: return 'c'

@f.register

def _(arg: B) -> str: return 'b'

def test_singledispatch() -> None:

assert f(A()) == 'a'

assert f(B()) == 'b'

assert f(C()) == 'c'

assert f(1) == 'default'

[case testRegisteredImplementationsInDifferentFiles]

from other_a import f, A, B, C

@f.register

def a(arg: A) -> int:

return 2

@f.register

def _(arg: C) -> int:

return 3

def test_singledispatch() -> None:

assert f(B()) == 1

assert f(A()) == 2

assert f(C()) == 3

assert f(1) == 0

[file other_a.py]

from functools import singledispatch

class A: pass

class B(A): pass

class C(B): pass

@singledispatch

def f(arg: object) -> int:

return 0

@f.register

def g(arg: B) -> int:

return 1

[case testOrderCanOnlyBeDeterminedFromMRONotIsinstanceChecks]

from mypy_extensions import trait

from functools import singledispatch

from typing import Any

@trait

class A: pass

@trait

class B: pass

class AB(A, B): pass

class BA(B, A): pass

@singledispatch

def f(arg: Any) -> str:

return "default"

@f.register

def fa(arg: A) -> str:

return "a"

@f.register

def fb(arg: B) -> str:

return "b"

def test_singledispatch() -> None:

assert f(AB()) == "a"

assert f(BA()) == "b"

[case testCallingFunctionBeforeAllImplementationsRegistered]

from functools import singledispatch

from typing import Any

class A: pass

class B(A): pass

@singledispatch

def f(arg: Any) -> str:

return 'default'

assert f(A()) == 'default'

assert f(B()) == 'default'

assert f(1) == 'default'

@f.register

def g(arg: A) -> str:

return 'a'

assert f(A()) == 'a'

assert f(B()) == 'a'

assert f(1) == 'default'

@f.register

def _(arg: B) -> str:

return 'b'

# TODO: Move whole testcase to a function when mypyc#1118 is fixed.

def test_final() -> None:

assert f(A()) == 'a'

assert f(B()) == 'b'

assert f(1) == 'default'

[case testDynamicallyRegisteringFunctionFromInterpretedCode]

from functools import singledispatch

from typing import Any

class A: pass

class B(A): pass

class C(B): pass

class D(C): pass

@singledispatch

def f(arg: Any) -> str:

return "default"

@f.register

def _(arg: B) -> str:

return 'b'

[file register_impl.py]

from native import f, A, B, C

@f.register(A)

def a(arg) -> str:

return 'a'

@f.register

def c(arg: C) -> str:

return 'c'

[file driver.py]

from native import f, A, B, C

from register_impl import a, c

# We need a custom driver here because register_impl has to be run before we test this (so that the

# additional implementations are registered)

assert f(C()) == 'c'

assert f(A()) == 'a'

assert f(B()) == 'b'

assert a(C()) == 'a'

assert c(A()) == 'c'

[case testMalformedDynamicRegisterCall]

from functools import singledispatch

from typing import Any

@singledispatch

def f(arg: Any) -> None:

pass

[file register.py]

from native import f

from testutil import assertRaises

with assertRaises(TypeError, 'Invalid first argument to `register()`'):

@f.register

def _():

pass

[file driver.py]

import register

[case testCacheClearedWhenNewFunctionRegistered]

from functools import singledispatch

@singledispatch

def f(arg: object) -> str:

return 'default'

[file register.py]

from native import f

class A: pass

class B: pass

class C: pass

# annotated function

assert f(A()) == 'default'

@f.register

def _(arg: A) -> str:

return 'a'

assert f(A()) == 'a'

# type passed as argument

assert f(B()) == 'default'

@f.register(B)

def _(arg: B) -> str:

return 'b'

assert f(B()) == 'b'

# 2 argument form

assert f(C()) == 'default'

def c(arg) -> str:

return 'c'

f.register(C, c)

assert f(C()) == 'c'

[file driver.py]

import register