Annotations guidelines

[crusaderky]

There's been some discordant opinions re. annotations recently, so I'd like to reach a consensus on project-wide guidelines.
The below is a first draft and represents my personal opinion only - please comment if you would like to add/amend.

General guidelines

• Dask code is annotated according to PEP 484 and later revisions.
  This means that the style for annotating dask can change over time as new PEPs are created or amended upstream.
• In addition to the PEPs, dask complies with typeshed and mypy guidelines. Notably:
  • __init__ return type is not annotated
  • return types should not be unions to avoid burdening the caller with type checking, unless you want the caller to go through type checking at runtime.
    e.g. def f() -> int | str should be replaced with def f() -> Any.
    When designing new API, you should avoid having multiple return types unless you want the user to explicitly check for them; consider instead using TypeVars or breaking your function into two separate ones.
• All new code should be annotated, if sensible. Code should not be annotated if adding annotations would make it unreasonably cumbersome to read.
• When tweaking existing code, it is recommended (but not required) to spend some time adding annotations to the updated functions, unless it would substantially blow up the scope of the PR and/or make it much harder to review.
• This guidelines document should not be enforced in code review. If you are reviewing a PR and you see a violation, before asking it to be amended please ask yourself how much friction you'd introduce by demanding it to be compliant. Consider using a github suggestion so that the developer just has to click on "accept" instead of doing the work themselves.

mypy

• mypy version and settings are hardcoded project-wide. For the sake of coherence and reproducibility, you should call mypy exclusively through pre-commit.
• You should use # type: ignore only when working around it would be unreasonable.
• If you use # type: ignore to work around a bug in mypy, you should add a FIXME with a cross-reference to the mypy issue on github.

Specific vs. generic

Specific is better than generic, unless it unreasonably harms readability.

For example, this is bad:

d: dict = {}

This is good:

d: dict[str, Callable[[str], bool]] = {}

However, this is bad:

d: dict[
    str | MyClass1 | MyClass2 | tuple[int, str, MyClass3],
    Callable[[str, SomethingElse, AndOneMore], bool]
    | Callable[[int, numpy.ndarray | dask.array.Array], bool]
] = {}

in the above case, it is better to be more generic for the sake of readability:

d: dict[Any, Callable[..., bool]] = {}

Frequently, components of very complex signatures are used repeatedly across a module; you should consider using TypeAlias (in the example above, there could be a TypeAlias for the dict keys).

You should use @overload, but only when the added verbosity is outweighed by the benefit.

Parameters and return types

Parameter types should be as abstract as possible - anything that is valid.
Return types should be as concrete as possible (however, do not use unions in return types - as explained above).
Prefer immutable collections in parameters whenever possible.

e.g.

def f(d: Mapping[str, int | None]) -> dict[str, int]:
    return {k: v for k, v in d.items() if v is not None}

Backporting

You may use annotations from the very latest version of Python, as follows:

from __future__ import annotations

from typing import TYPE_CHECKING

if TYPE_CHECKING:
    # TODO import from typing (requires Python >=3.10)
    from typing_extensions import TypeAlias

    UserID: TypeAlias = int | str

The TYPE_CHECKING guard is necessary as typing_extensions is not a runtime dependency.
The TODO comment is strongly recommended to make future upgrade work easier.

TYPE_CHECKING should be used only when actually necessary.

Delayed annotations

At the moment of writing, dask supports Python 3.8+ at runtime, but uses Python 3.9+ annotations.

from __future__ import annotations

x: int | dict[str, int]

Don't use Union, Optional, Dict, etc. or quoted annotations unless necessary.
Notable exception are cast and subclassing (see next paragraph), which are interpreted at runtime:

x = cast("int | dict[str, int]", x)

In this case, quoted annotations are preferred to Union etc.

typing vs. collections.abc

Always import everything that's possible from collections.abc.
Import from typing only when an object doesn't exist elsewhere.
Again, this requires from __future__ import annotations to work in Python 3.8:

from __future__ import annotations

