added type hints to lazy_property

randolf-scholz · randolf-scholz · commit ef2329f5cc9d · 2025-01-02T23:05:05.000+01:00
diff --git a/torch/distributions/utils.py b/torch/distributions/utils.py
@@ -1,7 +1,7 @@
 # mypy: allow-untyped-defs
 from functools import update_wrapper
 from numbers import Number
-from typing import Any, Dict
+from typing import Any, Callable, Dict, Generic, overload, TypeVar
 
 import torch
 import torch.nn.functional as F
@@ -130,19 +130,34 @@ def probs_to_logits(probs, is_binary=False):
     return torch.log(ps_clamped)
 
 
-class lazy_property:
+T = TypeVar("T", covariant=True)
+
+
+class lazy_property(Generic[T]):
     r"""
     Used as a decorator for lazy loading of class attributes. This uses a
     non-data descriptor that calls the wrapped method to compute the property on
     first call; thereafter replacing the wrapped method into an instance
     attribute.
     """
 
-    def __init__(self, wrapped):
-        self.wrapped = wrapped
+    def __init__(self, wrapped: Callable[..., T]) -> None:
+        self.wrapped: Callable[..., T] = wrapped
         update_wrapper(self, wrapped)  # type:ignore[arg-type]
 
-    def __get__(self, instance, obj_type=None):
+    @overload
+    def __get__(
+        self, instance: None, obj_type: Any = None
+    ) -> "_lazy_property_and_property[T]":
+        ...
+
+    @overload
+    def __get__(self, instance: object, obj_type: Any = None) -> T:
+        ...
+
+    def __get__(
+        self, instance: object, obj_type: Any = None
+    ) -> "T | _lazy_property_and_property[T]":
         if instance is None:
             return _lazy_property_and_property(self.wrapped)
         with torch.enable_grad():
@@ -151,14 +166,14 @@ def __get__(self, instance, obj_type=None):
         return value
 
 
-class _lazy_property_and_property(lazy_property, property):
+class _lazy_property_and_property(lazy_property[T], property):
     """We want lazy properties to look like multiple things.
 
     * property when Sphinx autodoc looks
     * lazy_property when Distribution validate_args looks
     """
 
-    def __init__(self, wrapped):
+    def __init__(self, wrapped: Callable[..., T]) -> None:
         property.__init__(self, wrapped)
 
 
