The library is not typed

[RomainBrault]

Describe the bug

MR #211 add py.typed to enable type checking support. However the library is not fully typed (at least not the public interface, see for instance https://github.com/arogozhnikov/einops/blob/master/einops/layers/torch.py.

Now mypy is complaining on each einops.layers.torch class call (Reduce, Rearrange, etc.).

-> error: Call to untyped function "Rearrange" in typed context [no-untyped-call]

The library should not have a py.typed file until its public interface is fully typed.

Reproduction steps
run mypy on any source file containing, for instance, from einops.layers.torch import Rearrange

Expected behavior
mypy should not raise [no-untyped-call] errors

Your platform
Ubuntu 22.04

[RomainBrault]

Full example:

Run mypy on

from einops.layers.torch import Rearrange

Rearrange("a b -> b a")

yields

> mypy test.py                                                                                                                                                                                        
test.py:6:1: error: Call to untyped function "Rearrange" in typed context  [no-untyped-call]
    Rearrange("a b -> b a")
    ^~~~~~~~~~~~~~~~~~~~~~~
Found 1 error in 1 file (checked 1 source file)

and mypy should be configured (if it is not by default) to raise error on untyped call (with the following pyproject.toml):

[tool.mypy]
disallow_untyped_calls = true

[RomainBrault]

I can see two quick and dirty solutions:

• remove the py.typed
• type any untyped public functions with Any

[arogozhnikov]

for pytorch, is it only constructors that need to be annotated?
does mypy demand to annotate __call__ or some other functions?

If it's only constructors, may worth just annotating that (potentially, with Any).

[RomainBrault]

@arogozhnikov Yes in my case it seems that only the constructors are problematic (so far I tested Rearrange and Reduce in pytorch). The functional api does not seeems to raise (I only tried the rearrange function from pytorch).

[arogozhnikov]

@RomainBrault
I've annotated the rest of public API in this branch:
https://github.com/arogozhnikov/einops/tree/type-hinting

Check if that branch works with your setup

[RomainBrault]

Thanks! I'll have time to test this next weekeend. But just looking at the code, it looks fine.