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
• nn
  • Type annotate pyro.nn.module #3337
  • Type annotate pyro.nn.dense_nn and pyro.nn.auto_reg_nn #3342
• ops
• optim Type Hints for Optim module and Setup MyPy as part of CI #2853
• params Add param_store.py type hints #3271
• poutine
  • Annotate handlers & add py.typed #3321
  • Mypy warn_unreachable=True #3312
  • Type annotate messengers #3309
  • Type annotate messengers #3308
  • Type annotate poutines #3306
  • Type annotate Trace, TraceMessenger, & pyro.poutine.guide #3299
  • Type annotate broadcast, condition, do, enum messengers #3295
  • Type annotate pyro.primitives & poutine.block_messenger #3292
  • Type annotate pyro.poutine.messenger #3290
  • Type annotate pyro.poutine.runtime #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