Use a Protocol for TRANSFORMER to ensure common arg names

Also cleans up some of the internals of the transformer decorator and
simplifies the types.

Follow-up to #4797

Author: maffoo

cirq-core/cirq/transformers/transformer_api.py

@@ -18,17 +18,17 @@
 import functools
 from typing import (
     Any,
-    Callable,
     Tuple,
     Hashable,
     List,
-    Type,
     overload,
+    Protocol,
+    Type,
     TYPE_CHECKING,
+    TypeVar,
 )
 import dataclasses
 import enum
-from cirq.circuits.circuit import CIRCUIT_TYPE
 
 if TYPE_CHECKING:
     import cirq
@@ -218,77 +218,41 @@ class TransformerContext:
     ignore_tags: Tuple[Hashable, ...] = ()
 
 
-TRANSFORMER = Callable[['cirq.AbstractCircuit', TransformerContext], 'cirq.AbstractCircuit']
-_TRANSFORMER_TYPE = Callable[['cirq.AbstractCircuit', TransformerContext], CIRCUIT_TYPE]
-
-
-def _transform_and_log(
-    func: _TRANSFORMER_TYPE[CIRCUIT_TYPE],
-    transformer_name: str,
-    circuit: 'cirq.AbstractCircuit',
-    context: TransformerContext,
-) -> CIRCUIT_TYPE:
-    """Helper to log initial and final circuits before and after calling the transformer."""
-
-    context.logger.register_initial(circuit, transformer_name)
-    transformed_circuit = func(circuit, context)
-    context.logger.register_final(transformed_circuit, transformer_name)
-    return transformed_circuit
-
-
-def _transformer_class(
-    cls: Type[_TRANSFORMER_TYPE[CIRCUIT_TYPE]],
-) -> Type[_TRANSFORMER_TYPE[CIRCUIT_TYPE]]:
-    old_func = cls.__call__
-
-    def transformer_with_logging_cls(
-        self: Type[_TRANSFORMER_TYPE[CIRCUIT_TYPE]],
-        circuit: 'cirq.AbstractCircuit',
-        context: TransformerContext,
-    ) -> CIRCUIT_TYPE:
-        def call_old_func(c: 'cirq.AbstractCircuit', ct: TransformerContext) -> CIRCUIT_TYPE:
-            return old_func(self, c, ct)
-
-        return _transform_and_log(call_old_func, cls.__name__, circuit, context)
+class TRANSFORMER(Protocol):
+    def __call__(
+        self, circuit: 'cirq.AbstractCircuit', context: TransformerContext
+    ) -> 'cirq.AbstractCircuit':
+        ...
 
-    setattr(cls, '__call__', transformer_with_logging_cls)
-    return cls
 
-
-def _transformer_func(func: _TRANSFORMER_TYPE[CIRCUIT_TYPE]) -> _TRANSFORMER_TYPE[CIRCUIT_TYPE]:
-    @functools.wraps(func)
-    def transformer_with_logging_func(
-        circuit: 'cirq.AbstractCircuit',
-        context: TransformerContext,
-    ) -> CIRCUIT_TYPE:
-        return _transform_and_log(func, func.__name__, circuit, context)
-
-    return transformer_with_logging_func
+_TRANSFORMER_T = TypeVar('_TRANSFORMER_T', bound=TRANSFORMER)
+_TRANSFORMER_CLS_T = TypeVar('_TRANSFORMER_CLS_T', bound=Type[TRANSFORMER])
 
 
 @overload
-def transformer(cls_or_func: _TRANSFORMER_TYPE[CIRCUIT_TYPE]) -> _TRANSFORMER_TYPE[CIRCUIT_TYPE]:
+def transformer(cls_or_func: _TRANSFORMER_T) -> _TRANSFORMER_T:
     pass
 
 
 @overload
-def transformer(
-    cls_or_func: Type[_TRANSFORMER_TYPE[CIRCUIT_TYPE]],
-) -> Type[_TRANSFORMER_TYPE[CIRCUIT_TYPE]]:
+def transformer(cls_or_func: _TRANSFORMER_CLS_T) -> _TRANSFORMER_CLS_T:
     pass
 
 
 def transformer(cls_or_func: Any) -> Any:
     """Decorator to verify API and append logging functionality to transformer functions & classes.
 
-    The decorated function or class must satisfy
-    `Callable[[cirq.Circuit, cirq.TransformerContext], cirq.Circuit]` API. For Example:
+    A transformer is a callable that takes as inputs a cirq.AbstractCircuit and
+    cirq.TransformerContext, and returns another cirq.AbstractCircuit without
+    modifying the input circuit. A transformer could be a function, for example:
 
     >>> @cirq.transformer
-    >>> def convert_to_cz(circuit: cirq.Circuit, context: cirq.TransformerContext) -> cirq.Circuit:
+    >>> def convert_to_cz(
+    >>>    circuit: cirq.AbstractCircuit, context: cirq.TransformerContext
+    >>> ) -> cirq.Circuit:
     >>>    ...
 
-    The decorated class must implement the `__call__` method to satisfy the above API.
+    Or it could be a class that implements `__call__` with the same API, for example:
 
     >>> @cirq.transformer
     >>> class ConvertToSqrtISwaps:
@@ -300,14 +264,45 @@ def transformer(cls_or_func: Any) -> Any:
     >>>        ...
 
     Args:
-        cls_or_func: The callable class or method to be decorated.
+        cls_or_func: The callable class or function to be decorated.
 
     Returns:
-        Decorated class / method which includes additional logging boilerplate. The decorated
-        callable always receives a copy of the input circuit so that the input is never mutated.
+        Decorated class / function which includes additional logging boilerplate.
     """
     if isinstance(cls_or_func, type):
-        return _transformer_class(cls_or_func)
+        cls = cls_or_func
+        method = cls.__call__
+
+        @functools.wraps(method)
+        def method_with_logging(self, circuit, context):
+            return _transform_and_log(
+                lambda circuit, context: method(self, circuit, context),
+                cls.__name__,
+                circuit,
+                context,
+            )
+
+        setattr(cls, '__call__', method_with_logging)
+        return cls
     else:
         assert callable(cls_or_func)
-        return _transformer_func(cls_or_func)
+        func = cls_or_func
+
+        @functools.wraps(func)
+        def func_with_logging(circuit, context):
+            return _transform_and_log(func, func.__name__, circuit, context)
+
+        return func_with_logging
+
+
+def _transform_and_log(
+    func: TRANSFORMER,
+    transformer_name: str,
+    circuit: 'cirq.AbstractCircuit',
+    context: TransformerContext,
+) -> 'cirq.AbstractCircuit':
+    """Helper to log initial and final circuits before and after calling the transformer."""
+    context.logger.register_initial(circuit, transformer_name)
+    transformed_circuit = func(circuit, context)
+    context.logger.register_final(transformed_circuit, transformer_name)
+    return transformed_circuit