Type annotate pyro.poutine.messenger (#3290)

pyro/poutine/guide.py

@@ -36,7 +36,7 @@ def __getstate__(self):
         state.pop("trace")
         return state
 
-    def __call__(self, *args, **kwargs) -> Dict[str, torch.Tensor]:
+    def __call__(self, *args, **kwargs) -> Dict[str, torch.Tensor]:  # type: ignore[override]
         """
         Draws posterior samples from the guide and replays the model against
         those samples.

pyro/poutine/messenger.py

@@ -1,13 +1,24 @@
 # Copyright (c) 2017-2019 Uber Technologies, Inc.
 # SPDX-License-Identifier: Apache-2.0
 
+from __future__ import annotations
+
 from contextlib import contextmanager
 from functools import partial
+from types import TracebackType
+from typing import Any, Callable, Iterator, List, Optional, Type, TypeVar, cast
+
+from .runtime import _PYRO_STACK, Message
 
-from .runtime import _PYRO_STACK
+_F = TypeVar("_F", bound=Callable)
 
 
-def _context_wrap(context, fn, *args, **kwargs):
+def _context_wrap(
+    context: Messenger,
+    fn: Callable,
+    *args: Any,
+    **kwargs: Any,
+) -> Any:
     with context:
         return fn(*args, **kwargs)
 
@@ -26,13 +37,17 @@ class _bound_partial(partial):
     def __init__(self, func):
         self.func = func
 
-    def __get__(self, instance, owner):
+    def __get__(
+        self,
+        instance: Optional[object],
+        owner: Optional[Type[object]] = None,
+    ) -> object:
         if instance is None:
             return self
         return partial(self.func, instance)
 
 
-def unwrap(fn):
+def unwrap(fn: Callable) -> Callable:
     """
     Recursively unwraps poutines.
     """
@@ -61,17 +76,15 @@ class Messenger:
     Most inference operations are implemented in subclasses of this.
     """
 
-    def __call__(self, fn):
+    def __call__(self, fn: _F) -> _F:
         if not callable(fn):
             raise ValueError(
-                "{} is not callable, did you mean to pass it as a keyword arg?".format(
-                    fn
-                )
+                f"{fn!r} is not callable, did you mean to pass it as a keyword arg?"
             )
         wraps = _bound_partial(partial(_context_wrap, self, fn))
-        return wraps
+        return cast(_F, wraps)
 
-    def __enter__(self):
+    def __enter__(self) -> Messenger:
         """
         :returns: self
         :rtype: pyro.poutine.Messenger
@@ -103,7 +116,12 @@ def __enter__(self):
             # but it could in principle be enabled...
             raise ValueError("cannot install a Messenger instance twice")
 
-    def __exit__(self, exc_type, exc_value, traceback):
+    def __exit__(
+        self,
+        exc_type: Optional[Type[Exception]],
+        exc_value: Optional[Exception],
+        traceback: Optional[TracebackType],
+    ) -> None:
         """
         :param exc_type: exception type, e.g. ValueError
         :param exc_value: exception instance?
@@ -146,30 +164,33 @@ def __exit__(self, exc_type, exc_value, traceback):
                 for i in range(loc, len(_PYRO_STACK)):
                     _PYRO_STACK.pop()
 
-    def _reset(self):
+    def _reset(self) -> None:
         pass
 
-    def _process_message(self, msg):
+    def _process_message(self, msg: Message) -> None:
         """
         :param msg: current message at a trace site
         :returns: None
 
         Process the message by calling appropriate method of itself based
         on message type. The message is updated in place.
         """
-        method = getattr(self, "_pyro_{}".format(msg["type"]), None)
+        method = getattr(self, f"_pyro_{msg['type']}", None)
         if method is not None:
-            return method(msg)
-        return None
+            method(msg)
 
-    def _postprocess_message(self, msg):
-        method = getattr(self, "_pyro_post_{}".format(msg["type"]), None)
+    def _postprocess_message(self, msg: Message) -> None:
+        method = getattr(self, f"_pyro_post_{msg['type']}", None)
         if method is not None:
-            return method(msg)
-        return None
+            method(msg)
 
     @classmethod
-    def register(cls, fn=None, type=None, post=None):
+    def register(
+        cls,
+        fn: Optional[Callable] = None,
+        type: Optional[str] = None,
+        post: Optional[bool] = None,
+    ) -> Callable:
         """
         :param fn: function implementing operation
         :param str type: name of the operation
@@ -197,7 +218,11 @@ def some_function(msg)
         return fn
 
     @classmethod
-    def unregister(cls, fn=None, type=None):
+    def unregister(
+        cls,
+        fn: Optional[Callable] = None,
+        type: Optional[str] = None,
+    ) -> Optional[Callable]:
         """
         :param fn: function implementing operation
         :param str type: name of the operation
@@ -227,7 +252,9 @@ def unregister(cls, fn=None, type=None):
 
 
 @contextmanager
-def block_messengers(predicate):
+def block_messengers(
+    predicate: Callable[[Messenger], bool]
+) -> Iterator[List[Messenger]]:
     """
     EXPERIMENTAL Context manager to temporarily remove matching messengers from
     the _PYRO_STACK. Note this does not call the ``.__exit__()`` and