text: # Copyright 2022, The TensorFlow Federated Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #      http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """Defines an abstract interface for representing a federated context."""
  
  import abc
  from typing import Any
  
  from tensorflow_federated.python.common_libs import py_typecheck
  from tensorflow_federated.python.core.api import computation_base
  from tensorflow_federated.python.core.impl.context_stack import context_base
  from tensorflow_federated.python.core.impl.context_stack import get_context_stack
  from tensorflow_federated.python.core.impl.types import computation_types
  from tensorflow_federated.python.core.impl.types import placements
  from tensorflow_federated.python.core.impl.types import type_analysis
  
  
  def contains_only_server_placed_data(
      type_signature: computation_types.Type) -> bool:
    """Determines if `type_signature` contains only server-placed data.
  
    Determines if `type_signature` contains only:
    * `tff.StructType`s
    * `tff.SequenceType`s
    * server-placed `tff.FederatedType`s
    * `tff.TensorType`s
  
    Args:
        type_signature: The `tff.Type` to test.
  
    Returns:
      `True` if `type_signature` contains only server-placed data, otherwise
      `False`.
    """
    py_typecheck.check_type(type_signature, computation_types.Type)
  
    def predicate(type_spec: computation_types.Type) -> bool:
      return (type_spec.is_struct() or
              (type_spec.is_federated() and
               type_spec.placement is placements.SERVER) or
              type_spec.is_sequence() or type_spec.is_tensor())
  
    return type_analysis.contains_only(type_signature, predicate)
  
  
  class FederatedContext(context_base.Context):
    """An abstract interface representing a federated context.
  
    A federated context supports invoking a limited set of `tff.Computation`s,
    making guarantees about what a `tff.Computation` can accept as an argument and
    what it returns when invoked.
  
    ## Restrictions on the TensorFlow Federated Type
  
    Arguments can be nested structures of values corresponding to the TensorFlow
    Federated type signature of the `tff.Computation`:
  
    *   Server-placed values must be represented by `numpy` arrays, `tf.Tensors`,
        or `tff.program.MaterializableValueReference`
    *   Client-placed values must be represented by structures of values returned
        by a `tff.program.FederatedDataSourceIterator`.
  
    Return values can be structures of `tff.program.MaterializableValueReference`s
    or a single `tff.program.MaterializableValueReference`, where a reference
    corresponds to the tensor-type of the TensorFlow Federated type signature in
    the return value of the invoked `tff.Computation`.
  
    ## TensorFlow Federated Type to Python Representation
  
    In order to interact with the value returned by a `tff.Computation`, it is
    helpful to be able to reason about the Python type of this value. In some way
    this Python type must depend on the TensorFlow Federated type signature of the
    associated value. To provide uniformity of experience and ease of reasoning,
    we specify the Python representation of values in a manner that can be stated
    entirely in the TensorFlow Federated typesystem.
  
    We have choosen to limit the TensorFlow Federated type signatures of invoked
    `tff.Computation`s to disallow the returning of client-placed values,
    `tff.SequenceTypes`, and `tff.FunctionTypes`, in order to reduced the area
    which needs to be supported by federated programs. Below we describe the
    mapping between TensorFlow Federated type signatures and Python
    representations of values that can be passed as arguments to or returned as
    results from `tff.Computation`s.
  
    Python representations of values that can be *accepted as an arguments to* or
    *returned as a value from* a `tff.Computation`:
  
    | TensorFlow Federated Type  | Python Representation                      |
    | -------------------------- | ------------------------------------------ |
    | `tff.TensorType`           | `tff.program.MaterializableValueReference` |
    | `tff.SequenceType`         | `tff.program.MaterializableValueReference` |
    | `tff.StructWithPythonType` | Python container of the                    |
    :                            : `tff.StructWithPythonType`                 :
    | `tff.StructType` (with no  | `collections.OrderedDict`                  |
    : Python type, all fields    :                                            :
    : named)                     :                                            :
    | `tff.StructType` (with no  | `tuple`                                    |
    : Python type, no fields     :                                            :
    : named)                     :                                            :
    | `tff.FederatedType`        | Python representation of the `member` of   |
    : (server-placed)            : the `tff.FederatedType`                    :
  
    Python representations of values that can be only be *accepted as an arguments
    to* a `tff.Computation`:
  
    | TFF Type            | Python Representation                   |
    | ------------------- | --------------------------------------- |
    | `tff.FederatedType` | Opaque object returned by               |
    : (client-placed)     : `tff.program.DataSourceIterator.select` :
    | `tff.FunctionType`  | `tff.Computation`                       |
    """
  
    @abc.abstractmethod
    def invoke(self, comp: computation_base.Computation, arg: Any) -> Any:
      """Invokes the `comp` with the argument `arg`.
  
      Args:
        comp: The `tff.Computation` being invoked.
        arg: The optional argument of `comp` corresponding to the TensorFlow
          Federated type signature of the `tff.Computation`; server-placed values
          must be represented by `numpy` arrays, `tf.Tensors`, or
          `tff.program.MaterializableValueReference`, and client-placed values
          must be represented by structures of values returned by a
          `tff.program.FederatedDataSourceIterator`.
  
      Returns:
        The result of invocation; will only contain structures of
        `tff.program.MaterializableValueReference`s or a single
        `tff.program.MaterializableValueReference`.
  
      Raises:
        ValueError: If the result type of the invoked comptuation does not contain
        only structures, server-placed values, or tensors.
      """
      raise NotImplementedError
  
  
  def check_in_federated_context():
    """Checks if the current context is a `tff.program.FederatedContext`."""
    context_stack = get_context_stack.get_context_stack()
    if not isinstance(context_stack.current, FederatedContext):
      raise ValueError(
          'Expected the current context to be a `tff.program.FederatedContext`, '
          f'found \'{type(context_stack.current)}\'.')
  