from collections.abc import Mapping

x: Mapping[str, int]

The only exception is when subclassing a collection. This is the only way to make it work in Python 3.8:

# TODO: import from collections.abc (requires Python >=3.9)
from typing import Mapping

class MyDict(Mapping[str, int]):
    ...

In this case, you should import the Mapping class only from typing. Don't import it from collections.abc as well.

Class annotations

You should always annotate all instance variables, both public and private, in the class header (C++ style).
Class variables should be explicitly marked with ClassVar.

mypy implements some intelligence to infer the type of instance variables from the __init__ method; it should not be relied upon.

Don't:

class C:
    x = 1
    def __init__(self):
        self.d: dict[str, int] = {}

Do:

class C:
    x: ClassVar[int] = 1
    d: dict[str, int]

    def __init__(self):
        self.d = {}

It is a good idea to keep Sphinx documentation together with annotations.
It is a good idea to use annotations to avoid explicitly declaring slots.

e.g.

class C:
    #: This is rendered by sphinx; don't forget the `:`!
    #: Line 2
    x: int

    __slots__ = tuple(__annotations__)

In sphinx:

.. autoclass:: mymodule.C
   :members:

Redundant annotations

Redundancy should be avoided.

Type specifications should also be avoided in Sphinx documentation when redundant with the annotations in the code.

For example:

class C:
    d: dict[str, int]

    def f(self, x: str):
        """Some description

        Parameters
        ----------
        x: int  # BAD - it's redundant with the signature; just use `x: `
            Some description
        """
        # BAD - the type of y can be fully inferred by the declaration of the instance variable
        y: int = self.d[x]

None

Defaults to None should be explicit, e.g. x: int | None = None.
You should always prefer | None to Optional or | NoneType.

[jacobtomlinson]

Thanks for starting this discussion @crusaderky. I'm very much in favour of adding more guidance around good practices for type annotations generally.

I'm less excited about having a style guide like this tucked away in the documentation somewhere that we nitpick PRs over. Especially community contributions by folks who don't work on Dask full time. I don't necessarily see the added friction both in terms of contributor time and reviewer time worth the tradeoff.

I like the way we enforce mypy, isort, black, etc via pre-commit and CI. Especially when things like black make consistency adjustments automatically. Is there any way we can tighten the mypy config to enforce some of this automatically?

Your guidelines above also focus a lot on "what" we should be doing, but not "why" is it beneficial for us to do it. Personally I like type annotations because VSCode gives me more helpful hinting. Other projects like typer depend on them directly for functionality. What benefits do you see enforcing tighter annotation rules having for the community?

[hendrikmakait]

Thanks for starting a discussion/guidelines document on this. Another thing I have noticed is that when fully following the numpydoc docstring style, we create redundant type information which (to the best of my knowledge) is never checked for consistency with the actual type hinting. This feels prone to errors resulting from updates to the function's signature without updating the docstring. See https://github.com/dask/distributed/blob/bd98e66039a38851ea33295a927255dcd5319601/distributed/scheduler.py#L3274-L3279 for an example of a typed docstring.

To avoid this, I suggest adjusting the docstring guidelines with a rule to avoid typing in docstrings.

[ian-r-rose]

Thanks for starting this discussion @crusaderky, it's useful to air a lot of this out. Most of what you have here is sensible to me, and I'd be happy to have it be made quasi-official. But I also agree with @jacobtomlinson that having these stylistic/linting guidelines be enforced with CI tools is far superior to nit-picking PR-by-PR. Unfortunately, I'm not aware of any configuration options for mypy that really address most of the stuff here. So I think I'd be in favor of making something like this a recommendation rather than a requirement.

A few specific thoughts:

Parameter types should be as abstract as possible - anything that is valid.

