Add support for extensions that are only able to wrap some values (#274)

Author: Strilanc

cirq/extension/extensions.py

@@ -32,7 +32,9 @@ def __init__(
             self,
             desired_to_actual_to_wrapper: Optional[Dict[
                 Type[T_DESIRED],
-                Dict[Type[T_ACTUAL], Callable[[T_ACTUAL], T_DESIRED]]]]=None
+                Dict[Type[T_ACTUAL],
+                     Callable[[T_ACTUAL],
+                              Optional[T_DESIRED]]]]]=None
             ) -> None:
         """Specifies extensions.
 
@@ -48,12 +50,12 @@ def __init__(
             {}
             if desired_to_actual_to_wrapper is None
             else desired_to_actual_to_wrapper
-        )  # type: Dict[Type[Any], Dict[Any, Callable[[Any], Any]]]
+        )  # type: Dict[Type[Any], Dict[Any, Callable[[Any], Optional[Any]]]]
 
     def add_cast(self,
                  desired_type: Type[T_DESIRED],
                  actual_type: Type[T_ACTUAL],
-                 conversion: Callable[[T_ACTUAL], T_DESIRED],
+                 conversion: Callable[[T_ACTUAL], Optional[T_DESIRED]],
                  also_add_inherited_conversions: bool = True,
                  overwrite_existing: bool = False) -> None:
         """Adds a way to turn one type of thing into another.
@@ -139,7 +141,9 @@ def try_cast(self,
             for actual_type in inspect.getmro(type(actual_value)):
                 wrapper = actual_to_wrapper.get(actual_type)
                 if wrapper:
-                    return wrapper(actual_value)
+                    wrapped = wrapper(actual_value)
+                    if wrapped is not None:
+                        return wrapped
 
         if isinstance(actual_value, desired_type):
             return actual_value

cirq/extension/extensions_test.py

@@ -263,3 +263,17 @@ def test_add_cast_redundant_including_subtypes():
                overwrite_existing=True)
     assert e.try_cast(Child(), Aunt) is o3
     assert e.try_cast(Child(), Cousin) is o3
+
+
+def test_add_potential_cast():
+    a = Aunt()
+    c1 = Child()
+    c2 = Child()
+
+    e = extension.Extensions()
+    e.add_cast(desired_type=Aunt,
+               actual_type=Child,
+               conversion=lambda e: a if e is c1 else None)
+
+    assert e.try_cast(c1, Aunt) is a
+    assert e.try_cast(c2, Aunt) is None