# In this unit (and most future units), we will add static type
# annotations to our programs, to be checked by ‹mypy›. Annotations
# can be attached to variables, function arguments and return types.
# In ‹--strict› mode (which we will be using), ‹mypy› requires that
# each function header (arguments and return type) is annotated.
# e.g. the function ‹divisor_count› takes a single ‹int› parameter
# and returns another ‹int›:

def divisor_count( n: int ) -> int:

    # Notice that variables, in most cases, do not need annotations:
    # types are inferred from the right-hand side of the initial
    # assignment.

    count = 0

    for i in range( 1, n + 1 ):
        if n % i == 0:
            count += 1
    return count

def test_divcount() -> None: # demo
    assert divisor_count( 5 )  == 2 # 1 and 5
    assert divisor_count( 6 )  == 4 # 1, 2, 3 and 6
    assert divisor_count( 12 ) == 6 # 1, 2, 3, 4, 6 and 12

# For built-in types, including compound types, we will use the
# actual builtins for annotations (for simple types, like ‹int› and
# ‹str›, this has worked for a long time; for compound, types, like
# ‹list› or ‹dict›, it works from Python 3.9 onwards).

# Compound types are «generic», i.e. they have one or more «type
# parameters». You know these from Haskell (they are everywhere) or
# perhaps C++/Java/C# (templates and generics, respectively). Like
# in Haskell but unlike in C++, generic types have no effect on the
# code itself – they are just annotations.  Type parameters are
# given in square brackets after the generic type.

def divisors( n: int ) -> list[ int ]:

    # As mentioned above, for variables, ‹mypy› can usually deduce
    # types automatically, even when they are of a generic type.
    # However, sometimes this fails, a prominent example being the
    # empty list – it's impossible to find the type parameter, since
    # there are no values to look at. Annotations of local variables
    # can be combined with initialization.

    res: list[ int ] = []

    for i in range( 1, n + 1 ):
        if n % i == 0:
            res.append( i )

    return res

def test_divisors() -> None: # demo
    assert divisors( 5 )  == [ 1, 5 ]
    assert divisors( 6 )  == [ 1, 2, 3, 6 ]
    assert divisors( 12 ) == [ 1, 2, 3, 4, 6, 12 ]

# Finally, it is quite common in Python that a particular name
# (variable, function parameter) can accept values of different
# types. For these cases, ‹mypy› supports «union types» (in a direct
# reference to the ‘types are sets’ idea, those are literally unions
# of their constituent types, when understood as sets).

# Before Python 3.10, the preferred way to write unions was to use
# helper classes from module ‹typing›: the more general
# ‹Union[ S, T ]› denotes the union of arbitrary two types (e.g.
# ‹Union[ int, str ]›). However, there is one very common union
# type, namely ‹Union[ T, NoneType ]› which can either take values
# from ‹T› or it can be ‹None›. Since it is so common, it can be
# written as ‹Optional[ T ]›. However, the new (Python 3.10) syntax
# is considerably nicer, and doesn't need extra imports:¹ ‹S | T› for
# a generic union and ‹None | T› for optional.

# For example:

def maybe_push( stack: list[ int ], value: None | int ) -> None:

    if value is not None:

        # Notice how ‹mypy› accepts this code even though it is
        # ostensibly ill-typed: on the face of it, ‹value› is not an
        # ‹int› (the annotation above says it could be ‹None›), but
        # ‹stack› only accepts elements of type ‹int›.
        #
        # The code is accepted because the following line is
        # «guarded»: the condition of the above ‹if› statement
        # means that the branch is only taken if ‹value› is an
        # actual ‹int›. In addition to conditionals, ‹mypy› will
        # also understand ‹assert› statements of this sort.

        stack.append( value )


def push_either( int_stack: list[ int ], str_stack: list[ str ],
                 value: int | str ) -> None:

    # Of course ‹is None› is a pretty special case: normally, we
    # will not want (or be able) to enumerate all possible values of
    # a given type. However, ‹mypy› also understand ‹isinstance›:

    if isinstance( value, int ):

        # In this branch, ‹mypy› takes ‹value› to be of type ‹int›,
        # since it is guarded by an ‹isinstance›.

        int_stack.append( value )
    else:

        # Of course, ‹else› branches are understood as well. In this
        # case, ‹value› can be anything in ‹int | str› that isn't
        # ‹int›, which just leaves ‹str›:

        str_stack.append( value )


def push_if_int( stack: list[ int ],
                 value: int | float | str ) -> None:

    # To make things more intuitive, ‹isinstance› actually lies
    # about types. The original meaning of ‹isinstance( x, c )› is
    # ‘is object ‹x› an instance of class ‹c›?’. But check this out:

    if isinstance( value, str | float ):
        pass

    # Clearly, ‹value› is «not» an instance of whatever type that
    # union is, because ‹str | list[ int ]› does «not» evaluate to
    # an actual superclass of ‹str›: there is only one, and that's
    # ‹object›. And ‹object› is also a superclass of ‹int›, so
    # things would go haywire there.

    # What on Earth is going on? Well, «metaclasses», that's what.
    # When you write ‹isinstance( x, c )›, the metaclass of ‹c› is
    # consulted about the matter, and can ‘claim’ ‹x› as an instance
    # of ‹c› even if they are completely unrelated (in the sense of
    # inheritance). We will get back to this in a later chapter.

    else:

        # There is one last thing to illustrate: the ‹else› branch
        # excludes both ‹str› and ‹list[ int ]›, leaving only ‹int›:

        stack.append( value )

def test_pushes() -> None: # demo
    int_stack: list[ int ] = []
    str_stack: list[ str ] = []

    push_if_int( int_stack, 3 )
    push_if_int( int_stack, 'xoxo' )
    assert int_stack == [ 3 ]

    push_either( int_stack, str_stack, 'xo' )
    assert int_stack == [ 3 ]
    assert str_stack == [ 'xo' ]

# Before we conclude this demo, there is one other case where
# variables need annotations, and it is when they are actually of
# union types:

def find_min( values: list[ int ] ) -> None | int:

    min_val: None | int = None

    for v in values:

        # Notice that the second operand of ‹or› is also treated as
        # a proper branch by ‹mypy›: if it is ever evaluated, we
        # already know that ‹min_val› is «not» ‹None› and hence is
        # an ‹int›, which means it can be compared with ‹v› (which
        # is also an ‹int›).

        if min_val is None or v < min_val:
            min_val = v

    return min_val

# ¹ To make the same syntax work with Python 3.9, you can use ‹from
#   __future__ import annotations›. This has the additional benefit
#   of making «forward references» possible. The mechanism here is
#   that with this ‹__future__› import, annotations are stored as
#   strings and they must be ‹eval›'d upon inspection to get the
#   actual annotations. Since Python 3.10, the ‹inspect› module
#   provides a helper to do that, ‹get_annotations›.

if __name__ == '__main__':
    test_divcount()
    test_divisors()
    test_pushes()