From fb4aa4b8e2a17223188ce3c3c04e929a179fda10 Mon Sep 17 00:00:00 2001
From: Mike Dusenberry <dusenberrymw@google.com>
Date: Tue, 12 Dec 2023 19:52:23 -0800
Subject: [PATCH] Fixes the type signature of retry_exception_types.

code was failing previously when passed in a list of Exception types for
`retry_exception_types`. the type `list[Exception]` indicates a list of
Exception *objects*, while `list[type[Exception]]` indicates a list of
Exception [sub]classes.

https://docs.python.org/3/library/typing.html#the-type-of-class-objects

PiperOrigin-RevId: 590419864
---
 edward2/maps.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/edward2/maps.py b/edward2/maps.py
index f7e7b75a..d29fe917 100644
--- a/edward2/maps.py
+++ b/edward2/maps.py
@@ -38,6 +38,7 @@ def robust_map(
     max_retries: int | None = ...,
     max_workers: int | None = ...,
     raise_error: Literal[False] = ...,
+    retry_exception_types: list[type[Exception]] | None = ...,
 ) -> Sequence[U | V]:
   ...
 
@@ -52,6 +53,7 @@ def robust_map(
     max_retries: int | None = ...,
     max_workers: int | None = ...,
     raise_error: Literal[True] = ...,
+    retry_exception_types: list[type[Exception]] | None = ...,
 ) -> Sequence[U]:
   ...
 
@@ -66,7 +68,7 @@ def robust_map(
     max_retries: int | None = None,
     max_workers: int | None = None,
     raise_error: bool = False,
-    retry_exception_types: list[Exception] | None = None,
+    retry_exception_types: list[type[Exception]] | None = None,
 ) -> Sequence[U | V]:
   """Maps a function to inputs using a threadpool.