Using Python type hints

[jpchen]

Since Pyro supports exclusively Python 3, we should look into adding type hints. This has a few advantages:

• pep8 linters can check for type errors
• easier for a reader to understand function signatures. useful for eg funsor where multiple dispatching is prevalent.
• sphinx docs can autogenerate types of args
• IDE static analysis tools will work (better)

e.g.

from typing import List, Optional

def foo(bar: List, baz: Optional[List] = None) -> List:
    return bar

Add type hints to modules

• [ ] contrib
• [ ] distributions
• [ ] distributions.transforms
• [ ] distributions.constraints
• [ ] infer
• [x] nn
  • #3337
  • #3342
• [ ] ops
• [x] optim #2853
• [x] params #3271
• [x] poutine
  • #3321
  • #3312
  • #3309
  • #3308
  • #3306
  • #3299
  • #3295
  • #3292
  • #3290
  • #3288
• [ ] scripts
• [ ] tests.test_primitives
• [ ] tests.test_generic (low priority)
• [ ] tests.poutine
• [ ] tests.ops
• [ ] tests.optim
• [ ] tests.perf
• [ ] tests.nn
• [ ] tests.infer
• [ ] tests.distributions
• [ ] tests.contrib

Additional functionality for type hints

• [ ] Give warnings on #type ignore in codebase Grey might be a good idea for finding #type ignores
• [ ] Give warnings on modules skipped by mypy

[fritzo]

I guess we should add a mypy test stage

[fritzo]

This would also allow us to remove sphinx type annotations, if we were to use https://pypi.org/project/sphinx-autodoc-typehints/

[kamathhrishi]

Can I work on this issue? Adding type hints to the existing codebase? Maybe not the entire codebase in a PR, but at least some sections of it. Then subsequent PR's could cover the codebase.

[fritzo]

Hi @kamathhrishi, sure, we'd appreciate any help! I especially like the idea of a bunch of small easy-to-review PRs that add type hints to one module at a time.

[kamathhrishi]

Thank you @fritzo. I have made a PR for optim module. Review it and let me know. I am new to the codebase and would like to take this as a starting step.

[fritzo]

It looks like we should also add py.types stub files to various directories to avoid errors like

infer.py:4: error: Skipping analyzing "pyro": module is installed, but missing library stubs or py.typed marker
infer.py:4: note: See https://mypy.readthedocs.io/en/stable/running_mypy.html#missing-imports
infer.py:5: error: Skipping analyzing "pyro.distributions": module is installed, but missing library stubs or py.typed marker
infer.py:8: error: Skipping analyzing "pyro.infer": module is installed, but missing library stubs or py.typed marker
infer.py:9: error: Skipping analyzing "pyro.infer.autoguide": module is installed, but missing library stubs or py.typed marker
infer.py:10: error: Skipping analyzing "pyro.infer.reparam": module is installed, but missing library stubs or py.typed marker
infer.py:11: error: Skipping analyzing "pyro.optim": module is installed, but missing library stubs or py.typed marker

[willtai]

Can I take a look at parts of this? I can open a PR for params