Add and update type annotations for BuildingBlocks.

* Added missing type annotations.
* Fixed existing type annotations.
* Removed type annotations for `cls`.
* Removed `Any` type annotations.
* Fixed lint errors.

PiperOrigin-RevId: 556914851

tensorflow_federated/python/core/impl/compiler/building_blocks.py

@@ -16,7 +16,7 @@
 import abc
 from collections.abc import Iterable, Iterator
 import enum
-from typing import Any, Optional
+from typing import Optional
 import zlib
 
 from tensorflow_federated.proto.v0 import computation_pb2 as pb
@@ -242,17 +242,14 @@ class Reference(ComputationBuildingBlock):
  """
 
  @classmethod
- def from_proto(
- cls: type['Reference'],
- computation_proto: pb.Computation,
- ) -> 'Reference':
+ def from_proto(cls, computation_proto: pb.Computation) -> 'Reference':
  _check_computation_oneof(computation_proto, 'reference')
  return cls(
  str(computation_proto.reference.name),
  type_serialization.deserialize_type(computation_proto.type),
  )
 
- def __init__(self, name, type_spec, context=None):
+ def __init__(self, name: str, type_spec: object, context=None):
  """Creates a reference to 'name' of type 'type_spec' in context 'context'.
 
  Args:
@@ -282,7 +279,7 @@ def children(self) -> Iterator[ComputationBuildingBlock]:
  return iter(())
 
  @property
- def name(self):
+ def name(self) -> str:
  return self._name
 
  @property
@@ -319,7 +316,12 @@ def from_proto(
  )
  return cls(selection, index=computation_proto.selection.index)
 
- def __init__(self, source, name=None, index=None):
+ def __init__(
+ self,
+ source: ComputationBuildingBlock,
+ name: Optional[str] = None,
+ index: Optional[int] = None,
+ ):
  """A selection from 'source' by a string or numeric 'name_or_index'.
 
  Exactly one of 'name' or 'index' must be specified (not None).
@@ -435,15 +437,19 @@ def from_proto(
  computation_proto: pb.Computation,
  ) -> 'Struct':
  _check_computation_oneof(computation_proto, 'struct')
- return cls(
- [
- (
- str(e.name) if e.name else None,
- ComputationBuildingBlock.from_proto(e.value),
- )
- for e in computation_proto.struct.element
- ]
- )
+
+ def _named_element(
+ proto: pb.Struct.Element,
+ ) -> tuple[Optional[str], ComputationBuildingBlock]:
+ if proto.name:
+ name = str(proto.name)
+ else:
+ name = None
+ element = ComputationBuildingBlock.from_proto(proto.value)
+ return (name, element)
+
+ elements = [_named_element(x) for x in computation_proto.struct.element]
+ return cls(elements)
 
  def __init__(self, elements, container_type=None):
  """Constructs a struct from the given list of elements.
@@ -552,7 +558,11 @@ def from_proto(
  arg = None
  return cls(fn, arg)
 
- def __init__(self, fn, arg=None):
+ def __init__(
+ self,
+ fn: ComputationBuildingBlock,
+ arg: Optional[ComputationBuildingBlock] = None,
+ ):
  """Creates a call to 'fn' with argument 'arg'.
 
  Args:
@@ -612,11 +622,11 @@ def children(self) -> Iterator[ComputationBuildingBlock]:
  yield self._argument
 
  @property
- def function(self):
+ def function(self) -> ComputationBuildingBlock:
  return self._function
 
  @property
- def argument(self):
+ def argument(self) -> Optional[ComputationBuildingBlock]:
  return self._argument
 
  def __repr__(self):
@@ -636,10 +646,7 @@ class Lambda(ComputationBuildingBlock):
  """
 
  @classmethod
- def from_proto(
- cls: type['Lambda'],
- computation_proto: pb.Computation,
- ) -> 'Lambda':
+ def from_proto(cls, computation_proto: pb.Computation) -> 'Lambda':
  _check_computation_oneof(computation_proto, 'lambda')
  the_lambda = getattr(computation_proto, 'lambda')
  if computation_proto.type.function.HasField('parameter'):
@@ -658,7 +665,7 @@ def from_proto(
  def __init__(
  self,
  parameter_name: Optional[str],
- parameter_type: Optional[Any],
+ parameter_type: Optional[object],
  result: ComputationBuildingBlock,
  ):
  """Creates a lambda expression.
@@ -780,10 +787,7 @@ class Block(ComputationBuildingBlock):
  """
 
  @classmethod
- def from_proto(
- cls: type['Block'],
- computation_proto: pb.Computation,
- ) -> 'Block':
+ def from_proto(cls, computation_proto: pb.Computation) -> 'Block':
  _check_computation_oneof(computation_proto, 'block')
  return cls(
  [
@@ -881,10 +885,7 @@ class Intrinsic(ComputationBuildingBlock):
  """
 
  @classmethod
- def from_proto(
- cls: type['Intrinsic'],
- computation_proto: pb.Computation,
- ) -> 'Intrinsic':
+ def from_proto(cls, computation_proto: pb.Computation) -> 'Intrinsic':
  _check_computation_oneof(computation_proto, 'intrinsic')
  return cls(
  computation_proto.intrinsic.uri,
@@ -952,17 +953,14 @@ class Data(ComputationBuildingBlock):
  """
 
  @classmethod
- def from_proto(
- cls: type['Data'],
- computation_proto: pb.Computation,
- ) -> 'Data':
+ def from_proto(cls, computation_proto: pb.Computation) -> 'Data':
  _check_computation_oneof(computation_proto, 'data')
  return cls(
  computation_proto.data.uri,
  type_serialization.deserialize_type(computation_proto.type),
  )
 
- def __init__(self, uri: str, type_spec: Any):
+ def __init__(self, uri: str, type_spec: object):
  """Creates a representation of data.
 
  Args:
@@ -1077,10 +1075,7 @@ class Placement(ComputationBuildingBlock):
  """
 
  @classmethod
- def from_proto(
- cls: type['Placement'],
- computation_proto: pb.Computation,
- ) -> 'Placement':
+ def from_proto(cls, computation_proto: pb.Computation) -> 'Placement':
  _check_computation_oneof(computation_proto, 'placement')
  return cls(
  placements.uri_to_placement_literal(
@@ -1423,7 +1418,6 @@ def _fit_with_inset(left, right, inset):
  for left_line, right_line in zip(left, right):
  if inset > 0:
  left_inset = 0
- right_inset = 0
  trailing_padding = _get_trailing_padding(left_line)
  if trailing_padding > 0:
  left_inset = min(trailing_padding, inset)