Allow custom definition for _sympy_pass_through (#3921)

Fixes #3916
Review: @maffoo @95-martin-orion

cirq/protocols/resolve_parameters.py

@@ -49,6 +49,16 @@ def _resolve_parameters_(self: Any, param_resolver: 'cirq.ParamResolver', recurs
         """Resolve the parameters in the effect."""
 
 
+class ResolvableValue(Protocol):
+    @doc_private
+    def _resolved_value_(self) -> Any:
+        """Returns a resolved value during parameter resolution.
+
+        Use this to mark a custom type as "resolved", instead of requiring
+        further parsing like we do with Sympy symbols.
+        """
+
+
 def is_parameterized(val: Any) -> bool:
     """Returns whether the object is parameterized with any Symbols.
 

cirq/study/resolver.py

@@ -15,6 +15,7 @@
 """Resolves ParameterValues to assigned values."""
 import numbers
 from typing import Any, Dict, Iterator, Optional, TYPE_CHECKING, Union, cast
+
 import numpy as np
 import sympy
 from sympy.core import numbers as sympy_numbers
@@ -34,6 +35,9 @@
     """Something that can be used to turn parameters into values.""",
 )
 
+# Used to mark values that are being resolved recursively to detect loops.
+_RecursionFlag = object()
+
 
 class ParamResolver:
     """Resolves parameters to actual values.
@@ -96,8 +100,8 @@ def value_of(
         """
 
         # Input is a pass through type, no resolution needed: return early
-        v = _sympy_pass_through(value)
-        if v is not None:
+        v = _resolve_value(value)
+        if v is not NotImplemented:
             return v
 
         # Handles 2 cases:
@@ -106,8 +110,8 @@ def value_of(
         # In both cases, return it directly.
         if value in self.param_dict:
             param_value = self.param_dict[value]
-            v = _sympy_pass_through(param_value)
-            if v is not None:
+            v = _resolve_value(param_value)
+            if v is not NotImplemented:
                 return v
 
         # Input is a string and is not in the dictionary.
@@ -121,8 +125,8 @@ def value_of(
         # in the dictionary ({'a': 1.0}).  Return it.
         if isinstance(value, sympy.Symbol) and value.name in self.param_dict:
             param_value = self.param_dict[value.name]
-            v = _sympy_pass_through(param_value)
-            if v is not None:
+            v = _resolve_value(param_value)
+            if v is not NotImplemented:
                 return v
 
         # The following resolves common sympy expressions
@@ -166,13 +170,13 @@ def value_of(
         # single symbol, since combinations are handled earlier in the method.
         if value in self._deep_eval_map:
             v = self._deep_eval_map[value]
-            if v is not None:
+            if v is not _RecursionFlag:
                 return v
             raise RecursionError('Evaluation of {value} indirectly contains itself.')
 
         # There isn't a full evaluation for 'value' yet. Until it's ready,
         # map value to None to identify loops in component evaluation.
-        self._deep_eval_map[value] = None
+        self._deep_eval_map[value] = _RecursionFlag
 
         v = self.value_of(value, recursive=False)
         if v == value:
@@ -235,7 +239,7 @@ def _from_json_dict_(cls, param_dict, **kwargs):
         return cls(dict(param_dict))
 
 
-def _sympy_pass_through(val: Any) -> Optional[Any]:
+def _resolve_value(val: Any) -> Any:
     if isinstance(val, numbers.Number) and not isinstance(val, sympy.Basic):
         return val
     if isinstance(val, sympy_numbers.IntegerConstant):
@@ -244,4 +248,7 @@ def _sympy_pass_through(val: Any) -> Optional[Any]:
         return val.p / val.q
     if val == sympy.pi:
         return np.pi
-    return None
+
+    getter = getattr(val, '_resolved_value_', None)
+    result = NotImplemented if getter is None else getter()
+    return result

cirq/study/resolver_test.py

@@ -230,6 +230,32 @@ def test_resolve_unknown_type():
     assert r.value_of(cirq.X) == cirq.X
 
 
+def test_custom_resolved_value():
+    class Foo:
+        def _resolved_value_(self):
+            return self
+
+    class Bar:
+        def _resolved_value_(self):
+            return NotImplemented
+
+    class Baz:
+        def _resolved_value_(self):
+            return 'Baz'
+
+    foo = Foo()
+    bar = Bar()
+    baz = Baz()
+
+    a = sympy.Symbol('a')
+    b = sympy.Symbol('b')
+    c = sympy.Symbol('c')
+    r = cirq.ParamResolver({a: foo, b: bar, c: baz})
+    assert r.value_of(a) is foo
+    assert r.value_of(b) is b
+    assert r.value_of(c) == 'Baz'
+
+
 def test_compose():
     """
     Calling cirq.resolve_paramters on a ParamResolver composes that resolver