How do I annotate this function with abs() and division?

[jwodder]

I'm trying to add type annotations to the following function, but I can't figure out anything that works:

def signum(x):
    """
    Returns the sign of ``x``: 1 if positive, -1 if negative, 0 if zero.  For
    complex numbers, returns the number with the same phase angle and magnitude
    1.
    """
    if x == 0:
        return 0
    else:
        return x / abs(x)

All of the following mypy runs were performed with mypy 0.961 on Python 3.9.13 on macOS 11.6.6 with the following config:

[mypy]
allow_incomplete_defs = False
allow_untyped_defs = False
ignore_missing_imports = True
no_implicit_optional = True
implicit_reexport = False
local_partial_types = True
pretty = True
show_error_codes = True
show_traceback = True
strict_equality = True
warn_redundant_casts = True
warn_return_any = True
warn_unreachable = True

---

First attempt:

from numbers import Number


def signum(x: Number) -> Number:
    ...

Results:

src/pywodlib/math/signum.py:10: error: Non-overlapping equality check (left
operand type: "Number", right operand type: "Literal[0]")  [comparison-overlap]
        if x == 0:
           ^
src/pywodlib/math/signum.py:11: error: Incompatible return value type (got
"int", expected "Number")  [return-value]
            return 0
                   ^
src/pywodlib/math/signum.py:13: error: Returning Any from function declared to
return "Number"  [no-any-return]
            return x / abs(x)
            ^
src/pywodlib/math/signum.py:13: error: Unsupported left operand type for /
("Number")  [operator]
            return x / abs(x)
                   ^
src/pywodlib/math/signum.py:13: error: Argument 1 to "abs" has incompatible
type "Number"; expected "SupportsAbs[<nothing>]"  [arg-type]
            return x / abs(x)
                           ^
Found 5 errors in 1 file (checked 64 source files)

---

Second attempt:

from typing import SupportsAbs, TypeVar

T = TypeVar("T", bound=SupportsAbs)

def signum(x: T) -> T:
    ...

Results:

src/pywodlib/math/signum.py:12: error: Incompatible return value type (got
"int", expected "T")  [return-value]
            return 0
                   ^
src/pywodlib/math/signum.py:14: error: Returning Any from function declared to
return "T"  [no-any-return]
            return x / abs(x)

---

Third attempt:

from typing import Protocol, SupportsAbs, TypeVar

T = TypeVar("T", bound="SupportsSignum")

class SupportsSignum(SupportsAbs[T]):
    def __truediv__(self, other: T) -> T:
        ...

def signum(x: T) -> T:
    ...

Results:

src/pywodlib/math/signum.py:15: error: Non-overlapping equality check (left
operand type: "T", right operand type: "Literal[0]")  [comparison-overlap]
        if x == 0:
           ^
src/pywodlib/math/signum.py:16: error: Incompatible return value type (got
"int", expected "T")  [return-value]
            return 0
                   ^
src/pywodlib/math/signum.py:18: error: Returning Any from function declared to
return "T"  [no-any-return]
            return x / abs(x)
            ^

---

Fourth attempt:

from typing import Protocol, TypeVar

T = TypeVar("T", bound="SupportsSignum")


class SupportsSignum(Protocol):
    def __abs__(self: T) -> T:
        ...

    def __truediv__(self: T, other: T) -> T:
        ...


def signum(x: T) -> T:
    """
    Returns the sign of ``x``: 1 if positive, -1 if negative, 0 if zero.  For
    complex numbers, returns the number with the same phase angle and magnitude
    1.
    """
    if x == 0:
        return x
    else:
        return x / abs(x)

Results:

src/pywodlib/math/signum.py:23: error: Unsupported operand types for / ("T" and
"SupportsSignum")  [operator]
            return x / abs(x)
                       ^

Now I'm out of ideas. Help?

[Akuli]

Sorry for the late reply. Looking at the docstring, the function is intended to be used only with floats and complex numbers, not with arbitrary types that happen to support absolute value and division (e.g. a 3D vector from some library). So you have two cases, either a float is passed in and the result is a float, or a complex is passed in and the result is a complex.

from __future__ import annotations  # for "|" syntax
from typing import overload

@overload
def signum(x: float) -> float: ...
@overload
def signum(x: complex) -> complex: ...

def signum(x: float | complex) -> float | complex:
    if x == 0:
        return 0
    else:
        return x / abs(x)

reveal_type(signum(5.2))  # float
reveal_type(signum(0))  # float
reveal_type(signum(1+2j))  # complex
reveal_type(signum(0+0j))  # complex

Note that it's fine to return float or complex instead of float | int or complex | int. Type checkers implicitly treat all ints as floats, and all floats as complexes.

This can be shortened slightly with a TypeVar that can be either float or complex.

from typing import TypeVar

T = TypeVar("T", float, complex)

def signum(x: T) -> T:
    if x == 0:
        return 0
    else:
        return x / abs(x)

reveal_type(signum(5.2))  # float
reveal_type(signum(0))  # float
reveal_type(signum(1+2j))  # complex
reveal_type(signum(0+0j))  # complex

This is pretty close to your second attempt, but your TypeVar bound was just more general than mine. Not all code must be written (or typed) in the most general way possible :)