Pass sweep_transformer recursively (#6951)

* fix bug

* rm print

Author: senecameeks

cirq-google/cirq_google/api/v2/sweeps.py

@@ -82,19 +82,27 @@ def sweep_to_proto(
     elif isinstance(sweep, cirq.Product):
         out.sweep_function.function_type = run_context_pb2.SweepFunction.PRODUCT
         for factor in sweep.factors:
-            sweep_to_proto(factor, out=out.sweep_function.sweeps.add())
+            sweep_to_proto(
+                factor, out=out.sweep_function.sweeps.add(), sweep_transformer=sweep_transformer
+            )
     elif isinstance(sweep, cirq.ZipLongest):
         out.sweep_function.function_type = run_context_pb2.SweepFunction.ZIP_LONGEST
         for s in sweep.sweeps:
-            sweep_to_proto(s, out=out.sweep_function.sweeps.add())
+            sweep_to_proto(
+                s, out=out.sweep_function.sweeps.add(), sweep_transformer=sweep_transformer
+            )
     elif isinstance(sweep, cirq.Zip):
         out.sweep_function.function_type = run_context_pb2.SweepFunction.ZIP
         for s in sweep.sweeps:
-            sweep_to_proto(s, out=out.sweep_function.sweeps.add())
+            sweep_to_proto(
+                s, out=out.sweep_function.sweeps.add(), sweep_transformer=sweep_transformer
+            )
     elif isinstance(sweep, cirq.Concat):
         out.sweep_function.function_type = run_context_pb2.SweepFunction.CONCAT
         for s in sweep.sweeps:
-            sweep_to_proto(s, out=out.sweep_function.sweeps.add())
+            sweep_to_proto(
+                s, out=out.sweep_function.sweeps.add(), sweep_transformer=sweep_transformer
+            )
     elif isinstance(sweep, cirq.Linspace) and not isinstance(sweep.key, sympy.Expr):
         sweep = cast(cirq.Linspace, sweep_transformer(sweep))
         out.single_sweep.parameter_key = sweep.key
@@ -143,7 +151,11 @@ def sweep_to_proto(
                 sweep_dict[cast(str, key)].append(cast(float, param_resolver.value_of(key)))
         out.sweep_function.function_type = run_context_pb2.SweepFunction.ZIP
         for key in sweep_dict:
-            sweep_to_proto(cirq.Points(key, sweep_dict[key]), out=out.sweep_function.sweeps.add())
+            sweep_to_proto(
+                cirq.Points(key, sweep_dict[key]),
+                out=out.sweep_function.sweeps.add(),
+                sweep_transformer=sweep_transformer,
+            )
     else:
         raise ValueError(f'cannot convert to v2 Sweep proto: {sweep}')
     return out
@@ -164,7 +176,7 @@ def sweep_from_proto(
     if which is None:
         return cirq.UnitSweep
     if which == 'sweep_function':
-        factors = [sweep_from_proto(m) for m in msg.sweep_function.sweeps]
+        factors = [sweep_from_proto(m, sweep_transformer) for m in msg.sweep_function.sweeps]
         func_type = msg.sweep_function.function_type
         if func_type == run_context_pb2.SweepFunction.PRODUCT:
             return cirq.Product(*factors)

cirq-google/cirq_google/api/v2/sweeps_test.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import math
+from copy import deepcopy
 from typing import Iterator
 
 import pytest
@@ -275,6 +276,64 @@ def strip_tunit_func(sweep: sweeps.SingleSweep):
     assert list(sweep.points)[0] == 1.0
 
 
+@pytest.mark.parametrize(
+    'sweep',
+    [
+        cirq.Concat(cirq.Points('a', [1, 2, 3]), cirq.Points('a', [4])),
+        cirq.Points('a', [1, 2, 3]) * cirq.Points('b', [4, 5, 6]),
+        cirq.ZipLongest(cirq.Points('a', [1, 2, 3]), cirq.Points('b', [1])),
+        cirq.Zip(cirq.Points('a', [1, 2, 3]), cirq.Points('b', [4, 5, 6])),
+    ],
+)
+def test_sweep_to_proto_with_func_on_resursive_sweep_succeeds(sweep):
+    def add_tunit_func(sweep: sweeps.SingleSweep):
+        if isinstance(sweep, cirq.Points):
+            sweep.points = [point * tunits.ns for point in sweep.points]  # type: ignore[misc]
+
+        return sweep
+
+    msg = v2.sweep_to_proto(sweep, sweep_transformer=add_tunit_func)
+
+    assert msg.sweep_function.sweeps[0].single_sweep.points.unit == tunits.ns.to_proto()
+
+
+@pytest.mark.parametrize(
+    'expected_sweep',
+    [
+        cirq.Concat(cirq.Points('a', [1.0, 2.0, 3.0]), cirq.Points('a', [4.0])),
+        cirq.Points('a', [1.0, 2.0, 3.0]) * cirq.Points('b', [4.0, 5.0, 6.0]),
+        cirq.ZipLongest(cirq.Points('a', [1.0, 2.0, 3.0]), cirq.Points('b', [1.0])),
+        cirq.Zip(cirq.Points('a', [1.0, 2.0, 3.0]), cirq.Points('b', [4.0, 5.0, 6.0])),
+        cirq.Points('a', [1, 2, 3])
+        + cirq.Points(
+            'b',
+            [4, 5, 6],
+            metadata=DeviceParameter(path=['path', 'to', 'parameter'], idx=2, units='GHz'),
+        ),
+    ],
+)
+def test_sweep_from_proto_with_func_on_resursive_sweep_succeeds(expected_sweep):
+    def add_tunit_func(sweep_to_transform: sweeps.SingleSweep):
+        sweep = deepcopy(sweep_to_transform)
+        if isinstance(sweep, cirq.Points):
+            sweep.points = [point * tunits.ns for point in sweep.points]  # type: ignore[misc]
+
+        return sweep
+
+    def strip_tunit_func(sweep_to_transform: sweeps.SingleSweep):
+        sweep = deepcopy(sweep_to_transform)
+        if isinstance(sweep, cirq.Points):
+            if isinstance(sweep.points[0], tunits.Value):
+                sweep.points = [point[point.unit] for point in sweep.points]
+
+        return sweep
+
+    msg = v2.sweep_to_proto(expected_sweep, sweep_transformer=add_tunit_func)
+    round_trip_sweep = v2.sweep_from_proto(msg, strip_tunit_func)
+
+    assert round_trip_sweep == expected_sweep
+
+
 def test_sweep_with_list_sweep():
     ls = cirq.study.to_sweep([{'a': 1, 'b': 2}, {'a': 3, 'b': 4}])
     proto = v2.sweep_to_proto(ls)