This is a really nice rule that is hard to enforce (a lot of contributors don't know about the reasoning for this, and digressing into type variance in a PR review is a tough ask). It would be super valuable if a tool like mypy could suggest "hey, it looks like you aren't mutating this parameter, consider using Sequence", but I don't think that exists today.

• return types should never be unions to avoid burdening the caller with type checking.
  e.g. def f() -> int | str must be replaced with def f() -> Any.

This was surprising to me, and I disagree with it. I see that this is listed in the typeshed style guide, which itself defers to this issue, which I don't find particularly convincing. The argument seems to be something like "it's annoying to check for different return types", which is true! But to me that reflects either (1) an underlying problem with the API, or (2) something that should be checked. To look at a special case: a huge fraction of bugs in a generic Python project (and I don't think dask is an exception) come from not checking for optional None in a return value. We should be doing that in many more places.

You should use @overload, but only when the added verbosity is outweighed by the benefit.

This is similar to the above: I sort of agree with this, but with the caveat that a lot of overloads probably meant we wrote a confusing API.

I wonder if there is a distinction to be made between annotating old code and annotating new code. If we have some old code that would involve a lot of overloads, for instance, it makes sense to me that we avoid that. But if new code would involve a lot of overloads, I think it would make sense to take a step back and reconsider the signature of a function.

[douglasdavis]

I'd like to +1 the idea of making CI enforcement the one and only requirement. That is, if things are passing in mypy then that particular patch is pretty much good to go, but small recommendations could be useful in PR discussion if clear/easy to address. The "return types should never be unions" topic is a perfect example IMO. It's not configurable with mypy and things like this exist in the codebase (e.g. __dask_graph__ can return None and xarray depends on that behavior). If it could be configured with mypy, perhaps it would be a good way to find some rough spots and do some cleanups, but in general it seems like it would be a good recommendation for new code, and not a specific requirement.

Chiming in on this comment from @jacobtomlinson:

What benefits do you see enforcing tighter annotation rules having for the community?

I think good type annotations are great developer documentation. I really liked Łukasz Langa's talk from PyCon this year, where inline docs is something he specifically mentions. Everything else in the talk is great too

[crusaderky]

when fully following the numpydoc docstring style, we create redundant type information

Good point. I updated the opening post.

I'm less excited about having a style guide like this tucked away in the documentation somewhere that we nitpick PRs over. Especially community contributions by folks who don't work on Dask full time. I don't necessarily see the added friction both in terms of contributor time and reviewer time worth the tradeoff.

I agree - we should not reject a PR based on subtle transgressions to these guidelines.

What benefits do you see enforcing tighter annotation rules having for the community?

A more consistent code style across the project - one that is shared with typeshed and hopefully more projects - makes it easier for newcomers to read the code.

having these stylistic/linting guidelines be enforced with CI tools is far superior to nit-picking PR-by-PR. Unfortunately, I'm not aware of any configuration options for mypy that really address most of the stuff here.

I just enabled

allow_incomplete_defs = false
allow_untyped_decorators = false
ignore_missing_imports = true
no_implicit_optional = true
show_error_codes = true
warn_redundant_casts = true
warn_unused_ignores = true
warn_unreachable = true

in dask/distributed and I plan to do the same in dask/dask.

So I think I'd be in favor of making something like this a recommendation rather than a requirement.

I agree.

This is a really nice rule that is hard to enforce (a lot of contributors don't know about the reasoning for this, and digressing into type variance in a PR review is a tough ask).

I agree. I've added a paragraph above stating that these guidelines should not slow down review.

It would be super valuable if a tool like mypy could suggest "hey, it looks like you aren't mutating this parameter, consider using Sequence", but I don't think that exists today.

Not that I'm aware of I'm afraid.

"it's annoying to check for different return types", which is true! But to me that reflects either (1) an underlying problem with the API, or (2) something that should be checked.

It's 99% of the times (1). An egregious case we have in house is distributed.Client. Most of its public methods return XYZ | Awaitable[XYZ] depending on the synchronous flag in the init. It would be catastrophic to write that down though, as it would make user code needlessly cumbersome.
To clarify: the way those methods work was a very good design when it was first implemented - but it's a design that predates annotations, and I can't see anything we can do now to make it annotations-friendly without completely breaking the public API.

I've added a paragraph stating that, when writing new API, you should make an effort to design it in a way that avoids unions in the return type.

[mkitti]

Could you clarify None vs NoneType, reintroduced in 3.10?

[crusaderky]

Could you clarify None vs NoneType, reintroduced in 3.10?

Could you articulate? I've never seen NoneType used in annotations

[mkitti]

Could you articulate? I've never seen NoneType used in annotations

None is a special case in PEP484. It was reintroduced in Python 3.10 via python/cpython#85976 (comment) . Outside of type annotations, None is clearly not a type.

In some cases, you can avoid mentioning either by using typing.Optional.

In [1]: from types import NoneType

In [2]: type(None)
Out[2]: NoneType

In [3]: isinstance(None, type)
Out[3]: False

In [4]: isinstance(NoneType, type)
Out[4]: True

In [5]: isinstance(None, NoneType)
Out[5]: True

In [6]: isinstance(None, None)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Input In [6], in <cell line: 1>()
----> 1 isinstance(None, None)

TypeError: isinstance() arg 2 must be a type, a tuple of types, or a union

In [7]: int | None == typing.Optional[int]
Out[7]: True

In [8]: int | NoneType == typing.Optional[int]
Out[8]: True

The conservative answer would be to stick with None rather than NoneType until a new PEP that supersedes PEP484 says otherwise. A programmer who prefers strict semantics might prefer to use NoneType. What is the guideline here between None, Optional, and NoneType?

[crusaderky]

I think | None is the most readable by far. I'll write it in the doc.

[multimeric]

Copying this from dask/dask#9233. I'm happy to help with the typing push. I just had two issues I wanted to discuss:

Dispatch

Basically we have an object containing a dictionary that maps types to functions that apply to those types, e.g. if we pass in a tuple to make_meta we get back a pd.Series. I'm sympathetic to this idea (having experience with R and its multiple dispatch methods), but unfortunately it's totally incompatible with Python's current type annotations, because we are effectively adding new method overloads at runtime. Actually the only way I could devise to type this system without entirely rewriting it, is to have each dispatch instance have its own Dispatch subclass with a ton of @overloads. e.g.

make_meta_dispatch = Dispatch("make_meta_dispatch")

becomes

class MakeMetaDispatch(Dispatch):
    @overload
    def __call__(self, arg: tuple[str, np.dtype], *args, **kwargs) -> pd.Series:
        ...
    @overload
    def __call__(self, arg: dict[str, np.dtype], *args, **kwargs) -> pd.DataFrame:
        ...
    def __call__(self, arg, *args, **kwargs):
        super().__call__(arg, *args, **kwargs)

make_meta_dispatch = MakeMetaDispatch("make_meta_dispatch")

I think this would work, but it's a bit unfortunate because the dispatch system is no longer decentralised like it was before: all the overload types have to be specified at the time of dispatch creation, even if the actual method body does not. I'm not sure this would be acceptable to the designer of this multiple dispatch system in the first place, but maybe it would be an improvement?

Docstrings

As far as I can tell, dask is using numpy docstrings, but I don't believe that system supports type annotations. What this effectively means is that you have to repeat yourself when you define the type signature and in the first line of the docstring. I'm not sure there's any easy solution that doesn't just involve using a different docstring style (like Google, which seems to handle this).

[crusaderky]

@multimeric

dispatch

Do I understand correctly that this paragraph is strictly about dask.dispatch.make_meta? If so, why can't we just annotate the function, as it is, with @overload + TypeVar?
If this one specific function cannot be accurately annotated without great effort and/or without making the code unreadable, it should be left unannotated or with just Any's.

docstrings

The paragraph in the opening comment, Redundant annotations`, addresses this.

[multimeric]

No, it's not just about make_meta. It's about type checking any function that calls a dispatch. make_meta is an example of that, but we can't check the correctness of make_meta itself, nor any other function that calls a dispatch without working out how to type these dispatch functions.